Custom script to convert from Supervisely label to YOLO label format for Machine Learning training

supervisely_json_to_yolo_txt.py

#!/usr/bin/env python3
"""
supervisely_json_to_yolo_txt.py

This script is intended to convert the Javascript Object notation based Supervisely labels from the FSOCO dataset to a txt-based yolo format we can use for training

The Supervisely label format specifies the top-left and bottom-right corners of the object in number of pixels from the origin (at the top left of the image)

The yolo label format specifies the object class as an integer, followed by  <x_center> <y_center> <width> <height> where each are a propotion (float values from 0 to 1, exclusive) of the total image width or height, measured from the origin at the top left corner

"""
import sys
import os

import math
import numpy as np

# from PIL import *
import json

def main():
    cwd = os.getcwd()
    input_directory = None
    output_directory = None
    for i in range(len(sys.argv)):
        if i == 0:
            continue
        if i < len(sys.argv) and os.path.exists(sys.argv[i]) and os.path.isdir(sys.argv[i]):
            input_directory = os.path.expanduser(sys.argv[i])
            input_directory = os.path.expandvars(input_directory)
            input_directory = os.path.abspath(input_directory)
            if i < len(sys.argv) - 1 and sys.argv[i + 1].strip() not in {"", None}:
                output_directory = os.path.expanduser(sys.argv[i + 1])
                output_directory = os.path.expandvars(output_directory)
                output_directory = os.path.abspath(output_directory)
                if not os.path.exists(output_directory):
                    os.mkdir(output_directory)
                break
    #

    while input_directory is None:
        print("Please enter the input directory, either as a relative or absolute path: ")
        input_directory = input("").strip()
        try:
            input_directory = os.path.expanduser(input_directory)
            input_directory = os.path.expandvars(input_directory)
            input_directory = os.path.abspath(input_directory)
        except:
            print(f"ERROR: input directory {input_directory} invalid, please try again")
            input_directory = None
            continue
        if not (os.path.exists(input_directory) and os.path.isdir(input_directory)):
            print(f"ERROR: input directory {input_directory} invalid, please try again")
            input_directory = None
            continue

    while output_directory is None:
        print("Please enter the output directory, either as a relative or absolute path: ")
        output_directory = input("").strip()
        try:
            output_directory = os.path.expanduser(output_directory)
            output_directory = os.path.expandvars(output_directory)
            output_directory = os.path.abspath(output_directory)
        except:
            print(f"ERROR: output directory {output_directory} invalid, please try again")
            output_directory = None
            continue
        if not (os.path.exists(output_directory) and os.path.isdir(output_directory)):
            print(f"ERROR: output directory {output_directory} invalid, please try again")
            output_directory = None
            continue

    print(f"input directory: {input_directory}")
    print(f"output directory: {output_directory}")
    docontinue = ' '
    while docontinue == ' ':
        docontinue = input(f"Is this correct? Y/n?: ").lower().strip()
        if docontinue not in {'y', 'n'}:
            print(f"ERROR: invalid input '{docontinue}', expected value 'y' or 'n'")
            docontinue = ' '

    if  docontinue != 'y':
        sys.exit("Ok, exiting...")

    supervisely_classnames = []
    input_classname = ""
    print("Type 'done' to finish input")
    while input_classname.lower() != 'done':
        print(f'supervisely classnames: {supervisely_classnames}')
        input_classname = str(input("Enter a target classname: "))
        input_classname.strip()
        if input_classname.lower() != 'done':
            supervisely_classnames.append(input_classname)

    classes_filename = None
    while classes_filename == None:
        classes_filename = input("Please enter the path to a yolo classes.txt file: ").strip()
        classes_filename = os.path.expanduser(classes_filename)
        classes_filename = os.path.expandvars(classes_filename)
        classes_filename = os.path.abspath(classes_filename)
        if (not os.path.isfile(classes_filename)) or classes_filename.split('.')[-1].lower() != 'txt':
            print(f"ERROR: invalid filename '{classes_filename}'. Please try again...")
            classes_filename = None
            continue

    classes_file = open(classes_filename)
    class_names = classes_file.readlines()
    classes_file.close()

    i = 0
    for line in class_names:
        line = line.strip()
        print(f"{i}: {line}")
        i += 1

    i -= 1
    class_id = None
    while class_id is None:
        class_id = input(f"Please select a class id for yolo label output, 0 to {i}: ")
        try:
            class_id = int(class_id)
        except:
            print(f"ERROR: {class_id} is not an in integer in range 0 to {i}. Please try again...")
            class_id = None
            continue

    print(f"Okay, using class id '{class_id}', representing a '{class_names[class_id]}'")

    for root, dirs, files in os.walk(input_directory, onerror=None):
        for aFile in files:
            final_ext = aFile.split('.')[-1].lower()
            if final_ext != "json":
                continue
            penultimate_ext = aFile.split('.')[-2].lower()
            if penultimate_ext not in {"png", "jpg", "jpeg"}:
                continue
            basename_no_ext = "".join(aFile.split('.')[0:-2])
            infile = os.path.join(root, aFile)
            infile = open(infile)
            ma_dict = json.load(infile)
            in_objs = ma_dict['objects'] # type : list of dicts, one for each object

            outfile = os.path.join(output_directory, basename_no_ext + ".txt")
            outfile = open(outfile, "w")

            for an_obj in in_objs:
                if an_obj['classTitle'] not in supervisely_classnames:
                    continue
                else:
                    in_box_pts = an_obj['points']['exterior']
                    topleft_pt = in_box_pts[0]
                    bottomright_pt = in_box_pts[1]
                    x_center_pxls = round((topleft_pt[0] + bottomright_pt[0]) / 2)
                    y_center_pxls = round((topleft_pt[1] + bottomright_pt[1]) / 2)
                    box_width_pxls = abs(bottomright_pt[0] - topleft_pt[0])
                    box_height_pxls = abs(bottomright_pt[1] - topleft_pt[1])
                    img_width_pxls = ma_dict['size']['width']
                    img_height_pxls = ma_dict['size']['height']
                    yolo_x_center = round((x_center_pxls / img_width_pxls),6)
                    yolo_y_center = round((y_center_pxls / img_height_pxls),6)
                    yolo_box_width = round((box_width_pxls / img_width_pxls),6)
                    yolo_box_height = round((box_height_pxls / img_height_pxls),6)
                    out_str = str(class_id) + " " + str(yolo_x_center) + " " + str(yolo_y_center)
                    out_str = out_str + " " + str(yolo_box_width) + " " + str(yolo_box_height) + "\n"
                    outfile.write(out_str)
            outfile.close()
            infile.close()
    print("Script completed succesfully without errors")




if __name__ == "__main__":
    main()