Skip to content
Snippets Groups Projects
Select Git revision
  • 3614416c7bb8fa3a0b0b537918a585beb4925367
  • main default protected
2 results

3-count_stats.py

Blame
  • 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()