Skip label_name which is None

wkentaro · Apr 29, 2020 · 86eab26 · 86eab26
1 parent 17f3ab9
commit 86eab26
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 11 deletions.
diff --git a/examples/label2rgb.py b/examples/label2rgb.py
@@ -18,7 +18,7 @@ def label2rgb():
         label, label_names=label_names, font_size=25
     )
     labelviz_withname2 = imgviz.label2rgb(
-        label, label_names=label_names, font_size=25, loc="lt"
+        label, label_names=label_names, font_size=25, loc="rb"
     )
     img = imgviz.color.rgb2gray(rgb)
     labelviz_withimg = imgviz.label2rgb(label, img=img)
@@ -38,7 +38,7 @@ def label2rgb():
     plt.axis("off")
 
     plt.subplot(133)
-    plt.title("loc=lt")
+    plt.title("loc=rb")
     plt.imshow(labelviz_withname2)
     plt.axis("off")
 

diff --git a/imgviz/label.py b/imgviz/label.py
@@ -109,11 +109,14 @@ def label2rgb(
     if label_names is None:
         return res
 
-    if loc == "centroid":
-        for label_i in np.unique(label):
-            if label_i == -1:
-                continue  # unlabeled
+    unique_labels = np.unique(label)
+    unique_labels = unique_labels[unique_labels != -1]
+    unique_labels = [l for l in unique_labels if label_names[l] is not None]
+    if len(unique_labels) == 0:
+        return res
 
+    if loc == "centroid":
+        for label_i in unique_labels:
             mask = label == label_i
             if 1.0 * mask.sum() / mask.size < thresh_suppress:
                 continue
@@ -138,11 +141,6 @@ def label2rgb(
                 font_path=font_path,
             )
     elif loc in ["rb", "lt"]:
-        unique_labels = np.unique(label)
-        unique_labels = unique_labels[unique_labels != -1]
-        if not unique_labels.size:
-            return res
-
         text_sizes = np.array(
             [
                 draw_module.text_size(