import gradio as gr
from ultralytics import YOLO
import cv2

MODELOS = {
    "best.pt": "model/best.pt",
    "bestV2.pt": "model/bestV2.pt",
    "model_merged1.pt": "model/model_merged1.pt"
}

# Carrega todos os modelos uma vez
modelos_carregados = {
    nome: YOLO(caminho)
    for nome, caminho in MODELOS.items()
}

def predict_image(img, modelo_escolhido):
    if img is None:
        return None

    model = modelos_carregados[modelo_escolhido]

    results = model(img)

    annotated_img = results[0].plot()
    annotated_img_rgb = cv2.cvtColor(annotated_img, cv2.COLOR_BGR2RGB)

    return annotated_img_rgb

demo = gr.Interface(
    fn=predict_image,
    inputs=[
        gr.Image(
            sources=["upload", "webcam", "clipboard"],
            type="numpy",
            label="Tirar foto"
        ),
        gr.Dropdown(
            choices=list(MODELOS.keys()),
            value="model_merged1.pt",
            label="Escolha o modelo"
        )
    ],
    outputs=gr.Image(type="numpy", label="Resultado do modelo"),
    title="Detecção de Objetos com YOLO",
    description="Escolha um modelo e envie uma imagem para detecção.",
    live=True
)

if __name__ == "__main__":
    demo.launch()