diff --git a/tools/image/tensor_to_image.py b/tools/image/tensor_to_image.py
index e04dde7e6bf1b1c9564f4198122a1fa5cf4db015..a4d7cda7a79497d38e3c3612fa5cde0a2760bc02 100644
--- a/tools/image/tensor_to_image.py
+++ b/tools/image/tensor_to_image.py
@@ -37,6 +37,7 @@ def tensors_to_images(input_files, image_shape):
         output = tf.placeholder(tf.string, name='output_file')
         # use the second channel if it is gray image
         if image_shape[2] == 2:
+            input = tf.nn.softmax(input)
             _, input = tf.split(input, 2, axis=2)
         tensor_data = tf.image.convert_image_dtype(input,
                                                    tf.uint8,