shuoshuoshuo
/
labelme2yolo
9
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

reformat code

This commit is contained in:
Wang Xin 2022-08-01 15:43:19 +08:00
parent 2e5b3fec07
commit 845afa0328
1 changed files with 134 additions and 126 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

124
src/labelme2yolo/l2y.py
View File

@ -1,12 +1,10 @@
''' """
Created on Aug 18, 2021 Created on Aug 18, 2021
@author: xiaosonh @author: xiaosonh
@author: GreatV(Wang Xin) @author: GreatV(Wang Xin)
''' """
import os import os
import sys
import argparse
import shutil import shutil
import math import math
import base64 import base64
@ -123,12 +121,40 @@ def apply_exif_orientation(image):
else: else:
return image return image
def _get_label_id_map(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 _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'))
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
f.write(yolo_obj_line)
class Labelme2YOLO(object): class Labelme2YOLO(object):
def __init__(self, json_dir): def __init__(self, json_dir):
self._json_dir = json_dir self._json_dir = json_dir
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): 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,
@ -147,57 +173,46 @@ class Labelme2YOLO(object):
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) \
for label_id, label in enumerate(label_set)])
def _train_test_split(self, folders, json_names, val_size, test_size): 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: 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_folder = os.path.join(self._json_dir, 'train/')
train_json_names = [train_sample_name + '.json' \ train_json_names = [train_sample_name + '.json'
for train_sample_name in os.listdir(train_folder) \ for train_sample_name in os.listdir(train_folder)
if os.path.isdir(os.path.join(train_folder, train_sample_name))] if os.path.isdir(os.path.join(train_folder, train_sample_name))]
val_folder = os.path.join(self._json_dir, 'val/') val_folder = os.path.join(self._json_dir, 'val/')
val_json_names = [val_sample_name + '.json' \ val_json_names = [val_sample_name + '.json'
for val_sample_name in os.listdir(val_folder) \ 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))]
test_folder = os.path.join(self._json_dir, 'test/') test_folder = os.path.join(self._json_dir, 'test/')
test_json_names = [test_sample_name + '.json' \ test_json_names = [test_sample_name + '.json'
for test_sample_name in os.listdir(test_folder) \ for test_sample_name in os.listdir(test_folder)
if os.path.isdir(os.path.join(test_folder, test_sample_name))] if os.path.isdir(os.path.join(test_folder, test_sample_name))]
return train_json_names, val_json_names, test_json_names return train_json_names, val_json_names, test_json_names
train_idxs, val_idxs = train_test_split(range(len(json_names)), train_indexes, val_indexes = train_test_split(range(len(json_names)),
test_size=val_size) test_size=val_size)
tmp_train_len = len(train_idxs) tmp_train_len = len(train_indexes)
test_idxs = [] test_indexes = []
if test_size > 1e-8: if test_size > 1e-8:
train_idxs, test_idxs = train_test_split(range(tmp_train_len), test_size=test_size / (1 - val_size)) train_indexes, test_indexes = train_test_split(
train_json_names = [json_names[train_idx] for train_idx in train_idxs] range(tmp_train_len), test_size=test_size / (1 - val_size))
val_json_names = [json_names[val_idx] for val_idx in val_idxs] train_json_names = [json_names[train_idx] for train_idx in train_indexes]
test_json_names = [json_names[test_idx] for test_idx in test_idxs] val_json_names = [json_names[val_idx] for val_idx in val_indexes]
test_json_names = [json_names[test_idx] for test_idx in test_indexes]
return train_json_names, val_json_names, test_json_names return train_json_names, val_json_names, test_json_names
def convert(self, val_size, test_size): def convert(self, val_size, test_size):
json_names = [file_name for file_name in os.listdir(self._json_dir) \ 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 \ if os.path.isfile(os.path.join(self._json_dir, file_name)) and
file_name.endswith('.json')] file_name.endswith('.json')]
folders = [file_name for file_name in os.listdir(self._json_dir) \ folders = [file_name for file_name in os.listdir(self._json_dir)
if os.path.isdir(os.path.join(self._json_dir, file_name))] 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) 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() self._make_train_val_dir()
@ -207,7 +222,8 @@ class Labelme2YOLO(object):
(train_json_names, val_json_names, test_json_names)): (train_json_names, val_json_names, test_json_names)):
pool = Pool(os.cpu_count() - 1) pool = Pool(os.cpu_count() - 1)
for json_name in json_names: 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.close()
pool.join() pool.join()
@ -218,7 +234,8 @@ class Labelme2YOLO(object):
json_path = os.path.join(self._json_dir, json_name) json_path = os.path.join(self._json_dir, json_name)
json_data = json.load(open(json_path)) json_data = json.load(open(json_path))
print('Converting %s for %s ...' % (json_name, target_dir.replace('/', ''))) print('Converting %s for %s ...' %
(json_name, target_dir.replace('/', '')))
img_path = self._save_yolo_image(json_data, img_path = self._save_yolo_image(json_data,
json_name, json_name,
@ -226,7 +243,7 @@ class Labelme2YOLO(object):
target_dir) target_dir)
yolo_obj_list = self._get_yolo_object_list(json_data, img_path) yolo_obj_list = self._get_yolo_object_list(json_data, img_path)
self._save_yolo_label(json_name, _save_yolo_label(json_name,
self._label_dir_path, self._label_dir_path,
target_dir, target_dir,
yolo_obj_list) yolo_obj_list)
@ -241,7 +258,7 @@ class Labelme2YOLO(object):
self._json_dir, '') self._json_dir, '')
yolo_obj_list = self._get_yolo_object_list(json_data, img_path) yolo_obj_list = self._get_yolo_object_list(json_data, img_path)
self._save_yolo_label(json_name, self._json_dir, _save_yolo_label(json_name, self._json_dir,
'', yolo_obj_list) '', yolo_obj_list)
def _get_yolo_object_list(self, json_data, img_path): def _get_yolo_object_list(self, json_data, img_path):
@ -252,9 +269,11 @@ class Labelme2YOLO(object):
# labelme circle shape is different from others # labelme circle shape is different from others
# it only has 2 points, 1st is circle center, 2nd is drag end point # it only has 2 points, 1st is circle center, 2nd is drag end point
if shape['shape_type'] == 'circle': 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: 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) yolo_obj_list.append(yolo_obj)
@ -279,7 +298,7 @@ class Labelme2YOLO(object):
def _get_other_shape_yolo_object(self, shape, img_h, img_w): def _get_other_shape_yolo_object(self, shape, img_h, img_w):
def __get_object_desc(obj_port_list): def __get_object_desc(obj_port_list):
__get_dist = lambda int_list: max(int_list) - min(int_list) def __get_dist(int_list): return max(int_list) - min(int_list)
x_lists = [port[0] for port in obj_port_list] x_lists = [port[0] for port in obj_port_list]
y_lists = [port[1] for port in obj_port_list] y_lists = [port[1] for port in obj_port_list]
@ -297,18 +316,6 @@ class Labelme2YOLO(object):
return label_id, yolo_center_x, yolo_center_y, yolo_w, yolo_h return label_id, yolo_center_x, yolo_center_y, yolo_w, yolo_h
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_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, image_dir_path, target_dir): def _save_yolo_image(self, json_data, json_name, image_dir_path, target_dir):
img_name = json_name.replace('.json', '.png') 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)
@ -320,14 +327,15 @@ class Labelme2YOLO(object):
return img_path return img_path
def _save_dataset_yaml(self): 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: 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/')) os.path.join(self._image_dir_path, 'train/'))
yaml_file.write('val: %s\n\n' % \ yaml_file.write('val: %s\n\n' %
os.path.join(self._image_dir_path, 'val/')) os.path.join(self._image_dir_path, 'val/'))
yaml_file.write('test: %s\n\n' % \ yaml_file.write('test: %s\n\n' %
os.path.join(self._image_dir_path, 'test/')) 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\n' % len(self._label_id_map))