Repository navigation
feat(keypoint-detection): enable ViTPose config/build/perf - #905
Merged
Merged
Conversation
Hyunsik Jeon (jeon185)
force-pushed
the
feat/keypoint-detection-enablement
branch
from
June 23, 2026 18:35
2007caf to
540ab05
Compare
Zhenchao Ni (zhenchaoni)
approved these changes
Jun 25, 2026
ViTPose keypoint-detection models could not pass the wmk pipeline:
1. Task resolution: Optimum registers the ViTPose ONNX export config but has
no task-to-class entry for keypoint-detection, and transformers'
AutoModelForKeypointDetection only recognizes SuperPoint. Add
MODEL_CLASS_MAPPING[(vitpose, keypoint-detection)] = VitPoseForPoseEstimation
(models/hf/vitpose.py) so the resolver loads the correct class.
2. MoE export: the vitpose-plus checkpoints use a Mixture-of-Experts backbone
whose patcher injects a constant dataset_index at export time. Optimum's
patch_model_for_export defaults model_kwargs to None, so the patcher crashed
on init. Pass an explicit model_kwargs={} in _get_optimum_patcher. Also wrap
the Step 3 hierarchy trace in the same patcher context (it previously ran
the model forward without the injected dataset_index, failing before export).
Verified config -> build -> perf on all 6 acceptance models in #284
(base-simple, plus-{small,base,large,huge}, synthpose-vitpose-huge-hf).
TasksManager can't infer a task from the ViTPose architecture, so config and build needed an explicit --task keypoint-detection. Add the (vitpose, None) sentinel that declares keypoint-detection as the default task - the same way SAM declares mask-generation - so the resolver reverse-looks-up keypoint-detection and --task becomes optional. Verified: winml config resolves keypoint-detection with no --task flag; added unit tests for the sentinel default and the mapping entry.
Hyunsik Jeon (jeon185)
force-pushed
the
feat/keypoint-detection-enablement
branch
from
June 26, 2026 22:59
24a31c3 to
e777754
Compare
vortex-captain
approved these changes
Jun 30, 2026
Zhenchao Ni (zhenchaoni)
pushed a commit
that referenced
this pull request
Jul 1, 2026
Adds the eval stage for keypoint-detection (ViTPose), so the COCO-keypoint models from #284 now go through `config -> build -> perf -> eval`. Stacked on #905 (the config/build/perf enablement) - that one should go in first. What's here: - `metrics/keypoint.py` - KeypointAPMetric. Computes the COCO keypoint score (OKS-based AP over 0.50:0.95) with pycocotools COCOeval, the same way object-detection already reuses the COCO mAP protocol. - `keypoint_detection_evaluator.py` - top-down evaluator. transformers has no keypoint-detection pipeline, so it runs the image processor and ONNX model directly: for each ground-truth person box it does preprocess -> model -> post_process_pose_estimation and scores against the GT keypoints. ViTPose is exported with a static batch of 1, so each person crop runs separately and the heatmaps are stacked back together for post-processing. It uses the GT person boxes (standard COCO top-down protocol - keeps the score about pose accuracy, not detection). - `scripts/build_coco_keypoints.py` - builds a local COCO val keypoints dataset. COCO has no script-free HF mirror for person keypoints, so this downloads the annotations once and fetches images individually, which means a small subset doesn't need the full image zip. - Schema, evaluator registry, default dataset, and unit tests for the metric and evaluator. Verified on the five COCO 17-keypoint models (vitpose-base-simple and vitpose-plus-{small,base,large,huge}): config -> build -> perf -> eval all pass and return COCO AP/AR. AP rises with model size as you'd expect. Absolute numbers are on the low side right now because the build quantizes with random calibration data, but relative comparison holds. ## synthpose-vitpose-huge-hf - not covered yet This is the one model from #284 that this PR does not evaluate. It predicts 52 anatomical keypoints instead of COCO's 17, so it can't be scored against COCO ground truth - the keypoint sets don't line up and OKS is only defined when they do. How it's handled for now: the metric checks the keypoint count up front and raises a clear, actionable error instead of failing with a numpy broadcast error deep inside pycocotools. Idea for finishing it: KeypointAPMetric already takes `sigmas` and `keypoint_names` as arguments, so the main missing piece is a dataset with SynthPose's 52-keypoint ground truth plus the matching OKS sigmas. I'd rather agree on the dataset and sigmas in review before adding that - happy to land it in this PR or as a follow-up, whichever you prefer. Refs #284.
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.
Enables the 6 ViTPose keypoint-detection models from #284 to pass
wmk config -> build -> perf(CPU and OpenVINO). Eval isn't included here - I'll do the accuracy side in a follow-up since it needs a couple of design decisions first.Two things were blocking all 6 models:
wmk configfailed with "Task 'keypoint-detection' not supported by TasksManager". Optimum has the ViTPose ONNX export config but no task->class entry for keypoint-detection, andAutoModelForKeypointDetectiononly covers SuperPoint. Added the(vitpose, keypoint-detection) -> VitPoseForPoseEstimationmapping, same way we already do it for CLIP/SAM.The plus checkpoints (MoE backbone) crashed during export with "dataset_index must be provided when using multiple experts". Optimum's VitPoseModelPatcher injects a constant dataset_index, but
patch_model_for_exportdefaults model_kwargs to None so it crashed on init. Passing an explicitmodel_kwargs={}fixes that. The trace step (Step 3) was also running the model outside the patcher context, so I wrapped it the same way the export step already is.The exporter change isn't ViTPose-specific - it helps any MoE model whose patcher injects forward args.
Verified config/build/perf on all 6: vitpose-base-simple, vitpose-plus-small/base/large/huge, and synthpose-vitpose-huge-hf. Added unit tests for the mapping and the patcher model_kwargs handling.
One note: you still need to pass
--task keypoint-detectionexplicitly for now - the task isn't auto-detected from the config yet. I left auto-detection out of this PR to keep it small; can add it here or as a follow-up if you'd prefer.Refs #284.