From 53b3dacc31e6eb35d16288e4ebe567222eeb2591 Mon Sep 17 00:00:00 2001 From: rooneysh Date: Thu, 19 Aug 2021 15:11:14 +0800 Subject: [PATCH] Update labelme2yolo.py --- labelme2yolo.py | 131 ++++++++++++++++++++++++++++++++++-------------- 1 file changed, 93 insertions(+), 38 deletions(-) diff --git a/labelme2yolo.py b/labelme2yolo.py index 10c65fa..2cb7604 100644 --- a/labelme2yolo.py +++ b/labelme2yolo.py @@ -7,6 +7,7 @@ import os import sys import argparse import shutil +import math import json import cv2 @@ -22,12 +23,12 @@ class Labelme2YOLO(object): self._label_id_map = self._get_label_id_map(self._json_dir) - self._label_dir_path = os.path.join(json_dir, 'YOLODataset/labels/') - self._image_dir_path = os.path.join(json_dir, 'YOLODataset/images/') - - self._make_train_val_dir() - def _make_train_val_dir(self): + self._label_dir_path = os.path.join(self._json_dir, + 'YOLODataset/labels/') + self._image_dir_path = os.path.join(self._json_dir, + 'YOLODataset/images/') + for yolo_path in (os.path.join(self._label_dir_path + 'train/'), os.path.join(self._label_dir_path + 'val/'), os.path.join(self._image_dir_path + 'train/'), @@ -59,7 +60,7 @@ class Labelme2YOLO(object): 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) \ - if os.path.isdir(os.path.join(val_folder, val_sample_name))] + if os.path.isdir(os.path.join(val_folder, val_sample_name))] return train_json_names, val_json_names @@ -70,14 +71,15 @@ class Labelme2YOLO(object): return train_json_names, val_json_names - def convert(self, val_size=0.2): + def convert(self, val_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 \ 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 = self._train_test_split(folders, json_names, val_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 @@ -89,12 +91,64 @@ class Labelme2YOLO(object): print('Converting %s for %s ...' % (json_name, target_dir.replace('/', ''))) - img_path = self._save_yolo_image(json_data, json_name, target_dir) + 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, 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, + 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) 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) + else: + 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 + + (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_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) @@ -102,41 +156,37 @@ class Labelme2YOLO(object): 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) - - yolo_obj_list = [] - for shape in json_data['shapes']: - obj_x_min, obj_w, obj_y_min, obj_h = __get_object_desc(shape['points']) + obj_x_min, obj_w, obj_y_min, obj_h = __get_object_desc(shape['points']) - img_h, img_w, _ = cv2.imread(img_path).shape - - 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) + 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) - label_id = self._label_id_map[shape['label']] - - yolo_obj_list.append((label_id, yolo_center_x, yolo_center_y, yolo__w, yolo_h)) - - return yolo_obj_list + label_id = self._label_id_map[shape['label']] + + return label_id, yolo_center_x, yolo_center_y, yolo__w, yolo_h - def _save_yolo_label(self, json_name, target_dir, yolo_obj_list): - txt_path = os.path.join(self._label_dir_path, + 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', '.text')) + json_name.replace('.json', '.txt')) with open(txt_path, 'w+') as f: - for yolo_obj in yolo_obj_list: - f.write('%s %s %s %s %s\n' % yolo_obj) + 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 + f.write(yolo_obj_line) - def _save_yolo_image(self, json_data, json_name, target_dir): - img = utils.img_b64_to_arr(json_data['imageData']) - + 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(self._image_dir_path, target_dir,img_name ) + img_path = os.path.join(image_dir_path, target_dir,img_name) - PIL.Image.fromarray(img).save(img_path) + if not os.path.exists(img_path): + img = utils.img_b64_to_arr(json_data['imageData']) + PIL.Image.fromarray(img).save(img_path) return img_path @@ -144,10 +194,15 @@ if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--json_dir',type=str, help='Please input the path of the labelme json files.') - parser.add_argument('--val_size',type=float, + parser.add_argument('--val_size',type=float, nargs='?', default=None, help='Please input the validation dataset size, for example 0.1 ') + parser.add_argument('--json_name',type=str, nargs='?', default=None, + help='If you put json name, it would convert only one json file to YOLO.') args = parser.parse_args(sys.argv[1:]) - + convertor = Labelme2YOLO(args.json_dir) - convertor.convert(val_size=args.val_size) + if args.json_name is None: + convertor.convert(val_size=args.val_size) + else: + convertor.convert_one(args.json_name)