Fix visualization when palette is None (NVIDIA#1177)

The palette may be `None`when working with grayscale labels. Fix NVIDIA#1147
ethantang95 · Oct 18, 2016 · ec17b6f · ec17b6f
1 parent 92a613e
commit ec17b6f
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/digits/dataset/generic/views.py b/digits/dataset/generic/views.py
@@ -147,7 +147,8 @@ def explore():
     db_path = job.path(db)
     labels = []
 
-    if COLOR_PALETTE_ATTRIBUTE in job.extension_userdata:
+    if COLOR_PALETTE_ATTRIBUTE in job.extension_userdata \
+            and job.extension_userdata[COLOR_PALETTE_ATTRIBUTE]:
         # assume single-channel 8-bit palette
         palette = job.extension_userdata[COLOR_PALETTE_ATTRIBUTE]
         palette = np.array(palette).reshape((len(palette)/3,3)) / 255.

diff --git a/digits/extensions/view/imageSegmentation/view.py b/digits/extensions/view/imageSegmentation/view.py
@@ -36,8 +36,9 @@ def __init__(self, dataset, **kwargs):
 
         # view options
         if kwargs['colormap'] == 'dataset':
-            if not COLOR_PALETTE_ATTRIBUTE in dataset.extension_userdata:
-                raise ValueError("Palette not found in dataset")
+            if not COLOR_PALETTE_ATTRIBUTE in dataset.extension_userdata or \
+                   not dataset.extension_userdata[COLOR_PALETTE_ATTRIBUTE]:
+                raise ValueError("No palette found in dataset - choose other colormap")
             palette = dataset.extension_userdata[COLOR_PALETTE_ATTRIBUTE]
             # assume 8-bit RGB palette and convert to N*3 numpy array
             palette = np.array(palette).reshape((len(palette)/3,3)) / 255.