CODEWITHSUNDEEP
pythonpytorchtensorboard
CommandoMan
CommandoMan

Reputation: 25

Setting Picture Colormap in Tensorboard with PyTorch

I'm pretty new to using Python, PyTorch and Tensorboard - moved from MATLAB due to the lacking automatic differentiation.

I am trying to use the above stated tools since I'm running an optimization problem - simple gradient descent for reconstructing distorted images. No machine learning or deep learning.

The point is that I need to see the images every few iterations of the algorithm, so I was told that tensorboard would be great for it. The only problem is that these images are shown in grayscale and I need to see it in a different colormap. Is there any way to change the colormap in tensorboard? Thanks!

Upvotes: 1

Views: 968

Answers (1)

Abuzar G
Abuzar G

Reputation: 111

You can colorize your tensor shape using tensorflow gather function. Following is a simple script for doing this. You may use other maps rather than 'Spectral':

import matplotlib
import matplotlib.cm
import tensorflow as tf

def colormap(shape):
    min = tf.reduce_min(shape)
    max = tf.reduce_max(shape)
    shape = (shape - min) / (max - min)
    gatherIndices = tf.to_int32(tf.round(shape * 255))
    map = matplotlib.cm.get_cmap('cividis')
    colors = tf.constant(map.colors, dtype=tf.float32)
    return tf.gather(colors, gatherIndices)

Upvotes: 1

Related Questions