Quellcode durchsuchen

Fix image orienation in instance_segmentation/labelme2voc.py

Kentaro Wada vor 5 Jahren
Ursprung
Commit
487e14b80b
1 geänderte Dateien mit 65 neuen und 68 gelöschten Zeilen
  1. 65 68
      examples/instance_segmentation/labelme2voc.py

+ 65 - 68
examples/instance_segmentation/labelme2voc.py

@@ -4,14 +4,12 @@ from __future__ import print_function
 
 import argparse
 import glob
-import json
 import os
 import os.path as osp
 import sys
 
 import imgviz
 import numpy as np
-import PIL.Image
 
 import labelme
 
@@ -66,73 +64,72 @@ def main():
         f.writelines('\n'.join(class_names))
     print('Saved class_names:', out_class_names_file)
 
-    for label_file in glob.glob(osp.join(args.input_dir, '*.json')):
-        print('Generating dataset from:', label_file)
-        with open(label_file) as f:
-            base = osp.splitext(osp.basename(label_file))[0]
-            out_img_file = osp.join(
-                args.output_dir, 'JPEGImages', base + '.jpg')
-            out_cls_file = osp.join(
-                args.output_dir, 'SegmentationClass', base + '.npy')
-            out_clsp_file = osp.join(
-                args.output_dir, 'SegmentationClassPNG', base + '.png')
-            if not args.noviz:
-                out_clsv_file = osp.join(
-                    args.output_dir,
-                    'SegmentationClassVisualization',
-                    base + '.jpg',
-                )
-            out_ins_file = osp.join(
-                args.output_dir, 'SegmentationObject', base + '.npy')
-            out_insp_file = osp.join(
-                args.output_dir, 'SegmentationObjectPNG', base + '.png')
-            if not args.noviz:
-                out_insv_file = osp.join(
-                    args.output_dir,
-                    'SegmentationObjectVisualization',
-                    base + '.jpg',
-                )
-
-            data = json.load(f)
-
-            img_file = osp.join(osp.dirname(label_file), data['imagePath'])
-            img = np.asarray(PIL.Image.open(img_file))
-            PIL.Image.fromarray(img).save(out_img_file)
-
-            cls, ins = labelme.utils.shapes_to_label(
-                img_shape=img.shape,
-                shapes=data['shapes'],
-                label_name_to_value=class_name_to_id,
+    for filename in glob.glob(osp.join(args.input_dir, '*.json')):
+        print('Generating dataset from:', filename)
+
+        label_file = labelme.LabelFile(filename=filename)
+
+        base = osp.splitext(osp.basename(filename))[0]
+        out_img_file = osp.join(
+            args.output_dir, 'JPEGImages', base + '.jpg')
+        out_cls_file = osp.join(
+            args.output_dir, 'SegmentationClass', base + '.npy')
+        out_clsp_file = osp.join(
+            args.output_dir, 'SegmentationClassPNG', base + '.png')
+        if not args.noviz:
+            out_clsv_file = osp.join(
+                args.output_dir,
+                'SegmentationClassVisualization',
+                base + '.jpg',
             )
-            ins[cls == -1] = 0  # ignore it.
-
-            # class label
-            labelme.utils.lblsave(out_clsp_file, cls)
-            np.save(out_cls_file, cls)
-            if not args.noviz:
-                clsv = imgviz.label2rgb(
-                    label=cls,
-                    img=imgviz.rgb2gray(img),
-                    label_names=class_names,
-                    font_size=15,
-                    loc='rb',
-                )
-                imgviz.io.imsave(out_clsv_file, clsv)
-
-            # instance label
-            labelme.utils.lblsave(out_insp_file, ins)
-            np.save(out_ins_file, ins)
-            if not args.noviz:
-                instance_ids = np.unique(ins)
-                instance_names = [str(i) for i in range(max(instance_ids) + 1)]
-                insv = imgviz.label2rgb(
-                    label=ins,
-                    img=imgviz.rgb2gray(img),
-                    label_names=instance_names,
-                    font_size=15,
-                    loc='rb',
-                )
-                imgviz.io.imsave(out_insv_file, insv)
+        out_ins_file = osp.join(
+            args.output_dir, 'SegmentationObject', base + '.npy')
+        out_insp_file = osp.join(
+            args.output_dir, 'SegmentationObjectPNG', base + '.png')
+        if not args.noviz:
+            out_insv_file = osp.join(
+                args.output_dir,
+                'SegmentationObjectVisualization',
+                base + '.jpg',
+            )
+
+        img = labelme.utils.img_data_to_arr(label_file.imageData)
+        imgviz.io.imsave(out_img_file, img)
+
+        cls, ins = labelme.utils.shapes_to_label(
+            img_shape=img.shape,
+            shapes=label_file.shapes,
+            label_name_to_value=class_name_to_id,
+        )
+        ins[cls == -1] = 0  # ignore it.
+
+        # class label
+        labelme.utils.lblsave(out_clsp_file, cls)
+        np.save(out_cls_file, cls)
+        if not args.noviz:
+            clsv = imgviz.label2rgb(
+                label=cls,
+                img=imgviz.rgb2gray(img),
+                label_names=class_names,
+                font_size=15,
+                loc='rb',
+            )
+            imgviz.io.imsave(out_clsv_file, clsv)
+
+        # instance label
+        labelme.utils.lblsave(out_insp_file, ins)
+        np.save(out_ins_file, ins)
+        if not args.noviz:
+            instance_ids = np.unique(ins)
+            instance_names = [str(i) for i in range(max(instance_ids) + 1)]
+            insv = imgviz.label2rgb(
+                label=ins,
+                img=imgviz.rgb2gray(img),
+                label_names=instance_names,
+                font_size=15,
+                loc='rb',
+            )
+            imgviz.io.imsave(out_insv_file, insv)
 
 
 if __name__ == '__main__':