Keypoint Detection & Pose Estimation

Keypoint tasks hide under half a dozen different names: human pose (17 body joints), face landmarks (68 or 478 points), hand pose (21 points), animal pose, robotic object pose, medical anatomy landmarks. Strip the domain label off any of them and the task is identical: detect K discrete points on an object and report their (x, y) coordinates.

That single structure underwrites motion capture, fitness apps, sports analytics, gesture control, animation, AR try-on, and robotic grasping. The 2D case is solved at production quality. 3D pose, recovering joint positions in world coordinates from a single camera, is still where the open research problems live.

The engineering question that actually matters here is scale, not accuracy. A single-image, single-person pose is a 20ms problem you can solve with a small model. Multi-person pose in a crowd at 30 fps is a different problem, with different architectures, even though both get filed under "pose estimation."

Top-down vs bottom-up

• Top-down: detect people first, then run a per-person keypoint model on each crop. Highest accuracy, and cost scales linearly with the number of people in frame.
• Bottom-up: one forward pass predicts every keypoint plus an association field, then groups them into instances. Constant time no matter how crowded the scene gets.

Top-down (HRNet, ViTPose) wins on accuracy. Bottom-up (OpenPose, HigherHRNet) wins on throughput once the scene fills up with people.

Heatmap regression

Instead of regressing (x, y) directly, predict an H x W heatmap per keypoint with a Gaussian blob centred at the true location.

target[k, y, x] = exp(-((x - cx_k)^2 + (y - cy_k)^2) / (2 sigma^2))

At inference, the argmax of each heatmap is the predicted keypoint location.

Heatmaps beat direct regression because a conv feature map's spatial structure already lines up with a spatial output; asking the network to collapse that structure into two raw numbers throws information away. A Gaussian target also regularises the loss: a small localisation error costs a small loss, not zero.

Sub-pixel localisation

Argmax only gives integer coordinates. For sub-pixel precision, fit a parabola to the argmax and its neighbours, or use the well-known offset (dx, dy) = 0.25 * (heatmap[y, x+1] - heatmap[y, x-1], ...) direction.

Part Affinity Fields (PAFs)

OpenPose's trick for bottom-up association. For each pair of connected keypoints (left shoulder to left elbow, say), predict a 2-channel field encoding the unit vector that points from one to the other. To decide whether a candidate shoulder and a candidate elbow belong to the same person, integrate the PAF along the line connecting them: the pair with the highest integral wins the match.

For each connection (limb):
  PAF channels: 2 (unit vector x, y)
  Line integral: sum over sample points of (PAF . line_direction)
  Higher integral = stronger match

Simple idea, and it scales to arbitrary crowd sizes without ever cropping a single person out of the frame.

COCO keypoints

The standard body-pose dataset annotates 17 keypoints per person and reports two metrics: PCK (Percentage of Correct Keypoints) and OKS (Object Keypoint Similarity). OKS is the keypoint analogue of IoU, and it's what COCO's mAP@OKS number is actually measuring.

2D vs 3D

• 2D pose: image coordinates, solved at production quality (MediaPipe, HRNet, ViTPose).
• 3D pose: world or camera coordinates, still active research. Three common approaches:
  • Lift 2D predictions to 3D with a small MLP (VideoPose3D).
  • Regress 3D directly from the image (PyMAF, MHFormer).
  • Use a multi-view capture rig (CMU Panoptic) to generate ground truth.

Step 1: Gaussian heatmap target

import numpy as np
import torch
 
def gaussian_heatmap(size, cx, cy, sigma=2.0):
    yy, xx = np.meshgrid(np.arange(size), np.arange(size), indexing="ij")
    return np.exp(-((xx - cx) ** 2 + (yy - cy) ** 2) / (2 * sigma ** 2)).astype(np.float32)
 
hm = gaussian_heatmap(64, 32, 32, sigma=2.0)
print(f"peak: {hm.max():.3f} at ({hm.argmax() % 64}, {hm.argmax() // 64})")

Per-keypoint heatmaps stacked along a channel axis give the full target tensor.

Step 2: Tiny keypoint head

A U-Net-style model that outputs K heatmap channels.

import torch.nn as nn
import torch.nn.functional as F
 
class TinyKeypointNet(nn.Module):
    def __init__(self, num_keypoints=4, base=16):
        super().__init__()
        self.down1 = nn.Sequential(nn.Conv2d(3, base, 3, 2, 1), nn.ReLU(inplace=True))
        self.down2 = nn.Sequential(nn.Conv2d(base, base * 2, 3, 2, 1), nn.ReLU(inplace=True))
        self.mid = nn.Sequential(nn.Conv2d(base * 2, base * 2, 3, 1, 1), nn.ReLU(inplace=True))
        self.up1 = nn.ConvTranspose2d(base * 2, base, 2, 2)
        self.up2 = nn.ConvTranspose2d(base, num_keypoints, 2, 2)
 
    def forward(self, x):
        h1 = self.down1(x)
        h2 = self.down2(h1)
        h3 = self.mid(h2)
        u1 = self.up1(h3)
        return self.up2(u1)

Input (N, 3, H, W), output (N, K, H, W). Loss is per-pixel MSE against Gaussian targets.

Step 3: Keypoint coordinates at inference

def heatmap_to_coords(heatmaps):
    """
    heatmaps: (N, K, H, W)
    returns:  (N, K, 2) float coordinates in image pixels
    """
    N, K, H, W = heatmaps.shape
    hm = heatmaps.reshape(N, K, -1)
    idx = hm.argmax(dim=-1)
    ys = (idx // W).float()
    xs = (idx % W).float()
    return torch.stack([xs, ys], dim=-1)
 
coords = heatmap_to_coords(torch.randn(2, 4, 32, 32))
print(f"coords: {coords.shape}")  # (2, 4, 2)

That's the whole inference step, one line. For sub-pixel refinement, interpolate around the argmax instead of taking it raw.

Step 4: Synthetic keypoint dataset

Simple: draw four points on a white canvas and learn to predict them.

def make_synthetic_sample(size=64):
    img = np.ones((3, size, size), dtype=np.float32)
    rng = np.random.default_rng()
    kps = rng.integers(8, size - 8, size=(4, 2))
    for cx, cy in kps:
        img[:, cy - 2:cy + 2, cx - 2:cx + 2] = 0.0
    hms = np.stack([gaussian_heatmap(size, cx, cy) for cx, cy in kps])
    return img, hms, kps

Easy enough for a tiny model to learn in a minute.

Step 5: Training

model = TinyKeypointNet(num_keypoints=4)
opt = torch.optim.Adam(model.parameters(), lr=3e-3)
 
for step in range(200):
    batch = [make_synthetic_sample() for _ in range(16)]
    imgs = torch.from_numpy(np.stack([b[0] for b in batch]))
    hms = torch.from_numpy(np.stack([b[1] for b in batch]))
    pred = model(imgs)
    # Upsample pred to full resolution
    pred = F.interpolate(pred, size=hms.shape[-2:], mode="bilinear", align_corners=False)
    loss = F.mse_loss(pred, hms)
    opt.zero_grad(); loss.backward(); opt.step()

• MediaPipe Pose: Google's production pose estimator. Ships WebGL and mobile runtimes with sub-10ms latency. Output is 33 landmarks per person, each an (x, y, z, visibility) tuple: x and y normalised to image size, z a relative depth estimate, visibility a confidence score that the point is actually in frame.
• MMPose (OpenMMLab): the comprehensive research codebase, every SOTA architecture with pretrained weights included. Output is a (num_people, K, 3) array, one (x, y, score) triple per keypoint, with K and the keypoint order set by whichever dataset format the model trained on (17 for COCO body).
• YOLOv8-pose: the fastest real-time multi-person pose option, a single forward pass and nothing else.
• transformers HumanDPT / PoseAnything: newer vision-language approaches for open-vocabulary pose, any object, any keypoint set.

This lesson produces:

• outputs/prompt-pose-stack-picker.md: a prompt that picks MediaPipe, YOLOv8-pose, HRNet, or ViTPose given latency, crowd size, and whether you need 2D or 3D.
• outputs/skill-heatmap-to-coords.md: a skill that writes the sub-pixel heatmap-to-coordinate routine every production pose model relies on.

1. (Easy) Train the tiny keypoint model on the synthetic 4-point dataset. Report mean L2 error between predicted and true keypoints after 200 steps.
2. (Medium) Add sub-pixel refinement: given the argmax position, fit a 1D parabola along x and y from the neighbouring pixels. Report the accuracy gain vs integer argmax.
3. (Hard) Build a 2-person synthetic dataset where each image shows two instances of the 4-keypoint pattern. Train a bottom-up pipeline with PAFs that predict which keypoint belongs to which instance, and evaluate OKS.

• OpenPose (Cao et al., 2017): bottom-up with PAFs, still the best writeup of the approach
• HRNet (Sun et al., 2019): the top-down reference architecture
• ViTPose (Xu et al., 2022): a plain ViT as pose backbone, current SOTA on many benchmarks
• MediaPipe Pose: production real-time pose, the fastest deployed stack in 2026