Select Git revision
labelme2yolo.py 2.92 KiB
"""
Converts LabelMe annotations to annotations compatible with YOLO.
The script does the following:
- cleans (!) the output directory and prepare it for training,
- splits the dataset on validation and training,
- converts all LabelMe annoations (*.json) to YOLO annoations (*.txt) and
- creates YOLO metadata (`.data`, `.names`, `train.txt` and `valid.txt`)
"""
import os
import json
import random
import base64
import shutil
import argparse
from pathlib import Path
from glob import glob
# Parse arguments
def arg_directory(path):
if os.path.isdir(path):
return path
else:
raise argparse.ArgumentTypeError(f'`{path}` is not a valid path')
parser = argparse.ArgumentParser(description='Convert LabelMe annotations to YOLO compatible')
parser.add_argument('-p','--path_to_data',type=arg_directory,help='Path to LabelMe annotations')
parser.add_argument('-d','--direction',type=arg_directory,help='Directory to which YOLO annotations will be stored')
args = parser.parse_args()
# YOLO metadata and files
yolo_names = []
yolo_names_path = os.path.join(args.direction, 'custom.names')
# Prepare output directory
os.mkdir(os.path.join(args.direction,'labels'))
# Convert image annotations
for index, labelme_annotation_path in enumerate(glob(f'{args.path_to_data}/*.json')):
image_id = os.path.basename(labelme_annotation_path).rstrip('.json')
labelme_annotation_file = open(labelme_annotation_path, 'r')
labelme_annotation = json.load(labelme_annotation_file)
yolo_annotation_path = os.path.join(args.direction,'labels', image_id + '.txt')
yolo_annotation_file = open(yolo_annotation_path, 'w')
yolo_image = base64.decodebytes(labelme_annotation['imageData'].encode())
yolo_image_path = os.path.join(args.direction, 'labels', image_id + '.jpg')
# Write YOLO image (and it to the list)
yolo_image_file = open(yolo_image_path, 'wb')
yolo_image_file.write(yolo_image)
yolo_image_file.close()
# Write YOLO image annotation
for shape in labelme_annotation['shapes']:
if shape['shape_type'] != 'rectangle':
print(
f'Invalid type `{shape["shape_type"]}` in annotation `annotation_path`')
continue
if shape['label'] not in yolo_names:
yolo_names.append(shape['label'])
points = shape['points']
scale_width = 1.0 / labelme_annotation['imageWidth']
scale_height = 1.0 / labelme_annotation['imageHeight']
width = abs(points[1][0] - points[0][0]) * scale_width
height = abs(points[1][1] - points[0][1]) * scale_height
x = (abs(points[1][0] + points[0][0]) / 2) * scale_width
y = (abs(points[1][1] + points[0][1]) / 2) * scale_height
object_class = shape['label']
yolo_annotation_file.write(f'{object_class} {x} {y} {width} {height}\n')
# Write YOLO names
yolo_names_file = open(yolo_names_path, 'w')
yolo_names_file.write(os.linesep.join(yolo_names))
yolo_names_file.close()