From 38882c72d5ecdb20d7034d26a44daa884afca578 Mon Sep 17 00:00:00 2001 From: Kentaro Wada Date: Thu, 14 Sep 2023 09:52:36 +0900 Subject: [PATCH] Support --labels to be comma separated text --- examples/instance_segmentation/labelme2voc.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/examples/instance_segmentation/labelme2voc.py b/examples/instance_segmentation/labelme2voc.py index aed97f751..ddf9539a3 100755 --- a/examples/instance_segmentation/labelme2voc.py +++ b/examples/instance_segmentation/labelme2voc.py @@ -20,7 +20,9 @@ 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( + "--labels", help="Labels file or comma separated text", required=True + ) parser.add_argument( "--noobject", help="Flag not to generate object label", action="store_true" ) @@ -50,11 +52,17 @@ def main(): os.makedirs(osp.join(args.output_dir, "SegmentationObjectVisualization")) print("Creating dataset:", args.output_dir) + if osp.exists(args.labels): + with open(args.labels) as f: + labels = [label.strip() for label in f if label] + else: + labels = [label.strip() for label in args.labels.split(",")] + class_names = [] class_name_to_id = {} - for i, line in enumerate(open(args.labels).readlines()): + for i, label in enumerate(labels): class_id = i - 1 # starts with -1 - class_name = line.strip() + class_name = label.strip() class_name_to_id[class_name] = class_id if class_id == -1: assert class_name == "__ignore__"