Plot only large objects

label_prettify.py

@@ -71,14 +71,15 @@ def seg_prettify(rgb_path, file_name):
 
     for i in np.unique(seg_labels):
         obj_idx_all = np.where(seg_labels == i)
-        obj_idx = random.randint(0, len(obj_idx_all[0]) - 1)
-        x, y = obj_idx_all[1][obj_idx], obj_idx_all[0][obj_idx]
-        obj_name = coco_label_map[int(i * 255)]
-        obj_name = obj_name.split(',')[0]
-        if islight(seg_map[int(y), int(x)]):
-            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(obj_idx_all[0]) > 20:  # only plot the label with its number of labelled pixel more than 20
+            obj_idx = random.randint(0,  len(obj_idx_all[0]) - 1)
+            x, y = obj_idx_all[1][obj_idx], obj_idx_all[0][obj_idx]
+            obj_name = coco_label_map[int(i * 255)]
+            obj_name = obj_name.split(',')[0]
+            if islight(seg_map[int(y), int(x)]):
+                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(file_name, bbox_inches='tight', transparent=True, pad_inches=0)