import os
import cv2
import glob

# === CONFIGURACAO ===
RESOLUCOES = {
    "640x640": (640, 640),
    "1280x720": (1280, 720),
}

FONTES = ["original", "augmented"]
PASTA_BASE = "dataset"

for fonte in FONTES:
    for nome_res, (larg, alt) in RESOLUCOES.items():
        pasta_imgs_saida = os.path.join(PASTA_BASE, nome_res, "images")
        pasta_lbls_saida = os.path.join(PASTA_BASE, nome_res, "labels")
        os.makedirs(pasta_imgs_saida, exist_ok=True)
        os.makedirs(pasta_lbls_saida, exist_ok=True)

        pasta_imgs_origem = os.path.join(PASTA_BASE, fonte, "images")
        pasta_lbls_origem = os.path.join(PASTA_BASE, fonte, "labels")

        imagens = sorted(glob.glob(os.path.join(pasta_imgs_origem, "*.jpg")))

        for i, caminho_img in enumerate(imagens, 1):
            nome_base = os.path.splitext(os.path.basename(caminho_img))[0]
            caminho_lbl = os.path.join(pasta_lbls_origem, f"{nome_base}.txt")
            if not os.path.exists(caminho_lbl):
                print(f"[!] Label ausente para {nome_base}, pulando...")
                continue

            # Carrega imagem
            img = cv2.imread(caminho_img)
            if img is None:
                print(f"[!] Erro ao ler imagem {caminho_img}")
                continue

            h_orig, w_orig = img.shape[:2]
            img_resized = cv2.resize(img, (larg, alt), interpolation=cv2.INTER_AREA)

            # Salva imagem redimensionada
            caminho_saida_img = os.path.join(pasta_imgs_saida, f"{nome_base}.jpg")
            cv2.imwrite(caminho_saida_img, img_resized)

            # Recalcula labels
            with open(caminho_lbl, 'r') as f:
                linhas = f.readlines()

            with open(os.path.join(pasta_lbls_saida, f"{nome_base}.txt"), 'w') as saida_lbl:
                for linha in linhas:
                    partes = linha.strip().split()
                    if len(partes) != 5:
                        continue
                    cls, x, y, w, h = map(float, partes)
                    # Conversao relativa antiga -> relativa nova
                    x_pix = x * w_orig
                    y_pix = y * h_orig
                    w_pix = w * w_orig
                    h_pix = h * h_orig

                    x_novo = x_pix / larg
                    y_novo = y_pix / alt
                    w_novo = w_pix / larg
                    h_novo = h_pix / alt

                    saida_lbl.write(f"{int(cls)} {x_novo:.6f} {y_novo:.6f} {w_novo:.6f} {h_novo:.6f}\n")

            print(f"[{fonte}] [{nome_res}] {i}/{len(imagens)} normalizado: {nome_base}")

print("\n✅ Normalizacao YOLO concluida!")