import os
import shutil
import xml.etree.cElementTree as et
from pathlib import Path

labels_dir = 'labels'
images_dir = 'images'

tag_size = 'size'
tag_object = 'object'
tag_name = 'name'
tag_bndbox = 'bndbox'
name_class_map = {'Truck': 1, 'Car': 2, 'Taxi': 2}  # 1 Вантажівка, 2 Машина легкова
forbidden_classes = ['Motorcycle']
default_class = 1


def convert_xml(folder):
    os.makedirs(images_dir, exist_ok=True)
    os.makedirs(labels_dir, exist_ok=True)

    for f in os.listdir(folder):
        if not f.endswith('.jpg'):
            continue
        label = f'{Path(f).stem}.xml'
        lines = read_xml(folder, label)
        if not lines:
            print(f'Image {f} has only forbidden classes in annotations')
            continue

        shutil.copy(os.path.join(folder, f), os.path.join(images_dir, f))
        with open(os.path.join(labels_dir, f'{Path(label).stem}.txt'), 'w') as label_file:
            label_file.writelines(lines)
            label_file.close()
        print(f'Image {f} has been processed successfully')


def read_xml(folder, label):
    tree = et.parse(os.path.join(folder, label))
    root = tree.getroot()
    lines = []
    size_dict = {size_ch.tag: size_ch.text for size_ch in root.findall(f'{tag_size}/*')}
    width = int(size_dict['width'])
    height = int(size_dict['height'])
    for node_object in tree.findall(tag_object):
        class_num = default_class
        c_x = c_y = c_w = c_h = 0
        for node_object_ch in node_object:
            if node_object_ch.tag == tag_name:
                key = node_object_ch.text
                if key in name_class_map:
                    class_num = name_class_map[key]
                else:
                    if key in forbidden_classes:
                        class_num = -1
                        continue
                    else:
                        class_num = default_class
            if node_object_ch.tag == tag_bndbox:
                bbox_dict = {bbox_ch.tag: bbox_ch.text for bbox_ch in node_object_ch}
                xmin = int(bbox_dict['xmin'])
                xmax = int(bbox_dict['xmax'])
                ymin = int(bbox_dict['ymin'])
                ymax = int(bbox_dict['ymax'])
                c_w = (xmax - xmin) / width
                c_h = (ymax - ymin) / height
                c_x = xmin / width + c_w / 2
                c_y = ymin / height + c_h / 2
        if class_num == -1:
            continue
        if c_x != 0 and c_y != 0 and c_w != 0 and c_h != 0:
            lines.append(f'{class_num} {round(c_x, 5)} {round(c_y, 5)} {round(c_w, 5)} {round(c_h, 5)}\n')
    return lines


if __name__ == '__main__':
    convert_xml('datasets/others/UAVimages')