# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import os
import xml.etree.cElementTree as ET
import numpy as np
import tqdm
import cv2 as cv
from aisscv import __RepoPath__


# # Convert Images and Fix Labels (after augmentations)
# This notebook does label-postprocessing. 
# It does:
#  - Read in all annotations (in pascal VOC .xml files)
#  - Resize images
#  - Fix labels

PATH = os.path.join(__RepoPath__.__RepoPath__, "data/dataset_png/annotations")
files = os.listdir(PATH)

files = sorted([f for f in files if f.endswith('.xml')])

# Sanity chec if all labels are read in correctly
len(files)


def box_on_image(img: np.ndarray, min_point:tuple, max_point:tuple, window_name:str = 'Demo') -> None:
    """
    For debugging purposes: Print a bounding box on an image and show the image.
    Press any key to continue.
    """
    image = cv.rectangle(img, min_point, max_point, color=(255,0,0), thickness = int(img.shape[0]/50))
    cv.imshow(window_name, image)
    cv.waitKey(0)
    cv.destroyAllWindows()


# ## Let the magic begin
# Now we 
#  - iterate over the label files 
#  - read in the image of each label 
#  - resize the image, if it is large
#  - resize the bounding boxes accordingly
#  - check that bounding boxes are non-negative ints
#  - make sure that min-coordinates are smaller than max-coordinates
#
# we also split up the data to a train and a validation dataset, by getting a random value u. If this values is >.15 the sample belongs to train, else to validation. Write the image and the label to the according directory. 

# +
NEW_PATH = os.path.join(__RepoPath__.__RepoPath__, "data/dataset_final_02")
if not os.path.exists(NEW_PATH):
    #os.mkdir(NEW_PATH)
    os.makedirs(os.path.join(NEW_PATH,'train', 'annotations'))
    os.makedirs(os.path.join(NEW_PATH,'validation', 'annotations'))
for ind, file in tqdm.tqdm(enumerate(files[:]), total = len(files[:])):
    
    # ob = root.find('object')
    #print(ind)
    try:
        tree = ET.parse(os.path.join(PATH, file))
        root = tree.getroot()
        image_path = os.path.join(PATH, '..', root.find('filename').text)
        if not os.path.exists(image_path):
            continue
        image = cv.imread(image_path)


        #cv.imshow('Test', image)

        (orig_height,orig_width,c) = image.shape
        #print('Shape: {}'.format((v,u,c)))
        if min(orig_height,orig_width) >= 2000:
            factor = 0.25
        elif min(orig_height,orig_width) >= 1000:
            factor = 0.5
        else:
            factor = 1
        #print(f'Orig width: {orig_width} and height {orig_height}, c:{c}, at file: {file}')
        
        resized = cv.resize(image, (int(orig_width*factor), int(orig_height*factor)))
        #print(f'New width: {resized.shape[1]} and height {resized.shape[0]}')
        size = root.find('size')
        new_width = int(np.floor(resized.shape[1]))
        new_height = int(np.floor(resized.shape[0]))
        assert new_width>0 and new_height >0, "Image dimensions must be positive"

        # Now, set the new values for the xml file.
        width_node = size.find('width')
        width_node.text = str(new_width)

        height_node = size.find('height')
        height_node.text = str(new_height)

        new_image_name = os.path.split(image_path)[-1].split('.')[0]+'.jpg'

        filename_node = root.find('filename')
        filename_node.text = new_image_name
        
        obj_to_del = []
        for ob in root.iter('object'):
            bbox = ob.find('bndbox')
            xmin_node = bbox.find('xmin')
            ymin_node = bbox.find('ymin')
            xmax_node = bbox.find('xmax')
            ymax_node = bbox.find('ymax')
            #print('Read values: min: {} max: {}'.format((xmin_node.text, ymin_node.text), (xmax_node.text, ymax_node.text)))
            #box_on_image(image, (int(float(xmin_node.text)), int(float(ymin_node.text))), (int(float(xmax_node.text)), int(float(ymax_node.text))), 'Original Size')
            
            xmin_value = int(np.floor(float(xmin_node.text)*factor))
            ymin_value = int(np.floor(float(ymin_node.text)*factor))
            xmax_value = int(np.ceil(float(xmax_node.text)*factor))
            ymax_value = int(np.ceil(float(ymax_node.text)*factor))
            #print('After factor: min: {} max: {}'.format((xmin_value, ymin_value), (xmax_value, ymax_value)))

            # Step1: Clamp values to image size
            xmin_value = max(min(xmin_value, new_width-1), 1)
            xmax_value = max(min(xmax_value, new_width-1), 1)
            ymin_value = max(min(ymin_value, new_height-1), 1)
            ymax_value = max(min(ymax_value, new_height-1), 1)

            # Step2: switch in order to get positive heights/widths
            xmin = min(xmin_value, xmax_value)
            xmax = max(xmin_value, xmax_value)
            ymin = min(ymin_value, ymax_value)
            ymax = max(ymin_value, ymax_value)

            # check if dmg is no longer in the image and add object in xml file to list for later deletion
            if not (0<xmin<xmax<new_width and 0<ymin<ymax<new_height):
                obj_to_del.append(ob)
                #root.remove(ob)
                continue
            assert 0<xmin<xmax<new_width and 0<ymin<ymax<new_height, 'Problem with coordinates: min{}, max:{}, image {} at file {}'.format((xmin,ymin), (xmax,ymax), (new_width, new_height), file)
            #box_on_image(resized, (xmin, ymin), (xmax, ymax), 'New Size')
            xmin_node.text = str(xmin)
            ymin_node.text = str(ymin)
            xmax_node.text = str(xmax)
            ymax_node.text = str(ymax)

        # delete objects that are no longer in the image (not possible during the iteration)
        [root.remove(ob) for ob in obj_to_del]
    except Exception as e:
        print('Exception: ', e.with_traceback())
        print(file)
        continue
    u = np.random.uniform()
    if u > 0.15:
        cv.imwrite(os.path.join(NEW_PATH, 'train', new_image_name), resized)
        tree.write(os.path.join(NEW_PATH, 'train', 'annotations', file))
    else:
        cv.imwrite(os.path.join(NEW_PATH, 'validation', new_image_name), resized)
        tree.write(os.path.join(NEW_PATH, 'validation', 'annotations', file))

    
# -



StarcoderdataJupyterScriptsDedupFiltered