Pose Estimation

Module: Computer Vision | Difficulty: Intermediate

2D Pose Estimation

Heatmap-Based Detection

Predict heatmaps for keypoints:

OKS (Object Keypoint Similarity)

AP@OKS

HRNet Architecture

Maintains high-resolution representations throughout:

import torch
import torch.nn as nn

class PoseHead(nn.Module):
    def __init__(self, in_channels=256, num_keypoints=17):
        super().__init__()
        self.head = nn.Sequential(
            nn.Conv2d(in_channels, in_channels, 3, padding=1),
            nn.BatchNorm2d(in_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels, in_channels, 3, padding=1),
            nn.BatchNorm2d(in_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels, in_channels, 3, padding=1),
            nn.BatchNorm2d(in_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels, num_keypoints, 1),
        )
    
    def forward(self, x):
        return self.head(x)

def generate_heatmaps(keypoints, img_size=64, sigma=2):
    B, K, _ = keypoints.shape
    heatmaps = torch.zeros(B, K, img_size, img_size)
    for k in range(K):
        x, y = keypoints[:, k, 0], keypoints[:, k, 1]
        xx, yy = torch.meshgrid(torch.arange(img_size), torch.arange(img_size), indexing='ij')
        heatmaps[:, k] = torch.exp(-((xx - x.view(-1, 1, 1))**2 + (yy - y.view(-1, 1, 1))**2) / (2 * sigma**2))
    return heatmaps

Key Takeaways

• Heatmap-based methods dominate 2D pose estimation
• OKS is the standard evaluation metric (analogous to IoU for boxes)
• HRNet preserves spatial detail for accurate keypoint localization