From 639cbcad0db0bfad2115f8d1e86906e9fd2473d7 Mon Sep 17 00:00:00 2001 From: greatx Date: Fri, 25 Nov 2022 10:14:06 +0800 Subject: [PATCH] update export data format: yolo polygon format --- README.md | 5 + src/labelme2yolo/__about__.py | 2 +- src/labelme2yolo/l2y.py | 213 +++++++++++++++++----------------- 3 files changed, 111 insertions(+), 109 deletions(-) diff --git a/README.md b/README.md index 0792faf..87dc795 100644 --- a/README.md +++ b/README.md @@ -8,6 +8,11 @@ Help converting LabelMe Annotation Tool JSON format to YOLO text file format. If you've already marked your segmentation dataset by LabelMe, it's easy to use this tool to help converting to YOLO format dataset. +--------- +## New + +- export data as yolo polygon annotation (for YOLOv5 v7.0 segmentation) + ## Parameters Explain **--json_dir** LabelMe JSON files folder path. diff --git a/src/labelme2yolo/__about__.py b/src/labelme2yolo/__about__.py index 4f7bd11..a64614b 100644 --- a/src/labelme2yolo/__about__.py +++ b/src/labelme2yolo/__about__.py @@ -1,4 +1,4 @@ # SPDX-FileCopyrightText: 2022-present Wang Xin # # SPDX-License-Identifier: MIT -__version__ = '0.0.1' +__version__ = '0.0.5' diff --git a/src/labelme2yolo/l2y.py b/src/labelme2yolo/l2y.py index 5d823a1..40797f9 100644 --- a/src/labelme2yolo/l2y.py +++ b/src/labelme2yolo/l2y.py @@ -123,214 +123,211 @@ def apply_exif_orientation(image): else: return image + class Labelme2YOLO(object): - + def __init__(self, json_dir): self._json_dir = json_dir - + self._label_id_map = self._get_label_id_map(self._json_dir) - + def _make_train_val_dir(self): - self._label_dir_path = os.path.join(self._json_dir, + self._label_dir_path = os.path.join(self._json_dir, 'YOLODataset/labels/') - self._image_dir_path = os.path.join(self._json_dir, + self._image_dir_path = os.path.join(self._json_dir, 'YOLODataset/images/') - - for yolo_path in (os.path.join(self._label_dir_path + 'train/'), + + for yolo_path in (os.path.join(self._label_dir_path + 'train/'), os.path.join(self._label_dir_path + 'val/'), os.path.join(self._label_dir_path + 'test/'), - os.path.join(self._image_dir_path + 'train/'), + os.path.join(self._image_dir_path + 'train/'), os.path.join(self._image_dir_path + 'val/'), os.path.join(self._image_dir_path + 'test/')): if os.path.exists(yolo_path): shutil.rmtree(yolo_path) - - os.makedirs(yolo_path) - + + os.makedirs(yolo_path) + 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) \ + + 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/') - train_json_names = [train_sample_name + '.json' \ - for train_sample_name in os.listdir(train_folder) \ + train_json_names = [train_sample_name + '.json' + for train_sample_name in os.listdir(train_folder) if os.path.isdir(os.path.join(train_folder, train_sample_name))] - + val_folder = os.path.join(self._json_dir, 'val/') - val_json_names = [val_sample_name + '.json' \ - for val_sample_name in os.listdir(val_folder) \ + val_json_names = [val_sample_name + '.json' + for val_sample_name in os.listdir(val_folder) if os.path.isdir(os.path.join(val_folder, val_sample_name))] - + test_folder = os.path.join(self._json_dir, 'test/') - test_json_names = [test_sample_name + '.json' \ - for test_sample_name in os.listdir(test_folder) \ - if os.path.isdir(os.path.join(test_folder, test_sample_name))] - + test_json_names = [test_sample_name + '.json' + for test_sample_name in os.listdir(test_folder) + if os.path.isdir(os.path.join(test_folder, test_sample_name))] + return train_json_names, val_json_names, test_json_names - - train_idxs, val_idxs = train_test_split(range(len(json_names)), + + train_idxs, val_idxs = train_test_split(range(len(json_names)), test_size=val_size) tmp_train_len = len(train_idxs) test_idxs = [] if test_size > 1e-8: - train_idxs, test_idxs = train_test_split(range(tmp_train_len), test_size=test_size / (1 - val_size)) + train_idxs, test_idxs = train_test_split( + range(tmp_train_len), test_size=test_size / (1 - val_size)) train_json_names = [json_names[train_idx] for train_idx in train_idxs] val_json_names = [json_names[val_idx] for val_idx in val_idxs] test_json_names = [json_names[test_idx] for test_idx in test_idxs] - + return train_json_names, val_json_names, test_json_names - + def convert(self, val_size, test_size): - json_names = [file_name for file_name in os.listdir(self._json_dir) \ - if os.path.isfile(os.path.join(self._json_dir, file_name)) and \ + json_names = [file_name for file_name in os.listdir(self._json_dir) + if os.path.isfile(os.path.join(self._json_dir, file_name)) and file_name.endswith('.json')] - folders = [file_name for file_name in os.listdir(self._json_dir) \ - if os.path.isdir(os.path.join(self._json_dir, file_name))] - train_json_names, val_json_names, test_json_names = self._train_test_split(folders, json_names, val_size, test_size) - + folders = [file_name for file_name in os.listdir(self._json_dir) + if os.path.isdir(os.path.join(self._json_dir, file_name))] + train_json_names, val_json_names, test_json_names = self._train_test_split( + folders, json_names, val_size, test_size) + self._make_train_val_dir() - + # convert labelme object to yolo format object, and save them to files # also get image from labelme json file and save them under images folder - for target_dir, json_names in zip(('train/', 'val/', 'test/'), + for target_dir, json_names in zip(('train/', 'val/', 'test/'), (train_json_names, val_json_names, test_json_names)): pool = Pool(os.cpu_count() - 1) for json_name in json_names: - pool.apply_async(self.covert_json_to_text, args=(target_dir, json_name)) + pool.apply_async(self.covert_json_to_text, + args=(target_dir, json_name)) pool.close() pool.join() - + print('Generating dataset.yaml file ...') self._save_dataset_yaml() def covert_json_to_text(self, target_dir, json_name): json_path = os.path.join(self._json_dir, json_name) json_data = json.load(open(json_path)) - - 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) - + + 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) 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) - + self._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) json_data = json.load(open(json_path)) - + print('Converting %s ...' % json_name) - - img_path = self._save_yolo_image(json_data, json_name, + + img_path = self._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, + self._save_yolo_label(json_name, self._json_dir, '', yolo_obj_list) - + def _get_yolo_object_list(self, json_data, img_path): yolo_obj_list = [] - + img_h, img_w, _ = cv2.imread(img_path).shape for shape in json_data['shapes']: # labelme circle shape is different from others # it only has 2 points, 1st is circle center, 2nd is drag end point if shape['shape_type'] == 'circle': - yolo_obj = self._get_circle_shape_yolo_object(shape, img_h, img_w) + yolo_obj = self._get_circle_shape_yolo_object( + shape, img_h, img_w) else: - yolo_obj = self._get_other_shape_yolo_object(shape, img_h, img_w) - + yolo_obj = self._get_other_shape_yolo_object( + shape, img_h, img_w) + yolo_obj_list.append(yolo_obj) - + return yolo_obj_list - + def _get_circle_shape_yolo_object(self, shape, img_h, img_w): obj_center_x, obj_center_y = shape['points'][0] - - radius = math.sqrt((obj_center_x - shape['points'][1][0]) ** 2 + + + radius = math.sqrt((obj_center_x - shape['points'][1][0]) ** 2 + (obj_center_y - shape['points'][1][1]) ** 2) obj_w = 2 * radius obj_h = 2 * radius - - yolo_center_x= round(float(obj_center_x / img_w), 6) + + yolo_center_x = round(float(obj_center_x / img_w), 6) yolo_center_y = round(float(obj_center_y / img_h), 6) yolo_w = round(float(obj_w / img_w), 6) yolo_h = round(float(obj_h / img_h), 6) - + label_id = self._label_id_map[shape['label']] - + return label_id, yolo_center_x, yolo_center_y, yolo_w, yolo_h - + def _get_other_shape_yolo_object(self, shape, img_h, img_w): - def __get_object_desc(obj_port_list): - __get_dist = lambda int_list: max(int_list) - min(int_list) - - x_lists = [port[0] for port in obj_port_list] - y_lists = [port[1] for port in obj_port_list] - - return min(x_lists), __get_dist(x_lists), min(y_lists), __get_dist(y_lists) - - obj_x_min, obj_w, obj_y_min, obj_h = __get_object_desc(shape['points']) - - yolo_center_x= round(float((obj_x_min + obj_w / 2.0) / img_w), 6) - yolo_center_y = round(float((obj_y_min + obj_h / 2.0) / img_h), 6) - yolo_w = round(float(obj_w / img_w), 6) - yolo_h = round(float(obj_h / img_h), 6) - + + point_list = shape['points'] + points = np.zeros(2 * len(point_list)) + 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] label_id = self._label_id_map[shape['label']] - - return label_id, yolo_center_x, yolo_center_y, yolo_w, yolo_h - + + return (label_id, points.tolist()) + 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, + 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) - + 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) - + 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') - + yaml_path = os.path.join( + self._json_dir, 'YOLODataset/', 'dataset.yaml') + with open(yaml_path, 'w+') as yaml_file: - yaml_file.write('train: %s\n' % \ + yaml_file.write('train: %s\n' % os.path.join(self._image_dir_path, 'train/')) - yaml_file.write('val: %s\n\n' % \ + yaml_file.write('val: %s\n' % os.path.join(self._image_dir_path, 'val/')) - yaml_file.write('test: %s\n\n' % \ + yaml_file.write('test: %s\n' % os.path.join(self._image_dir_path, 'test/')) - yaml_file.write('nc: %i\n\n' % len(self._label_id_map)) - + yaml_file.write('nc: %i\n' % len(self._label_id_map)) + names_str = '' for label, _ in self._label_id_map.items(): names_str += "'%s', " % label