Converting semantic JSON structure in RICO dataset to Yolo txt format

rico_json_to_yolo_txt.py

"""
This script, which converts ui elements in json to yolo txt format, was written in line with the need for the project
I was working on. Possible to make improvements on performance.
"""

import os
import shutil
import json

#bounds => [left, top, right, bottom]

input_path = "/path/to/rico/json/file"
output_path = "/path/to/rico/txt/file"

for filename in os.listdir(input_path):
    with open(os.path.join(input_path, filename)) as f:
        data = json.load(f)

        def get_child_components(data, labels=[]):
            if 'children' in list(data.keys()):
                for c in data['children']:

                    label_map = {
                        'Image': 0,
                        'Text': 1,
                        'Text Button': 2,
                        'Input': 3,
                        'Icon': 4,
                        'List Item': 5,
                        'Another': 6,
                        # more mappings if needed
                    }

                    #Convert these bounds to YOLO format
                    width = c['bounds'][2] - c['bounds'][0]
                    height = c['bounds'][3] - c['bounds'][1]

                    x_center = c['bounds'][0] + width/2
                    y_center = c['bounds'][1] + height/2
                    
                    #base = data['bounds']
                    
                    image_width = 1440 #int(base[2])
                    image_height = 2560 #int(base[3])
                    
                    x_center_norm = x_center/image_width
                    y_center_norm = y_center/image_height
                    width_norm = width/image_width
                    height_norm = height/image_height

                    component_index = label_map.get(c['componentLabel'], 0)

                    yolo_format = [component_index, x_center_norm, y_center_norm, width_norm, height_norm]

                    labels.append({c['componentLabel']: yolo_format})

                    get_child_components(c, labels)
            return labels

        try:
            os.mkdir(output_path)
        except:
            pass
        
        output_file_path = os.path.join(output_path, filename[:-5] + '.txt')
        
        with open(output_file_path, 'w') as f:
            for i in get_child_components(data):
                for k in i.keys():
                    f.write('{} {} {} {} {}\n'.format(int(i[k][0]), float(i[k][1]), float(i[k][2]), float(i[k][3]), float(i[k][4])))