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

add output_format option

This commit is contained in:
Wang Xin 2023-03-16 13:48:24 +08:00
parent 14b716e173
commit 1f31b4bae9
5 changed files with 49 additions and 70 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
.gitignore vendored
View File

@ -1,3 +1,5 @@
dataset/
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]

4
README.md
View File

@ -14,7 +14,7 @@ If you've already marked your segmentation dataset by LabelMe, it's easy to use
## New
- export data as yolo polygon annotation (for YOLOv5 v7.0 segmentation)
- Now you can choose the output format of the label text. The available options are `plygon` and `bbox`.
## Installation
```console
@ -30,6 +30,8 @@ pip install labelme2yolo
**--json_name (Optional)** Convert single LabelMe JSON file.
**--output_format (Optional)** The output format of label.
## How to Use
### 1. Convert JSON files, split training, validation and test dataset by --val_size and --test_size

4
src/labelme2yolo/__about__.py
View File

@ -1,5 +1,5 @@
# SPDX-FileCopyrightText: 2022-present Wang Xin <xinwang614@gmail.com>
# SPDX-FileCopyrightText: 2023-present Wang Xin <xinwang614@gmail.com>
#
# SPDX-License-Identifier: MIT
__version__ = '0.0.8'
__version__ = '0.0.9'

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

@ -16,14 +16,14 @@ def run():
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=None,
help="Please input the validation dataset size, for example 0.1 ",
help="Please input the test dataset size, for example 0.1.",
)
parser.add_argument(
"--json_name",
@ -32,13 +32,21 @@ def run():
default=None,
help="If you put json name, it would convert only one json file to YOLO.",
)
parser.add_argument(
"--output_format",
type=str,
default="polygon",
help='The default output format for labelme2yolo is "polygon".'
' However, you can choose to output in bbox format by specifying the "bbox" option.',
)
args = parser.parse_args()
if not args.json_dir:
parser.print_help()
return 0
convertor = Labelme2YOLO(args.json_dir)
convertor = Labelme2YOLO(args.json_dir, args.output_format)
if args.json_name is None:
convertor.convert(val_size=args.val_size, test_size=args.test_size)

75
src/labelme2yolo/l2y.py
View File

@ -93,27 +93,28 @@ def get_label_id_map(json_dir):
return OrderedDict([(label, label_id) for label_id, label in enumerate(label_set)])
def extend_point_list(point_list):
def extend_point_list(point_list, format="polygon"):
xmin = min([float(point) for point in point_list[::2]])
xmax = max([float(point) for point in point_list[::2]])
ymin = min([float(point) for point in point_list[1::2]])
ymax = max([float(point) for point in point_list[1::2]])
if (format == "polygon"):
return np.array([xmin, ymin, xmax, ymin, xmax, ymax, xmin, ymax])
if (format == "bbox"):
return np.array([xmin, ymin, xmax - xmin, ymax - ymin])
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")
)
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
)
for yolo_obj in yolo_obj_list:
label, points = yolo_obj
points = [str(item) for item in points]
yolo_obj_line = f"{label} {' '.join(points)}\n"
f.write(yolo_obj_line)
@ -130,10 +131,11 @@ def save_yolo_image(json_data, json_name, image_dir_path, target_dir):
class Labelme2YOLO(object):
def __init__(self, json_dir):
def __init__(self, json_dir, output_format):
self._json_dir = json_dir
self._output_format = output_format
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):
self._label_dir_path = os.path.join(self._json_dir,
@ -152,20 +154,6 @@ class Labelme2YOLO(object):
os.makedirs(yolo_path)
@staticmethod
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):
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/')
@ -233,12 +221,12 @@ class Labelme2YOLO(object):
print('Converting %s for %s ...' %
(json_name, target_dir.replace('/', '')))
img_path = self._save_yolo_image(json_data,
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)
self._save_yolo_label(json_name,
save_yolo_label(json_name,
self._label_dir_path,
target_dir,
yolo_obj_list)
@ -249,11 +237,11 @@ class Labelme2YOLO(object):
print('Converting %s ...' % json_name)
img_path = self._save_yolo_image(json_data, json_name,
img_path = 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,
save_yolo_label(json_name, self._json_dir,
'', yolo_obj_list)
def _get_yolo_object_list(self, json_data, img_path):
@ -298,35 +286,14 @@ class Labelme2YOLO(object):
points[::2] = [float(point[0]) / img_w for point in point_list]
points[1::2] = [float(point[1]) / img_h for point in point_list]
if len(points) == 4:
if self._output_format == "polygon":
points = extend_point_list(points)
if self._output_format == "bbox":
points = extend_point_list(points, "bbox")
label_id = self._label_id_map[shape['label']]
return label_id, points.tolist()
@staticmethod
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 in yolo_obj_list:
label, points = yolo_obj
points = [str(item) for item in points]
yolo_obj_line = f"{label} {' '.join(points)}\n"
f.write(yolo_obj_line)
@staticmethod
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(image_dir_path, target_dir, img_name)
if not os.path.exists(img_path):
img = img_b64_to_arr(json_data['imageData'])
PIL.Image.fromarray(img).save(img_path)
return img_path
def _save_dataset_yaml(self):
yaml_path = os.path.join(
self._json_dir, 'YOLODataset/', 'dataset.yaml')