import gradio as gr
from ultralytics import YOLO
import numpy as np
from PIL import Image

# 1. Carregar o seu modelo treinado
# O Hugging Face baixa automaticamente do seu repositório de modelos
# Substitua 'PericlesRodrigues01' pelo seu usuário se for diferente
model = YOLO("https://huggingface.co/PericlesRodrigues01/player-detector/resolve/main/best.pt")

def detect_objects(image):
    # 2. Fazer a predição
    # conf=0.25: Confiança mínima
    # iou=0.4: Eliminar caixas duplicadas
    results = model.predict(image, conf=0.25, iou=0.4)
    
    # 3. Gerar a imagem com as caixas desenhadas
    # O YOLO retorna a imagem em formato BGR (padrão OpenCV), 
    # mas o Gradio precisa de RGB. Vamos converter.
    result_array = results[0].plot() # Cria o array com os desenhos
    result_rgb = result_array[..., ::-1] # Inverte de BGR para RGB
    
    return Image.fromarray(result_rgb)

# 4. Criar a Interface do Gradio
interface = gr.Interface(
    fn=detect_objects,                    # Função que será chamada
    inputs=gr.Image(type="pil"),          # Entrada: Imagem enviada pelo usuário
    outputs=gr.Image(type="pil"),         # Saída: Imagem com as detecções
    title="⚽ Detecção de Jogadores - YOLOv8",
    description="Faça upload de uma imagem de futebol ou tênis para detectar Jogadores, Bolas e Equipamentos.",
)

# 5. Iniciar o App
if __name__ == "__main__":
    interface.launch()