Fix no obj error

label_prettify.py

@@ -34,22 +34,27 @@ def obj_detection_prettify(rgb_path, path_name):
 
     plt.imshow(rgb)
 
-    num_objs = np.unique(obj_labels)[:-1].max()
-    plt.imshow(obj_labels, cmap='terrain', vmax=num_objs + 1 / 255., alpha=0.8)
-    cmap = matplotlib.colormaps.get_cmap('terrain')
-    for i in np.unique(obj_labels)[:-1]:
-        obj_idx_all = np.where(obj_labels == i)
-        x, y = obj_idx_all[1].mean(), obj_idx_all[0].mean()
-        obj_name = obj_label_map[obj_labels_dict[str(int(i * 255))]]
-        obj_name = obj_name.split(',')[0]
-        if islight([c*255 for c in cmap(i / num_objs)[:3]]):
-            plt.text(x, y, obj_name, c='black', horizontalalignment='center', verticalalignment='center', clip_on=True)
-        else:
-            plt.text(x, y, obj_name, c='white', horizontalalignment='center', verticalalignment='center', clip_on=True)
+    if len(np.unique(obj_labels)) == 1:
+        plt.axis('off')
+        plt.savefig(path_name, bbox_inches='tight', transparent=True, pad_inches=0)
+        plt.close()
+    else:
+        num_objs = np.unique(obj_labels)[:-1].max()
+        plt.imshow(obj_labels, cmap='terrain', vmax=num_objs + 1 / 255., alpha=0.8)
+        cmap = matplotlib.colormaps.get_cmap('terrain')
+        for i in np.unique(obj_labels)[:-1]:
+            obj_idx_all = np.where(obj_labels == i)
+            x, y = obj_idx_all[1].mean(), obj_idx_all[0].mean()
+            obj_name = obj_label_map[obj_labels_dict[str(int(i * 255))]]
+            obj_name = obj_name.split(',')[0]
+            if islight([c*255 for c in cmap(i / num_objs)[:3]]):
+                plt.text(x, y, obj_name, c='black', horizontalalignment='center', verticalalignment='center', clip_on=True)
+            else:
+                plt.text(x, y, obj_name, c='white', horizontalalignment='center', verticalalignment='center', clip_on=True)
 
-    plt.axis('off')
-    plt.savefig(path_name, bbox_inches='tight', transparent=True, pad_inches=0)
-    plt.close()
+        plt.axis('off')
+        plt.savefig(path_name, bbox_inches='tight', transparent=True, pad_inches=0)
+        plt.close()
 
 
 def seg_prettify(rgb_path, file_name):