detector-uvas / app.py
rugarce's picture
Update app.py
6b0d0df verified
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()