Feat/embedding extraction - #1327
Open
unaxEtxeberriaBieleDigital wants to merge 29 commits into
Open
Conversation
…xeberriaBieleDigital/rf-detr into feat/embedding-extraction
unaxEtxeberriaBieleDigital
requested review from
Borda,
SkalskiP,
isaacrob and
probicheaux
as code owners
August 11, 2026 06:33
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## develop #1327 +/- ##
=======================================
Coverage 86% 86%
=======================================
Files 112 112
Lines 14520 14563 +43
=======================================
+ Hits 12479 12520 +41
- Misses 2041 2043 +2 🚀 New features to boost your workflow:
|
…xeberriaBieleDigital/rf-detr into feat/embedding-extraction
Contributor
There was a problem hiding this comment.
Pull request overview
Adds per-detection query embedding extraction across eager and optimized inference paths.
Changes:
- Adds
return_embeddingsto model forwarding, optimization, prediction, and postprocessing. - Exposes embeddings through detection and keypoint result data.
- Adds tests, documentation, and changelog coverage.
Reviewed changes
Copilot reviewed 14 out of 15 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
.gitignore |
Adds local tooling exclusions. |
CHANGELOG.md |
Records embedding extraction support. |
docs/learn/run/detection.md |
Documents embedding extraction. |
docs/learn/run/keypoints.md |
Documents keypoint result embeddings. |
docs/learn/run/segmentation.md |
Documents segmentation result embeddings. |
src/rfdetr/detr.py |
Integrates embeddings into inference and prediction APIs. |
src/rfdetr/evaluation/coco_eval.py |
Normalizes sigma values to Python floats. |
src/rfdetr/models/lwdetr.py |
Returns decoder query embeddings. |
src/rfdetr/models/postprocess.py |
Attaches selected embeddings to results. |
tests/inference/helpers.py |
Extends inference test doubles. |
tests/inference/test_model_inference.py |
Updates optimization test contracts. |
tests/inference/test_predict.py |
Tests eager and optimized embedding prediction. |
tests/inference/test_predict_eval_mode.py |
Updates the mocked forward signature. |
tests/models/test_lwdetr_embeddings.py |
Tests eager and exported model embeddings. |
tests/models/test_postprocess.py |
Tests embedding gathering and attachment. |
Suppressed comments (2)
tests/models/test_lwdetr_embeddings.py:41
- This fixture-builder helper has a descriptive docstring but no
Examplesdoctest. Test helpers are required to exercise their behavior directly via doctest; add a small construction example.
"""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()``.
"""
tests/models/test_lwdetr_embeddings.py:86
- This fixture-builder helper lacks the mandatory direct
Examplesdoctest for helpers intests/. Add a minimal example that constructs the export-ready model and verifies its state.
"""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]``).
"""
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+96
to
+97
| if out_embeddings is not None: | ||
| self._attach_embeddings(results, out_embeddings, topk_boxes) |
Comment on lines
+616
to
+618
| 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] |
Comment on lines
+17
to
+18
| def _build_feature_batch(batch_size: int, hidden_dim: int) -> list[NestedTensor]: | ||
| return [ |
Comment on lines
+212
to
+213
| 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)``.""" |
_postprocess_masks drops rows scoring at or below score_threshold before returning results, but _attach_embeddings still gathered embeddings for every unfiltered topk_boxes row. predict() then indexed the full embedding tensor with a boolean mask sized to the filtered scores, raising a shape mismatch whenever a detection fell below the threshold. _attach_embeddings now optionally accepts scores/score_threshold and reproduces the same per-image filtering _postprocess_masks applies, so embeddings line up 1:1 with the (possibly filtered) segmentation rows. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Repository guidance requires every non-test helper under tests/ to carry a docstring with a direct Examples doctest. _build_feature_batch, _make_detection_model, and _make_export_ready_model were missing them. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Repository guidance requires every non-test helper under tests/ to carry a docstring with a direct Examples doctest. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…xeberriaBieleDigital/rf-detr into feat/embedding-extraction
Adds a Colab-runnable cookbook demonstrating the value of the new predict(..., return_embeddings=True) interface: it 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 (ROC-AUC, precision@k versus random review), surfacing confident detections with no matching annotation, and nearest-neighbour object retrieval. Registers the notebook in cards.yaml/NOTES.md, cross-links it from the detection docs, and notes it in the changelog. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…xeberriaBieleDigital/rf-detr into feat/embedding-extraction
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Add the ability to extract the embeddings of each query after the decoder layers. It is related to the issue #1272.
Type of Change
Testing
I made sure all the tests passed and added a few new tests for the new feature.
Test details:
All tests follow TDD (failed before implementation) and pass in the current CPU test suite; no tests marked gpu were added.
Checklist
Additional Context