diff --git a/.gitignore b/.gitignore index 2dc53ca..89d4981 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,5 @@ +dataset/ + # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] diff --git a/README.md b/README.md index 23fcba4..d377d99 100644 --- a/README.md +++ b/README.md @@ -14,7 +14,7 @@ If you've already marked your segmentation dataset by LabelMe, it's easy to use ## New - export data as yolo polygon annotation (for YOLOv5 v7.0 segmentation) - +- Now you can choose the output format of the label text. The available options are `plygon` and `bbox`. ## Installation ```console @@ -30,6 +30,8 @@ pip install labelme2yolo **--json_name (Optional)** Convert single LabelMe JSON file. +**--output_format (Optional)** The output format of label. + ## How to Use ### 1. Convert JSON files, split training, validation and test dataset by --val_size and --test_size diff --git a/src/labelme2yolo/__about__.py b/src/labelme2yolo/__about__.py index 6de232a..c4ee3a4 100644 --- a/src/labelme2yolo/__about__.py +++ b/src/labelme2yolo/__about__.py @@ -1,5 +1,5 @@ -# SPDX-FileCopyrightText: 2022-present Wang Xin +# SPDX-FileCopyrightText: 2023-present Wang Xin # # SPDX-License-Identifier: MIT -__version__ = '0.0.8' +__version__ = '0.0.9' diff --git a/src/labelme2yolo/cli/__init__.py b/src/labelme2yolo/cli/__init__.py index 64f96e8..85a9d9b 100644 --- a/src/labelme2yolo/cli/__init__.py +++ b/src/labelme2yolo/cli/__init__.py @@ -16,14 +16,14 @@ def run(): type=float, nargs="?", default=None, - help="Please input the validation dataset size, for example 0.1 ", + help="Please input the validation dataset size, for example 0.1.", ) parser.add_argument( "--test_size", type=float, nargs="?", default=None, - help="Please input the validation dataset size, for example 0.1 ", + help="Please input the test dataset size, for example 0.1.", ) parser.add_argument( "--json_name", @@ -32,13 +32,21 @@ def run(): default=None, help="If you put json name, it would convert only one json file to YOLO.", ) + parser.add_argument( + "--output_format", + type=str, + default="polygon", + help='The default output format for labelme2yolo is "polygon".' + ' However, you can choose to output in bbox format by specifying the "bbox" option.', + ) + args = parser.parse_args() if not args.json_dir: parser.print_help() return 0 - convertor = Labelme2YOLO(args.json_dir) + convertor = Labelme2YOLO(args.json_dir, args.output_format) if args.json_name is None: convertor.convert(val_size=args.val_size, test_size=args.test_size) diff --git a/src/labelme2yolo/l2y.py b/src/labelme2yolo/l2y.py index f677715..d4b6aff 100644 --- a/src/labelme2yolo/l2y.py +++ b/src/labelme2yolo/l2y.py @@ -93,27 +93,28 @@ def get_label_id_map(json_dir): return OrderedDict([(label, label_id) for label_id, label in enumerate(label_set)]) -def extend_point_list(point_list): +def extend_point_list(point_list, format="polygon"): xmin = min([float(point) for point in point_list[::2]]) xmax = max([float(point) for point in point_list[::2]]) ymin = min([float(point) for point in point_list[1::2]]) ymax = max([float(point) for point in point_list[1::2]]) - return np.array([xmin, ymin, xmax, ymin, xmax, ymax, xmin, ymax]) + if (format == "polygon"): + return np.array([xmin, ymin, xmax, ymin, xmax, ymax, xmin, ymax]) + if (format == "bbox"): + return np.array([xmin, ymin, xmax - xmin, ymax - ymin]) def save_yolo_label(json_name, label_dir_path, target_dir, yolo_obj_list): - txt_path = os.path.join( - label_dir_path, target_dir, json_name.replace(".json", ".txt") - ) + txt_path = os.path.join(label_dir_path, + target_dir, + json_name.replace(".json", ".txt")) with open(txt_path, "w+") as f: - for yolo_obj_idx, yolo_obj in enumerate(yolo_obj_list): - yolo_obj_line = ( - "%s %s %s %s %s\n" % yolo_obj - if yolo_obj_idx + 1 != len(yolo_obj_list) - else "%s %s %s %s %s" % yolo_obj - ) + for yolo_obj in yolo_obj_list: + label, points = yolo_obj + points = [str(item) for item in points] + yolo_obj_line = f"{label} {' '.join(points)}\n" f.write(yolo_obj_line) @@ -130,10 +131,11 @@ def save_yolo_image(json_data, json_name, image_dir_path, target_dir): class Labelme2YOLO(object): - def __init__(self, json_dir): + def __init__(self, json_dir, output_format): self._json_dir = json_dir + self._output_format = output_format - self._label_id_map = self._get_label_id_map(self._json_dir) + self._label_id_map = get_label_id_map(self._json_dir) def _make_train_val_dir(self): self._label_dir_path = os.path.join(self._json_dir, @@ -152,20 +154,6 @@ class Labelme2YOLO(object): os.makedirs(yolo_path) - @staticmethod - def _get_label_id_map(self, json_dir): - label_set = set() - - for file_name in os.listdir(json_dir): - if file_name.endswith('json'): - json_path = os.path.join(json_dir, file_name) - data = json.load(open(json_path)) - for shape in data['shapes']: - label_set.add(shape['label']) - - return OrderedDict([(label, label_id) - for label_id, label in enumerate(label_set)]) - def _train_test_split(self, folders, json_names, val_size, test_size): if len(folders) > 0 and 'train' in folders and 'val' in folders and 'test' in folders: train_folder = os.path.join(self._json_dir, 'train/') @@ -233,15 +221,15 @@ class Labelme2YOLO(object): print('Converting %s for %s ...' % (json_name, target_dir.replace('/', ''))) - img_path = self._save_yolo_image(json_data, - json_name, - self._image_dir_path, - target_dir) + img_path = save_yolo_image(json_data, + json_name, + self._image_dir_path, + target_dir) yolo_obj_list = self._get_yolo_object_list(json_data, img_path) - self._save_yolo_label(json_name, - self._label_dir_path, - target_dir, - yolo_obj_list) + save_yolo_label(json_name, + self._label_dir_path, + target_dir, + yolo_obj_list) def convert_one(self, json_name): json_path = os.path.join(self._json_dir, json_name) @@ -249,12 +237,12 @@ class Labelme2YOLO(object): print('Converting %s ...' % json_name) - img_path = self._save_yolo_image(json_data, json_name, - self._json_dir, '') + img_path = save_yolo_image(json_data, json_name, + self._json_dir, '') yolo_obj_list = self._get_yolo_object_list(json_data, img_path) - self._save_yolo_label(json_name, self._json_dir, - '', yolo_obj_list) + save_yolo_label(json_name, self._json_dir, + '', yolo_obj_list) def _get_yolo_object_list(self, json_data, img_path): yolo_obj_list = [] @@ -298,35 +286,14 @@ class Labelme2YOLO(object): points[::2] = [float(point[0]) / img_w for point in point_list] points[1::2] = [float(point[1]) / img_h for point in point_list] if len(points) == 4: - points = extend_point_list(points) + if self._output_format == "polygon": + points = extend_point_list(points) + if self._output_format == "bbox": + points = extend_point_list(points, "bbox") label_id = self._label_id_map[shape['label']] return label_id, points.tolist() - @staticmethod - def _save_yolo_label(self, json_name, label_dir_path, target_dir, yolo_obj_list): - txt_path = os.path.join(label_dir_path, - target_dir, - json_name.replace('.json', '.txt')) - - with open(txt_path, 'w+') as f: - for yolo_obj in yolo_obj_list: - label, points = yolo_obj - points = [str(item) for item in points] - yolo_obj_line = f"{label} {' '.join(points)}\n" - f.write(yolo_obj_line) - - @staticmethod - def _save_yolo_image(self, json_data, json_name, image_dir_path, target_dir): - img_name = json_name.replace('.json', '.png') - img_path = os.path.join(image_dir_path, target_dir, img_name) - - if not os.path.exists(img_path): - img = img_b64_to_arr(json_data['imageData']) - PIL.Image.fromarray(img).save(img_path) - - return img_path - def _save_dataset_yaml(self): yaml_path = os.path.join( self._json_dir, 'YOLODataset/', 'dataset.yaml')