"""Detectors: closed-set RF-DETR and open-vocabulary models driven by a Vocabulary. A Vocabulary maps each class name to the text prompts that stand for it. Open-vocabulary models score every prompt; a hit on any prompt counts for its class, and overlapping hits from synonyms of one class collapse to the best-scoring box. """ import json from dataclasses import dataclass from pathlib import Path import cv2 import numpy as np import supervision as sv import torch import transformers from PIL import Image from transformers.image_transforms import center_to_corners_format from transformers import ( AutoImageProcessor, AutoProcessor, GroundingDinoForObjectDetection, MMGroundingDinoForObjectDetection, OmDetTurboForObjectDetection, Owlv2ForObjectDetection, RfDetrForInstanceSegmentation, ) from hub import model_path # Task -> display name -> {folder under models/, Hub repo that fills it, default threshold}. # RF-DETR has a fixed COCO head. Open-vocabulary models score the prompts in a Vocabulary. TASKS = { "Detection (COCO)": { "RF-DETR Nano": {"dir": "rf-detr-nano", "repo": "Roboflow/rf-detr-nano", "conf": 0.5}, "RF-DETR Small": {"dir": "rf-detr-small", "repo": "Roboflow/rf-detr-small", "conf": 0.5}, "RF-DETR Medium": {"dir": "rf-detr-medium", "repo": "Roboflow/rf-detr-medium", "conf": 0.5}, "RF-DETR Base": {"dir": "rf-detr-base", "repo": "Roboflow/rf-detr-base", "conf": 0.5}, "RF-DETR Large": {"dir": "rf-detr-large", "repo": "Roboflow/rf-detr-large", "conf": 0.5}, }, "Segmentation (COCO)": { "RF-DETR Seg Nano": {"dir": "rf-detr-seg-nano", "repo": "Roboflow/rf-detr-seg-nano", "conf": 0.5}, "RF-DETR Seg Small": {"dir": "rf-detr-seg-small", "repo": "Roboflow/rf-detr-seg-small", "conf": 0.5}, "RF-DETR Seg Medium": {"dir": "rf-detr-seg-medium", "repo": "Roboflow/rf-detr-seg-medium", "conf": 0.5}, "RF-DETR Seg Large": {"dir": "rf-detr-seg-large", "repo": "Roboflow/rf-detr-seg-large", "conf": 0.5}, "RF-DETR Seg XLarge": {"dir": "rf-detr-seg-xlarge", "repo": "Roboflow/rf-detr-seg-xlarge", "conf": 0.5}, "RF-DETR Seg XXLarge": {"dir": "rf-detr-seg-xxlarge", "repo": "Roboflow/rf-detr-seg-xxlarge", "conf": 0.5}, }, "Zero-shot detection": { "OWLv2 Base": {"dir": "owlv2-base-patch16-ensemble", "repo": "google/owlv2-base-patch16-ensemble", "conf": 0.2}, "Grounding DINO Tiny": {"dir": "grounding-dino-tiny", "repo": "IDEA-Research/grounding-dino-tiny", "conf": 0.3}, "MM Grounding DINO Tiny": {"dir": "mm-grounding-dino-tiny", "repo": "openmmlab-community/mm_grounding_dino_tiny_o365v1_goldg_v3det", "conf": 0.3}, "LLMDet Tiny": {"dir": "llmdet-tiny", "repo": "iSEE-Laboratory/llmdet_tiny", "conf": 0.3}, "LLMDet Base": {"dir": "llmdet-base", "repo": "iSEE-Laboratory/llmdet_base", "conf": 0.3}, "OmDet-Turbo Tiny": {"dir": "omdet-turbo-swin-tiny", "repo": "omlab/omdet-turbo-swin-tiny-hf", "conf": 0.3}, }, } ZERO_SHOT_TASK = "Zero-shot detection" DETECTOR_SPECS = [spec for models in TASKS.values() for spec in models.values()] # IoU above which two hits of the same class are the same object (synonym prompts overlap). SYNONYM_NMS_IOU = 0.5 # Grounding DINO grounds every salient object to *some* phrase in its text, so when the # vocabulary spans several passes the same object can come back under a different class # in each. Hits this overlapping are one object; the best-scoring class wins. SAME_OBJECT_IOU = 0.7 ARTICLES = ("a ", "an ", "the ") _CACHE: dict[str, tuple] = {} # ----------------------------- # Vocabulary # ----------------------------- @dataclass(frozen=True) class Vocabulary: """Class names in display order and the text prompts that stand for each one.""" prompts_by_class: dict[str, list[str]] @classmethod def from_json(cls, path: str | Path) -> "Vocabulary": """Read `{"class": ["prompt", ...], ...}`. Any other shape raises ValueError.""" data = json.loads(Path(path).read_text(encoding="utf-8")) if not isinstance(data, dict) or not data: raise ValueError(f"{path}: expected a non-empty object of class -> list of prompts") for name, prompts in data.items(): if not isinstance(prompts, list) or not prompts or not all(isinstance(p, str) and p.strip() for p in prompts): raise ValueError(f"{path}: class {name!r} needs a non-empty list of prompt strings") return cls({name.strip().lower(): [p.strip() for p in prompts] for name, prompts in data.items()}) @classmethod def from_names(cls, text: str) -> "Vocabulary": """Build from a comma-separated list; each name becomes its own prompt with an article.""" names = [] for raw in (text or "").split(","): name = raw.strip().lower() if name and name not in names: names.append(name) return cls({name: [with_article(name)] for name in names}) @property def classes(self) -> list[str]: return list(self.prompts_by_class) @property def prompts(self) -> list[str]: return [p for prompts in self.prompts_by_class.values() for p in prompts] @property def class_of_prompt(self) -> np.ndarray: """Class index for each position in `prompts`.""" return np.array([i for i, prompts in enumerate(self.prompts_by_class.values()) for _ in prompts], dtype=int) def clip_queries(self) -> list[str]: """Prompts as OWLv2 saw them in training: `a photo of a person`.""" return [f"a photo of {p}" for p in self.prompts] def category_names(self, class_ids: list[int] | None = None) -> list[tuple[int, str]]: """(class index, bare category) pairs as Grounding DINO and OmDet-Turbo saw them: `cartoon character`.""" ids = range(len(self.prompts_by_class)) if class_ids is None else class_ids classes = self.classes return [(i, strip_article(p)) for i in ids for p in self.prompts_by_class[classes[i]]] def class_chunks(self, tokenizer, max_tokens: int) -> list[list[int]]: """Group class indices so each group's prompts, joined with ". ", fit the text encoder. A class's synonyms always share a pass, so its score is the max over all of them. """ budget = max_tokens - 2 chunks, current, used = [], [], 0 for i, prompts in enumerate(self.prompts_by_class.values()): cost = sum(len(tokenizer(strip_article(p) + ".", add_special_tokens=False).input_ids) for p in prompts) if current and used + cost > budget: chunks.append(current) current, used = [], 0 current.append(i) used += cost if current: chunks.append(current) return chunks def strip_article(text: str) -> str: text = text.strip().lower() for article in ARTICLES: if text.startswith(article): return text[len(article):] return text def with_article(name: str) -> str: if name.startswith(ARTICLES) or name.endswith("s"): return name return ("an " if name[0] in "aeiou" else "a ") + name # ----------------------------- # Loading # ----------------------------- def load_detector(spec: dict): """Load (processor, model) once per folder; kept on CPU until a GPU call. The model class comes from `architectures` in the folder's config.json, so a new folder needs no code change beyond its `TASKS` entry. """ folder = spec["dir"] if folder not in _CACHE: path = model_path(spec) arch = json.loads((path / "config.json").read_text())["architectures"][0] cls = getattr(transformers, arch) processor_cls = AutoImageProcessor if arch.startswith("RfDetr") else AutoProcessor _CACHE[folder] = (processor_cls.from_pretrained(path), cls.from_pretrained(path).eval()) return _CACHE[folder] def is_open_vocabulary(model) -> bool: return isinstance(model, (Owlv2ForObjectDetection, OmDetTurboForObjectDetection, GroundingDinoForObjectDetection, MMGroundingDinoForObjectDetection)) # ----------------------------- # Inference # ----------------------------- @torch.inference_mode() def detect(processor, model, rgb: np.ndarray, conf: float, device: str, vocabulary: Vocabulary | None = None) -> sv.Detections: """Run one RGB frame through any detector; open-vocabulary models need a Vocabulary.""" image = Image.fromarray(rgb) if is_open_vocabulary(model): if vocabulary is None or not vocabulary.classes: return sv.Detections.empty() return detect_open_vocabulary(processor, model, image, conf, device, vocabulary) return detect_closed_set(processor, model, image, conf, device) def detect_at_resolution(processor, model, rgb: np.ndarray, conf: float, device: str, vocabulary: Vocabulary | None, max_side: int) -> sv.Detections: """Detect on a frame downscaled to `max_side` (0 = source), mapping boxes and masks back.""" h, w = rgb.shape[:2] if max_side <= 0 or max(h, w) <= max_side: return detect(processor, model, rgb, conf, device, vocabulary) scale = max_side / max(h, w) small = np.array(Image.fromarray(rgb).resize((round(w * scale), round(h * scale)), Image.BILINEAR)) det = detect(processor, model, small, conf, device, vocabulary) if len(det) == 0: return det det.xyxy = det.xyxy / scale if det.mask is not None: det.mask = np.stack([cv2.resize(m.astype(np.uint8), (w, h), interpolation=cv2.INTER_NEAREST).astype(bool) for m in det.mask]) return det def detect_closed_set(processor, model, image: Image.Image, conf: float, device: str) -> sv.Detections: inputs = processor(images=image, return_tensors="pt").to(device) outputs = model(**inputs) target_sizes = [(image.height, image.width)] if isinstance(model, RfDetrForInstanceSegmentation): results = processor.post_process_instance_segmentation( outputs, target_sizes=target_sizes, threshold=conf, return_binary_maps=True )[0] return segmentation_to_detections(results, model.config.id2label) results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=conf)[0] return sv.Detections.from_transformers(results, id2label=model.config.id2label) def detect_open_vocabulary(processor, model, image: Image.Image, conf: float, device: str, vocabulary: Vocabulary) -> sv.Detections: target_sizes = [(image.height, image.width)] if isinstance(model, Owlv2ForObjectDetection): queries = vocabulary.clip_queries() inputs = processor(text=[queries], images=image, return_tensors="pt").to(device) results = processor.post_process_grounded_object_detection( model(**inputs), threshold=conf, target_sizes=target_sizes, text_labels=[queries] )[0] det = detections_from_prompt_hits(results, vocabulary) elif isinstance(model, OmDetTurboForObjectDetection): categories = [name for _, name in vocabulary.category_names()] inputs = processor(images=image, text=categories, return_tensors="pt").to(device) results = processor.post_process_grounded_object_detection( model(**inputs), text_labels=[categories], target_sizes=target_sizes, threshold=conf, nms_threshold=1.0 )[0] det = detections_from_prompt_hits(results, vocabulary) else: # Grounding DINO family parts = [ detect_grounding_dino_pass(processor, model, image, conf, device, vocabulary, chunk) for chunk in vocabulary.class_chunks(processor.tokenizer, model.config.max_text_len) ] det = sv.Detections.merge(parts) if len(parts) > 1: det = det.with_nms(threshold=SAME_OBJECT_IOU, class_agnostic=True) return det.with_nms(threshold=SYNONYM_NMS_IOU, class_agnostic=False) def detect_grounding_dino_pass(processor, model, image: Image.Image, conf: float, device: str, vocabulary: Vocabulary, class_ids: list[int]) -> sv.Detections: """One Grounding DINO pass over the prompts of `class_ids`, scored per class from token logits. Each box gets, per class, the max sigmoid logit over that class's prompt tokens; the box takes its best class and keeps that score. This replaces phrase decoding, which merges neighbouring prompts and needs string matching to recover the class. """ prompts = vocabulary.category_names(class_ids) text = ". ".join(p for _, p in prompts) + "." inputs = processor(images=image, text=text, return_tensors="pt").to(device) outputs = model(**inputs) offsets = processor.tokenizer(text, return_offsets_mapping=True)["offset_mapping"] token_class = np.full(len(offsets), -1, dtype=int) cursor = 0 for ci, prompt in prompts: start = text.index(prompt, cursor) end = start + len(prompt) cursor = end for t, (s, e) in enumerate(offsets): if s >= start and e <= end and e > s: token_class[t] = ci probs = outputs.logits[0, :, : len(offsets)].sigmoid() # (queries, tokens) class_scores = torch.stack([probs[:, torch.from_numpy(token_class == ci).to(probs.device)].max(dim=1).values for ci in class_ids], dim=1) best_score, best_col = class_scores.max(dim=1) keep = best_score > conf if not keep.any(): return sv.Detections.empty() boxes = center_to_corners_format(outputs.pred_boxes[0][keep]) * torch.tensor([image.width, image.height] * 2, device=probs.device) class_id = np.array(class_ids, dtype=int)[best_col[keep].cpu().numpy()] return labelled_detections({"boxes": boxes, "scores": best_score[keep]}, class_id, vocabulary) # ----------------------------- # Result conversion # ----------------------------- def detections_from_prompt_hits(results: dict, vocabulary: Vocabulary) -> sv.Detections: """For models whose `labels` index into the prompt list (OWLv2, OmDet-Turbo).""" prompt_index = results["labels"].detach().cpu().numpy().astype(int).reshape(-1) if len(prompt_index) == 0: return sv.Detections.empty() return labelled_detections(results, vocabulary.class_of_prompt[prompt_index], vocabulary) def labelled_detections(results: dict, class_id: np.ndarray, vocabulary: Vocabulary) -> sv.Detections: boxes, scores = results["boxes"], results["scores"] boxes = boxes.detach().cpu().numpy() if torch.is_tensor(boxes) else boxes scores = scores.detach().cpu().numpy() if torch.is_tensor(scores) else scores classes = vocabulary.classes return sv.Detections( xyxy=np.asarray(boxes, dtype=np.float32).reshape(-1, 4), confidence=np.asarray(scores, dtype=np.float32).reshape(-1), class_id=class_id, data={"class_name": np.array([classes[c] for c in class_id])}, ) def segmentation_to_detections(results: dict, id2label: dict) -> sv.Detections: """Convert a transformers instance-segmentation result to sv.Detections. Handles both the (N, H, W) binary-map stack and the (H, W) id-map layout, without assuming segment ids line up with array indices. """ infos = results.get("segments_info") or [] seg = results.get("segmentation") if seg is None or len(infos) == 0: return sv.Detections.empty() seg = seg.detach().cpu().numpy() if torch.is_tensor(seg) else np.asarray(seg) if seg.ndim == 3: n = seg.shape[0] ids = np.array([int(i["id"]) for i in infos]) if len(infos) == n and ids.min() >= 0 and ids.max() < n: idx = ids elif len(infos) == n and ids.min() >= 1 and ids.max() <= n: idx = ids - 1 else: idx = np.arange(min(n, len(infos))) infos = infos[: len(idx)] masks = seg[idx].astype(bool) else: masks = np.stack([seg == int(i["id"]) for i in infos]).astype(bool) keep = masks.reshape(len(masks), -1).any(axis=1) masks, infos = masks[keep], [i for i, k in zip(infos, keep) if k] if len(infos) == 0: return sv.Detections.empty() class_id = np.array([int(i["label_id"]) for i in infos]) confidence = np.array([float(i.get("score", 1.0)) for i in infos], dtype=np.float32) return sv.Detections( xyxy=sv.mask_to_xyxy(masks).astype(np.float32), mask=masks, confidence=confidence, class_id=class_id, data={"class_name": np.array([id2label.get(c, str(c)) for c in class_id])}, )