import albumentations as a
import numpy as np
#from ultralytics import YOLO
from os import listdir as ls
import cv2 as cv
#from pybboxes import BoundingBox as bb

#import torch

def b_c(b, s = (1920, 1080)):
    return [
        round(b[0] * s[0] - b[2] * s[0] / 2),
        round(b[0] * s[0] + b[2] * s[0] / 2),
        round(b[1] * s[1] - b[3] * s[1] / 2),
        round(b[1] * s[1] + b[3] * s[1] / 2),
    ]

def sq(img, bbox):
    for i in bbox:
        b = b_c(i, s = (img.shape[1], img.shape[0]))
        #print(b)
        cv.rectangle(img, (b[0], b[2]), (b[1], b[3]), (255, 0, 0), 2)

base_path = '../videos/total'
tagt_path = '../videos/total-exp'

t = a.Compose([
    a.BBoxSafeRandomCrop(),
    a.HorizontalFlip(p=0.5),
    a.RandomBrightnessContrast(p=0.2)
], bbox_params=a.BboxParams(format='yolo', label_fields=['class_labels']))

fs = set([i.rsplit(".", 1)[0] for i in ls(base_path) if i != "classes.txt"])

for i in list(fs)[:1]:
    img = cv.cvtColor(cv.imread(f"{base_path}/{i}.jpg"), cv.COLOR_BGR2RGB)
    box = [tuple(map(float, j.split()[1:])) for j in open(f"{base_path}/{i}.txt").read().split("\n") if j]
    
    #for j in open(f"{base_path}/{i}.txt").read().split("\n"):
    #    box.append(tuple(map(float, j.split()[1:])))

    print(box)

    lbl = np.array(['ch'] * len(box))
    box = np.array(box)

    #sq(img, box)

    r = t(image = img, bboxes = box, class_labels = lbl)

    out_img = r['image']
    out_box = r['bboxes']
    out_lbl = r['class_labels']

    sq(out_img, out_box)
    
    from matplotlib import pyplot as plt
    plt.imshow(out_img)
    plt.show()

#y = YOLO("m.yaml")
#y.info()

#o = y.model(torch.randn(1, 3, 640, 640))