labelme2yolo.py

"""
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()