Skip to content

enhancement(pose): add Sapiens / ViTPose++ as an additional selectable backend #3

Description

@kmch4n

Summary

Integrate a state-of-the-art pose estimator as an optional backend for users with GPU resources who need maximum landmark accuracy on heavily occluded or extreme skiing postures. Primary candidates are Sapiens (Meta, ECCV 2024 Oral) and ViTPose++ (TPAMI 2023). Both clearly outperform MediaPipe and YOLO11-Pose on public benchmarks.

Motivation

  • Even YOLO11-Pose, while stronger than MediaPipe, is not at the top of published accuracy for pose estimation.
  • Sapiens sets a new state of the art on Humans-5K (+7.6 mAP over prior SOTA) and is trained on 300M in-the-wild human images — closer to our domain than COCO-only models. It natively supports 1K-resolution inference.
  • ViTPose++ achieves 81.1 AP on COCO keypoint test-dev with multiple practical ONNX export paths and community runtimes.
  • This opens a path to "best-shot quality" analyses for key reference images (portfolio / marketing shots), while keeping lightweight backends for real-time video.

Scope

Additive to the backend abstraction introduced in issue #2. This issue is blocked by that one.

  • Add a new backend selectable via SKISENSE_POSE_BACKEND=sapiens or vitpose.
  • CUDA GPU only; raise a clear error on CPU/MPS.
  • Use ONNX-based inference to avoid heavy PyTorch weights and keep the install footprint reasonable.

Implementation Plan

1. Pre-work: evaluate Sapiens vs ViTPose++

  • Run both on a curated set of 10–20 representative skiing frames and compare:
    • Right-leg detection success rate on deeply angulated turns
    • Landmark consistency across consecutive frames (for the video flow)
    • Inference time per frame on RTX 30/40 class GPU
    • Model file size
  • Pick the winner as the default but keep the code path extensible enough to keep both.

2. ONNX model preparation

  • ViTPose++: export from ViTAE-Transformer/ViTPose, or reuse pre-exported weights from Pukei-Pukei/ViTPose-ONNX or onnx-community/vitpose-base-simple.
  • Sapiens: export from facebookresearch/sapiens (sapiens_0.6b as the default VRAM-friendly variant; allow sapiens_1b for users with more VRAM).
  • Store exported weights in models/ and add a download script mirroring the current MediaPipe asset-download flow.

3. Dependencies

  • Add onnxruntime-gpu (CUDA) to a new requirements-sota.txt installable as an optional extra.
  • Do not force existing users onto ONNX Runtime.

4. Backend implementations (src/skisense/backends/)

  • sapiens_backend.py — uses the Sapiens COCO-17 head.
  • vitpose_backend.py — uses the ViTPose++ COCO-17 output.
  • Both conform to the backend interface defined in the YOLO11-Pose issue.
  • Preprocessing: resize ROI to model input (e.g., 256×192 for ViTPose, 1024×768 or higher for Sapiens), pad to aspect ratio, normalize to ImageNet mean/std.
  • Post-processing: heatmap → keypoint decoding with sub-pixel refinement (DARK decoding recommended for sub-pixel accuracy).

5. Configuration

  • POSE_BACKEND accepts sapiens / vitpose in addition to mediapipe / yolo11.
  • SAPIENS_MODEL_VARIANT=0.3b|0.6b|1b|2b (default 0.6b).
  • VITPOSE_MODEL_VARIANT=base-simple|huge-simple (default base-simple).
  • Fail fast with a clear error if CUDA is unavailable when a SOTA backend is selected.

6. Drawing and analyzer integration

  • Reuse the topology adapter from the YOLO11-Pose issue. Both Sapiens and ViTPose++ emit COCO-17.
  • No changes to pose_analyzer.py beyond what the prior issue already introduced.

7. Documentation

  • New README section: Backend accuracy tiers.
  • Table comparing MediaPipe / YOLO11-Pose / SOTA backend on (accuracy, speed, VRAM, install size).
  • Example shots produced with each backend, included as images/backend_comparison.png.

8. Benchmark / validation

  • New script scripts/benchmark_backends.py:
    • Runs all enabled backends on the same set of test frames.
    • Writes a CSV with per-frame score, per-joint angles, and inference time.
    • Generates a side-by-side comparison image.

Files Touched

  • src/skisense/config.py
  • .env.example
  • requirements-sota.txt (new)
  • src/skisense/backends/sapiens_backend.py (new)
  • src/skisense/backends/vitpose_backend.py (new)
  • scripts/benchmark_backends.py (new)
  • README.md
  • images/backend_comparison.png (new)

Acceptance Criteria

  • With SKISENSE_POSE_BACKEND=sapiens (or vitpose) and onnxruntime-gpu installed, python run.py skier.jpg --image produces a visibly better pose on the reference heavy-angulation test image (right-leg landmarks detected, non-zero score on previously failing frames).
  • On CPU or MPS, the SOTA backend raises a clear RuntimeError naming the required environment.
  • Default behavior (SKISENSE_POSE_BACKEND=mediapipe) is unchanged.
  • Benchmark script runs end-to-end and produces the comparison CSV/image.

Blocked By

References

Metadata

Metadata

Assignees

No one assigned

    Labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions