Feat/embedding extraction

.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

CHANGELOG.md

@@ -8,6 +8,8 @@ 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`.
+
 - `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))
@@ -23,6 +25,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 ### Fixed
 
 - The default (torchvision-native) training pipeline no longer silently corrupts keypoint annotations when `keypoint_flip_pairs` is empty on a schema that has genuine left/right pairs. `RandomHorizontalFlip` on this backend always mirrored keypoint x-coordinates when a flip was drawn, but only relabeled left/right joints `if self.keypoint_flip_pairs:` — with an empty list (the pydantic default, and one possible outcome when automatic flip-pair inference from dataset metadata doesn't match an asymmetric schema), affected training samples got their keypoints mirrored in position while keeping their original left/right label, with no warning. `_build_torchvision_pipeline` now drops the flip entirely for an empty-but-not-`None` `keypoint_flip_pairs`, logging the same warning the Albumentations backend already emits via `filter_keypoint_hflip_augmentations` (worded for this backend's lack of an editable `aug_config`), matching the annotation-safety behavior that backend has had since #1122. An empty list can also legitimately mean the schema has no left/right pairs at all (e.g. a single midline keypoint) — the unpatched flip was already harmless there since nothing needed relabeling; this fix disables it there too, for consistency with the Albumentations backend's existing contract, at the cost of a now-unavailable-by-default augmentation for that narrower case. Detection-only pipelines (`keypoint_flip_pairs=None`) and keypoint pipelines with real pairs are unaffected.
+
 - `BestModelCallback` no longer treats PyTorch Lightning's pre-training sanity-check validation pass as a real epoch's result. Its EMA-checkpoint tracking and the `smooth_alpha` smoothing accumulator are custom bookkeeping that sit outside `ModelCheckpoint`'s own `trainer.sanity_checking` guard (the guard the regular-checkpoint path already inherits), so a positive sanity-check score — common when starting a new training run initialized with `pretrain_weights` from a checkpoint pretrained on a different dataset — could get written out as the permanent "best" `checkpoint_best_ema.pth` before a single real epoch ran, and real training could then never surpass it. Note this is distinct from PTL's own `resume`/`ckpt_path` restart, which PTL itself skips the sanity check for (`not val_loop.restarting`). ([#1348](https://github.com/roboflow/rf-detr/issues/1348))
 
 ## [1.9.3] — 2026-08-17

docs/learn/run/detection.md

@@ -70,6 +70,31 @@ 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`.
+
 ## Run on video, webcam, or RTSP stream
 
 These examples use OpenCV for decoding and display. Replace `<SOURCE_VIDEO_PATH>`, `<WEBCAM_INDEX>`, and `<RTSP_STREAM_URL>` with your inputs. `<WEBCAM_INDEX>` is usually `0` for the default camera.

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 `<SOURCE_VIDEO_PATH>`, `<WEBCAM_INDEX>`, and `<RTSP_STREAM_URL>` with your inputs. `<WEBCAM_INDEX>` is usually `0` for the default camera.

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 `<SOURCE_VIDEO_PATH>`, `<WEBCAM_INDEX>`, and `<RTSP_STREAM_URL>` with your inputs. `<WEBCAM_INDEX>` is usually `0` for the default camera.

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:
@@ -1207,6 +1211,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.
 
@@ -1238,6 +1243,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
@@ -1255,7 +1265,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):
@@ -1324,7 +1334,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)
 
@@ -1355,6 +1365,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):
@@ -1401,7 +1412,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):
@@ -1440,6 +1451,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:
@@ -1457,7 +1469,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):
@@ -2246,6 +2258,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.
@@ -2275,6 +2288,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.
 
@@ -2515,14 +2536,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],
@@ -2533,6 +2572,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)
@@ -2592,6 +2633,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"]
@@ -2616,6 +2660,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.

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: