app.py

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

MODEL_ID = "CarlosAraEspinosa/yolo_finetuned_raccoon"

device = "cuda" if torch.cuda.is_available() else "cpu"

processor = AutoImageProcessor.from_pretrained(MODEL_ID)
model = AutoModelForObjectDetection.from_pretrained(MODEL_ID).to(device)
model.eval()

id2label = model.config.id2label

def predict(image: Image.Image, threshold: float):
    image = image.convert("RGB")

    inputs = processor(images=image, return_tensors="pt").to(device)

    with torch.no_grad():
        outputs = model(**inputs)

    target_sizes = torch.tensor([image.size[::-1]], device=device)  # (h, w)
    results = processor.post_process_object_detection(
        outputs, threshold=threshold, target_sizes=target_sizes
    )[0]

    drawn = image.copy()
    draw = ImageDraw.Draw(drawn)

    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
        box = [round(x, 2) for x in box.tolist()]
        xmin, ymin, xmax, ymax = box
        class_name = id2label[int(label)]
        text = f"{class_name}: {float(score):.2f}"

        draw.rectangle([xmin, ymin, xmax, ymax], outline="red", width=3)
        draw.text((xmin, max(0, ymin - 12)), text, fill="red")

    return drawn

demo = gr.Interface(
    fn=predict,
    inputs=[
        gr.Image(type="pil", label="Sube una imagen"),
        gr.Slider(0.0, 1.0, value=0.5, step=0.05, label="Threshold"),
    ],
    outputs=gr.Image(type="pil", label="Detecciones"),
    examples=[
        ["raccoon1.jpg", 0.5],
        ["raccoon2.jpg", 0.5],
    ],
    title="Detección de mapaches",
)

if __name__ == "__main__":
    demo.launch()