import gradio as gr
import torch
from transformers import AutoImageProcessor, AutoModelForObjectDetection
from PIL import Image, ImageDraw, ImageFont

# --- CONFIGURACIÓN ---
REPO_ID = "rugarce/detr_resnet_50_finetuned_grape" 

# Cargar modelo y procesador
print(f"Cargando modelo desde {REPO_ID}...")
image_processor = AutoImageProcessor.from_pretrained(REPO_ID)
model = AutoModelForObjectDetection.from_pretrained(REPO_ID)

def detect_grapes(image):
    if image is None:
        return None
    
    # 1. Preprocesar
    inputs = image_processor(images=image, return_tensors="pt")
    
    # 2. Inferencia
    with torch.no_grad():
        outputs = model(**inputs)
    
    # 3. Post-procesar
    target_sizes = torch.tensor([image.size[::-1]])
    results = image_processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.5)[0]
    
    # 4. Dibujar
    draw = ImageDraw.Draw(image)
    
    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
        box = [round(i, 2) for i in box.tolist()]
        label_name = model.config.id2label[label.item()]
        confidence = round(score.item(), 2)
        
        # Coordenadas
        xmin, ymin, xmax, ymax = box
        
        # Dibujar rectángulo (Rojo y grueso)
        draw.rectangle(box, outline="red", width=4)
        
        # Dibujar fondo del texto para que se lea mejor
        text = f"{label_name}: {confidence}"
        
        # Calculamos tamaño del texto (aproximación si no hay fuente cargada)
        # Usamos textbbox si está disponible (Pillow reciente) o estimación simple
        try:
            text_bbox = draw.textbbox((xmin, ymin), text)
            draw.rectangle(text_bbox, fill="red")
            draw.text((xmin, ymin), text, fill="white")
        except:
            # Fallback para versiones antiguas de Pillow
            draw.rectangle((xmin, ymin, xmin + 100, ymin + 15), fill="red")
            draw.text((xmin, ymin), text, fill="white")
        
    return image

# --- INTERFAZ GRADIO ---
description_text = """
# 🍇 Detector de Uvas con IA
Sube una imagen de un viñedo y el modelo (DETR-ResNet-50) detectará los racimos de uvas.
"""

# --- AQUÍ ESTÁN LOS EJEMPLOS ---
# Asegúrate de que los archivos existan en la pestaña "Files" de tu Space.
example_images = [
    ["CDY_2018.jpg"], 
    ["SVB_1956.jpg"]
]

iface = gr.Interface(
    fn=detect_grapes,
    inputs=gr.Image(type="pil", label="Imagen Original"),
    outputs=gr.Image(type="pil", label="Detección"),
    title="Detector de Frutas WGISD",
    description=description_text,
    examples=example_images,  # <--- Aquí cargamos la lista
    cache_examples=False      # Ponlo en True si quieres que pre-calcule los resultados al iniciar
)

iface.launch()