diff --git a/python/tensorboard/summary.py b/python/tensorboard/summary.py index 08cdf4b..d80d253 100644 --- a/python/tensorboard/summary.py +++ b/python/tensorboard/summary.py @@ -36,7 +36,7 @@ import logging import re as _re import bisect -from six import StringIO +from six import BytesIO from six.moves import range from PIL import Image import numpy as np @@ -172,25 +172,33 @@ def image(tag, tensor): if not isinstance(tensor, np.ndarray): # try conversion, if failed then need handle by user. tensor = np.ndarray(tensor, dtype=np.float32) - shape = tensor.shape - height, width, channel = shape[0], shape[1], shape[2] - if channel == 1: - # walk around. PIL's setting on dimension. - tensor = np.reshape(tensor, (height, width)) - image = make_image(tensor, height, width, channel) + if not 1 < tensor.ndim < 4: + raise ValueError("Expecting images to be 2 or 3 dimensional arrays." + "The provided is " + str(tensor.ndim)) + if tensor.dtype in (np.float16, np.float32, np.float64): + tensor = np.asarray(tensor * 255, np.uint8) + if tensor.ndim == 2: + height, width = tensor.shape + else: + height, width, channel = tensor.shape + if channel == 1: + # walk around. PIL's setting on dimension. + tensor = np.reshape(tensor, (height, width)) + image = make_image(tensor, height, width) return Summary(value=[Summary.Value(tag=tag, image=image)]) -def make_image(tensor, height, width, channel): +def make_image(tensor, height, width): """Convert an numpy representation image to Image protobuf""" image = Image.fromarray(tensor) - output = StringIO() + if image.mode != "RGB": + image = image.convert("RGB") + output = BytesIO() image.save(output, format='PNG') image_string = output.getvalue() output.close() return Summary.Image(height=height, width=width, - colorspace=channel, encoded_image_string=image_string)