Segmentation - force image into [0-255] before viz

NVIDIA · Oct 14, 2016 · d546fe8 · d546fe8
1 parent 1cb8310
commit d546fe8
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/digits/extensions/view/imageSegmentation/view.py b/digits/extensions/view/imageSegmentation/view.py
@@ -206,7 +206,14 @@ def process_data(self, input_id, input_data, output_data):
                                        fill_data[:, :, x] * line_mask)
 
         # Input image with outlines
-        input_image = PIL.Image.fromarray(input_data)
+        input_max = input_data.max()
+        input_min = input_data.min()
+        input_range = input_max - input_min
+        if input_range > 255:
+            input_data = (input_data - input_min) * 255.0 / input_range
+        elif input_min < 0:
+            input_data -= input_min
+        input_image = PIL.Image.fromarray(input_data.astype('uint8'))
         input_image.format = 'png'
 
         # Fill image