Bläddra i källkod

Add flag --no-vis to semantic_segmentation/labelme2voc.py

BenjaminMueggenburg 5 år sedan
förälder
incheckning
5b24ff3888

+ 3 - 1
examples/semantic_segmentation/README.md

@@ -16,7 +16,9 @@ labelme data_annotated --labels labels.txt --nodata
 #   - data_dataset_voc/JPEGImages
 #   - data_dataset_voc/SegmentationClass
 #   - data_dataset_voc/SegmentationClassVisualization
-./labelme2voc.py data_annotated data_dataset_voc --labels labels.txt
+# Supports flags:
+#   --no-vis  -  No visualization results produced (SegementationClassVisualization)
+python ./labelme2voc.py data_annotated data_dataset_voc --labels labels.txt
 ```
 
 <img src="data_dataset_voc/JPEGImages/2011_000003.jpg" width="33%" /> <img src="data_dataset_voc/SegmentationClassPNG/2011_000003.png" width="33%" /> <img src="data_dataset_voc/SegmentationClassVisualization/2011_000003.jpg" width="33%" />

+ 14 - 9
examples/semantic_segmentation/labelme2voc.py

@@ -22,6 +22,8 @@ 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('--no-vis', help='no visualization results to be created', action='store_false')
+
     args = parser.parse_args()
 
     if osp.exists(args.output_dir):
@@ -31,7 +33,8 @@ def main():
     os.makedirs(osp.join(args.output_dir, 'JPEGImages'))
     os.makedirs(osp.join(args.output_dir, 'SegmentationClass'))
     os.makedirs(osp.join(args.output_dir, 'SegmentationClassPNG'))
-    os.makedirs(osp.join(args.output_dir, 'SegmentationClassVisualization'))
+    if not args.no_vis:
+        os.makedirs(osp.join(args.output_dir, 'SegmentationClassVisualization'))
     print('Creating dataset:', args.output_dir)
 
     class_names = []
@@ -65,11 +68,12 @@ def main():
                 args.output_dir, 'SegmentationClass', base + '.npy')
             out_png_file = osp.join(
                 args.output_dir, 'SegmentationClassPNG', base + '.png')
-            out_viz_file = osp.join(
-                args.output_dir,
-                'SegmentationClassVisualization',
-                base + '.jpg',
-            )
+            if not args.no_vis:
+                out_viz_file = osp.join(
+                    args.output_dir,
+                    'SegmentationClassVisualization',
+                    base + '.jpg',
+                )
 
             data = json.load(f)
 
@@ -86,9 +90,10 @@ def main():
 
             np.save(out_lbl_file, lbl)
 
-            viz = labelme.utils.draw_label(
-                lbl, img, class_names, colormap=colormap)
-            PIL.Image.fromarray(viz).save(out_viz_file)
+            if not args.no_vis:
+                viz = labelme.utils.draw_label(
+                    lbl, img, class_names, colormap=colormap)
+                PIL.Image.fromarray(viz).save(out_viz_file)
 
 
 if __name__ == '__main__':