# start od 0 
import cv2
from labels_list import labels_list
from tqdm import tqdm

import sys


dataset = sys.argv[1]


def to_yolo_format(entity, y_total, x_total):
    label_str, ref = entity.split(':')

    label_id = labels_list.index(label_str)
    x0,y0,x1,y1 = [float(i) for i in ref.split(',')]
    x_middle = (x0 + x1) / 2
    y_middle = (y0 + y1) / 2
    x_w = x1 - x0
    y_w = y1 - y0

    x_middle /= x_total
    y_middle /= y_total
    x_w /= x_total
    y_w /= y_total


    items_to_write = (label_id, x_middle, y_middle, x_w, y_w)
    items_to_write = [str(i) for i in items_to_write]
    str_to_write = ' '.join(items_to_write) + '\n'
    return str_to_write


with open(f'{dataset}/expected.tsv','r') as f_exp, open(f'{dataset}/in.tsv') as f_in:
    for line_in, line_exp in zip(f_in, f_exp):
        img_name = line_in.rstrip()
        img_references = line_exp.rstrip().split(' ')
        img = cv2.imread('images/' + img_name)
        img_dimensions = img.shape
        y_total, x_total, _ = img_dimensions
        with open(f'{dataset}-dataset/labels/' +  img_name.replace('.jpg','.txt'), 'w') as f_out:
            for entity in img_references:
                str_to_write = to_yolo_format(entity, y_total, x_total)
                f_out.write(str_to_write)