[Feature Proposal] Implementation of Keypoint Detection Head (Architecture & Loss)

[Kagura-Ahad]

Search before asking

• I have searched the RF-DETR issues and found no similar feature requests.

Description

Following up on discussions in #418 and #169, and recent conversations regarding the roadmap for Keypoint support.
I understand that full pre-training for a pose model is computationally heavy and currently a lower priority than OBB (#56). However, I would like to contribute by implementing the architectural components as of right now, so the model is "pose-ready" when the team decides to allocate compute resources for training.

Proposed Implementation Plan:

1. Model Architecture: Extend the Transformer Decoder to support a Keypoint Head (likely an MLP predicting K points per query).
2. Matcher: Update the Hungarian Matcher to include Keypoint cost (using OKS - Object Keypoint Similarity).
3. Loss Function: Implement the OKS Loss calculation for the training loop.
4. Data Pipeline: Ensure the dataloader structure can accept COCO-Keypoint format annotations.

Use case

The Keypoint support is on developers' roadmap

Additional

Questions for the Team

• Do you prefer I mirror the implementation style of RT-DETR / YOLOv8-Pose for the head?
• Should this live in a specific feature branch (e.g., feat/keypoints)?
  I am happy to open a PR with the structural changes + unit tests verifying the forward pass/loss calculation works on dummy data.

Are you willing to submit a PR?

• Yes I'd like to help by submitting a PR!

[isaacrob-roboflow]

I have a different architecture in mind and am planning to bang it all out once we have bandwidth :)

[Lucasmogsan]

@Kagura-Ahad (and @isaacrob-roboflow) let me know if you're looking in to this, I would potentially be interesting in contributing in order to speed up the process a bit.

For reference the DETRPose and obviously SAM3DB which could be interesting to consider.

[michaelmohamed]

I have a working Keypoint/Pose implementation here.

[isaacrob-roboflow]

@michaelmohamed have you benchmarked your approach on any common datasets?

Haven't had bandwidth to work on this myself yet but don't think it'll take all that long once we get started

[michaelmohamed]

I’m currently testing the implementation on private datasets with comparable results to what I have seen using other frameworks. Although it’s still early, the intent is for this to act as a starting point for future development, or others who need something today. I’m happy to align with whatever the team rolls out, when available. Docs here.