import gradio as gr
import numpy as np
from PIL import Image
import torch
import matplotlib.pyplot as plt
import io

# Modelo
model = torch.hub.load('ultralytics/yolov5', 'custom', path='bestyolo5.pt', trust_repo=True)

def detect(img):
    if img is None:
        return None, 0
    
    img_arr = np.array(img)
    results = model(img_arr)
    
    # Criar figura
    fig, ax = plt.subplots(figsize=(10, 8))
    ax.imshow(img_arr)
    
    cattle_count = 0
    
    # Processar detecções
    for *xyxy, conf, cls in results.xyxy[0].cpu().numpy():
        x1, y1, x2, y2 = map(int, xyxy)
        label = model.names[int(cls)]
        
        if label == 'cattle':
            cattle_count += 1
            
        # Desenhar bounding box
        ax.add_patch(plt.Rectangle((x1, y1), x2-x1, y2-y1, 
                                 fill=False, color='red', linewidth=2))
        
        # Adicionar label
        ax.text(x1, y1-10, f'{label} {conf:.2f}', 
               color='white', fontsize=10, 
               bbox={'facecolor': 'red', 'alpha': 0.7})
    
    ax.set_title(f'Detecções: {cattle_count} gado encontrado')
    plt.axis('off')
    plt.tight_layout()
    
    # Converter para PIL Image
    buf = io.BytesIO()
    plt.savefig(buf, format='png', bbox_inches='tight', dpi=150)
    buf.seek(0)
    pil_img = Image.open(buf)
    plt.close(fig)
    
    return pil_img, cattle_count

# Interface Gradio
iface = gr.Interface(
    fn=detect,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Image(type="pil", label="Imagem com Detecções"), 
        gr.Number(label="Número de Gado Detectado")
    ],
    title="🐄 YOLOv5 Contador de Gado",
    description="Detector de objetos treinado para contar gado usando YOLOv5.",
    examples=[["example1.jpg"]] if False else None,  # Remova ou adicione imagens de exemplo
    theme=gr.themes.Soft()
)

if __name__ == "__main__":
    iface.launch(share=True, debug=True)