Skip visualization for images with no annotation in labelme2coco.py

anmsajedulalam · Oct 15, 2021 · c7ba122 · c7ba122
1 parent 4aa785f
commit c7ba122
Showing 1 changed file with 17 additions and 15 deletions.
diff --git a/examples/instance_segmentation/labelme2coco.py b/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"
  )