Update labelme2yolo.py
This commit is contained in:
rooneysh
GitHub
parent
eda5192d35
commit
53b3dacc31
131
labelme2yolo.py
131
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)
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue