瀏覽代碼

Skip visualization for images with no annotation in labelme2coco.py

Kentaro Wada 3 年之前
父節點
當前提交
c7ba122e1b
共有 1 個文件被更改,包括 17 次插入15 次删除
  1. 17 15
      examples/instance_segmentation/labelme2coco.py

+ 17 - 15
examples/instance_segmentation/labelme2coco.py

@@ -159,21 +159,23 @@ def main():
             )
 
         if not args.noviz:
-            labels, captions, masks = zip(
-                *[
-                    (class_name_to_id[cnm], cnm, msk)
-                    for (cnm, gid), msk in masks.items()
-                    if cnm in class_name_to_id
-                ]
-            )
-            viz = imgviz.instances2rgb(
-                image=img,
-                labels=labels,
-                masks=masks,
-                captions=captions,
-                font_size=15,
-                line_width=2,
-            )
+            viz = img
+            if masks:
+                labels, captions, masks = zip(
+                    *[
+                        (class_name_to_id[cnm], cnm, msk)
+                        for (cnm, gid), msk in masks.items()
+                        if cnm in class_name_to_id
+                    ]
+                )
+                viz = imgviz.instances2rgb(
+                    image=img,
+                    labels=labels,
+                    masks=masks,
+                    captions=captions,
+                    font_size=15,
+                    line_width=2,
+                )
             out_viz_file = osp.join(
                 args.output_dir, "Visualization", base + ".jpg"
             )