import os
import json


def generate_labels(data_path, save_path):
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    for file_name in os.listdir(os.path.join(data_path)):
        if os.path.isdir(os.path.join(data_path, file_name)):
            continue
        if file_name.endswith('.json'):
            with open(os.path.join(data_path, file_name), 'r') as f:
                data = f.read()
            data_dict = json.loads(data)
            save_name = str(data_dict["imagePath"]).split('.')
            save_name = save_name[0] + '.txt'
            save_file = os.path.join(save_path, save_name)
            width = data_dict["imageWidth"]
            height = data_dict["imageHeight"]
            result_list = []
            for i in range(len(data_dict["shapes"])):
                label_list = []
                label = int(data_dict["shapes"][i]["label"])
                label_list.append(label)
                x1 = data_dict["shapes"][i]["points"][0][0]
                y1 = data_dict["shapes"][i]["points"][0][1]
                x2 = data_dict["shapes"][i]["points"][1][0]
                y2 = data_dict["shapes"][i]["points"][1][1]
                center_x = ((x1 + x2) / 2) / width
                center_y = ((y1 + y2) / 2) / height
                w = (x2 - x1) / width
                h = (y2 - y1) / height
                label_list.append(round(center_x, 6))
                label_list.append(round(center_y, 6))
                label_list.append(round(w, 6))
                label_list.append(round(h, 6))
                result_list.append(label_list)
            file_w = open(save_file, 'w')
            for i in range(len(result_list)):
                line = result_list[i]
                line = str(line).strip('[').strip(']').replace(',', '')
                file_w.write(line + '\n')
            file_w.close()


if __name__ == "__main__":
    input_path = 'dataset'
    output_path = 'labels'
    generate_labels(input_path, output_path)