Feat/embedding extraction

[unaxEtxeberriaBieleDigital]

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

• New feature (non-breaking change that adds functionality)
• Documentation update

Testing

I made sure all the tests passed and added a few new tests for the new feature.

• [ X ] I have tested this change locally
• [ X ] I have added/updated tests for this change

Test details:

• tests/models/test_lwdetr_embeddings.py (new): covers LWDETR.forward(return_embeddings=...) on the eager path (toggles the "embeddings" key, verifies shape [B, Q, H] taken from the last decoder layer only — consistent with the exported/traced path, not a concatenation of all layers — and confirms it doesn't alter pred_logits / pred_boxes ) and the exported/traced path via export(return_embeddings=...) + forward_export (embeddings absent by default, present as the last tuple element—also with segmentation masks—and with shape [B, Q, H] from the last decoder layer).
• tests/models/test_postprocess.py  (new  TestAttachEmbeddings  class): covers  PostProcess._attach_embeddings  and its integration into  forward()  — no  embeddings  in outputs means the key is not added; with  embeddings , it performs  gather  by exact  topk_boxes  indices (order and batch respected); correct coexistence with segmentation masks in the same call.
• tests/inference/test_predict.py  (new  TestPredictReturnEmbeddings  class) and  tests/inference/helpers.py  (updated): covers  RFDETR.predict(return_embeddings=...)  on eager model (by default doesn't add  detections.data["embeddings"] ; with  True  adds them with shape  (K, H)  and forwards the kwarg to the underlying model; also propagated to  key_points.data["embeddings"]  on keypoint outputs) and on optimized model ( inference(return_embeddings=...) ): correct match between  predict()  and the optimized model, and  RuntimeError  when the  return_embeddings  value in  predict()  doesn't match the one used in  inference() .

All tests follow TDD (failed before implementation) and pass in the current CPU test suite; no tests marked  gpu  were added.

Checklist

• [ X ] My code follows the style guidelines of this project
• [ X ] I have performed a self-review of my own code
• [ X ] I have commented my code where necessary, particularly in hard-to-understand areas
• [ X ] My changes generate no new warnings or errors
• [ X ] I have updated the documentation accordingly (if applicable)

Additional Context

[codecov]

Codecov Report

❌ Patch coverage is 92.98246% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 86%. Comparing base (3467c55) to head (9c0833e).
⚠️ Report is 2 commits behind head on develop.

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:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Copilot]

Pull request overview

Adds per-detection query embedding extraction across eager and optimized inference paths.

Changes:

• Adds return_embeddings to 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 Examples doctest. 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 Examples doctest for helpers in tests/. 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.