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

View File

@ -2,8 +2,6 @@
#
# SPDX-License-Identifier: MIT
import argparse
import sys
# from labelme2yolo.__about__ import version
from labelme2yolo.l2y import Labelme2YOLO
@ -13,15 +11,15 @@ def run():
help='Please input the path of the labelme json files.')
parser.add_argument('--val_size', type=float, nargs='?', default=None,
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 ')
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()
if not args.json_dir:
parser.print_help(sys.stderr)
sys.exit(1)
parser.print_help()
return 0
convertor = Labelme2YOLO(args.json_dir)
@ -29,3 +27,5 @@ def run():
convertor.convert(val_size=args.val_size, test_size=args.test_size)
else:
convertor.convert_one(args.json_name)
return 0

View File

@ -76,7 +76,7 @@ def img_data_to_png_data(img_data):
return f_in.read()
def _get_label_id_map(json_dir):
def get_label_id_map(json_dir):
label_set = set()
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)])
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,
target_dir,
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)
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_path = os.path.join(image_dir_path, target_dir, img_name)
@ -119,7 +119,7 @@ class Labelme2YOLO(object):
def __init__(self, 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):
self._label_dir_path = os.path.join(self._json_dir,
@ -150,10 +150,11 @@ class Labelme2YOLO(object):
test_size=val_size)
tmp_train_len = len(train_indexes)
test_indexes = []
if test_size > 1e-8:
if test_size:
train_indexes, test_indexes = 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_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]
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 ...' %
(json_name, target_dir.replace('/', '')))
img_path = _save_yolo_image(json_data,
json_name,
self._image_dir_path,
target_dir)
img_path = 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)
_save_yolo_label(json_name,
self._label_dir_path,
target_dir,
yolo_obj_list)
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)
@ -215,12 +216,12 @@ class Labelme2YOLO(object):
print('Converting %s ...' % json_name)
img_path = _save_yolo_image(json_data, json_name,
self._json_dir, '')
img_path = save_yolo_image(json_data, json_name,
self._json_dir, '')
yolo_obj_list = self._get_yolo_object_list(json_data, img_path)
_save_yolo_label(json_name, self._json_dir,
'', yolo_obj_list)
save_yolo_label(json_name, self._json_dir,
'', yolo_obj_list)
def _get_yolo_object_list(self, json_data, img_path):
yolo_obj_list = []