willalexm/vitpose-l-swimxyz

ViTPose-L finetuned on SwimXYZ (multi-view swimmer pose)

2D human pose estimation (17 COCO keypoints) for swimmers, across five camera views. This is a ViTPose-L model finetuned on the synthetic SwimXYZ dataset. The SwimXYZ paper reports finetuning ViTPose but released no weights; this fills that gap.

Off-the-shelf COCO pose models degrade badly on swimmers (small/distant figures, prone orientations, water occlusion and refraction). This model is trained specifically for that regime.

Results — PCK@0.05 on held-out synthetic validation

View	Base ViTPose-L (COCO)	This model
Aerial	0.818	0.993
Front	0.375	0.962
Side above-water	0.754	0.984
Side underwater	0.613	0.989
Side water-level	0.700	0.983
Overall	0.653	0.982

Trained on 25,000 freestyle frames (5,000 per view) for 5 epochs, AdamW LR 1e-4, input 256×192 → heatmaps 64×48, on a single RTX 5090.

Files

2D pose (ViTPose-L):

• vitpose-l-swimxyz-mv.pth — PyTorch checkpoint ({"state_dict", "epoch", "val_pck", "per_view_pck"})
• vitpose-l-swimxyz-mv.onnx — ONNX export (opset 17, dynamic batch), torch-parity verified

3D lift (temporal 2D→3D):

• lifter-swimxyz.pth / .onnx — trained on GT 2D (30.3 mm MPJPE)
• lifter-swimxyz-realistic.pth / .onnx — trained with the finetuned ViTPose's measured 2D noise injected — the end-to-end lifter (33.9 mm MPJPE)

SMPL mesh (temporal 3D→SMPL):

• smpl-regressor-swimxyz.pth / .onnx — from SMPL-24 joints (1.5 mm recon)
• smpl-regressor-coco.pth / .onnx — from the 12 COCO body joints the lifter outputs (canonicalized), the piece that closes the loop (5.3 mm recon)

Stroke classifier (per-frame):

• stroke-classifier.pth / .onnx — temporal 4-way classifier (freestyle / backstroke / breaststroke / butterfly) over the 12 COCO body joints. 100% val accuracy (all four strokes separate cleanly in pose space). Input (N,27,12,3) canonical body joints → (N,4) stroke logits; gives a per-frame stroke label for the whole pipeline.

End-to-end pipeline

video → ViTPose-L 2D (PCK 0.982) → temporal lifter → 3D joints → SMPL body → mesh

The realistic lifter is trained on the finetuned detector's own 2D-error distribution (5.8 px front … 22 px underwater, ~16 px overall), so 33.9 mm MPJPE is an honest video→3D estimate, not a GT-2D upper bound. The temporal model absorbs detector noise well (+3.6 mm vs clean). The smpl-regressor-coco model consumes the lifter's COCO body joints (canonicalized to remove camera orientation) and outputs SMPL body pose for mesh recovery.

3D pose (2D→3D lifting)

A temporal convolutional lifter turns the 2D keypoint sequence above into root-relative 3D joint positions. Trained on SwimXYZ's ground-truth 3D (3D_cam.txt, camera-space, pelvis-centered) across all five views.

MPJPE on held-out synthetic validation: 30.3 mm (mean-pose baseline: 670 mm).

View	MPJPE
Aerial	35 mm
Front	41 mm
Side above-water	32 mm
Side underwater	29 mm
Side water-level	31 mm

• Input: (N, 27, 17, 2) — 27-frame window of screen-normalized 2D keypoints (VideoPose3D convention: x/w*2-1, y/w*2-h/w), COCO-17 order.
• Output: (N, 17, 3) — root-relative 3D joints (metres) for the centre frame.

import torch
from lift_model import TemporalModel   # training/src/lift_model.py
lift = TemporalModel()
lift.load_state_dict(torch.load("lifter-swimxyz.pth")["state_dict"]); lift.eval()
xyz = lift(win)   # win: (N,27,17,2) -> (N,17,3) metres, root-relative

Model I/O

• Input: (N, 3, 256, 192) float32, ImageNet-normalized RGB person crop (top-down — supply a person box from a detector, e.g. YOLO).
• Output: (N, 17, 64, 48) keypoint heatmaps, COCO-17 joint order.

Usage

import torch
# model definition: https://github.com/JunkyByte/easy_ViTPose  (ViTPose class)
from vit_models.model import ViTPose
from configs.train_configs.ViTPose_large_coco_256x192 import model as cfg

net = ViTPose(cfg)
ckpt = torch.load("vitpose-l-swimxyz-mv.pth", map_location="cpu", weights_only=False)
net.load_state_dict(ckpt["state_dict"])
net.eval()
# heatmaps = net(crop)   # crop: (N,3,256,192) ImageNet-normalized

ONNX:

import onnxruntime as ort
sess = ort.InferenceSession("vitpose-l-swimxyz-mv.onnx")
heatmaps = sess.run(None, {"input": crop_np})[0]  # (N,17,64,48)

SMPL mesh recovery (3D→SMPL)

A temporal regressor maps a window of 3D joints to SMPL body parameters (6D-rotation pose + shape), trained on SwimXYZ's 60 freestyle SMPL motions with 6D-param + joint-MPJPE + vertex-PVE losses through a differentiable SMPL layer. Feed the output to the SMPL body model to recover a full mesh.

Joint-reconstruction error on held-out synthetic validation: 1.5 mm — note this is measured with clean SMPL-derived 24-joint input, so it reflects the regressor inverting the joint→SMPL map, not end-to-end mesh-from-video accuracy. End-to-end (video → 2D → lift → SMPL) is dominated by the lifter (~30 mm) plus a COCO-17 → SMPL-24 joint-definition gap; wiring that bridge and fine-tuning on predicted joints is the natural next step.

Requires the licensed SMPL body model (https://smpl.is.tue.mpg.de) at inference to turn parameters into a mesh; not redistributed here.

Provenance & training details

• Finetuned from the COCO ViTPose-L weights in JunkyByte/easy_ViTPose.
• SwimXYZ labels are OpenPose COCO-18 order (not the 25-joint file header), bottom-origin (y_img = H - y); mapped to COCO-17. Validated by overlay.
• Full pipeline, code, and writeup: SwimLab training/ (build → finetune → export).

Citation

If you use this model, please cite SwimXYZ and ViTPose:

@inproceedings{fiche2023swimxyz, title={SwimXYZ: A large-scale dataset of synthetic swimming motions and videos}, author={Fiche, Gu{\'e}nol{\'e} and others}, year={2023}}
@article{xu2022vitpose, title={ViTPose: Simple Vision Transformer Baselines for Human Pose Estimation}, author={Xu, Yufei and others}, journal={NeurIPS}, year={2022}}