Fine-Tune ViTPose++ for Keypoint Detection with Transformers

TL;DR

The old SuperGradients/YOLO-NAS-Pose workflow is replaced with a current transformers 5.x workflow built around VitPoseForPoseEstimation.
The recommended starting checkpoint is usyd-community/vitpose-plus-base: it is Transformers-native, Apache-2.0, actively used on the Hub, and much more current than the original YOLO-NAS-Pose setup.
VitPoseForPoseEstimation does not implement a built-in training loss yet, so this tutorial fine-tunes the model with a small PyTorch loop that generates target heatmaps from COCO-style keypoints.

Abstract

Pose estimation tutorials age quickly because the model, dataset, and training framework often move together. The older version of this tutorial depended on SuperGradients and YOLO-NAS-Pose, which made the workflow harder to keep current with the Hugging Face ecosystem. This rewrite shows how to fine-tune a modern ViTPose++ keypoint detector with transformers, a COCO-style annotation file, and a minimal heatmap loss. By the end, you will know how to train the pose head, run inference with RT-DETR person boxes, and decide when you should switch to a dedicated pose framework instead.

Requirements

Python 3.10+
torch
transformers>=5.9.0
pillow
numpy
tqdm
A COCO-style keypoint dataset with image paths, bounding boxes, and 17 keypoints per person

pip install -U "transformers>=5.9.0" torch pillow numpy tqdm

Optional packages for richer evaluation and visualization:

pip install -U opencv-python pycocotools supervision

Prerequisites

You are comfortable with PyTorch training loops.
You understand COCO keypoint annotations: each keypoint is stored as (x, y, visibility).
Your training data uses the same 17-keypoint COCO human-pose schema as the checkpoint used below.

If your dataset uses a different keypoint layout, do not silently reuse this code. You need a compatible checkpoint or a resized/reinitialized prediction head, plus a new skeleton definition and evaluation protocol.

Why ViTPose++ Now

ViTPose is a top-down pose estimator. It does not find people by itself. Instead, a detector first produces person bounding boxes, and ViTPose predicts keypoints inside each box.

For a current Transformers-native tutorial, usyd-community/vitpose-plus-base is the most practical default:

Model	Why it matters	Tradeoff
`usyd-community/vitpose-plus-base`	Strong default, Apache-2.0, high Hub usage, works with `VitPoseForPoseEstimation`	Requires a detector for boxes
`usyd-community/vitpose-plus-huge`	Better capacity for inference or large hardware	Much heavier to fine-tune
`facebook/sapiens2-pose-*`	Newer Sapiens-family pose checkpoints on the Hub	Larger and not the cleanest starting point for this Transformers fine-tuning tutorial
Ultralytics YOLO pose models	Great if you want detector and pose in one YOLO-style workflow	Not the Transformers-native path shown here

There is no useful official Hugging Face leaderboard for this exact fine-tuning task, so the choice is based on current Hub availability, Transformers support, license, and practical reproducibility.

The Honest Transformers Caveat

The current VitPoseForPoseEstimation API returns predicted heatmaps, but its built-in loss is not implemented. Passing labels is not enough.

That means a working fine-tuning recipe needs to do three things explicitly:

Convert annotations into instance crops.
Generate target heatmaps aligned with the model output.
Compute the heatmap loss in your own PyTorch loop.

This is less convenient than Trainer, but it has one advantage: you can see every assumption in the training path.

Approach

This tutorial uses a simple top-down training setup:

Load a pretrained ViTPose++ checkpoint.
Read COCO-style keypoint annotations.
Crop each annotated person instance.
Generate one Gaussian heatmap per keypoint.
Freeze the backbone and fine-tune the pose head.
Save the model and run inference with RT-DETR person detections.

The code below is intentionally small. It is a baseline you can extend with stronger augmentation, OKS evaluation, distributed training, and model-card publishing after the core loop works.

Load the Model

import torch
from transformers import AutoProcessor, VitPoseForPoseEstimation

MODEL_ID = "usyd-community/vitpose-plus-base"
DATASET_INDEX = 0  # 0 is the COCO human-pose expert for ViTPose++.

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

processor = AutoProcessor.from_pretrained(MODEL_ID)
model = VitPoseForPoseEstimation.from_pretrained(MODEL_ID).to(device)

For the current checkpoint, the model predicts 17 heatmaps:

print(model.config.num_labels)
# 17

The output heatmap size for the default processor is 64 x 48:

from PIL import Image

dummy = Image.new("RGB", (192, 256), color="white")
inputs = processor(dummy, boxes=[[[0, 0, 192, 256]]], return_tensors="pt").to(device)

with torch.no_grad():
    outputs = model(**inputs, dataset_index=torch.tensor([DATASET_INDEX], device=device))

print(outputs.heatmaps.shape)
# torch.Size([1, 17, 64, 48])

Prepare COCO-Style Keypoints

This tutorial expects a COCO-style annotation file:

dataset/
  images/
    000000000001.jpg
    000000000002.jpg
  annotations/
    person_keypoints_train.json
    person_keypoints_val.json

Each annotation should include:

image_id
bbox in COCO format: [x, y, width, height]
keypoints as a flat list of 17 * 3 values
num_keypoints

The dataset class below turns each person annotation into one crop. That keeps the target heatmap math straightforward: the crop itself is the model input, and the model sees one box that covers the crop.

import json
from pathlib import Path

import numpy as np
import torch
from PIL import Image
from torch.utils.data import Dataset


def padded_xyxy_from_xywh(box, image_width, image_height, padding=0.15):
    x, y, w, h = box
    pad_x = w * padding
    pad_y = h * padding

    x1 = max(0, int(round(x - pad_x)))
    y1 = max(0, int(round(y - pad_y)))
    x2 = min(image_width, int(round(x + w + pad_x)))
    y2 = min(image_height, int(round(y + h + pad_y)))

    return x1, y1, x2, y2


def make_target_heatmaps(keypoints, crop_size, heatmap_size, sigma=2.0):
    crop_width, crop_height = crop_size
    heatmap_height, heatmap_width = heatmap_size

    targets = np.zeros((len(keypoints), heatmap_height, heatmap_width), dtype=np.float32)
    weights = np.zeros((len(keypoints), 1), dtype=np.float32)

    xs = np.arange(heatmap_width, dtype=np.float32)
    ys = np.arange(heatmap_height, dtype=np.float32)
    grid_x, grid_y = np.meshgrid(xs, ys)

    for joint_id, (x, y, visibility) in enumerate(keypoints):
        if visibility <= 0:
            continue
        if x < 0 or y < 0 or x >= crop_width or y >= crop_height:
            continue

        heatmap_x = x / max(crop_width, 1) * heatmap_width
        heatmap_y = y / max(crop_height, 1) * heatmap_height

        targets[joint_id] = np.exp(
            -((grid_x - heatmap_x) ** 2 + (grid_y - heatmap_y) ** 2) / (2 * sigma**2)
        )
        weights[joint_id, 0] = 1.0

    return torch.from_numpy(targets), torch.from_numpy(weights)


class CocoPoseCropDataset(Dataset):
    def __init__(
        self,
        image_dir,
        annotation_file,
        processor,
        heatmap_size=(64, 48),
        num_keypoints=17,
        padding=0.15,
        max_samples=None,
    ):
        self.image_dir = Path(image_dir)
        self.processor = processor
        self.heatmap_size = heatmap_size
        self.num_keypoints = num_keypoints
        self.padding = padding

        with open(annotation_file, "r", encoding="utf-8") as f:
            coco = json.load(f)

        self.images = {image["id"]: image for image in coco["images"]}
        self.annotations = [
            ann
            for ann in coco["annotations"]
            if ann.get("num_keypoints", 0) > 0
            and len(ann.get("keypoints", [])) == num_keypoints * 3
            and ann.get("bbox", [0, 0, 0, 0])[2] > 1
            and ann.get("bbox", [0, 0, 0, 0])[3] > 1
        ]

        if max_samples is not None:
            self.annotations = self.annotations[:max_samples]

    def __len__(self):
        return len(self.annotations)

    def __getitem__(self, index):
        annotation = self.annotations[index]
        image_info = self.images[annotation["image_id"]]
        image = Image.open(self.image_dir / image_info["file_name"]).convert("RGB")

        x1, y1, x2, y2 = padded_xyxy_from_xywh(
            annotation["bbox"],
            image_width=image.width,
            image_height=image.height,
            padding=self.padding,
        )

        crop = image.crop((x1, y1, x2, y2))
        crop_width, crop_height = crop.size

        keypoints = np.asarray(annotation["keypoints"], dtype=np.float32).reshape(
            self.num_keypoints, 3
        )
        keypoints[:, 0] -= x1
        keypoints[:, 1] -= y1

        target_heatmaps, target_weights = make_target_heatmaps(
            keypoints,
            crop_size=(crop_width, crop_height),
            heatmap_size=self.heatmap_size,
        )

        inputs = self.processor(
            crop,
            boxes=[[[0, 0, crop_width, crop_height]]],
            return_tensors="pt",
        )

        return {
            "pixel_values": inputs["pixel_values"][0],
            "target_heatmaps": target_heatmaps,
            "target_weights": target_weights,
        }

Build DataLoaders

from torch.utils.data import DataLoader


def collate_pose_batch(batch):
    return {
        "pixel_values": torch.stack([item["pixel_values"] for item in batch]),
        "target_heatmaps": torch.stack([item["target_heatmaps"] for item in batch]),
        "target_weights": torch.stack([item["target_weights"] for item in batch]),
    }


train_dataset = CocoPoseCropDataset(
    image_dir="dataset/images",
    annotation_file="dataset/annotations/person_keypoints_train.json",
    processor=processor,
    heatmap_size=(64, 48),
    max_samples=None,
)

val_dataset = CocoPoseCropDataset(
    image_dir="dataset/images",
    annotation_file="dataset/annotations/person_keypoints_val.json",
    processor=processor,
    heatmap_size=(64, 48),
    max_samples=None,
)

train_loader = DataLoader(
    train_dataset,
    batch_size=16,
    shuffle=True,
    num_workers=4,
    collate_fn=collate_pose_batch,
    pin_memory=True,
)

val_loader = DataLoader(
    val_dataset,
    batch_size=16,
    shuffle=False,
    num_workers=4,
    collate_fn=collate_pose_batch,
    pin_memory=True,
)

Fine-Tune the Pose Head

Start conservatively by freezing the backbone and training only the heatmap head. This is a useful first pass because it verifies the dataset, target generation, and loss before you spend GPU time updating the full model.

from tqdm.auto import tqdm


for parameter in model.parameters():
    parameter.requires_grad = False

for parameter in model.head.parameters():
    parameter.requires_grad = True

optimizer = torch.optim.AdamW(model.head.parameters(), lr=5e-4, weight_decay=1e-4)


def heatmap_mse_loss(predicted, target, weights):
    per_joint_loss = (predicted - target).pow(2).mean(dim=(-1, -2), keepdim=True)
    weighted_loss = per_joint_loss * weights[:, :, None, :]
    return weighted_loss.sum() / weights.sum().clamp_min(1.0)


def train_one_epoch(model, loader, optimizer, device):
    model.train()
    total_loss = 0.0

    for batch in tqdm(loader, desc="train"):
        pixel_values = batch["pixel_values"].to(device)
        target_heatmaps = batch["target_heatmaps"].to(device)
        target_weights = batch["target_weights"].to(device)
        dataset_index = torch.full(
            (pixel_values.shape[0],),
            DATASET_INDEX,
            dtype=torch.long,
            device=device,
        )

        optimizer.zero_grad(set_to_none=True)
        outputs = model(pixel_values=pixel_values, dataset_index=dataset_index)
        loss = heatmap_mse_loss(outputs.heatmaps, target_heatmaps, target_weights)
        loss.backward()
        optimizer.step()

        total_loss += loss.item()

    return total_loss / max(len(loader), 1)


@torch.no_grad()
def evaluate_loss(model, loader, device):
    model.eval()
    total_loss = 0.0

    for batch in tqdm(loader, desc="val"):
        pixel_values = batch["pixel_values"].to(device)
        target_heatmaps = batch["target_heatmaps"].to(device)
        target_weights = batch["target_weights"].to(device)
        dataset_index = torch.full(
            (pixel_values.shape[0],),
            DATASET_INDEX,
            dtype=torch.long,
            device=device,
        )

        outputs = model(pixel_values=pixel_values, dataset_index=dataset_index)
        loss = heatmap_mse_loss(outputs.heatmaps, target_heatmaps, target_weights)
        total_loss += loss.item()

    return total_loss / max(len(loader), 1)


for epoch in range(5):
    train_loss = train_one_epoch(model, train_loader, optimizer, device)
    val_loss = evaluate_loss(model, val_loader, device)
    print(f"epoch={epoch + 1} train_loss={train_loss:.5f} val_loss={val_loss:.5f}")

Once the head-only run behaves sensibly, unfreeze the last few backbone blocks and lower the learning rate:

for layer in model.backbone.encoder.layer[-2:]:
    for parameter in layer.parameters():
        parameter.requires_grad = True

optimizer = torch.optim.AdamW(
    [parameter for parameter in model.parameters() if parameter.requires_grad],
    lr=5e-5,
    weight_decay=1e-4,
)

Do this only after you have verified that the heatmaps and validation loss are moving in the right direction. If the targets are misaligned, unfreezing more layers will only make the model better at learning the wrong geometry.

Save the Fine-Tuned Model

OUTPUT_DIR = "vitpose-plus-base-coco-keypoints-finetuned"

model.save_pretrained(OUTPUT_DIR)
processor.save_pretrained(OUTPUT_DIR)

You can push the directory to the Hub later with huggingface_hub or the hf CLI.

Run Inference with RT-DETR Boxes

For inference, use a detector to find people, convert the detector boxes to COCO [x, y, width, height] format, and pass those boxes into ViTPose++.

import torch
from PIL import Image
from transformers import AutoProcessor, RTDetrForObjectDetection, VitPoseForPoseEstimation

DETECTOR_ID = "PekingU/rtdetr_r50vd_coco_o365"
POSE_MODEL_DIR = "vitpose-plus-base-coco-keypoints-finetuned"

image = Image.open("example.jpg").convert("RGB")

detector_processor = AutoProcessor.from_pretrained(DETECTOR_ID)
detector = RTDetrForObjectDetection.from_pretrained(DETECTOR_ID).to(device)

pose_processor = AutoProcessor.from_pretrained(POSE_MODEL_DIR)
pose_model = VitPoseForPoseEstimation.from_pretrained(POSE_MODEL_DIR).to(device)

detector_inputs = detector_processor(images=image, return_tensors="pt").to(device)

with torch.no_grad():
    detector_outputs = detector(**detector_inputs)

detections = detector_processor.post_process_object_detection(
    detector_outputs,
    target_sizes=torch.tensor([(image.height, image.width)], device=device),
    threshold=0.3,
)[0]

person_boxes_xyxy = detections["boxes"][detections["labels"] == 0].detach().cpu().numpy()

if len(person_boxes_xyxy) == 0:
    raise ValueError("No person detections found. Try lowering the detector threshold.")

person_boxes_xywh = person_boxes_xyxy.copy()
person_boxes_xywh[:, 2] = person_boxes_xyxy[:, 2] - person_boxes_xyxy[:, 0]
person_boxes_xywh[:, 3] = person_boxes_xyxy[:, 3] - person_boxes_xyxy[:, 1]

pose_inputs = pose_processor(
    image,
    boxes=[person_boxes_xywh],
    return_tensors="pt",
).to(device)

dataset_index = torch.full(
    (pose_inputs["pixel_values"].shape[0],),
    DATASET_INDEX,
    dtype=torch.long,
    device=device,
)

with torch.no_grad():
    pose_outputs = pose_model(**pose_inputs, dataset_index=dataset_index)

pose_results = pose_processor.post_process_pose_estimation(
    pose_outputs,
    boxes=[person_boxes_xywh],
)

print(pose_results[0][0].keys())
# dict_keys(['keypoints', 'scores', 'labels', 'bbox'])

Sanity Checks

Before running a long training job, check these items:

A single batch returns pixel_values with shape (batch, 3, 256, 192).
outputs.heatmaps has shape (batch, 17, 64, 48).
Target heatmaps have the same shape as model heatmaps.
Visible keypoints produce non-empty heatmaps.
A batch can overfit. If one or two batches cannot drive training loss down, inspect crop alignment before changing the optimizer.
Validation loss is computed on the same keypoint schema as training.

Failure Modes and Caveats

The keypoint schema does not match

vitpose-plus-base is a 17-keypoint model in this setup. If your dataset has 20 animal keypoints, 24 dog keypoints, or 133 whole-body keypoints, the head and target layout must change. Do not pad labels and hope the model figures it out.

The boxes are wrong

Top-down pose estimation depends on good boxes. If your detector boxes are loose, missing limbs, or aimed at the wrong class, the pose model will fail even if the keypoint head is well trained.

Heatmaps are misaligned

The simplest way to avoid processor geometry bugs is to train on person crops and pass a box covering the crop. If you train directly on full images with arbitrary boxes, you need to reproduce the processor’s affine transform when generating target heatmaps.

Head-only fine-tuning plateaus

That is normal. Head-only training is a diagnostic baseline. For a real domain shift, unfreeze the last few ViT blocks, lower the learning rate, and use stronger validation.

You need production-grade pose training

If you need full COCO OKS evaluation, multi-GPU recipes, advanced augmentation, and mature pose-specific losses, use a dedicated pose framework such as MMPose. Transformers is excellent for portable inference and lightweight adaptation, but ViTPose training still requires custom glue.

Practical Guidance

Use usyd-community/vitpose-plus-base for the first fine-tuning run. It is small enough to iterate on, but current enough to justify replacing the older YOLO-NAS-Pose tutorial.

Use usyd-community/vitpose-plus-huge when inference quality matters more than iteration speed. For fine-tuning, start with base, confirm the data path, then scale up.

Use an Ultralytics YOLO pose model if your real requirement is a single package that detects and estimates pose end to end. Use ViTPose++ if you want a Transformers-native pose model and are comfortable pairing it with a detector.

Summary

The modern Hugging Face path for pose estimation is not a drop-in replacement for the old SuperGradients tutorial. The model stack is better aligned with current Transformers APIs, but training requires an explicit heatmap target and custom loss. That tradeoff is worth it when you want a portable ViTPose++ workflow you can inspect, adapt, and save with standard Hugging Face model tooling.

Next Steps

Run the head-only loop on a tiny subset and confirm it overfits.
Add COCO OKS or PCK evaluation before comparing checkpoints.
Add augmentations only after the unaugmented crop geometry is correct.
Push the trained model and processor to the Hugging Face Hub with a model card that documents the keypoint schema.