Merge pull request #1167 from lukeyeager/fix-segmentation-extensions

Fixes for segmentation extensions
NVIDIA · Oct 14, 2016 · 48ba2ab · 48ba2ab
2 parents e1d31a4 + d546fe8
commit 48ba2ab
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/digits/extensions/data/imageSegmentation/data.py b/digits/extensions/data/imageSegmentation/data.py
@@ -35,6 +35,9 @@ def __init__(self, **kwargs):
             # choose random seed and add to userdata so it gets persisted
             self.userdata['seed'] = random.randint(0, 1000)
 
+        if 'colormap_method' not in self.userdata:
+            self.userdata['colormap_method'] = 'label'
+
         random.seed(self.userdata['seed'])
 
         if self.userdata['colormap_method'] == "label":

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