diff --git a/.gitignore b/.gitignore index 57e78f4f9..c66da35a6 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,8 @@ __pycache__/ *.py[cod] *$py.class +visualizer +.vscode # C extensions *.so @@ -227,5 +229,10 @@ docs/cookbooks/*.py docs/cookbooks/export_executorch/ docs/cookbooks/export_tensorrt/ +# deprecated code +visualizer + +# VSCode metadata directory +.vscode # MacOS files .DS_Store diff --git a/CHANGELOG.md b/CHANGELOG.md index 9e475e054..9b82c1538 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), ### Added +- `RFDETR.predict(..., return_embeddings=True)` now attaches a per-detection embedding vector as `detections.data["embeddings"]` (or `key_points.data["embeddings"]` for keypoint outputs), shape `(K, H)`, gathered with the same indices used for boxes/masks/keypoints — useful for downstream similarity search, clustering, or re-identification. On the eager (unoptimized) model this can be toggled per `predict()` call. On a model optimized via `model.inference(...)`, the exported/traced forward pass has fixed control flow, so `inference()` now also accepts `return_embeddings` and the value must be decided at optimization time and match the `return_embeddings` passed to `predict()`; a mismatch raises `RuntimeError`. + +- Added the `embedding-extraction` cookbook, a Colab-runnable notebook that extracts per-detection query embeddings on a public COCO 2017 validation subset (streamed from a Hugging Face parquet export, no API key) and applies them to dataset quality auditing: ranking mislabelled annotations against injected label noise (reported as ROC-AUC and precision@k versus random review), surfacing confident detections with no matching annotation, and nearest-neighbour object retrieval. + - `deploy_to_roboflow()`'s `version` argument is now optional: when omitted, the highest existing dataset version of the target project is resolved automatically via the Roboflow API (falling back to version `1` for a project with no generated versions, where the Roboflow SDK then raises its usual "Version number 1 is not found."). Passing an explicit `version` behaves exactly as before, with no extra API call. ([#1116](https://github.com/roboflow/rf-detr/issues/1116)) - Added a live opt-in end-to-end CI job (`roboflow-deploy-e2e`, `-m e2e_roboflow`) that generates a fresh dataset version in a dedicated Roboflow test project, deploys a real model with `version` omitted, and independently polls the server-side trained-model status — catching silent server-side upload failures that `deploy_to_roboflow()`'s return value cannot surface. ([#1116](https://github.com/roboflow/rf-detr/issues/1116)) diff --git a/docs/cookbooks/NOTES.md b/docs/cookbooks/NOTES.md index 5d267cdd3..ae9142358 100644 --- a/docs/cookbooks/NOTES.md +++ b/docs/cookbooks/NOTES.md @@ -36,7 +36,7 @@ If jupytext is not installed: `pip install jupytext` (or `uv add jupytext --dev` description: One sentence describing what the notebook demonstrates. ``` -Available labels (reuse these to keep tags standardised): `TRAINING`, `AUGMENTATION`, `EXPORT`, `TFLITE`, `PYTORCH LIGHTNING`, `INFERENCE`, `SEGMENTATION`, `DEPLOY`. Current tag colours are assigned dynamically by the docs UI, so they may change if cards or labels are added or reordered. +Available labels (reuse these to keep tags standardised): `TRAINING`, `AUGMENTATION`, `EXPORT`, `TFLITE`, `PYTORCH LIGHTNING`, `INFERENCE`, `SEGMENTATION`, `DEPLOY`, `EMBEDDINGS`. Current tag colours are assigned dynamically by the docs UI, so they may change if cards or labels are added or reordered. For newly added or updated notebooks, write markdown cells in plain, notebook-portable Markdown only — no MkDocs Material syntax (`!!! note`, `=== "Tab"`, admonition blocks). `mkdocs-jupyter` renders the site copy through the same plain renderer as a raw `.ipynb`, so MkDocs-only syntax shows up as literal text (e.g. `!!! warning "..."`) instead of a styled callout. Use a blockquote (`> **Note:** ...`) for callouts instead. @@ -51,6 +51,7 @@ For newly added or updated notebooks, write markdown cells in plain, notebook-po | ----------------------------------- | ----------------------------------------------- | ------- | | `custom-augmentations.ipynb` | Custom Augmentations and Live Training Progress | v1.5.0 | | `custom-optimizer-scheduler.ipynb` | Custom Optimizer and LR Scheduler | v1.9.0 | +| `embedding-extraction.ipynb` | Query Embeddings for Dataset Quality Auditing | v1.10.0 | | `export-coreml.ipynb` | Export to Native CoreML & Run Inference | v1.9.0 | | `export-tensorrt.ipynb` | Export to TensorRT & Run Inference | v1.9.0 | | `export-executorch.ipynb` | Export to ExecuTorch & Run Inference | v1.9.0 | diff --git a/docs/cookbooks/cards.yaml b/docs/cookbooks/cards.yaml index 242f6da78..8383d6067 100644 --- a/docs/cookbooks/cards.yaml +++ b/docs/cookbooks/cards.yaml @@ -1,4 +1,10 @@ cards: + - href: embedding-extraction/ + name: "Query Embeddings for Dataset Quality Auditing" + labels: [INFERENCE, EMBEDDINGS] + version: v1.10.0 + author: unaxEtxeberriaBieleDigital + description: "Extract per-detection query embeddings with predict(return_embeddings=True) and use them to rank mislabelled annotations, surface unannotated objects, and retrieve similar instances — measured on a public COCO subset." - href: export-coreml/ name: "Export to Native CoreML & Run Inference" labels: [EXPORT, INFERENCE, DEPLOY] diff --git a/docs/cookbooks/embedding-extraction.ipynb b/docs/cookbooks/embedding-extraction.ipynb new file mode 100644 index 000000000..51a037257 --- /dev/null +++ b/docs/cookbooks/embedding-extraction.ipynb @@ -0,0 +1,688 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "70944f0c", + "metadata": {}, + "outputs": [], + "source": [ + "# ------------------------------------------------------------------------\n", + "# RF-DETR\n", + "# Copyright (c) 2025 Roboflow. All Rights Reserved.\n", + "# Licensed under the Apache License, Version 2.0 [see LICENSE for details]\n", + "# ------------------------------------------------------------------------" + ] + }, + { + "cell_type": "markdown", + "id": "5a19b5e2", + "metadata": {}, + "source": [ + "# Query Embeddings for Dataset Quality Auditing\n", + "\n", + "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/roboflow/rf-detr/blob/develop/docs/cookbooks/embedding-extraction.ipynb)\n", + "\n", + "RF-DETR predicts objects with a set of *queries*. Each query carries a hidden-state vector from the last decoder\n", + "layer — the same vector the box and class heads read from. `predict(..., return_embeddings=True)` exposes that\n", + "vector for every returned detection, so you get one embedding per detected object instead of one embedding per\n", + "image.\n", + "\n", + "**Why per-object embeddings matter.** Image-level embeddings (CLIP, DINOv2, a classifier backbone) describe a\n", + "whole scene. Quality problems in a detection dataset are *per annotation*: a box labelled `truck` that actually\n", + "contains a `bus`, an object nobody annotated, a duplicated instance. Those live at the object level, and RF-DETR\n", + "already computes an object-level representation as part of normal inference — this cookbook just reads it out.\n", + "\n", + "**What this notebook demonstrates**, using a public COCO 2017 validation subset, no API key required:\n", + "\n", + "| Section | Task | What it shows |\n", + "|---|---|---|\n", + "| 3 | Extract embeddings | `predict(..., return_embeddings=True)` -> `detections.data[\"embeddings\"]`, shape `(K, H)` |\n", + "| 4 | Structure | Detections of the same category cluster together in a 2D projection |\n", + "| 5 | Probe | A k-NN probe recovers the ground-truth category from the embedding alone |\n", + "| 6 | **Mislabelled objects** | Label noise is injected into a known subset, then ranked back out (ROC-AUC, precision@k) |\n", + "| 7 | **Unannotated objects** | Confident detections with no ground-truth overlap are surfaced as missing labels |\n", + "| 8 | Similarity search | Nearest-neighbour retrieval of visually similar instances |\n", + "| 9 | Optimized models | How `inference(return_embeddings=True)` differs from the eager path |\n", + "\n", + "Section 6 is the industrial-inspection use case from the feature request: we *know* which annotations were\n", + "corrupted, so the ranking quality is measured rather than eyeballed.\n", + "\n", + "> **Runtime:** ~5 minutes on a Colab T4 with the default settings. It also runs on CPU (slower). Nothing here\n", + "> requires training." + ] + }, + { + "cell_type": "markdown", + "id": "71a326b9", + "metadata": {}, + "source": [ + "## 1. Setup\n", + "\n", + "`rfdetr` brings the model and `supervision` detections; `pyarrow` streams the dataset subset straight from the\n", + "Hugging Face parquet endpoint; `scikit-learn` provides the projection and the k-NN probe." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1a6404b2", + "metadata": {}, + "outputs": [], + "source": [ + "# Until 1.10.0 is on PyPI, install the branch that adds the embedding interface instead:\n", + "!pip install -q \"rfdetr @ git+https://github.com/roboflow/rf-detr.git@feat/embedding-extraction\" \n", + "\n", + "!uv pip install pyarrow fsspec aiohttp scikit-learn matplotlib" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6e5a90dd", + "metadata": {}, + "outputs": [], + "source": [ + "\"\"\"Extract per-detection RF-DETR query embeddings and use them to audit a detection dataset.\"\"\"\n", + "\n", + "import io\n", + "import warnings\n", + "from dataclasses import dataclass\n", + "\n", + "import fsspec\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import pyarrow.parquet as pq\n", + "import torch\n", + "from PIL import Image\n", + "from sklearn.decomposition import PCA\n", + "from sklearn.manifold import TSNE\n", + "from sklearn.metrics import roc_auc_score\n", + "from sklearn.neighbors import NearestNeighbors\n", + "\n", + "from rfdetr import RFDETRSmall\n", + "from rfdetr.assets.coco_classes import COCO_CLASSES\n", + "\n", + "warnings.filterwarnings(\"ignore\", category=UserWarning)\n", + "\n", + "SEED = 0\n", + "NUM_IMAGES = 300 # 100 images per parquet row group; raise for a denser embedding space\n", + "CONFIDENCE_THRESHOLD = 0.5\n", + "IOU_MATCH_THRESHOLD = 0.5 # detection <-> ground-truth association\n", + "IOU_UNANNOTATED_THRESHOLD = 0.3 # below this a detection covers no annotated object\n", + "NOISE_FRACTION = 0.10 # share of matched annotations whose label is corrupted in section 6\n", + "K_NEIGHBOURS = 10\n", + "\n", + "rng = np.random.default_rng(SEED)\n", + "torch.manual_seed(SEED)\n", + "\n", + "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", + "print(f\"device: {DEVICE}\")" + ] + }, + { + "cell_type": "markdown", + "id": "acfbbd1b", + "metadata": {}, + "source": [ + "## 2. Load a public COCO validation subset\n", + "\n", + "The subset is streamed from the Hugging Face parquet export of COCO 2017 (`rafaelpadilla/coco2017`). Only the\n", + "row groups we need are downloaded — 100 images each, roughly 16 MB — so this stays small and repeatable, and no\n", + "credentials are involved.\n", + "\n", + "Its `label` field indexes the 91-entry COCO category table, i.e. the same raw COCO category ids RF-DETR's\n", + "COCO-pretrained checkpoints return as `class_id`. Ground truth and predictions are therefore directly comparable\n", + "without a remapping table." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "821f8275", + "metadata": {}, + "outputs": [], + "source": [ + "PARQUET_URL = \"https://huggingface.co/api/datasets/rafaelpadilla/coco2017/parquet/default/val/0.parquet\"\n", + "\n", + "\n", + "@dataclass\n", + "class Sample:\n", + " \"\"\"One dataset image with its ground-truth annotations.\n", + "\n", + " Attributes:\n", + " image_id: COCO image id.\n", + " image: Decoded RGB image.\n", + " boxes: Ground-truth boxes in ``xyxy`` pixel coordinates, shape ``(N, 4)``.\n", + " labels: Raw COCO category id per ground-truth box, shape ``(N,)``.\n", + " \"\"\"\n", + "\n", + " image_id: int\n", + " image: Image.Image\n", + " boxes: np.ndarray\n", + " labels: np.ndarray\n", + "\n", + "\n", + "def load_coco_val_subset(url: str, num_images: int) -> list[Sample]:\n", + " \"\"\"Stream the first ``num_images`` COCO validation images from a parquet export.\n", + "\n", + " Only whole row groups are read, so the download stops as soon as enough images are collected.\n", + "\n", + " Args:\n", + " url: HTTP(S) location of the parquet shard.\n", + " num_images: Number of images to return.\n", + "\n", + " Returns:\n", + " Decoded samples with ground-truth boxes in ``xyxy`` pixel coordinates.\n", + " \"\"\"\n", + " samples: list[Sample] = []\n", + " with fsspec.open(url) as file:\n", + " parquet_file = pq.ParquetFile(file)\n", + " for group_index in range(parquet_file.metadata.num_row_groups):\n", + " if len(samples) >= num_images:\n", + " break\n", + " table = parquet_file.read_row_group(group_index)\n", + " for row in table.to_pylist():\n", + " if len(samples) >= num_images:\n", + " break\n", + " objects = row[\"objects\"]\n", + " keep = [i for i, crowd in enumerate(objects[\"iscrowd\"]) if not crowd]\n", + " if not keep:\n", + " continue\n", + " # COCO stores boxes as [x, y, width, height]; convert to xyxy.\n", + " xywh = np.asarray([objects[\"bbox\"][i] for i in keep], dtype=np.float32)\n", + " boxes = np.column_stack([xywh[:, 0], xywh[:, 1], xywh[:, 0] + xywh[:, 2], xywh[:, 1] + xywh[:, 3]])\n", + " samples.append(\n", + " Sample(\n", + " image_id=int(row[\"image_id\"]),\n", + " image=Image.open(io.BytesIO(row[\"image\"][\"bytes\"])).convert(\"RGB\"),\n", + " boxes=boxes,\n", + " labels=np.asarray([objects[\"label\"][i] for i in keep], dtype=np.int64),\n", + " )\n", + " )\n", + " return samples\n", + "\n", + "\n", + "samples = load_coco_val_subset(PARQUET_URL, NUM_IMAGES)\n", + "print(f\"images: {len(samples)}\")\n", + "print(f\"ground-truth objects: {sum(len(s.labels) for s in samples)}\")" + ] + }, + { + "cell_type": "markdown", + "id": "d68e5bca", + "metadata": {}, + "source": [ + "## 3. Extract per-detection embeddings\n", + "\n", + "This is the whole public interface: pass `return_embeddings=True` to `predict()`. The embeddings are gathered with\n", + "the same indices used for boxes, scores, and classes, so row `i` of `detections.data[\"embeddings\"]` belongs to\n", + "detection `i` — no bookkeeping required on your side." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0c24a952", + "metadata": { + "lines_to_next_cell": 1 + }, + "outputs": [], + "source": [ + "model = RFDETRSmall()\n", + "\n", + "detections_per_image = []\n", + "for start in range(0, len(samples), 8):\n", + " batch = [sample.image for sample in samples[start : start + 8]]\n", + " detections_per_image.extend(model.predict(batch, threshold=CONFIDENCE_THRESHOLD, return_embeddings=True))\n", + "\n", + "first_with_objects = next(det for det in detections_per_image if len(det) > 0)\n", + "embeddings = first_with_objects.data[\"embeddings\"]\n", + "print(f\"detections in this image: {len(first_with_objects)}\")\n", + "print(f\"embeddings shape: {embeddings.shape} (one row per detection, {embeddings.shape[1]}-d)\")\n", + "print(f\"embeddings dtype: {embeddings.dtype}\")" + ] + }, + { + "cell_type": "markdown", + "id": "50e0ca75", + "metadata": {}, + "source": [ + "## 4. Associate detections with annotations\n", + "\n", + "To audit *annotations*, each detection is matched to the ground-truth box it covers (greedy, class-agnostic, by\n", + "IoU). Matched detections inherit the annotation's label and become the audit set. Unmatched confident detections\n", + "are kept aside for section 7." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "172e77df", + "metadata": {}, + "outputs": [], + "source": [ + "def iou_matrix(boxes_a: np.ndarray, boxes_b: np.ndarray) -> np.ndarray:\n", + " \"\"\"Compute the pairwise IoU between two sets of ``xyxy`` boxes.\n", + "\n", + " Args:\n", + " boxes_a: Boxes with shape ``(N, 4)``.\n", + " boxes_b: Boxes with shape ``(M, 4)``.\n", + "\n", + " Returns:\n", + " IoU values with shape ``(N, M)``.\n", + " \"\"\"\n", + " if len(boxes_a) == 0 or len(boxes_b) == 0:\n", + " return np.zeros((len(boxes_a), len(boxes_b)), dtype=np.float32)\n", + " top_left = np.maximum(boxes_a[:, None, :2], boxes_b[None, :, :2])\n", + " bottom_right = np.minimum(boxes_a[:, None, 2:], boxes_b[None, :, 2:])\n", + " overlap = np.clip(bottom_right - top_left, 0.0, None).prod(axis=2)\n", + " area_a = (boxes_a[:, 2] - boxes_a[:, 0]) * (boxes_a[:, 3] - boxes_a[:, 1])\n", + " area_b = (boxes_b[:, 2] - boxes_b[:, 0]) * (boxes_b[:, 3] - boxes_b[:, 1])\n", + " return overlap / np.clip(area_a[:, None] + area_b[None, :] - overlap, 1e-9, None)\n", + "\n", + "\n", + "def match_to_ground_truth(detection_boxes: np.ndarray, gt_boxes: np.ndarray, iou_threshold: float) -> np.ndarray:\n", + " \"\"\"Greedily assign each detection to at most one ground-truth box.\n", + "\n", + " Pairs are consumed in descending IoU order, so a ground-truth box is claimed by its best detection only.\n", + "\n", + " Args:\n", + " detection_boxes: Predicted boxes with shape ``(N, 4)`` in ``xyxy``.\n", + " gt_boxes: Ground-truth boxes with shape ``(M, 4)`` in ``xyxy``.\n", + " iou_threshold: Minimum IoU for a valid match.\n", + "\n", + " Returns:\n", + " Index of the matched ground-truth box per detection, or ``-1`` where unmatched, shape ``(N,)``.\n", + " \"\"\"\n", + " ious = iou_matrix(detection_boxes, gt_boxes)\n", + " assignment = np.full(len(detection_boxes), -1, dtype=np.int64)\n", + " if ious.size == 0:\n", + " return assignment\n", + " taken_gt: set[int] = set()\n", + " order = np.dstack(np.unravel_index(np.argsort(ious, axis=None)[::-1], ious.shape))[0]\n", + " for detection_index, gt_index in order:\n", + " if ious[detection_index, gt_index] < iou_threshold:\n", + " break\n", + " if assignment[detection_index] != -1 or gt_index in taken_gt:\n", + " continue\n", + " assignment[detection_index] = gt_index\n", + " taken_gt.add(int(gt_index))\n", + " return assignment\n", + "\n", + "\n", + "matched_embeddings: list[np.ndarray] = []\n", + "matched_labels: list[int] = [] # ground-truth category id\n", + "matched_sources: list[tuple[int, int]] = [] # (sample index, detection index)\n", + "unannotated_sources: list[tuple[int, int]] = []\n", + "\n", + "for sample_index, (sample, detections) in enumerate(zip(samples, detections_per_image)):\n", + " if len(detections) == 0:\n", + " continue\n", + " assignment = match_to_ground_truth(detections.xyxy, sample.boxes, IOU_MATCH_THRESHOLD)\n", + " best_iou = iou_matrix(detections.xyxy, sample.boxes).max(axis=1, initial=0.0)\n", + " for detection_index, gt_index in enumerate(assignment):\n", + " if gt_index >= 0:\n", + " matched_embeddings.append(detections.data[\"embeddings\"][detection_index])\n", + " matched_labels.append(int(sample.labels[gt_index]))\n", + " matched_sources.append((sample_index, detection_index))\n", + " elif best_iou[detection_index] < IOU_UNANNOTATED_THRESHOLD:\n", + " unannotated_sources.append((sample_index, detection_index))\n", + "\n", + "embeddings = np.stack(matched_embeddings)\n", + "labels = np.asarray(matched_labels)\n", + "# Cosine similarity is the natural metric here, so work with unit-norm rows throughout.\n", + "unit_embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)\n", + "\n", + "print(f\"matched detections (audit set): {len(embeddings)}\")\n", + "print(f\"confident detections without any annotation: {len(unannotated_sources)}\")\n", + "print(f\"distinct categories: {len(np.unique(labels))}\")" + ] + }, + { + "cell_type": "markdown", + "id": "606235e2", + "metadata": {}, + "source": [ + "## 5. Do the embeddings carry object semantics?\n", + "\n", + "Two checks before relying on them. First a 2D projection (PCA to 50 dimensions, then t-SNE) coloured by the\n", + "ground-truth category: same-category detections should land together. Second a leave-one-out k-NN probe: if a\n", + "detection's neighbours predict its category well above the majority-class baseline, the geometry is meaningful,\n", + "which is exactly what the audit in section 6 depends on." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8852b33e", + "metadata": {}, + "outputs": [], + "source": [ + "top_categories = [int(c) for c, _ in sorted(zip(*np.unique(labels, return_counts=True)), key=lambda p: -p[1])[:8]]\n", + "plot_mask = np.isin(labels, top_categories)\n", + "\n", + "projection = PCA(n_components=min(50, unit_embeddings.shape[1]), random_state=SEED).fit_transform(unit_embeddings)\n", + "projected_2d = TSNE(n_components=2, init=\"pca\", perplexity=30, random_state=SEED).fit_transform(projection)\n", + "\n", + "plt.figure(figsize=(8, 6))\n", + "for category in top_categories:\n", + " category_mask = labels == category\n", + " plt.scatter(\n", + " projected_2d[category_mask, 0], projected_2d[category_mask, 1], s=14, alpha=0.75, label=COCO_CLASSES[category]\n", + " )\n", + "plt.legend(loc=\"best\", fontsize=8)\n", + "plt.title(f\"RF-DETR query embeddings ({int(plot_mask.sum())} detections, t-SNE)\")\n", + "plt.xticks([])\n", + "plt.yticks([])\n", + "plt.tight_layout()\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a5686b82", + "metadata": {}, + "outputs": [], + "source": [ + "neighbours = NearestNeighbors(n_neighbors=K_NEIGHBOURS + 1, metric=\"cosine\").fit(unit_embeddings)\n", + "_, neighbour_indices = neighbours.kneighbors(unit_embeddings)\n", + "neighbour_indices = neighbour_indices[:, 1:] # drop self-match\n", + "\n", + "neighbour_labels = labels[neighbour_indices]\n", + "knn_prediction = np.asarray([np.bincount(row).argmax() for row in neighbour_labels])\n", + "knn_accuracy = float((knn_prediction == labels).mean())\n", + "majority_baseline = float(np.bincount(labels).max() / len(labels))\n", + "\n", + "print(f\"{K_NEIGHBOURS}-NN leave-one-out accuracy: {knn_accuracy:.3f}\")\n", + "print(f\"majority-class baseline: {majority_baseline:.3f}\")" + ] + }, + { + "cell_type": "markdown", + "id": "f5543d09", + "metadata": {}, + "source": [ + "## 6. Find mislabelled annotations\n", + "\n", + "The industrial-inspection scenario: an operator mislabels some objects and you want the review queue ordered so\n", + "the worst annotations surface first.\n", + "\n", + "To *measure* this rather than guess, we corrupt a known 10% of the matched annotations by reassigning a random\n", + "different category. Each annotation is then scored by how much its own label disagrees with its neighbourhood in\n", + "embedding space — the neighbours are visually similar objects, so a label they do not share is suspicious:\n", + "\n", + "```\n", + "suspicion = 1 - (similarity-weighted share of the k nearest neighbours carrying the same label)\n", + "```\n", + "\n", + "Nothing about which annotations were corrupted enters the score; the corruption mask is only used afterwards to\n", + "grade the ranking with ROC-AUC and precision@k." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4c48cc99", + "metadata": {}, + "outputs": [], + "source": [ + "noisy_labels = labels.copy()\n", + "category_pool = np.unique(labels)\n", + "num_corrupted = int(NOISE_FRACTION * len(labels))\n", + "corrupted_indices = rng.choice(len(labels), size=num_corrupted, replace=False)\n", + "for index in corrupted_indices:\n", + " alternatives = category_pool[category_pool != labels[index]]\n", + " noisy_labels[index] = rng.choice(alternatives)\n", + "\n", + "is_corrupted = np.zeros(len(labels), dtype=bool)\n", + "is_corrupted[corrupted_indices] = True\n", + "print(f\"corrupted annotations: {is_corrupted.sum()} / {len(labels)}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f5b15a13", + "metadata": {}, + "outputs": [], + "source": [ + "neighbour_similarity = 1.0 - neighbours.kneighbors(unit_embeddings)[0][:, 1:]\n", + "agreement = (noisy_labels[neighbour_indices] == noisy_labels[:, None]).astype(np.float32)\n", + "suspicion = 1.0 - (agreement * neighbour_similarity).sum(axis=1) / np.clip(neighbour_similarity.sum(axis=1), 1e-9, None)\n", + "\n", + "ranking = np.argsort(suspicion)[::-1]\n", + "roc_auc = roc_auc_score(is_corrupted, suspicion)\n", + "review_budget = num_corrupted\n", + "precision_at_k = float(is_corrupted[ranking[:review_budget]].mean())\n", + "random_precision = float(is_corrupted.mean())\n", + "\n", + "print(f\"ROC-AUC of the suspicion score: {roc_auc:.3f}\")\n", + "print(f\"precision@{review_budget} (embedding ranking): {precision_at_k:.3f}\")\n", + "print(f\"precision@{review_budget} (random review): {random_precision:.3f}\")\n", + "print(f\"lift over random review: {precision_at_k / max(random_precision, 1e-9):.1f}x\")" + ] + }, + { + "cell_type": "markdown", + "id": "0226d4c1", + "metadata": {}, + "source": [ + "Reading the numbers: ROC-AUC is the probability that a corrupted annotation outranks a clean one, and\n", + "precision@k is the share of genuinely bad annotations in a review queue the size of the corruption. Compare it\n", + "against the random-review baseline — that gap is the labour the embeddings save.\n", + "\n", + "The plot below shows the score distributions, and the grid shows the top-ranked suspects with their (possibly\n", + "corrupted) label and the label their neighbourhood votes for." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "49e9e44e", + "metadata": {}, + "outputs": [], + "source": [ + "plt.figure(figsize=(8, 4))\n", + "plt.hist(suspicion[~is_corrupted], bins=40, alpha=0.7, density=True, label=\"clean annotations\")\n", + "plt.hist(suspicion[is_corrupted], bins=40, alpha=0.7, density=True, label=\"corrupted annotations\")\n", + "plt.xlabel(\"suspicion score\")\n", + "plt.ylabel(\"density\")\n", + "plt.title(f\"Label-error ranking (ROC-AUC {roc_auc:.3f})\")\n", + "plt.legend()\n", + "plt.tight_layout()\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2175eb8a", + "metadata": {}, + "outputs": [], + "source": [ + "def crop_detection(index: int, padding: float = 0.06) -> Image.Image:\n", + " \"\"\"Crop the image region of an audited detection.\n", + "\n", + " Args:\n", + " index: Row index into the matched-detection arrays.\n", + " padding: Fraction of the box size added on every side for context.\n", + "\n", + " Returns:\n", + " The cropped RGB region.\n", + " \"\"\"\n", + " sample_index, detection_index = matched_sources[index]\n", + " sample = samples[sample_index]\n", + " x1, y1, x2, y2 = detections_per_image[sample_index].xyxy[detection_index]\n", + " pad_x, pad_y = padding * (x2 - x1), padding * (y2 - y1)\n", + " box = (\n", + " max(0.0, x1 - pad_x),\n", + " max(0.0, y1 - pad_y),\n", + " min(float(sample.image.width), x2 + pad_x),\n", + " min(float(sample.image.height), y2 + pad_y),\n", + " )\n", + " return sample.image.crop(box)\n", + "\n", + "\n", + "neighbour_vote = np.asarray([np.bincount(noisy_labels[row]).argmax() for row in neighbour_indices])\n", + "\n", + "figure, axes = plt.subplots(2, 5, figsize=(15, 7))\n", + "for axis, index in zip(axes.ravel(), ranking[:10]):\n", + " axis.imshow(crop_detection(index))\n", + " verdict = \"injected error\" if is_corrupted[index] else \"flagged\"\n", + " axis.set_title(\n", + " f\"label: {COCO_CLASSES[noisy_labels[index]]}\\nneighbours: {COCO_CLASSES[neighbour_vote[index]]}\\n{verdict}\",\n", + " fontsize=9,\n", + " )\n", + " axis.axis(\"off\")\n", + "figure.suptitle(\"Top-ranked annotation suspects\")\n", + "figure.tight_layout()\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "936e0888", + "metadata": {}, + "source": [ + "Suspects that are *not* injected errors are worth a look rather than a dismissal: they are typically ambiguous\n", + "crops, near-duplicate categories (`car` vs `truck`), or annotations that were already questionable in the\n", + "original dataset." + ] + }, + { + "cell_type": "markdown", + "id": "86143dca", + "metadata": {}, + "source": [ + "## 7. Find unannotated objects\n", + "\n", + "The same pass also finds the opposite failure mode. A confident detection that overlaps no annotation is either a\n", + "false positive or an object nobody labelled — the second is a dataset bug that silently penalises training and\n", + "evaluation. Grouping the candidates by nearest neighbours among *audited* detections tells you what the model\n", + "thinks they are." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d5211395", + "metadata": {}, + "outputs": [], + "source": [ + "unannotated_sorted = sorted(\n", + " unannotated_sources,\n", + " key=lambda source: float(detections_per_image[source[0]].confidence[source[1]]),\n", + " reverse=True,\n", + ")\n", + "\n", + "figure, axes = plt.subplots(2, 5, figsize=(15, 7))\n", + "for axis, (sample_index, detection_index) in zip(axes.ravel(), unannotated_sorted[:10]):\n", + " sample = samples[sample_index]\n", + " detections = detections_per_image[sample_index]\n", + " x1, y1, x2, y2 = detections.xyxy[detection_index]\n", + " axis.imshow(sample.image.crop((float(x1), float(y1), float(x2), float(y2))))\n", + " axis.set_title(\n", + " f\"{COCO_CLASSES[int(detections.class_id[detection_index])]}\"\n", + " f\" ({detections.confidence[detection_index]:.2f})\\nno annotation\",\n", + " fontsize=9,\n", + " )\n", + " axis.axis(\"off\")\n", + "figure.suptitle(\"Confident detections with no matching annotation\")\n", + "figure.tight_layout()\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "ca1a69e9", + "metadata": {}, + "source": [ + "## 8. Similarity search over objects\n", + "\n", + "Because the embeddings live in a shared space, a single detection can be used as a query to retrieve visually\n", + "similar instances across the dataset — the building block for near-duplicate removal, re-identification, and\n", + "\"show me more objects like this one\" review tooling." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1a87a175", + "metadata": {}, + "outputs": [], + "source": [ + "query_index = int(ranking[0])\n", + "_, retrieved = neighbours.kneighbors(unit_embeddings[query_index : query_index + 1], n_neighbors=6)\n", + "\n", + "figure, axes = plt.subplots(1, 6, figsize=(16, 3.2))\n", + "for rank, (axis, index) in enumerate(zip(axes, retrieved[0])):\n", + " axis.imshow(crop_detection(int(index)))\n", + " axis.set_title(\"query\" if rank == 0 else f\"#{rank} · {COCO_CLASSES[labels[index]]}\", fontsize=9)\n", + " axis.axis(\"off\")\n", + "figure.suptitle(\"Nearest neighbours of a query object\")\n", + "figure.tight_layout()\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "381ce55d", + "metadata": {}, + "source": [ + "## 9. Embeddings from an optimized model\n", + "\n", + "`model.inference()` traces/compiles the forward pass, and a traced graph has fixed control flow: whether\n", + "embeddings are produced cannot be decided per call any more. Pass the flag to `inference()` and keep `predict()`\n", + "consistent with it — a mismatch raises `RuntimeError` instead of silently returning nothing.\n", + "\n", + "```python\n", + "model = RFDETRSmall()\n", + "model.inference(return_embeddings=True) # decided once, at optimization time\n", + "\n", + "detections = model.predict(image, threshold=0.5, return_embeddings=True)\n", + "detections.data[\"embeddings\"] # (K, H)\n", + "```\n", + "\n", + "The eager path used throughout this notebook has no such constraint — `return_embeddings` is a plain per-call\n", + "argument.\n", + "\n", + "Segmentation models behave identically. Keypoint models attach embeddings to `key_points.data[\"embeddings\"]`,\n", + "since they return `sv.KeyPoints` rather than `sv.Detections`." + ] + }, + { + "cell_type": "markdown", + "id": "aef1b3ba", + "metadata": {}, + "source": [ + "## 10. Takeaways\n", + "\n", + "- `predict(..., return_embeddings=True)` returns one embedding per detection in `detections.data[\"embeddings\"]`,\n", + " aligned row-for-row with boxes, scores, and classes.\n", + "- They are the decoder hidden states the detection heads already consume, so they cost nothing extra to obtain\n", + " and require no second model.\n", + "- Neighbourhood disagreement in that space ranks mislabelled annotations far above random review, and detections\n", + " without ground-truth overlap surface missing annotations — both measured above on a public dataset.\n", + "- The same vectors support similarity search, near-duplicate detection, clustering, and active-learning\n", + " selection.\n", + "\n", + "**Next steps:** run this over your own dataset with a fine-tuned checkpoint (`RFDETRSmall(pretrain_weights=...)`),\n", + "and route the top-ranked suspects into your labelling tool as a review queue." + ] + } + ], + "metadata": { + "jupytext": { + "cell_metadata_filter": "-all", + "main_language": "python", + "notebook_metadata_filter": "-all" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/learn/run/detection.md b/docs/learn/run/detection.md index 935d181a3..67057fbd2 100644 --- a/docs/learn/run/detection.md +++ b/docs/learn/run/detection.md @@ -70,6 +70,33 @@ For memory-constrained inference-only deployments with the `rfdetr` package, opt model.inference(compile=False, inplace=True, dtype="float16") ``` +## Extract Embeddings + +Pass `return_embeddings=True` to `predict()` to also get a per-detection embedding vector, gathered with the same indices used for boxes (and masks/keypoints, where applicable). This is useful for downstream tasks like similarity search, clustering, or re-identification. Embeddings are attached as `detections.data["embeddings"]` with shape `(K, H)` — one row per detection. + +```python +from rfdetr import RFDETRMedium + +model = RFDETRMedium() + +detections = model.predict("https://media.roboflow.com/dog.jpg", threshold=0.5, return_embeddings=True) + +embeddings = detections.data["embeddings"] # shape (K, H) +``` + +!!! note "Optimized models decide this at `inference()` time" + + On a model optimized with `model.inference(...)`, the exported/traced forward pass has fixed control flow, so whether embeddings are computed can't be toggled per `predict()` call — it must match the `return_embeddings` value passed to `inference()`: + + ```python + model.inference(compile=False, return_embeddings=True) + detections = model.predict(image, return_embeddings=True) + ``` + + Calling `predict(return_embeddings=...)` with a value that doesn't match the optimized model raises `RuntimeError`. + +For an end-to-end walkthrough — ranking mislabelled annotations, surfacing unannotated objects, and object-level similarity search — see the [Query Embeddings for Dataset Quality Auditing](../../cookbooks/embedding-extraction.ipynb) cookbook. + ## Run on video, webcam, or RTSP stream These examples use OpenCV for decoding and display. Replace ``, ``, and `` with your inputs. `` is usually `0` for the default camera. diff --git a/docs/learn/run/keypoints.md b/docs/learn/run/keypoints.md index faab1917e..e0d9a3c80 100644 --- a/docs/learn/run/keypoints.md +++ b/docs/learn/run/keypoints.md @@ -68,6 +68,8 @@ Keypoints with `visible=False` are skipped by supervision annotators. To hide lo For fine-tuning on a custom keypoint dataset, see [Keypoint preview custom datasets](../train/index.md#keypoint-preview-custom-datasets). +Pass `return_embeddings=True` to `predict()` to also get a per-instance embedding vector attached as `key_points.data["embeddings"]`. See [Extract Embeddings](detection.md#extract-embeddings) for details, including the constraints that apply to optimized models. + ## Run on video, webcam, or RTSP stream These examples use OpenCV for decoding and display. Replace ``, ``, and `` with your inputs. `` is usually `0` for the default camera. diff --git a/docs/learn/run/segmentation.md b/docs/learn/run/segmentation.md index 1e7f8a4b6..528a7c772 100644 --- a/docs/learn/run/segmentation.md +++ b/docs/learn/run/segmentation.md @@ -64,6 +64,8 @@ For memory-constrained inference-only deployments with the `rfdetr` package, opt model.inference(compile=False, inplace=True, dtype="float16") ``` +Pass `return_embeddings=True` to `predict()` to also get a per-detection embedding vector attached as `detections.data["embeddings"]`. See [Extract Embeddings](detection.md#extract-embeddings) for details, including the constraints that apply to optimized models. + ## Run on video, webcam, or RTSP stream These examples use OpenCV for decoding and display. Replace ``, ``, and `` with your inputs. `` is usually `0` for the default camera. diff --git a/src/rfdetr/detr.py b/src/rfdetr/detr.py index 21ce601e5..9b7987ea7 100644 --- a/src/rfdetr/detr.py +++ b/src/rfdetr/detr.py @@ -416,6 +416,10 @@ def __init__(self, *, trust_checkpoint: bool = False, **kwargs: Any) -> None: self._optimized_resolution: int | None = None self._optimized_dtype: torch.dtype | None = None self._optimized_inplace = False + # Whether the currently optimized `inference_model` was exported with embeddings enabled. + # Unlike the eager model, the exported/traced forward pass can't take `return_embeddings` + # as a runtime argument, so this is fixed at `inference()` time. + self._optimized_return_embeddings = False self._has_been_trained = False def maybe_download_pretrain_weights(self) -> None: @@ -1208,6 +1212,7 @@ def inference( dtype: torch.dtype | str = torch.float32, *, inplace: bool = False, + return_embeddings: bool = False, ) -> None: """Optimize the model for inference with optional JIT compilation and dtype casting. @@ -1239,6 +1244,11 @@ def inference( inference-only path because ``export()`` mutates the module and dtype casting mutates its parameters. Requires ``compile=False``. With the default ``dtype=torch.float32``, the dtype cast is a no-op, so memory savings come only from clearing the base model reference rather than from dtype reduction. + return_embeddings: If ``True``, the optimized model also returns per-query embeddings from + ``predict(..., return_embeddings=True)``. Unlike the unoptimized model, this cannot be toggled per + call: the exported/traced forward pass has fixed control flow, so whether embeddings are computed + must be decided here, at optimization time. Calling ``predict(return_embeddings=True)`` on a model + optimized with ``return_embeddings=False`` (or vice versa) raises ``RuntimeError``. Raises: TypeError: If ``dtype`` is not a ``torch.dtype``, or if ``dtype`` is a @@ -1256,7 +1266,7 @@ def inference( ... self.linear = torch.nn.Linear(1, 1) ... def forward(self, x): ... return {"pred_boxes": self.linear(x[:, :1, :1, :1].squeeze(-1).squeeze(-1))} - ... def export(self): + ... def export(self, return_embeddings=False): ... return None >>> class _TinyContext: ... def __init__(self): @@ -1325,7 +1335,7 @@ def inference( with cuda_ctx: inference_model: Any = self.model.model if inplace else deepcopy(self.model.model) inference_model.eval() - inference_model.export() + inference_model.export(return_embeddings=return_embeddings) inference_model = inference_model.to(dtype=dtype) @@ -1356,6 +1366,7 @@ def inference( self._optimized_resolution = self.model.resolution self._is_optimized_for_inference = True self._optimized_dtype = dtype + self._optimized_return_embeddings = return_embeddings except Exception: # Ensure the object is left in a consistent, unoptimized state if optimization fails. with contextlib.suppress(Exception): @@ -1402,7 +1413,7 @@ def remove_optimized_model(self) -> None: ... self.linear = torch.nn.Linear(1, 1) ... def forward(self, x): ... return {"pred_boxes": self.linear(x[:, :1, :1, :1].squeeze(-1).squeeze(-1))} - ... def export(self): + ... def export(self, return_embeddings=False): ... return None >>> class _TinyContext: ... def __init__(self): @@ -1441,6 +1452,7 @@ def remove_optimized_model(self) -> None: self._optimized_resolution = None self._optimized_dtype = None self._optimized_inplace = False + self._optimized_return_embeddings = False @property def is_optimized_inplace(self) -> bool: @@ -1458,7 +1470,7 @@ def is_optimized_inplace(self) -> bool: ... self.linear = torch.nn.Linear(1, 1) ... def forward(self, x): ... return {"pred_boxes": self.linear(x[:, :1, :1, :1].squeeze(-1).squeeze(-1))} - ... def export(self): + ... def export(self, return_embeddings=False): ... return None >>> class _TinyContext: ... def __init__(self): @@ -2247,6 +2259,7 @@ def predict( shape: tuple[int, int] | None = None, patch_size: int | None = None, include_source_image: bool = True, + return_embeddings: bool = False, **kwargs: Any, ) -> Detections | KeyPoints | list[Detections | KeyPoints]: """Performs model inference on the input images. @@ -2276,6 +2289,14 @@ def predict( ``key_points.data["source_image"]`` because Supervision ``KeyPoints`` currently has no collection-level metadata field. Defaults to ``True``. Set to ``False`` to reduce memory use when source images are not needed. + return_embeddings: + Whether to also return per-detection embeddings, one per selected query, gathered with the same + indices used for boxes/masks/keypoints. Embeddings are attached as + ``detections.data["embeddings"]`` (shape ``(K, H)``) for detection/segmentation outputs, or + ``key_points.data["embeddings"]`` for keypoint outputs. If the model has been optimized via + :meth:`inference`, this must match the ``return_embeddings`` value passed to that call — the + exported/traced forward pass has fixed control flow and cannot toggle this per call; mismatches + raise ``RuntimeError``. **kwargs: Additional keyword arguments. @@ -2516,14 +2537,32 @@ class and ``class_id=1`` is ``"__background__"``. ) if self._is_optimized_for_inference: + if return_embeddings != self._optimized_return_embeddings: + raise RuntimeError( + f"predict(return_embeddings={return_embeddings}) does not match the optimized model, which was " + f"prepared with inference(return_embeddings={self._optimized_return_embeddings}). The " + "exported/traced forward pass has fixed control flow, so this must be decided at " + "inference()/optimization time, not per predict() call. Call " + f"model.inference(..., return_embeddings={return_embeddings}) (optionally after " + "model.remove_optimized_model()) to change it.", + ) inference_model = self.model.inference_model assert inference_model is not None, "inference_model is set whenever _is_optimized_for_inference is True." predictions = inference_model(batch_tensor.to(dtype=self._optimized_dtype)) else: model = self.model.model assert model is not None, "self.model.model is only cleared when optimized for inference." - predictions = model(batch_tensor) + predictions = model(batch_tensor, return_embeddings=return_embeddings) if isinstance(predictions, tuple): + # Only the exported/traced (optimized) forward pass returns a plain tuple; its structure is: + # (pred_boxes, pred_logits[, pred_masks | pred_keypoints][, embeddings]). + # Embeddings are always the *last* element when `_optimized_return_embeddings` is True, regardless + # of whether masks/keypoints are also present, so we pop them off first rather than relying on + # `len(predictions)` alone to infer the full structure. + embeddings = None + if self._optimized_return_embeddings: + embeddings = predictions[-1] + predictions = predictions[:-1] return_predictions = { "pred_logits": predictions[1], "pred_boxes": predictions[0], @@ -2534,6 +2573,8 @@ class and ``class_id=1`` is ``"__background__"``. return_predictions["pred_keypoints"] = predictions[2] else: return_predictions["pred_masks"] = predictions[2] + if embeddings is not None: + return_predictions["embeddings"] = embeddings predictions = return_predictions target_sizes = torch.tensor(orig_sizes, device=self.model.device) results = self.model.postprocess(predictions, target_sizes=target_sizes, score_threshold=threshold) @@ -2593,6 +2634,9 @@ class and ``class_id=1`` is ``"__background__"``. keypoints = result["keypoints"][keep] keypoints_array = keypoints.float().cpu().numpy() has_keypoints = keypoints_array is not None + embeddings_array = None + if "embeddings" in result: + embeddings_array = result["embeddings"][keep].float().cpu().numpy() if "masks" in result: masks = result["masks"] @@ -2617,6 +2661,8 @@ class and ``class_id=1`` is ``"__background__"``. if include_source_image: detections.metadata["source_image"] = source_images[i] # type: ignore[index] detections.data["source_shape"] = np.tile(np.array(orig_sizes[i], dtype=np.int64), (len(detections), 1)) + if embeddings_array is not None: + detections.data["embeddings"] = embeddings_array # Attach class names so callers can map class_id → name without a # separate lookup. Always set data["class_name"] for a consistent interface. diff --git a/src/rfdetr/evaluation/coco_eval.py b/src/rfdetr/evaluation/coco_eval.py index 7a5f7189a..ca6c658a4 100644 --- a/src/rfdetr/evaluation/coco_eval.py +++ b/src/rfdetr/evaluation/coco_eval.py @@ -175,7 +175,9 @@ def _resolve_keypoint_oks_sigmas(coco_gt: COCO, keypoint_oks_sigmas: list[float] return None _warn_custom_keypoint_oks_sigma_once(keypoint_count) - return np.full(keypoint_count, _DEFAULT_CUSTOM_KEYPOINT_OKS_SIGMA, dtype=np.float32).tolist() + return [ + float(sigma) for sigma in np.full(keypoint_count, _DEFAULT_CUSTOM_KEYPOINT_OKS_SIGMA, dtype=np.float32).tolist() + ] def _resolve_group_keypoint_oks_sigmas( @@ -187,7 +189,10 @@ def _resolve_group_keypoint_oks_sigmas( if keypoint_count == len(_COCO_PERSON_KEYPOINT_SIGMAS): return None _warn_custom_keypoint_oks_sigma_once(keypoint_count) - return np.full(keypoint_count, _DEFAULT_CUSTOM_KEYPOINT_OKS_SIGMA, dtype=np.float32).tolist() + return [ + float(sigma) + for sigma in np.full(keypoint_count, _DEFAULT_CUSTOM_KEYPOINT_OKS_SIGMA, dtype=np.float32).tolist() + ] sigmas = np.asarray(keypoint_oks_sigmas, dtype=np.float32) if sigmas.ndim != 1 or sigmas.size == 0: diff --git a/src/rfdetr/models/lwdetr.py b/src/rfdetr/models/lwdetr.py index a4ed9d916..4cfec08b0 100644 --- a/src/rfdetr/models/lwdetr.py +++ b/src/rfdetr/models/lwdetr.py @@ -226,6 +226,7 @@ def __init__( ) self._export = False + self._export_return_embeddings = False def reinitialize_detection_head(self, num_classes: int) -> None: """Resize the detection classification head to *num_classes* outputs. @@ -347,8 +348,23 @@ def reset_keypoint_gaussian_parameters(self) -> None: for keypoint_embed in enc_keypoint_embed: _reset_keypoint_gaussian_output_rows(keypoint_embed) - def export(self) -> None: + def export(self, return_embeddings: bool = False) -> None: + """Switch ``forward`` to the traced/compiled ``forward_export`` path. + + Also recursively calls ``export()`` on any submodule that exposes it, so the whole model + (backbone, transformer, segmentation head, etc.) is prepared for tracing/compilation together. + + Args: + return_embeddings: If ``True``, ``forward_export`` also appends per-query embeddings as the last + element of its output tuple. Unlike the eager ``forward``, ``forward_export`` is traced/compiled + with fixed control flow, so this must be decided here, at export time, rather than passed as a + runtime argument to ``forward_export`` itself. + """ self._export = True + # `forward_export` is traced/compiled with a fixed control flow, so whether + # embeddings are appended to the output tuple must be decided at export time + # rather than passed as a runtime argument (unlike the eager `forward`). + self._export_return_embeddings = return_embeddings self._forward_origin = self.forward self.forward = self.forward_export # type: ignore[method-assign,assignment] for name, m in self.named_modules(): @@ -456,7 +472,12 @@ def _aggregate_keypoint_class_logits(self, keypoint_predictions: Tensor) -> Tens class_boost = class_boost[..., :detection_num_classes] return class_boost - def forward(self, samples: NestedTensor, targets: list[dict[str, Tensor]] | None = None) -> dict[str, Any]: + def forward( + self, + samples: NestedTensor, + targets: list[dict[str, Tensor]] | None = None, + return_embeddings: bool = False, + ) -> dict[str, Any]: """The forward expects a NestedTensor, which consists of: - samples.tensor: batched images, of shape [batch_size x 3 x H x W] @@ -471,6 +492,22 @@ def forward(self, samples: NestedTensor, targets: list[dict[str, Tensor]] | None information on how to retrieve the unnormalized bounding box. - "aux_outputs": Optional, only returned when auxiliary losses are activated. It is a list of dictionaries containing the two above keys for each decoder layer. + - "embeddings": Optional, only present when ``return_embeddings=True``. Per-query embeddings with shape + [batch_size x num_queries x hidden_dim], taken from the last decoder layer's hidden + state — consistent with the exported/traced path. + - "keypoint_embeddings": Optional, only present when ``return_embeddings=True`` and the model predicts + keypoints (GroupPose). Per-keypoint embeddings with shape + [batch_size x num_queries x num_keypoints x hidden_dim], taken from the last decoder layer. + + Args: + samples: Batched input images and their padding mask, or a list of image tensors to be + converted into a :class:`NestedTensor`. + targets: Optional ground-truth annotations, unused by the forward pass itself but accepted for + interface parity with training call sites. + return_embeddings: If ``True``, also computes and returns ``"embeddings"`` (and, for keypoint + models, ``"keypoint_embeddings"``) in the output dict. Unlike the exported/traced + ``forward_export``, this can be toggled per call since the eager forward has no fixed + control-flow constraint. """ if isinstance(samples, (list, Tensor)): samples = nested_tensor_from_tensor_list(samples) @@ -576,6 +613,12 @@ def forward(self, samples: NestedTensor, targets: list[dict[str, Tensor]] | None outputs_keypoints, ) + if return_embeddings and hs is not None: + # hs shape: [L, B, Q, H] — take only the last decoder layer to match the exported path. + out["embeddings"] = hs[-1] + if keypoint_hs is not None: + out["keypoint_embeddings"] = keypoint_hs[-1] + if self.two_stage: assert self.transformer.enc_out_class_embed is not None group_detr = self.group_detr if self.training else 1 @@ -621,6 +664,21 @@ def forward(self, samples: NestedTensor, targets: list[dict[str, Tensor]] | None return out def forward_export(self, tensors: Tensor) -> tuple[Tensor, ...]: + """Traced/compiled forward pass used after :meth:`export`. + + Unlike :meth:`forward`, this has fixed control flow (required for tracing/compilation), so it takes a + single image tensor (no ``NestedTensor``/mask) and returns a plain tuple instead of a dict. Whether + embeddings are included is decided once at :meth:`export` time via ``self._export_return_embeddings``, + not per call. + + Args: + tensors: Batched input images, shape ``[batch_size x 3 x H x W]``. + + Returns: + A tuple ``(pred_boxes, pred_logits[, pred_masks | pred_keypoints][, embeddings])``. Embeddings are + always the last element when present, so callers must gate on ``self._export_return_embeddings`` + (known at export time) rather than inferring structure from the tuple length alone. + """ srcs, _, poss, cross_attn_srcs = self.backbone(tensors) # only use one group in inference refpoint_embed_weight = self.refpoint_embed.weight[: self.num_queries] @@ -643,6 +701,7 @@ def forward_export(self, tensors: Tensor) -> tuple[Tensor, ...]: outputs_masks = None outputs_keypoints = None + outputs_embeddings: Tensor | None = None if hs is not None: if self.bbox_reparam: @@ -678,6 +737,10 @@ def forward_export(self, tensors: Tensor) -> tuple[Tensor, ...]: ], tensors.shape[-2:], )[0] + if self._export_return_embeddings: + # Unlike eager `forward`, the exported/traced decoder only returns the last + # layer's hidden state (shape [B, Q, H]), so no multi-layer reshape is needed here. + outputs_embeddings = hs else: assert self.two_stage, "if not using decoder, two_stage must be True" assert self.transformer.enc_out_class_embed is not None @@ -699,13 +762,25 @@ def forward_export(self, tensors: Tensor) -> tuple[Tensor, ...]: tensors.shape[-2:], skip_blocks=True, )[0] - + if self._export_return_embeddings: + outputs_embeddings = hs_enc + + # Embeddings are always appended as the *last* element of the output tuple (when + # present) so their position doesn't depend on whether masks/keypoints are also + # returned. Callers must gate on `self._export_return_embeddings` (known at export + # time) rather than inferring structure from `len(...)`, since that would be + # ambiguous once embeddings coexist with masks or keypoints. if outputs_masks is not None: - return outputs_coord, outputs_class, outputs_masks - if outputs_keypoints is not None: - return outputs_coord, outputs_class, outputs_keypoints + base_predictions: tuple[Tensor, ...] = (outputs_coord, outputs_class, outputs_masks) + elif outputs_keypoints is not None: + base_predictions = (outputs_coord, outputs_class, outputs_keypoints) else: - return outputs_coord, outputs_class + base_predictions = (outputs_coord, outputs_class) + + if self._export_return_embeddings: + assert outputs_embeddings is not None, "outputs_embeddings must be set when _export_return_embeddings." + return base_predictions + (outputs_embeddings,) + return base_predictions @torch.jit.unused def _set_aux_loss( diff --git a/src/rfdetr/models/postprocess.py b/src/rfdetr/models/postprocess.py index ec78864d1..b12bc4028 100644 --- a/src/rfdetr/models/postprocess.py +++ b/src/rfdetr/models/postprocess.py @@ -51,7 +51,7 @@ def forward( Args: outputs: Model output dictionary containing ``pred_logits`` and ``pred_boxes`` plus optional - ``pred_masks`` or ``pred_keypoints``. + ``pred_masks``, ``pred_keypoints``, or ``embeddings``. target_sizes: Per-image ``(height, width)`` tensor. For inference and evaluation this should be the original image size so normalized boxes and keypoints are returned in source-image pixel coordinates. score_threshold: Optional confidence threshold already known to the caller. Detections scoring at or @@ -71,13 +71,14 @@ def forward( out_logits, out_bbox = outputs["pred_logits"], outputs["pred_boxes"] out_masks = outputs.get("pred_masks") out_keypoints = outputs.get("pred_keypoints") + out_embeddings = outputs.get("embeddings") self._validate_outputs(out_logits, out_masks, out_keypoints, target_sizes) scores, labels, topk_boxes = self._select_topk(out_logits) boxes = self._gather_and_scale_boxes(out_bbox, topk_boxes, target_sizes) if out_masks is not None: - return self._postprocess_masks( + results = self._postprocess_masks( out_masks, scores, labels, @@ -87,9 +88,60 @@ def forward( self.upsample_masks_to_image_size, score_threshold=score_threshold, ) - if out_keypoints is not None: - return self._postprocess_keypoints(out_keypoints, scores, labels, boxes, topk_boxes, target_sizes) - return self._postprocess_boxes(scores, labels, boxes) + elif out_keypoints is not None: + results = self._postprocess_keypoints(out_keypoints, scores, labels, boxes, topk_boxes, target_sizes) + else: + results = self._postprocess_boxes(scores, labels, boxes) + + if out_embeddings is not None: + # `_postprocess_masks` drops rows scoring at or below `score_threshold` before returning + # `results`, so segmentation embeddings must be filtered with the same per-image predicate + # or their row count would diverge from `results[i]["scores"]`/`"masks"`. The keypoint and + # box-only paths never filter on `score_threshold` (see the `forward` docstring), so no + # filtering is needed there either. + self._attach_embeddings( + results, + out_embeddings, + topk_boxes, + scores, + score_threshold=score_threshold if out_masks is not None else None, + ) + + return results + + @staticmethod + def _attach_embeddings( + results: list[dict[str, torch.Tensor]], + out_embeddings: torch.Tensor, + topk_boxes: torch.Tensor, + scores: torch.Tensor | None = None, + score_threshold: float | None = None, + ) -> None: + """Gather per-query embeddings for the selected top-k queries and attach them in-place. + + Args: + results: Per-image result dicts already populated by :meth:`forward`; mutated in-place. + out_embeddings: Raw per-query embeddings with shape ``(B, Q, H)``. + topk_boxes: Selected query indices with shape ``(B, K)``, same indices used to gather + boxes/masks/keypoints so embeddings line up 1:1 with the returned detections. + scores: Optional pre-filter object scores with shape ``(B, K)``, matching ``topk_boxes``. + Required (together with ``score_threshold``) to reproduce the same row filtering + ``_postprocess_masks`` already applied to ``results``; ``None`` skips filtering. + score_threshold: Optional confidence threshold already applied to ``results`` by + ``_postprocess_masks``. When set (together with ``scores``), rows scoring at or below + it are dropped here too, so embeddings line up 1:1 with each image's filtered rows. + ``None`` (the default) keeps every selected query, matching the unfiltered paths. + """ + for i, res_i in enumerate(results): + k_idx = topk_boxes[i] + if score_threshold is not None and scores is not None: + sel = (scores[i] > score_threshold).nonzero(as_tuple=True)[0] + k_idx = k_idx[sel] + res_i["embeddings"] = torch.gather( + out_embeddings[i], + 0, + k_idx.unsqueeze(-1).repeat(1, out_embeddings.shape[-1]), + ) # [K, H] @staticmethod def _validate_outputs( diff --git a/tests/inference/helpers.py b/tests/inference/helpers.py index 4fb5c3333..395e63b07 100644 --- a/tests/inference/helpers.py +++ b/tests/inference/helpers.py @@ -21,6 +21,32 @@ from rfdetr.detr import RFDETR +class _IdentityAcceptingReturnEmbeddings(torch.nn.Module): + """``torch.nn.Identity``-like stub whose ``forward`` also accepts (and ignores) ``return_embeddings``. + + ``predict()`` always calls the unoptimized model with ``return_embeddings=``, mirroring the real + ``LWDETR.forward`` signature. Plain ``torch.nn.Identity`` doesn't accept that kwarg, so test doubles standing + in for the base model use this instead. + + Examples: + >>> stub = _IdentityAcceptingReturnEmbeddings() + >>> x = torch.zeros(1) + >>> torch.equal(stub(x, return_embeddings=True), x) + True + """ + + def __init__(self, hidden_dim: int = 4) -> None: + """Initialise the stub, remembering ``hidden_dim`` for the synthetic embeddings tensor.""" + super().__init__() + self.hidden_dim = hidden_dim + self.last_return_embeddings: bool | None = None + + def forward(self, x: torch.Tensor, return_embeddings: bool = False) -> torch.Tensor: + """Return the input unchanged, recording ``return_embeddings`` for later assertions.""" + self.last_return_embeddings = return_embeddings + return x + + class _BaseFakeRFDETR(RFDETR): """RFDETR test double that skips weight downloads and returns a minimal model config. @@ -55,15 +81,19 @@ def __init__( labels: list[int] | None = None, include_keypoints: bool = False, num_keypoints: int = 17, + include_embeddings: bool = False, + embedding_dim: int = 4, ) -> None: - """Initialise stub with optional class names, label list, and keypoint flag.""" + """Initialise stub with optional class names, label list, keypoint flag, and embeddings flag.""" self.device = torch.device("cpu") self.resolution = 28 - self.model = torch.nn.Identity() + self.model = _IdentityAcceptingReturnEmbeddings(hidden_dim=embedding_dim) self.class_names = class_names self._labels = labels if labels is not None else [1] self._include_keypoints = include_keypoints self._num_keypoints = num_keypoints + self._include_embeddings = include_embeddings + self._embedding_dim = embedding_dim def postprocess( self, @@ -85,6 +115,11 @@ def postprocess( result["keypoint_precision_cholesky"] = torch.full( (len(self._labels), self._num_keypoints, 3), 0.25, dtype=torch.float32 ) + if self._include_embeddings: + # Identifiable per-label embeddings: row i is filled with value i, so tests can assert on content. + result["embeddings"] = torch.stack( + [torch.full((self._embedding_dim,), float(i)) for i in range(len(self._labels))] + ) results.append(result) return results diff --git a/tests/inference/test_model_inference.py b/tests/inference/test_model_inference.py index 99cd54b2f..776779f0a 100644 --- a/tests/inference/test_model_inference.py +++ b/tests/inference/test_model_inference.py @@ -25,7 +25,7 @@ def __init__(self) -> None: def forward(self, x: torch.Tensor) -> dict[str, torch.Tensor]: return {"pred_boxes": self.linear(x[:, :1, :1, :1].squeeze(-1).squeeze(-1))} - def export(self) -> None: + def export(self, return_embeddings: bool = False) -> None: pass @@ -531,7 +531,7 @@ def test_inplace_export_failure_module_mutations_are_not_undone(self) -> None: original_model = rfdetr.model.model mutated: dict[str, bool] = {"happened": False} - def _mutating_export() -> None: + def _mutating_export(return_embeddings: bool = False) -> None: mutated["happened"] = True raise RuntimeError("export failed mid-mutation") diff --git a/tests/inference/test_predict.py b/tests/inference/test_predict.py index 0fb8d24be..31ffe612c 100644 --- a/tests/inference/test_predict.py +++ b/tests/inference/test_predict.py @@ -160,9 +160,166 @@ def _make_optimized_keypoint_model() -> tuple[RFDETR, _TupleOutputModelContext]: model._optimized_resolution = stub.resolution model._optimized_has_been_compiled = False model._optimized_dtype = torch.float32 + model._optimized_return_embeddings = False return model, stub +class _TupleOutputEmbeddingsModelContext: + """Model context whose forward returns a tuple ending in embeddings, mirroring an optimized detection model. + + Used to test that ``predict()`` correctly pops embeddings off the end of the optimized tuple output (rather than + misreading them as ``pred_masks``/``pred_keypoints``) when ``_optimized_return_embeddings`` is True. + """ + + def __init__(self, embedding_dim: int = 4) -> None: + self.device = torch.device("cpu") + self.resolution = 28 + self.class_names = ["object"] + self.args = SimpleNamespace(use_grouppose_keypoints=False, num_keypoints_per_class=[]) + self.model = torch.nn.Identity() + self.inference_model = self._forward + self.embedding_dim = embedding_dim + self.captured_predictions: dict[str, torch.Tensor] | None = None + + def _forward(self, batch_tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + batch = batch_tensor.shape[0] + boxes = torch.tensor([[[0.5, 0.5, 0.2, 0.2]]] * batch) + logits = torch.full((batch, 1, 1), 10.0) + embeddings = torch.full((batch, 1, self.embedding_dim), 3.0) + return boxes, logits, embeddings + + def postprocess( + self, + predictions: dict[str, torch.Tensor], + target_sizes: torch.Tensor, + score_threshold: float | None = None, + ) -> list[dict[str, torch.Tensor]]: + self.captured_predictions = predictions + batch = target_sizes.shape[0] + results = [] + for _ in range(batch): + result: dict[str, torch.Tensor] = { + "scores": torch.tensor([0.9]), + "labels": torch.tensor([0]), + "boxes": torch.tensor([[0.0, 0.0, 1.0, 1.0]]), + } + if "embeddings" in predictions: + result["embeddings"] = torch.full((1, self.embedding_dim), 3.0) + results.append(result) + return results + + +def _make_optimized_embeddings_model(embedding_dim: int = 4) -> tuple[RFDETR, _TupleOutputEmbeddingsModelContext]: + """Build a ``_DummyRFDETR`` wired to look like it ran ``inference(return_embeddings=True)``. + + Examples: + >>> model, stub = _make_optimized_embeddings_model(embedding_dim=6) + >>> model._optimized_return_embeddings + True + >>> isinstance(stub, _TupleOutputEmbeddingsModelContext) + True + """ + model = _DummyRFDETR() + stub = _TupleOutputEmbeddingsModelContext(embedding_dim=embedding_dim) + model.model = stub + model._is_optimized_for_inference = True + model._optimized_resolution = stub.resolution + model._optimized_has_been_compiled = False + model._optimized_dtype = torch.float32 + model._optimized_return_embeddings = True + return model, stub + + +class TestPredictReturnEmbeddings: + """Tests for ``predict(return_embeddings=...)`` on both the eager and optimized paths.""" + + def test_eager_default_omits_embeddings(self) -> None: + """By default (return_embeddings=False), detections.data has no 'embeddings' key.""" + img = PIL.Image.new("RGB", (64, 48), color=(128, 128, 128)) + model = _DummyRFDETR() + model.model = _DummyModel(labels=[0, 1]) + + detections = model.predict(img) + + assert "embeddings" not in detections.data + + def test_eager_default_forwards_return_embeddings_false_to_model(self) -> None: + """By default, predict() calls the underlying model with return_embeddings=False.""" + img = PIL.Image.new("RGB", (64, 48), color=(128, 128, 128)) + model = _DummyRFDETR() + dummy = _DummyModel(labels=[0, 1]) + model.model = dummy + + model.predict(img) + + assert dummy.model.last_return_embeddings is False + + def test_eager_return_embeddings_true_attaches_data(self) -> None: + """return_embeddings=True attaches detections.data['embeddings'] with shape (K, H).""" + img = PIL.Image.new("RGB", (64, 48), color=(128, 128, 128)) + embedding_dim = 5 + model = _DummyRFDETR() + model.model = _DummyModel(labels=[0, 1], include_embeddings=True, embedding_dim=embedding_dim) + + detections = model.predict(img, return_embeddings=True) + + assert "embeddings" in detections.data + assert detections.data["embeddings"].shape == (2, embedding_dim) + assert np.array_equal(detections.data["embeddings"][0], np.zeros(embedding_dim)) + assert np.array_equal(detections.data["embeddings"][1], np.ones(embedding_dim)) + + def test_eager_return_embeddings_true_forwards_kwarg_to_model(self) -> None: + """predict(return_embeddings=True) calls the underlying model with return_embeddings=True.""" + img = PIL.Image.new("RGB", (64, 48), color=(128, 128, 128)) + model = _DummyRFDETR() + dummy = _DummyModel(labels=[0], include_embeddings=True) + model.model = dummy + + model.predict(img, return_embeddings=True) + + assert dummy.model.last_return_embeddings is True + + def test_eager_keypoints_propagate_embeddings_to_keypoint_data(self) -> None: + """Keypoint outputs also expose 'embeddings' via key_points.data (shared dict construction path).""" + img = PIL.Image.new("RGB", (64, 48), color=(128, 128, 128)) + model = _DummyRFDETR() + model.model = _DummyModel(labels=[0], include_keypoints=True, include_embeddings=True) + + result = model.predict(img, return_embeddings=True) + + assert isinstance(result, sv.KeyPoints) + assert "embeddings" in result.data + + def test_optimized_return_embeddings_true_matches_attaches_data(self) -> None: + """An optimized model prepared with return_embeddings=True must attach embeddings on predict().""" + img = PIL.Image.new("RGB", (64, 48), color=(128, 128, 128)) + embedding_dim = 6 + model, stub = _make_optimized_embeddings_model(embedding_dim=embedding_dim) + + detections = model.predict(img, return_embeddings=True) + + assert stub.captured_predictions is not None + assert "embeddings" in stub.captured_predictions + assert "embeddings" in detections.data + assert detections.data["embeddings"].shape == (1, embedding_dim) + + def test_optimized_mismatch_true_predict_false_optimized_raises(self) -> None: + """predict(return_embeddings=False) on a model optimized with return_embeddings=True raises RuntimeError.""" + img = PIL.Image.new("RGB", (64, 48), color=(128, 128, 128)) + model, _stub = _make_optimized_embeddings_model() + + with pytest.raises(RuntimeError, match="does not match the optimized model"): + model.predict(img, return_embeddings=False) + + def test_optimized_mismatch_false_predict_true_optimized_raises(self) -> None: + """predict(return_embeddings=True) on a model optimized with return_embeddings=False raises RuntimeError.""" + img = PIL.Image.new("RGB", (64, 48), color=(128, 128, 128)) + model, _stub = _make_optimized_keypoint_model() # optimized with return_embeddings=False + + with pytest.raises(RuntimeError, match="does not match the optimized model"): + model.predict(img, return_embeddings=True) + + class TestPredictOptimizedInferenceKeypoints: """Regression tests for GitHub #1208: inference() breaks keypoint predict().""" diff --git a/tests/inference/test_predict_eval_mode.py b/tests/inference/test_predict_eval_mode.py index 9953eb043..01ddb2b9d 100644 --- a/tests/inference/test_predict_eval_mode.py +++ b/tests/inference/test_predict_eval_mode.py @@ -120,7 +120,10 @@ def test_predict_puts_module_in_eval_mode(self, monkeypatch: pytest.MonkeyPatch) monkeypatch.setattr( rfdetr.model.model, "forward", - lambda batch: {"pred_logits": torch.zeros(1, 10, 81), "pred_boxes": torch.zeros(1, 10, 4)}, + lambda batch, return_embeddings=False: { + "pred_logits": torch.zeros(1, 10, 81), + "pred_boxes": torch.zeros(1, 10, 4), + }, ) monkeypatch.setattr( rfdetr.model, diff --git a/tests/models/test_lwdetr_embeddings.py b/tests/models/test_lwdetr_embeddings.py new file mode 100644 index 000000000..b9a797954 --- /dev/null +++ b/tests/models/test_lwdetr_embeddings.py @@ -0,0 +1,242 @@ +# ------------------------------------------------------------------------ +# RF-DETR +# Copyright (c) 2025 Roboflow. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ +"""Unit tests for ``return_embeddings`` support in LWDETR's eager and exported forward passes.""" + +from unittest.mock import MagicMock + +import torch +from torch import nn + +from rfdetr.models.lwdetr import LWDETR +from rfdetr.utilities.tensors import NestedTensor + + +def _build_feature_batch(batch_size: int, hidden_dim: int) -> list[NestedTensor]: + """Build a single-level backbone feature map wrapped in a ``NestedTensor``, batched. + + Examples: + >>> features = _build_feature_batch(batch_size=2, hidden_dim=8) + >>> len(features) + 1 + >>> features[0].tensors.shape + torch.Size([2, 8, 4, 4]) + """ + return [ + NestedTensor( + torch.zeros(batch_size, hidden_dim, 4, 4), + torch.zeros(batch_size, 4, 4, dtype=torch.bool), + ) + ] + + +def _make_detection_model( + *, + batch_size: int = 2, + num_queries: int = 3, + hidden_dim: int = 8, + num_classes: int = 5, + num_decoder_layers: int = 2, + two_stage: bool = False, + segmentation_head: nn.Module | None = None, +) -> tuple[LWDETR, MagicMock]: + """Build an LWDETR detection model whose backbone/transformer are mocked with fixed-shape outputs. + + The mock backbone's ``return_value`` is a plain 3-tuple ``(features, poss, cross_attn_features)``, matching the + eager forward's unpacking (``forward``). For the exported/traced path (``forward_export``), which unpacks a 4-tuple + ``(feats, masks, poss, cross_attn_feats)``, tests reconfigure ``backbone.return_value`` before calling ``export()``. + + Examples: + >>> model, transformer = _make_detection_model(batch_size=1, num_queries=2, hidden_dim=4, num_classes=3) + >>> transformer.d_model + 4 + >>> outputs = model(torch.ones(1, 3, 8, 8)) + >>> outputs["pred_logits"].shape + torch.Size([1, 2, 3]) + """ + features = _build_feature_batch(batch_size=batch_size, hidden_dim=hidden_dim) + poss = [torch.zeros(batch_size, hidden_dim, 4, 4)] + + backbone = MagicMock() + backbone.return_value = (features, poss, None) + + transformer = MagicMock() + transformer.d_model = hidden_dim + transformer.return_value = ( + torch.zeros(num_decoder_layers, batch_size, num_queries, hidden_dim), # hs (all decoder layers) + torch.zeros(num_decoder_layers, batch_size, num_queries, 4), # ref_unsigmoid + torch.zeros(batch_size, num_queries, hidden_dim), # hs_enc + torch.zeros(batch_size, num_queries, 4), # ref_enc + ) + + model = LWDETR( + backbone=backbone, + transformer=transformer, + segmentation_head=segmentation_head, + num_classes=num_classes, + num_queries=num_queries, + aux_loss=False, + group_detr=1, + two_stage=two_stage, + lite_refpoint_refine=False, + bbox_reparam=False, + ) + model.eval() + return model, transformer + + +def _make_export_ready_model( + *, + batch_size: int = 2, + num_queries: int = 3, + hidden_dim: int = 8, + num_classes: int = 5, + segmentation_head: nn.Module | None = None, +) -> LWDETR: + """Build an LWDETR model with mocked backbone/transformer shaped for ``forward_export`` (traced/optimized path). + + Unlike :func:`_make_detection_model`, the mock backbone here returns the export-time 4-tuple ``(feats, masks, poss, + cross_attn_feats)`` expected by ``forward_export``, and the mock transformer returns the export-time last-decoder- + layer-only shapes ``[B, Q, H]`` (not ``[L, B, Q, H]``). + + Examples: + >>> model = _make_export_ready_model(batch_size=1, num_queries=2, hidden_dim=4, num_classes=3) + >>> model.export() + >>> predictions = model(torch.ones(1, 3, 8, 8)) + >>> len(predictions) + 2 + """ + features = _build_feature_batch(batch_size=batch_size, hidden_dim=hidden_dim) + masks = [torch.zeros(batch_size, 4, 4, dtype=torch.bool)] + poss = [torch.zeros(batch_size, hidden_dim, 4, 4)] + + backbone = MagicMock() + backbone.return_value = (features, masks, poss, None) + + transformer = MagicMock() + transformer.d_model = hidden_dim + transformer.return_value = ( + torch.zeros(batch_size, num_queries, hidden_dim), # hs (last decoder layer only) + torch.zeros(batch_size, num_queries, 4), # ref_unsigmoid + torch.zeros(batch_size, num_queries, hidden_dim), # hs_enc + torch.zeros(batch_size, num_queries, 4), # ref_enc + ) + + model = LWDETR( + backbone=backbone, + transformer=transformer, + segmentation_head=segmentation_head, + num_classes=num_classes, + num_queries=num_queries, + aux_loss=False, + group_detr=1, + two_stage=False, + lite_refpoint_refine=False, + bbox_reparam=False, + ) + model.eval() + return model + + +class TestEagerForwardReturnEmbeddings: + """``LWDETR.forward(return_embeddings=...)`` (unoptimized/eager path).""" + + def test_return_embeddings_false_omits_embeddings_key(self) -> None: + """By default, no 'embeddings' key is added to the output dict.""" + model, _ = _make_detection_model() + outputs = model(torch.ones(2, 3, 8, 8), return_embeddings=False) + + assert "embeddings" not in outputs + + def test_return_embeddings_true_adds_embeddings_with_expected_shape(self) -> None: + """return_embeddings=True adds 'embeddings' with shape [B, Q, H] from the last decoder layer only.""" + batch_size, num_queries, hidden_dim, num_decoder_layers = 2, 3, 8, 2 + model, _ = _make_detection_model( + batch_size=batch_size, + num_queries=num_queries, + hidden_dim=hidden_dim, + num_decoder_layers=num_decoder_layers, + ) + + outputs = model(torch.ones(batch_size, 3, 8, 8), return_embeddings=True) + + assert "embeddings" in outputs + assert outputs["embeddings"].shape == (batch_size, num_queries, hidden_dim) + + def test_return_embeddings_does_not_affect_other_outputs(self) -> None: + """Turning on return_embeddings must not change pred_logits/pred_boxes shapes or values.""" + model, _ = _make_detection_model() + x = torch.ones(2, 3, 8, 8) + + torch.manual_seed(0) + out_without = model(x, return_embeddings=False) + torch.manual_seed(0) + out_with = model(x, return_embeddings=True) + + assert torch.equal(out_without["pred_logits"], out_with["pred_logits"]) + assert torch.equal(out_without["pred_boxes"], out_with["pred_boxes"]) + + +class TestExportForwardReturnEmbeddings: + """``LWDETR.export(return_embeddings=...)`` + traced ``forward_export`` (optimized path).""" + + def test_export_return_embeddings_false_appends_nothing(self) -> None: + """export(return_embeddings=False) keeps the base 2-tuple (coord, class) for detection-only models.""" + model = _make_export_ready_model() + model.export(return_embeddings=False) + + predictions = model(torch.ones(2, 3, 8, 8)) + + assert isinstance(predictions, tuple) + assert len(predictions) == 2 + + def test_export_return_embeddings_true_appends_embeddings_last(self) -> None: + """export(return_embeddings=True) appends embeddings as the last element of the tuple.""" + batch_size, num_queries, hidden_dim = 2, 3, 8 + model = _make_export_ready_model(batch_size=batch_size, num_queries=num_queries, hidden_dim=hidden_dim) + model.export(return_embeddings=True) + + predictions = model(torch.ones(batch_size, 3, 8, 8)) + + assert isinstance(predictions, tuple) + assert len(predictions) == 3 + coord, logits, embeddings = predictions + assert coord.shape == (batch_size, num_queries, 4) + assert logits.shape[:2] == (batch_size, num_queries) + # Exported decoder only returns the last layer's hidden state -> [B, Q, H], no L*H reshape. + assert embeddings.shape == (batch_size, num_queries, hidden_dim) + + def test_export_return_embeddings_true_with_segmentation_head_still_appends_last(self) -> None: + """When a mask head is present, embeddings remain the last tuple element (position 3, after masks).""" + batch_size, num_queries, hidden_dim = 2, 3, 8 + mask_h, mask_w = 4, 4 + seg_head = MagicMock() + seg_head.return_value = [torch.zeros(batch_size, num_queries, mask_h, mask_w)] + + model = _make_export_ready_model( + batch_size=batch_size, + num_queries=num_queries, + hidden_dim=hidden_dim, + segmentation_head=seg_head, + ) + model.export(return_embeddings=True) + + predictions = model(torch.ones(batch_size, 3, 8, 8)) + + assert isinstance(predictions, tuple) + assert len(predictions) == 4 + coord, logits, masks, embeddings = predictions + assert masks.shape == (batch_size, num_queries, mask_h, mask_w) + assert embeddings.shape == (batch_size, num_queries, hidden_dim) + + def test_export_return_embeddings_defaults_to_false_when_omitted(self) -> None: + """Export() with no arguments behaves exactly like export(return_embeddings=False).""" + model = _make_export_ready_model() + model.export() + + predictions = model(torch.ones(2, 3, 8, 8)) + + assert isinstance(predictions, tuple) + assert len(predictions) == 2 diff --git a/tests/models/test_postprocess.py b/tests/models/test_postprocess.py index 19e0544e7..a7c1c715a 100644 --- a/tests/models/test_postprocess.py +++ b/tests/models/test_postprocess.py @@ -723,3 +723,117 @@ def test_score_threshold_is_ignored_outside_the_mask_path(self, head): assert baseline.keys() == with_threshold.keys() for key in baseline: torch.testing.assert_close(baseline[key], with_threshold[key], rtol=0.0, atol=0.0, equal_nan=True) + + +class TestAttachEmbeddings: + """Tests for :meth:`PostProcess._attach_embeddings` and its wiring into ``forward``.""" + + def test_forward_omits_embeddings_key_when_absent_from_outputs(self) -> None: + """No 'embeddings' key is added to results when outputs has no 'embeddings' entry.""" + pp = PostProcess(num_select=2) + outputs = { + "pred_logits": torch.tensor([[[10.0, -10.0], [9.0, -10.0]]]), + "pred_boxes": torch.tensor([[[0.3, 0.3, 0.2, 0.2], [0.6, 0.6, 0.1, 0.1]]]), + } + target_sizes = torch.tensor([[480, 640]]) + + results = pp(outputs, target_sizes) + + assert "embeddings" not in results[0] + + def test_forward_attaches_embeddings_with_expected_shape(self) -> None: + """Embeddings in outputs is gathered per-selected-query and attached with shape (K, H).""" + hidden_dim = 4 + pp = PostProcess(num_select=2) + outputs = { + "pred_logits": torch.tensor([[[10.0, -10.0], [9.0, -10.0]]]), + "pred_boxes": torch.tensor([[[0.3, 0.3, 0.2, 0.2], [0.6, 0.6, 0.1, 0.1]]]), + "embeddings": torch.arange(2 * hidden_dim, dtype=torch.float32).reshape(1, 2, hidden_dim), + } + target_sizes = torch.tensor([[480, 640]]) + + results = pp(outputs, target_sizes) + + assert "embeddings" in results[0] + assert results[0]["embeddings"].shape == (2, hidden_dim) + + def test_attach_embeddings_gathers_by_topk_indices(self) -> None: + """_attach_embeddings must select embeddings for the exact queries in topk_boxes, in that order.""" + hidden_dim = 3 + # 4 queries per image, embeddings identifiable by their first column value == query index. + out_embeddings = torch.zeros(1, 4, hidden_dim) + for q in range(4): + out_embeddings[0, q, 0] = q + + # Select queries [2, 0] (reversed, non-contiguous) as the "top-k" selection. + topk_boxes = torch.tensor([[2, 0]]) + results: list[dict[str, torch.Tensor]] = [{}] + + PostProcess._attach_embeddings(results, out_embeddings, topk_boxes) + + gathered = results[0]["embeddings"] + assert gathered.shape == (2, hidden_dim) + assert gathered[0, 0].item() == 2 + assert gathered[1, 0].item() == 0 + + def test_attach_embeddings_batch_of_two_images(self) -> None: + """Embeddings are gathered independently per image in the batch.""" + out_embeddings = torch.tensor( + [ + [[0.0, 0.0], [1.0, 1.0], [2.0, 2.0]], # image 0 + [[10.0, 10.0], [11.0, 11.0], [12.0, 12.0]], # image 1 + ] + ) + topk_boxes = torch.tensor([[1, 0], [2, 1]]) + results: list[dict[str, torch.Tensor]] = [{}, {}] + + PostProcess._attach_embeddings(results, out_embeddings, topk_boxes) + + assert torch.equal(results[0]["embeddings"], torch.tensor([[1.0, 1.0], [0.0, 0.0]])) + assert torch.equal(results[1]["embeddings"], torch.tensor([[12.0, 12.0], [11.0, 11.0]])) + + def test_forward_attaches_embeddings_alongside_masks(self) -> None: + """Embeddings and masks can coexist in the same forward() call without interfering with each other.""" + batch, num_queries, mask_h, mask_w, hidden_dim, num_classes = 1, 4, 8, 8, 4, 2 + pp = PostProcess(num_select=4) + outputs = { + "pred_logits": torch.randn(batch, num_queries, num_classes), + "pred_boxes": torch.rand(batch, num_queries, 4), + "pred_masks": torch.randn(batch, num_queries, mask_h, mask_w), + "embeddings": torch.randn(batch, num_queries, hidden_dim), + } + target_sizes = torch.tensor([[256, 256]]) + + results = pp(outputs, target_sizes) + + assert "masks" in results[0] + assert "embeddings" in results[0] + assert results[0]["embeddings"].shape == (4, hidden_dim) + + def test_forward_filters_embeddings_by_score_threshold_with_masks(self) -> None: + """When masks are filtered by score_threshold, embeddings must be filtered the same way. + + ``_postprocess_masks`` drops rows scoring at or below ``score_threshold`` before returning results, so + ``scores``/``labels``/``boxes``/``masks`` all end up with fewer rows than the raw top-k selection. Embeddings + must be filtered with the same per-image predicate, or the row counts diverge and callers indexing the + embeddings tensor with a boolean mask sized to the filtered scores hit a shape mismatch. + """ + hidden_dim = 4 + pp = PostProcess(num_select=3) + # Deterministic sigmoid scores: query 0 -> ~1.0 (kept), query 1 -> ~0.0 (dropped), query 2 -> ~0.99 (kept). + logits = torch.tensor([[[10.0], [-10.0], [5.0]]]) + outputs = { + "pred_logits": logits, + "pred_boxes": torch.rand(1, 3, 4), + "pred_masks": torch.randn(1, 3, 8, 8), + "embeddings": torch.arange(3 * hidden_dim, dtype=torch.float32).reshape(1, 3, hidden_dim), + } + target_sizes = torch.tensor([[256, 256]]) + + results = pp(outputs, target_sizes, score_threshold=0.5) + + assert results[0]["embeddings"].shape[0] == results[0]["scores"].shape[0] == 2 + assert results[0]["masks"].shape[0] == 2 + # Query 1 (logit -10) is below threshold and must be dropped from embeddings too. + kept_first_values = {row[0].item() for row in results[0]["embeddings"]} + assert kept_first_values == {0.0, 2 * hidden_dim}