video-intelligence

Running on Zero

Format prompts per model: bare categories for Grounding DINO and OmDet-Turbo, 'a photo of' queries for OWLv2, articles on typed names

6d66381 verified 6 days ago

Raw History Blame Contribute Delete

16.6 kB

"""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])},
    )

	"""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])},
	)