Add data_dataset_coco/Visualization

examples/instance_segmentation/labelme2coco.py

@@ -10,8 +10,8 @@
 import sys
 import uuid
 
+import imgviz
 import numpy as np
-import PIL.Image
 
 import labelme
 
@@ -29,13 +29,18 @@ def main():
  parser.add_argument("input_dir", help="input annotated directory")
  parser.add_argument("output_dir", help="output dataset directory")
  parser.add_argument("--labels", help="labels file", required=True)
+ parser.add_argument(
+ "--noviz", help="no visualization", action="store_true"
+ )
  args = parser.parse_args()
 
  if osp.exists(args.output_dir):
  print("Output directory already exists:", args.output_dir)
  sys.exit(1)
  os.makedirs(args.output_dir)
  os.makedirs(osp.join(args.output_dir, "JPEGImages"))
+ if not args.noviz:
+ os.makedirs(osp.join(args.output_dir, "Visualization"))
  print("Creating dataset:", args.output_dir)
 
  now = datetime.datetime.now()
@@ -85,7 +90,7 @@ def main():
  out_img_file = osp.join(args.output_dir, "JPEGImages", base + ".jpg")
 
  img = labelme.utils.img_data_to_arr(label_file.imageData)
- PIL.Image.fromarray(img).convert("RGB").save(out_img_file)
+ imgviz.io.imsave(out_img_file, img)
  data["images"].append(
  dict(
  license=0,
@@ -153,6 +158,27 @@ 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,
+ )
+ out_viz_file = osp.join(
+ args.output_dir, "Visualization", base + ".jpg"
+ )
+ imgviz.io.imsave(out_viz_file, viz)
+
  with open(out_ann_file, "w") as f:
  json.dump(data, f)