Update app.py

app.py

@@ -1,20 +1,47 @@
-from huggingface_hub import from_pretrained_fastai
 import gradio as gr
-from fastai.vision.all import *
+import torch
+from transformers import AutoModelForObjectDetection, AutoImageProcessor
+from PIL import Image, ImageDraw
 
-
-
-# repo_id = "YOUR_USERNAME/YOUR_LEARNER_NAME"
+# Definir el repositorio en Hugging Face
 repo_id = "JeanCGuerrero/Practica2"
 
-learner = from_pretrained_fastai(repo_id)
-labels = learner.dls.vocab
+# Cargar el modelo de Hugging Face
+model = AutoModelForObjectDetection.from_pretrained(repo_id)
+image_processor = AutoImageProcessor.from_pretrained(repo_id)
 
-# Definimos una función que se encarga de llevar a cabo las predicciones
+# Función que se encarga de llevar a cabo las predicciones
 def predict(img):
-    #img = PILImage.create(img)
-    pred,pred_idx,probs = learner.predict(img)
-    return {labels[i]: float(probs[i]) for i in range(len(labels))}
-    
-# Creamos la interfaz y la lanzamos. 
-gr.Interface(fn=predict, inputs=gr.Image(), outputs=gr.Label(num_top_classes=3),examples=['raccoon-108.jpg','raccoon-133.jpg']).launch(share=False)
+    img = img.convert("RGB")  # Asegurar que la imagen esté en formato RGB
+    inputs = image_processor(images=img, return_tensors="pt")
+
+    with torch.no_grad():
+        outputs = model(**inputs)
+
+    # Procesar los resultados
+    target_sizes = torch.tensor([img.size[::-1]])
+    results = image_processor.post_process_object_detection(outputs, threshold=0.5, target_sizes=target_sizes)[0]
+
+    # Dibujar las detecciones en la imagen
+    draw = ImageDraw.Draw(img)
+    detecciones = []
+
+    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+        box = [round(i, 2) for i in box.tolist()]
+        x, y, x2, y2 = box
+        draw.rectangle([x, y, x2, y2], outline="red", width=3)
+        class_name = f"Clase {label.item()} - Confianza: {round(score.item(), 2)}"
+        draw.text((x, y), class_name, fill="red")
+        detecciones.append(class_name)
+
+    return img, "\n".join(detecciones)
+
+# Crear la interfaz y lanzarla con Gradio
+gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=[gr.Image(), gr.Text()],
+    title="Detector de Mapaches 🦝",
+    description="Este modelo detecta mapaches en imágenes usando DETR con el dataset Raccoon.",
+    examples=['raccoon-108.jpg', 'raccoon-108.jpg']
+).launch(share=False)