from transformers import pipeline
from pathlib import Path
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import colormaps
from matplotlib.patches import Patch
import cv2
from matplotlib import patches

detc_modelo = 'valentinafeve/yolos-fashionpedia'
detc = pipeline('object-detection', detc_modelo)

sgm_modelo = 'mattmdjaga/segformer_b0_clothes'
sgm = pipeline('image-segmentation', sgm_modelo)

def processar_imagem(image_path):
    # Carrega a imagem
    imagem = Image.open(Path(image_path))

    # Segmentação
    segmentacao = sgm(imagem)

    # Detecção de objetos
    deteccao = detc(imagem)

    return imagem, segmentacao, deteccao

# segmentacao = sgm(imagem) feito na função processar_imagem

def adicionar_mascara(dados_imagem, dados_mascara, cor):
    imagem_mascara = dados_imagem.copy()
    imagem_mascara = np.where(
        dados_mascara[:, :, np.newaxis] == 255,
        cor,
        imagem_mascara,
    ).astype(dados_imagem.dtype)
    return cv2.addWeighted(dados_imagem, 0.5, imagem_mascara, 0.5, 0)


def plotar_segmentos(imagem, segmentacao, nome_colormap):
    cmap = colormaps.get_cmap(nome_colormap)
    cores = [
        (np.array(cmap(x)[:3]) * 255).astype(int)
        for x in np.linspace(0, 1, len(segmentacao))
    ]

    imagem_final = np.array(imagem).copy()
    legendas = []
    for segmento, cor in zip(segmentacao, cores):
        dados_mascara = np.array(segmento['mask'])
        label_mascara = segmento['label']
        imagem_final = adicionar_mascara(imagem_final, dados_mascara, cor)
        legendas.append(Patch(facecolor=cor/255, edgecolor='black', label=label_mascara))

    fig, ax = plt.subplots(figsize=(16, 16))
    plt.legend(handles=legendas)

    ax.imshow(imagem_final)

#plotar_segmentos(imagem=imagem, segmentacao=segmentacao, nome_colormap='hsv') 

# deteccao = detc(imagem) feito na função processar_imagem

# Função para desenhar os retângulos
def plotar_deteccao(imagem, deteccao):
    fig, ax = plt.subplots(figsize=(16, 16))
    ax.imshow(imagem)

    for det in deteccao:
        if det['score'] >= 0.80:
            box = det['box']
            origem = (box['xmin'], box['ymin'])
            largura = box['xmax'] - box['xmin']
            altura = box['ymax'] - box['ymin']

            rect = patches.Rectangle(origem, largura, altura, linewidth=3, edgecolor='red', facecolor='none')
            ax.add_patch(rect)

            texto = f"{det['label']} {100 * det['score']:.2f}%"
            ax.text(box['xmin'] - 50, box['ymin'] + altura / 2, texto,
                    bbox={'facecolor': 'red', 'alpha': 0.8}, ha='right', va='center')

    plt.axis('off')
    plt.show()

# plotar_deteccao(imagem=imagem, deteccao=deteccao)