Create app.py

app.py

@@ -0,0 +1,119 @@
+import os
+import gradio as gr
+import torch
+from torchvision.ops import nms
+from PIL import Image, ImageDraw
+
+from ultralytics import YOLO
+from transformers import AutoImageProcessor, AutoModelForObjectDetection
+from huggingface_hub import hf_hub_download
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# ---- CAMBIA ESTO ----
+YOLO_REPO_ID = "adcelis/TU_REPO_YOLOV8"      # repo donde subiste best.pt
+YOLO_FILENAME = "best.pt"                   # nombre del archivo en el repo
+
+DETR_REPO_ID = "adcelis/detr_finetuned_raccoon"  # tu repo DETR
+# ----------------------
+
+# Si tu repo es privado, crea un secret HF_TOKEN en el Space y descomenta:
+# HF_TOKEN = os.environ.get("HF_TOKEN")
+
+yolo_path = hf_hub_download(repo_id=YOLO_REPO_ID, filename=YOLO_FILENAME)  # , token=HF_TOKEN
+yolo_model = YOLO(yolo_path)
+
+detr_processor = AutoImageProcessor.from_pretrained(DETR_REPO_ID)  # , token=HF_TOKEN
+detr_model = AutoModelForObjectDetection.from_pretrained(DETR_REPO_ID).to(DEVICE)  # , token=HF_TOKEN
+detr_model.eval()
+
+
+def yolo_predict(pil_img, conf=0.25):
+    res = yolo_model.predict(pil_img, conf=conf, verbose=False)[0]
+    boxes = res.boxes.xyxy.cpu()
+    scores = res.boxes.conf.cpu()
+    labels = res.boxes.cls.cpu().long()
+    names = res.names  # dict id->label
+    return boxes, scores, labels, names
+
+
+@torch.no_grad()
+def detr_predict(pil_img, conf=0.5):
+    inputs = detr_processor(images=[pil_img], return_tensors="pt").to(DEVICE)
+    outputs = detr_model(**inputs)
+    target_sizes = torch.tensor([[pil_img.size[1], pil_img.size[0]]], device=DEVICE)
+    results = detr_processor.post_process_object_detection(outputs, threshold=conf, target_sizes=target_sizes)[0]
+    return results["boxes"].cpu(), results["scores"].cpu(), results["labels"].cpu()
+
+
+def ensemble_union_nms(boxes1, scores1, labels1, boxes2, scores2, labels2,
+                       w2=0.8, iou_thr=0.5, score_thr=0.25):
+    boxes = torch.cat([boxes1, boxes2], dim=0)
+    scores = torch.cat([scores1, scores2 * w2], dim=0)
+    labels = torch.cat([labels1, labels2], dim=0)
+
+    keep = scores >= score_thr
+    boxes, scores, labels = boxes[keep], scores[keep], labels[keep]
+
+    if boxes.numel() == 0:
+        return boxes, scores, labels
+
+    keep_all = []
+    for cls in labels.unique():
+        idx = torch.where(labels == cls)[0]
+        k = nms(boxes[idx], scores[idx], iou_thr)
+        keep_all.append(idx[k])
+    keep_all = torch.cat(keep_all)
+    keep_all = keep_all[scores[keep_all].argsort(descending=True)]
+    return boxes[keep_all], scores[keep_all], labels[keep_all]
+
+
+def draw_boxes(pil_img, boxes, scores, labels, names):
+    img = pil_img.copy()
+    draw = ImageDraw.Draw(img)
+    for b, s, l in zip(boxes, scores, labels):
+        x1, y1, x2, y2 = [float(x) for x in b.tolist()]
+        draw.rectangle((x1, y1, x2, y2), outline="green", width=2)
+        label = names.get(int(l), str(int(l)))
+        draw.text((x1, y1), f"{label} {float(s):.2f}", fill="black")
+    return img
+
+
+def run(pil_img, yolo_conf, detr_conf, w2, iou_thr, score_thr):
+    pil_img = pil_img.convert("RGB")
+
+    b1, s1, l1, names = yolo_predict(pil_img, conf=yolo_conf)
+    b2, s2, l2 = detr_predict(pil_img, conf=detr_conf)
+
+    be, se, le = ensemble_union_nms(b1, s1, l1, b2, s2, l2,
+                                    w2=w2, iou_thr=iou_thr, score_thr=score_thr)
+
+    out_img = draw_boxes(pil_img, be, se, le, names)
+
+    rows = []
+    for b, s, l in zip(be, se, le):
+        x1, y1, x2, y2 = [round(float(x), 2) for x in b.tolist()]
+        rows.append([names.get(int(l), str(int(l))), round(float(s), 3), x1, y1, x2, y2])
+
+    return out_img, rows
+
+
+demo = gr.Interface(
+    fn=run,
+    inputs=[
+        gr.Image(type="pil", label="Imagen"),
+        gr.Slider(0.05, 0.9, value=0.25, step=0.05, label="YOLO conf"),
+        gr.Slider(0.05, 0.9, value=0.5, step=0.05, label="DETR conf"),
+        gr.Slider(0.1, 1.5, value=0.8, step=0.05, label="Peso DETR (w2)"),
+        gr.Slider(0.1, 0.9, value=0.5, step=0.05, label="NMS IoU"),
+        gr.Slider(0.05, 0.9, value=0.25, step=0.05, label="Score mínimo (post-ensemble)"),
+    ],
+    outputs=[
+        gr.Image(type="pil", label="Ensemble (NMS)"),
+        gr.Dataframe(headers=["label", "score", "x1", "y1", "x2", "y2"], label="Detecciones"),
+    ],
+    title="Ensemble YOLOv8 + DETR con Non-Maximum Suppression",
+)
+
+if __name__ == "__main__":
+    demo.launch()