From 26856479f908ce4dbc59329ede5ce33c6240c0da Mon Sep 17 00:00:00 2001
From: liyin <liyin@xiaomi.com>
Date: Mon, 5 Aug 2019 18:06:32 +0800
Subject: [PATCH] Add softmax

---
 tools/image/tensor_to_image.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tools/image/tensor_to_image.py b/tools/image/tensor_to_image.py
index e04dde7e..a4d7cda7 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,
-- 
GitLab