Select Git revision
oracle_parser_arc_eager.h
-
Alexis Nasr authoredAlexis Nasr authored
labelme2yolo.py 3.38 KiB
import os
import json
import base64
import shutil
import argparse
from glob import glob
from pathlib import Path
def arg_directory(path):
if os.path.isdir(path):
return path
else:
raise argparse.ArgumentTypeError(f'`{path}` is not a valid path')
def create_directory(directory):
try:
os.makedirs(os.path.join(directory,'labels'), exist_ok=True)
except Exception as e:
print(f'Error creating directory {directory}: {e}')
return (directory)
def convert_labelme_to_yolo(labelme_annotation_path, yolo_directory):
# Load LabelMe annotation
image_id = Path(labelme_annotation_path).stem
with open(labelme_annotation_path, 'r') as labelme_annotation_file:
labelme_annotation = json.load(labelme_annotation_file)
# YOLO annotation and image paths
yolo_annotation_path = os.path.join(yolo_directory, 'labels', f'{image_id}.txt')
yolo_image_path = os.path.join(yolo_directory, 'images', f'{image_id}.jpg')
with open(yolo_annotation_path, 'w') as yolo_annotation_file:
yolo_image_data = base64.b64decode(labelme_annotation['imageData'])
# Write YOLO image
with open(yolo_image_path, 'wb') as yolo_image_file:
yolo_image_file.write(yolo_image_data)
# Write YOLO image annotation
for shape in labelme_annotation['shapes']:
if shape['shape_type'] != 'rectangle':
print(f'Invalid type `{shape["shape_type"]}` in annotation `{labelme_annotation_path}`')
continue
label = shape['label']
# shape['points'] format : [[x1,y1],[x2,y2]...] #
scale_width = 1.0 / labelme_annotation['imageWidth']
scale_height = 1.0 / labelme_annotation['imageHeight']
width = abs(shape['points'][1][0] - shape['points'][0][0]) * scale_width
height = abs(shape['points'][1][1] - shape['points'][0][1]) * scale_height
x = min(shape['points'][0][0], shape['points'][1][0]) * scale_width + width / 2
y = min(shape['points'][0][1], shape['points'][1][1]) * scale_height + height / 2
if x+width/2 > 1 or y+height/2>1:
print(f'Error with bounding box values over 1 in file {yolo_image_file}')
annotation_line = f'{label} {x} {y} {width} {height}\n'
yolo_annotation_file.write(annotation_line)
def main(args):
yolo_names = set()
for labelme_annotation_path in glob(f'{args.path_to_data}/*.json'):
convert_labelme_to_yolo(labelme_annotation_path, args.directory)
with open(labelme_annotation_path, 'r') as labelme_annotation_file:
labelme_annotation = json.load(labelme_annotation_file)
for shape in labelme_annotation['shapes']:
yolo_names.add(shape['label'])
# Write YOLO names
yolo_names_path = os.path.join(args.directory, 'custom.names')
with open(yolo_names_path, 'w') as yolo_names_file:
yolo_names_file.write('\n'.join(yolo_names))
if __name__ == "__main__":
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', required=True)
parser.add_argument('-d', '--directory', type=arg_directory, help='Directory to which YOLO annotations will be stored', required=True)
args = parser.parse_args()
main(args)