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

bug fixed

This commit is contained in:
Wang Xin 2022-08-10 16:51:21 +08:00
parent a3a824b5ce
commit 8abb8e4bd4
2 changed files with 24 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
src/labelme2yolo/cli/__init__.py
View File

@ -2,8 +2,6 @@
# #
# SPDX-License-Identifier: MIT # SPDX-License-Identifier: MIT
import argparse import argparse
import sys
# from labelme2yolo.__about__ import version
from labelme2yolo.l2y import Labelme2YOLO from labelme2yolo.l2y import Labelme2YOLO
@ -13,15 +11,15 @@ def run():
help='Please input the path of the labelme json files.') help='Please input the path of the labelme json files.')
parser.add_argument('--val_size', type=float, nargs='?', default=None, parser.add_argument('--val_size', 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=0.0, 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 validation dataset size, for example 0.1 ')
parser.add_argument('--json_name', type=str, nargs='?', default=None, 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.') help='If you put json name, it would convert only one json file to YOLO.')
args = parser.parse_args() args = parser.parse_args()
if not args.json_dir: if not args.json_dir:
parser.print_help(sys.stderr) parser.print_help()
sys.exit(1) return 0
convertor = Labelme2YOLO(args.json_dir) convertor = Labelme2YOLO(args.json_dir)
@ -29,3 +27,5 @@ def run():
convertor.convert(val_size=args.val_size, test_size=args.test_size) convertor.convert(val_size=args.val_size, test_size=args.test_size)
else: else:
convertor.convert_one(args.json_name) convertor.convert_one(args.json_name)
return 0

37
src/labelme2yolo/l2y.py
View File

@ -76,7 +76,7 @@ def img_data_to_png_data(img_data):
return f_in.read() return f_in.read()
def _get_label_id_map(json_dir): def get_label_id_map(json_dir):
label_set = set() label_set = set()
for file_name in os.listdir(json_dir): for file_name in os.listdir(json_dir):
@ -90,7 +90,7 @@ def _get_label_id_map(json_dir):
for label_id, label in enumerate(label_set)]) for label_id, label in enumerate(label_set)])
def _save_yolo_label(json_name, label_dir_path, target_dir, yolo_obj_list): def save_yolo_label(json_name, label_dir_path, target_dir, yolo_obj_list):
txt_path = os.path.join(label_dir_path, txt_path = os.path.join(label_dir_path,
target_dir, target_dir,
json_name.replace('.json', '.txt')) json_name.replace('.json', '.txt'))
@ -103,7 +103,7 @@ def _save_yolo_label(json_name, label_dir_path, target_dir, yolo_obj_list):
f.write(yolo_obj_line) f.write(yolo_obj_line)
def _save_yolo_image(json_data, json_name, image_dir_path, target_dir): def save_yolo_image(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)
@ -119,7 +119,7 @@ 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 = _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,
@ -150,10 +150,11 @@ class Labelme2YOLO(object):
test_size=val_size) test_size=val_size)
tmp_train_len = len(train_indexes) tmp_train_len = len(train_indexes)
test_indexes = [] test_indexes = []
if test_size > 1e-8: if test_size:
train_indexes, test_indexes = train_test_split( train_indexes, test_indexes = train_test_split(
range(tmp_train_len), test_size=test_size / (1 - val_size)) range(tmp_train_len), test_size=test_size / (1 - val_size))
train_json_names = [json_names[train_idx] for train_idx in train_indexes] train_json_names = [json_names[train_idx]
for train_idx in train_indexes]
val_json_names = [json_names[val_idx] for val_idx in val_indexes] 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] test_json_names = [json_names[test_idx] for test_idx in test_indexes]
@ -198,16 +199,16 @@ class Labelme2YOLO(object):
print('Converting %s for %s ...' % print('Converting %s for %s ...' %
(json_name, target_dir.replace('/', ''))) (json_name, target_dir.replace('/', '')))
img_path = _save_yolo_image(json_data, img_path = save_yolo_image(json_data,
json_name, json_name,
self._image_dir_path, self._image_dir_path,
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)
_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)
def convert_one(self, json_name): def convert_one(self, json_name):
json_path = os.path.join(self._json_dir, json_name) json_path = os.path.join(self._json_dir, json_name)
@ -215,12 +216,12 @@ class Labelme2YOLO(object):
print('Converting %s ...' % json_name) print('Converting %s ...' % json_name)
img_path = _save_yolo_image(json_data, json_name, img_path = save_yolo_image(json_data, json_name,
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)
_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):
yolo_obj_list = [] yolo_obj_list = []