-
-
Save jeasinema/e58d405439bfab9ce9c37540c2470488 to your computer and use it in GitHub Desktop.
A utility function for TensorFlow that maps a grayscale image to a matplotlib colormap for use with TensorBoard image summaries.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#!/usr/bin/env python | |
# -*- coding:UTF-8 -*- | |
# File Name : colorize.py | |
# Purpose : | |
# Creation Date : 21-12-2017 | |
# Created By : Jeasine Ma [jeasinema[at]gmail[dot]com] | |
# ref: https://gist.github.com/jimfleming/c1adfdb0f526465c99409cc143dea97b | |
import matplotlib | |
import matplotlib.cm | |
import cv2 | |
import numpy as np | |
def colorize(value, factor=1, vmin=None, vmax=None): | |
""" | |
A utility function for TensorFlow that maps a grayscale image to a matplotlib | |
colormap for use with TensorBoard image summaries. | |
By default it will normalize the input value to the range 0..1 before mapping | |
to a grayscale colormap. | |
Arguments: | |
- value: 2D Tensor of shape [height, width] or 3D Tensor of shape | |
[height, width, 1]. | |
- factor: resize factor, scalar | |
- vmin: the minimum value of the range used for normalization. | |
(Default: value minimum) | |
- vmax: the maximum value of the range used for normalization. | |
(Default: value maximum) | |
Example usage: | |
``` | |
output = tf.random_uniform(shape=[256, 256, 1]) | |
output_color = colorize(output, vmin=0.0, vmax=1.0, cmap='viridis') | |
tf.summary.image('output', output_color) | |
``` | |
Returns a 3D tensor of shape [height, width, 3]. | |
""" | |
# normalize | |
value = np.sum(value, axis=-1) | |
vmin = np.min(value) if vmin is None else vmin | |
vmax = np.max(value) if vmax is None else vmax | |
value = (value - vmin) / (vmax - vmin) # vmin..vmax | |
value = (value * 255).astype(np.uint8) | |
value = cv2.applyColorMap(value, cv2.COLORMAP_PARULA) # or cv2.COLORMAP_JET | |
value = cv2.cvtColor(value, cv2.COLOR_BGR2RGB) | |
x, y, _ = value.shape | |
value = cv2.resize(value, (y * factor, x * factor)) | |
return value | |
def colorize_seg(value, factor=1, color='viridis'): | |
""" | |
Arguments: | |
- value: 2D Tensor of shape [height, width] or 3D Tensor of shape | |
[height, width, 1]., with discrete value, like segmentation gt mask | |
- factor: resize factor, scalar | |
- color: name of color map, i.e. viridis, jet, plasma | |
Returns a 3D tensor of shape [height, width, 3]. | |
""" | |
# normalize | |
n_max, n_min = np.max(value), np.min(value) | |
value = (value - n_min)/(n_max-n_min) | |
cmap = cm.get_cmap(color) | |
value = cmap(value)*255. | |
x, y, _ = value.shape | |
value = cv2.resize(value, (y * factor, x * factor)) | |
return value[..., :3] | |
def tf_colorize(value, factor=1, vmin=None, vmax=None, cmap=None): | |
""" | |
A utility function for TensorFlow that maps a grayscale image to a matplotlib | |
colormap for use with TensorBoard image summaries. | |
By default it will normalize the input value to the range 0..1 before mapping | |
to a grayscale colormap. | |
Arguments: | |
- value: 2D Tensor of shape [height, width] or 3D Tensor of shape | |
[height, width, 1]. | |
- factor: resize factor, scalar | |
- vmin: the minimum value of the range used for normalization. | |
(Default: value minimum) | |
- vmax: the maximum value of the range used for normalization. | |
(Default: value maximum) | |
- cmap: a valid cmap named for use with matplotlib's `get_cmap`. | |
(Default: 'gray') | |
Example usage: | |
``` | |
output = tf.random_uniform(shape=[256, 256, 1]) | |
output_color = colorize(output, vmin=0.0, vmax=1.0, cmap='viridis') | |
tf.summary.image('output', output_color) | |
``` | |
Returns a 3D tensor of shape [height, width, 3]. | |
""" | |
import tensorflow as tf | |
# normalize | |
vmin = tf.reduce_min(value) if vmin is None else vmin | |
vmax = tf.reduce_max(value) if vmax is None else vmax | |
value = (value - vmin) / (vmax - vmin) # vmin..vmax | |
# squeeze last dim if it exists | |
value = tf.squeeze(value) | |
# quantize | |
indices = tf.to_int32(tf.round(value * 255)) | |
# gather | |
cm = matplotlib.cm.get_cmap(cmap if cmap is not None else 'gray') | |
colors = tf.constant(cm.colors, dtype=tf.float32) | |
value = tf.gather(colors, indices) | |
return value |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment