如何从PyTorch姿态估计模型输出张量提取带置信度的关键点？

Hey there! Let's walk through how to convert that model output tensor into usable keypoint coordinates with confidence scores. I’ve worked with this exact HRNet-based pose estimation repo before, so I can break this down clearly.

First, understand your output tensor #

Your output shape [1, 16, 64, 56] translates to:

• 1: Batch size (you’re processing one image at a time)
• 16: Number of keypoints the model is predicting (this might differ from the 17 you mentioned if the model config was adjusted—double-check your setup if you expect 17!)
• 64x56: The size of the heatmaps for each keypoint. Each heatmap is a 2D grid where the peak value indicates the confidence of the keypoint being at that location, and the peak’s position maps to the keypoint’s coordinates.

Step 1: Process the heatmaps to get raw coordinates & confidence #

Each keypoint’s heatmap has a peak value (confidence) and a corresponding (y, x) position in the heatmap grid. Here’s how to extract those:

First, get your original image dimensions (we’ll need this to map back to pixel coordinates later):

# If your input `img` is a numpy array
original_h, original_w = img.shape[:2]
# If it's a PIL Image
original_h, original_w = img.size[1], img.size[0]

Then process the model output:

model.eval()
img_t = ToTensor()(img).unsqueeze(0)
# Always use torch.no_grad() during inference to save memory
with torch.no_grad():
    out = model(img_t)

# Remove the batch dimension since we're working with one image
heatmaps = out.squeeze(0)  # Shape becomes [16, 64, 56]
num_keypoints = heatmaps.shape[0]

Step 2: Convert heatmap positions to original image pixels #

The heatmap is a downscaled version of your input image (HRNet typically outputs heatmaps at 1/4 the input resolution). We’ll map the heatmap coordinates back to your original image’s pixel space, and capture confidence scores:

Basic method (integer coordinates)

This is simple and works well for most cases:

keypoints = []
for kp_idx in range(num_keypoints):
    heatmap = heatmaps[kp_idx]  # Get the heatmap for this keypoint
    # Find the maximum confidence value and its flattened index
    conf, flat_pos = torch.max(heatmap.view(-1), 0)
    # Convert flattened index to (y, x) in the heatmap grid
    hm_y = flat_pos // heatmap.shape[1]
    hm_x = flat_pos % heatmap.shape[1]
    # Map to original image coordinates
    orig_x = float(hm_x * original_w / heatmap.shape[1])
    orig_y = float(hm_y * original_h / heatmap.shape[0])
    # Store the keypoint data
    keypoints.append({
        "keypoint_id": kp_idx,
        "x": orig_x,
        "y": orig_y,
        "confidence": float(conf)
    })

Improved method (sub-pixel accuracy)

For more precise coordinates, you can use sub-pixel interpolation around the heatmap peak. This helps with fine-grained alignment:

def get_subpixel_keypoint(heatmap, original_h, original_w):
    hm_h, hm_w = heatmap.shape
    # Get initial peak position (integer coordinates)
    conf, flat_pos = torch.max(heatmap.view(-1), 0)
    hm_y = flat_pos // hm_w
    hm_x = flat_pos % hm_w
    # Convert to float for interpolation
    hm_x, hm_y = hm_x.float(), hm_y.float()

    # Calculate sub-pixel offset using surrounding pixels
    if 0 < hm_x < hm_w - 1 and 0 < hm_y < hm_h - 1:
        # Compute gradients in x and y directions
        dx = heatmap[hm_y, hm_x + 1] - heatmap[hm_y, hm_x - 1]
        dy = heatmap[hm_y + 1, hm_x] - heatmap[hm_y - 1, hm_x]
        # Update coordinates with sub-pixel adjustment
        hm_x += dx / (2 * (heatmap[hm_y, hm_x + 1] + heatmap[hm_y, hm_x - 1] - 2 * heatmap[hm_y, hm_x]))
        hm_y += dy / (2 * (heatmap[hm_y + 1, hm_x] + heatmap[hm_y - 1, hm_x] - 2 * heatmap[hm_y, hm_x]))

    # Map to original image coordinates
    orig_x = hm_x * original_w / hm_w
    orig_y = hm_y * original_h / hm_h
    return orig_x.item(), orig_y.item(), conf.item()

# Use this function for each keypoint
keypoints = []
for kp_idx in range(num_keypoints):
    x, y, conf = get_subpixel_keypoint(heatmaps[kp_idx], original_h, original_w)
    keypoints.append({
        "keypoint_id": kp_idx,
        "x": x,
        "y": y,
        "confidence": conf
    })

Handling multiple people #

Wait—if you’re trying to detect multiple people in an image, this model (HRNet from the repo) is a top-down model. That means you first need to run a person detector (like Faster R-CNN) to get bounding boxes for each person, then crop each person from the original image, resize the crop to the model’s input size, run inference on each crop, and map the keypoints back to the original image using the bounding box coordinates.

For example, if you have a bounding box with top-left corner (x1, y1), width bbox_w, height bbox_h, and you resized the crop to (256, 224) for the model, the mapped coordinates would be:

orig_x = x1 + (crop_keypoint_x * bbox_w / 224)
orig_y = y1 + (crop_keypoint_y * bbox_h / 256)

内容的提问来源于stack exchange，提问作者Ioanna