Spaces:
Running on Zero
Running on Zero
| """DEIMv2-S: Real-Time Object Detection Meets DINOv3. | |
| A Gradio demo for the DEIMv2-S checkpoint (DINOv3-vit_tiny backbone + STA adapter), | |
| the flagship model of the paper "Real-Time Object Detection Meets DINOv3" | |
| (arXiv:2509.20787): 50.9 COCO AP at 9.7M parameters. | |
| Pipeline (mirrors the authors' tools/inference/torch_inf.py + hf_models.ipynb): | |
| RGB image -> Resize(640x640) -> ToTensor -> ImageNet normalize | |
| -> DEIMv2 forward (backbone -> HybridEncoder -> DEIMTransformer -> PostProcessor) | |
| -> sigmoid scores, boxes scaled to the ORIGINAL image size. | |
| """ | |
| import colorsys | |
| import time | |
| import gradio as gr | |
| import spaces # must precede any torch/CUDA import | |
| import torch | |
| import torch.nn as nn | |
| from huggingface_hub import PyTorchModelHubMixin | |
| from PIL import Image, ImageDraw, ImageFont | |
| from torchvision import transforms | |
| from engine.backbone import DINOv3STAs, HGNetv2 | |
| from engine.deim import DFINETransformer, DEIMTransformer, HybridEncoder, LiteEncoder | |
| from engine.deim.postprocessor import PostProcessor | |
| MODEL_ID = "Intellindust/DEIMv2_DINOv3_S_COCO" | |
| IMAGE_SIZE = (640, 640) | |
| # COCO categories, 1-indexed as in the authors' hf_models.ipynb (model label i -> label_map[i+1]) | |
| LABEL_MAP = { | |
| 1: "person", 2: "bicycle", 3: "car", 4: "motorcycle", 5: "airplane", | |
| 6: "bus", 7: "train", 8: "truck", 9: "boat", 10: "traffic light", | |
| 11: "fire hydrant", 12: "stop sign", 13: "parking meter", 14: "bench", | |
| 15: "bird", 16: "cat", 17: "dog", 18: "horse", 19: "sheep", | |
| 20: "cow", 21: "elephant", 22: "bear", 23: "zebra", 24: "giraffe", | |
| 25: "backpack", 26: "umbrella", 27: "handbag", 28: "tie", 29: "suitcase", | |
| 30: "frisbee", 31: "skis", 32: "snowboard", 33: "sports ball", 34: "kite", | |
| 35: "baseball bat", 36: "baseball glove", 37: "skateboard", 38: "surfboard", 39: "tennis racket", | |
| 40: "bottle", 41: "wine glass", 42: "cup", 43: "fork", 44: "knife", | |
| 45: "spoon", 46: "bowl", 47: "banana", 48: "apple", 49: "sandwich", | |
| 50: "orange", 51: "broccoli", 52: "carrot", 53: "hot dog", 54: "pizza", | |
| 55: "donut", 56: "cake", 57: "chair", 58: "couch", 59: "potted plant", | |
| 60: "bed", 61: "dining table", 62: "toilet", 63: "tv", 64: "laptop", | |
| 65: "mouse", 66: "remote", 67: "keyboard", 68: "cell phone", 69: "microwave", | |
| 70: "oven", 71: "toaster", 72: "sink", 73: "refrigerator", 74: "book", | |
| 75: "clock", 76: "vase", 77: "scissors", 78: "teddy bear", 79: "hair drier", | |
| 80: "toothbrush", | |
| } | |
| NUM_CLASSES = len(LABEL_MAP) | |
| # Per-category colors exactly as the authors' notebook builds them (HSV sweep) | |
| CATEGORY_COLORS = {} | |
| for _i, _label in enumerate(LABEL_MAP.values()): | |
| _rgb = colorsys.hsv_to_rgb(_i / NUM_CLASSES, 0.8, 0.9) | |
| CATEGORY_COLORS[_label] = tuple(int(c * 255) for c in _rgb) | |
| class DEIMv2(nn.Module, PyTorchModelHubMixin): | |
| """DEIMv2 detector: backbone -> encoder -> decoder -> postprocessor. | |
| Matches the class in the authors' hf_models.ipynb; from_pretrained() passes | |
| config.json into __init__ and loads model.safetensors. | |
| """ | |
| def __init__(self, config): | |
| super().__init__() | |
| if "DINOv3STAs" in config: | |
| self.backbone = DINOv3STAs(**config["DINOv3STAs"]) | |
| else: | |
| self.backbone = HGNetv2(**config["HGNetv2"]) | |
| if "LiteEncoder" in config: | |
| self.encoder = LiteEncoder(**config["LiteEncoder"]) | |
| else: | |
| self.encoder = HybridEncoder(**config["HybridEncoder"]) | |
| if "DEIMTransformer" in config: | |
| self.decoder = DEIMTransformer(**config["DEIMTransformer"]) | |
| else: | |
| self.decoder = DFINETransformer(**config["DFINETransformer"]) | |
| self.postprocessor = PostProcessor(**config["PostProcessor"]) | |
| def forward(self, x, orig_target_sizes): | |
| x = self.backbone(x) | |
| x = self.encoder(x) | |
| x = self.decoder(x) | |
| x = self.postprocessor(x, orig_target_sizes) | |
| return x | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| print(f"[app] loading {MODEL_ID} ...", flush=True) | |
| MODEL = DEIMv2.from_pretrained(MODEL_ID).to(DEVICE).eval() | |
| print("[app] model loaded", flush=True) | |
| # ImageNet normalization (the DINOv3/ViT-backbone variants are trained with it) | |
| TRANSFORM = transforms.Compose([ | |
| transforms.Resize(IMAGE_SIZE), | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), | |
| ]) | |
| def _load_font(): | |
| try: | |
| return ImageFont.load_default(size=16) | |
| except TypeError: | |
| return ImageFont.load_default() | |
| FONT = _load_font() | |
| def _text_size(draw, text, font): | |
| if hasattr(draw, "textbbox"): | |
| box = draw.textbbox((0, 0), text, font=font) | |
| return box[2] - box[0], box[3] - box[1] | |
| return draw.textsize(text, font=font) | |
| def detect(image, conf_threshold=0.4, progress=None): | |
| """Run DEIMv2-S object detection on an image. | |
| Args: | |
| image: input PIL image (RGB or RGBA/grayscale, converted to RGB). | |
| conf_threshold: minimum sigmoid confidence for a detection to be kept. | |
| progress: Gradio progress tracker. | |
| Returns: | |
| (annotated image, detection table, summary string) | |
| """ | |
| if image is None: | |
| raise ValueError("Please upload or select an image first.") | |
| if progress: | |
| progress((1, 3), desc="Preprocessing") | |
| img = image.convert("RGB") if image.mode != "RGB" else image | |
| orig_size = torch.tensor([[img.width, img.height]], device=DEVICE) | |
| t0 = time.time() | |
| im_data = TRANSFORM(img).unsqueeze(0).to(DEVICE) | |
| with torch.no_grad(): | |
| outputs = MODEL(im_data, orig_size) | |
| inference_ms = (time.time() - t0) * 1000 | |
| if progress: | |
| progress((2, 3), desc="Drawing boxes") | |
| labels, boxes, scores = outputs[0]["labels"], outputs[0]["boxes"], outputs[0]["scores"] | |
| keep = scores > conf_threshold | |
| annotated = img.copy() | |
| draw = ImageDraw.Draw(annotated) | |
| lines = [] | |
| n = int(keep.sum()) | |
| for lbl, box, scr in zip(labels[keep], boxes[keep], scores[keep]): | |
| name = LABEL_MAP[int(lbl.item()) + 1] | |
| color = CATEGORY_COLORS.get(name, (255, 0, 0)) | |
| conf = float(scr.item()) | |
| x1, y1, x2, y2 = (float(v) for v in box.tolist()) | |
| draw.rectangle([x1, y1, x2, y2], outline=color, width=3) | |
| text = f"{name} {conf:.2f}" | |
| tw, th = _text_size(draw, text, FONT) | |
| ty = y1 - th - 4 if y1 - th - 4 > 0 else y1 + 4 | |
| draw.rectangle([x1, ty, x1 + tw + 6, ty + th + 4], fill=color) | |
| draw.text((x1 + 3, ty + 2), text, fill=(255, 255, 255), font=FONT) | |
| lines.append([ | |
| name, | |
| round(conf, 3), | |
| f"[{x1:.0f}, {y1:.0f}, {x2:.0f}, {y2:.0f}]", | |
| ]) | |
| if progress: | |
| progress((3, 3), desc="Done") | |
| summary = f"{n} object{'s' if n != 1 else ''} above {conf_threshold:.2f} · {inference_ms:.0f} ms" | |
| return annotated, lines, summary | |
| with gr.Blocks() as demo: | |
| gr.HTML( | |
| """ | |
| <div style="display:flex; align-items:center; gap:12px; flex-wrap:wrap;"> | |
| <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/ZeroGPU- | |
| on-orange-badge.svg" alt="ZeroGPU" width="140"> | |
| <h1 style="margin:0; font-size:1.6em;">DEIMv2-S · Real-Time Object Detection Meets DINOv3</h1> | |
| </div> | |
| <p style="margin-top:6px; color:var(--body-text-color-subdued);"> | |
| DEIMv2-S (DINOv3-vit_tiny backbone with a Spatial Tuning Adapter) hits | |
| <b>50.9 COCO AP at 9.7M parameters</b> — the paper's headline result. | |
| Upload an image, adjust the confidence threshold, and get bounding boxes, | |
| class labels and scores for all 80 COCO categories. | |
| <a href="https://arxiv.org/abs/2509.20787">Paper</a> · | |
| <a href="https://github.com/Intellindust-AI-Lab/DEIMv2">Code</a> · | |
| <a href="https://huggingface.co/Intellindust/DEIMv2_DINOv3_S_COCO">Weights</a> | |
| </p> | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| image_in = gr.Image(type="pil", label="Input image", height=420) | |
| conf_slider = gr.Slider(0.05, 0.95, value=0.4, step=0.05, | |
| label="Confidence threshold", info="Keep detections with sigmoid score above this") | |
| btn = gr.Button("Detect objects", variant="primary") | |
| with gr.Column(): | |
| image_out = gr.Image(type="pil", label="Detections", height=420) | |
| summary = gr.Markdown() | |
| detections_table = gr.Dataframe( | |
| headers=["class", "score", "box (x1, y1, x2, y2)"], | |
| datatype=["str", "number", "str"], | |
| label="Detections", | |
| interactive=False, | |
| wrap=True, | |
| ) | |
| gr.Examples( | |
| examples=[ | |
| ["examples/authors_example.jpg"], | |
| ["examples/motorcycle_street.jpg"], | |
| ["examples/elephant_asian.jpg"], | |
| ["examples/dog.jpg"], | |
| ["examples/cafe_interior.jpg"], | |
| ], | |
| inputs=[image_in], | |
| fn=detect, | |
| outputs=[image_out, detections_table, summary], | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| run_on_click=True, | |
| label="Examples — try one", | |
| ) | |
| gr.Markdown( | |
| """ | |
| **Notes** · Preprocessing matches the authors' inference script (Resize 640² → ImageNet | |
| normalize); boxes are scaled back to the original image resolution. Scores are sigmoid | |
| confidences (no NMS — DEIMv2 is end-to-end). Weights are released by Intellindust AI Lab | |
| under a non-commercial research license. | |
| """ | |
| ) | |
| btn.click(detect, inputs=[image_in, conf_slider], outputs=[image_out, detections_table, summary]) | |
| demo.launch(theme=gr.themes.Citrus(), mcp_server=True) | |