Add pose estimation keypoint preprocessing to Sapiens2ImageProcessor by Sainava · Pull Request #47199 · huggingface/transformers

Sainava · 2026-07-09T07:36:09Z

What does this PR do?

As discussed with @guarin in the previous Sapiens2 loss function PR (#46764), this PR implements the preprocessing data pipeline required for Sapiens2 pose estimation fine-tuning.

Specifically, it updates Sapiens2ImageProcessor (_preprocess_image_like_inputs) to handle raw keypoint coordinates. It implements:

Vectorized 2D Gaussian Heatmap Generation: Transforms raw (x, y) image coordinates into scaled (H, W) heatmap targets using inverse affine transformations based on the cropped bounding boxes.
Visibility Masking: Handles COCO-style visibility flags, zeroing out weights for invisible or out-of-bounds keypoints.
Memory-Optimized Weight Tensors: Formats the label_weights to (num_kps, 1, 1) to utilize PyTorch broadcasting during loss calculation, avoiding dense (num_kps, H, W) allocations.

Testing

I ran the full integration test suite for the Sapiens2 Image Processor locally to verify all shapes, visibility masking, and exceptions function as expected.
Command run: pytest tests/models/sapiens2/test_image_processing_sapiens2.py
Results: 22 passed, 7 skipped (All passing)

I confirm that this is not a pure code agent PR.

Before submitting

This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
Did you read the contributor guideline and the
Pull Request checks?
Was this discussed/approved via a Github issue or the forum? Please add a link
to it if that's the case.
Did you make sure to update the documentation with your changes according to the guidelines?
Did you write any new necessary tests?

Who can review?

Hi @guarin , following up on our chat in #46764! Here is the initial implementation for the Sapiens2 preprocessing pipeline we discussed.

…eration

guarin

Hey thanks a lot for the PR! This looks already pretty good, all the important pieces are here. Left some comments regarding naming, structure, and tests :)

guarin · 2026-07-10T10:29:53Z

+                    center = centers[person_idx]
+                    scale = scales[person_idx]
+
+                    hm_coords = (raw_coords - center + 0.5 * scale) / scale * heatmap_size


Suggested change

hm_coords = (raw_coords - center + 0.5 * scale) / scale * heatmap_size

hm_coords = ((raw_coords - center) / scale + 0.5) * heatmap_size

guarin · 2026-07-10T10:30:40Z

+                    person_kps = image_keypoints[person_idx]
+


Suggested change

person_kps = image_keypoints[person_idx]

person_keypoints = image_keypoints[person_idx]

guarin · 2026-07-10T10:31:15Z

+                        weights_list.append(torch.zeros((0, heatmap_h, heatmap_w), dtype=torch.float32))
+                        continue
+
+                    person_kps_tensor = torch.tensor(person_kps, dtype=torch.float32)


We prefer to not use abbreviations

Suggested change

person_kps_tensor = torch.tensor(person_kps, dtype=torch.float32)

person_keypoints_tensor = torch.tensor(person_kps, dtype=torch.float32)

Same for other variables below

guarin · 2026-07-10T10:31:27Z

+                    center = centers[person_idx]
+                    scale = scales[person_idx]
+
+                    hm_coords = (raw_coords - center + 0.5 * scale) / scale * heatmap_size


Suggested change

hm_coords = (raw_coords - center + 0.5 * scale) / scale * heatmap_size

heatmap_coords = (raw_coords - center + 0.5 * scale) / scale * heatmap_size

guarin · 2026-07-10T10:49:13Z

+            heatmap_h = self.size["height"] // 4
+            heatmap_w = self.size["width"] // 4


Let's pack the whole heatmap generation logic into a generate_upd_gaussian_heatmaps function. Could you also add some tests to test_processing_sapiens2.py that assert that the generated outputs match outputs from the original implementation in https://github.com/facebookresearch/sapiens2/blob/7e5bae88456ac418ff0e58e74106c9fe192055d4/sapiens/pose/src/datasets/codecs/utils/gaussian_heatmap.py#L155?

Something that seems suspicious is that the original uses a np.maximum that is missing from our implementation (might not be necessary): https://github.com/facebookresearch/sapiens2/blob/7e5bae88456ac418ff0e58e74106c9fe192055d4/sapiens/pose/src/datasets/codecs/utils/gaussian_heatmap.py#L229

HuggingFaceDocBuilderDev · 2026-07-10T12:24:30Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

github-actions · 2026-07-11T20:50:29Z

[For maintainers] Suggested jobs to run (before merge)

run-slow: sapiens2

github-actions · 2026-07-11T20:54:37Z

CI recap

Dashboard: View test results in Grafana
Latest run: 29166551086:2
Result: success | Jobs: 5 | Tests: 317 | Failures: 7 | Duration: 4m 8s

Sainava · 2026-07-11T21:03:30Z

Hey @guarin,

I've implemented the changes.

Regarding the np.maximum : The original codebase uses a single global (K, H, W) tensor for all people, which requires np.maximum to prevent overlapping keypoints from summing past 1.0. Because our implementation creates isolated (num_kps, H, W) tensors per individual person crop, cross-person overlap is not possible, making np.maximum unnecessary here I think.

Let me know if anything else is needed :)

guarin

Hey thanks for the updates! Sorry for the delay, was off for a while. I left some comments, my biggest concern is that the current implementation might not match the original one from sapiens2.

guarin · 2026-07-27T08:46:35Z

+        keypoint_heatmap_downscale_factor: int = 4,
+        keypoint_heatmap_sigma: float = 6.0,


Suggested change

keypoint_heatmap_downscale_factor: int = 4,

keypoint_heatmap_sigma: float = 6.0,

Not needed as they are already in **kwargs. Please add the default values above like this:

class Sapiens2ImageProcessor(BeitImageProcessor): valid_kwargs = Sapiens2ImageProcessorKwargs ... keypoint_heatmap_downscale_factor = 4 keypoint_heatmap_sigma = 6.0

guarin · 2026-07-27T08:47:46Z

+        kwargs["keypoint_heatmap_downscale_factor"] = keypoint_heatmap_downscale_factor
+        kwargs["keypoint_heatmap_sigma"] = keypoint_heatmap_sigma


Suggested change

kwargs["keypoint_heatmap_downscale_factor"] = keypoint_heatmap_downscale_factor

kwargs["keypoint_heatmap_sigma"] = keypoint_heatmap_sigma

guarin · 2026-07-27T08:49:57Z

+        input_data_format: ChannelDimension | None,
        return_tensors: str | TensorType | None,


Are you sure input_data_format can be None here? It is usually only ChannelDimension in the codebase.

guarin · 2026-07-27T08:57:18Z

+        data["pixel_values"] = self._preprocess(images, **images_kwargs)
+


Suggested change

data["pixel_values"] = self._preprocess(images, **images_kwargs)

data["pixel_values"] = self._preprocess(images, **images_kwargs, boxes=boxes)

guarin · 2026-07-27T08:57:30Z

+            processed_segmentation_maps = self._preprocess(
+                images=processed_segmentation_maps, **segmentation_maps_kwargs
+            )
+


Suggested change

processed_segmentation_maps = self._preprocess(

images=processed_segmentation_maps, **segmentation_maps_kwargs

)

processed_segmentation_maps = self._preprocess(

images=processed_segmentation_maps, **segmentation_maps_kwargs, boxes=boxes

)

guarin · 2026-07-27T09:05:11Z

+            dist_sq = (grid_x.unsqueeze(0) - xs) ** 2 + (grid_y.unsqueeze(0) - ys) ** 2
+            person_heatmaps = torch.exp(-dist_sq / (2 * (sigma**2)))


Suggested change

dist_sq = (grid_x.unsqueeze(0) - xs) ** 2 + (grid_y.unsqueeze(0) - ys) ** 2

person_heatmaps = torch.exp(-dist_sq / (2 * (sigma**2)))

distance_squared = (grid_x.unsqueeze(0) - xs) ** 2 + (grid_y.unsqueeze(0) - ys) ** 2

person_heatmaps = torch.exp(-distance_squared / (2 * (sigma**2)))

guarin · 2026-07-27T09:16:57Z

+    """Generates UDP Gaussian heatmaps and visibility weights from raw keypoint coordinates."""
+    heatmap_height = output_size[0] // downscale_factor


Please document the shape and content of the returned lists

guarin · 2026-07-27T09:39:22Z

+
+            valid_mask = mask * (~out_of_bounds).float()
+            person_heatmaps = person_heatmaps * valid_mask
+            person_weights = valid_mask.expand_as(person_heatmaps)


Hmm I think this is missing the 3 * sigma bounds from the original code. All weights that are outside of a 3 * sigma radius from the keypoint center must be set to 0: https://github.com/facebookresearch/sapiens2/blob/7e5bae88456ac418ff0e58e74106c9fe192055d4/sapiens/pose/src/datasets/codecs/utils/gaussian_heatmap.py#L204-L206

guarin · 2026-07-27T09:48:29Z

+                weights_list.append(
+                    torch.zeros((0, heatmap_height, heatmap_width), dtype=torch.float32, device=device)
+                )


Is this shape correct? Shouldn't the weights have one entry per keypoint instead of one entry for each heatmap pixel?

guarin · 2026-07-27T09:50:10Z

+    def test_generate_udp_gaussian_heatmaps_parity(self):
+        import numpy as np
+        import torch
+
+        from transformers.models.sapiens2.image_processing_sapiens2 import (
+            box_xywh_to_cxcywh,
+            boxes_to_crop_params,
+            generate_udp_gaussian_heatmaps,
+        )


I think the current test is not ideal as it basically re-implements the logic from the tested function. Could you instead copy the full original function from sapiens2 and write a test that checks that the output from our function matches the output of the original function?

Sainava added 2 commits July 9, 2026 12:50

Add pose estimation keypoint preprocessing via vectorized heatmap gen…

05bdebb

…eration

Fix ruff trailing whitespace errors in sapiens2 tests

e9f9761

guarin reviewed Jul 10, 2026

View reviewed changes

Sainava and others added 3 commits July 11, 2026 14:35

Merge branch 'main' into feat/sapiens2-pose-pipeline

a88d9ee

Refactor Sapiens2 pose preprocessing and add parity tests

a54ab41

Refactor Sapiens2 pose preprocessing and add parity tests

d0eb7d5

guarin reviewed Jul 27, 2026

View reviewed changes

	hm_coords = (raw_coords - center + 0.5 * scale) / scale * heatmap_size
	hm_coords = ((raw_coords - center) / scale + 0.5) * heatmap_size

	person_kps = image_keypoints[person_idx]
	person_keypoints = image_keypoints[person_idx]

	person_kps_tensor = torch.tensor(person_kps, dtype=torch.float32)
	person_keypoints_tensor = torch.tensor(person_kps, dtype=torch.float32)

	hm_coords = (raw_coords - center + 0.5 * scale) / scale * heatmap_size
	heatmap_coords = (raw_coords - center + 0.5 * scale) / scale * heatmap_size

		heatmap_h = self.size["height"] // 4
		heatmap_w = self.size["width"] // 4

		keypoint_heatmap_downscale_factor: int = 4,
		keypoint_heatmap_sigma: float = 6.0,

		kwargs["keypoint_heatmap_downscale_factor"] = keypoint_heatmap_downscale_factor
		kwargs["keypoint_heatmap_sigma"] = keypoint_heatmap_sigma

		input_data_format: ChannelDimension \| None,
		return_tensors: str \| TensorType \| None,

		data["pixel_values"] = self._preprocess(images, **images_kwargs)

		dist_sq = (grid_x.unsqueeze(0) - xs) 2 + (grid_y.unsqueeze(0) - ys) 2
		person_heatmaps = torch.exp(-dist_sq / (2 * (sigma**2)))

		"""Generates UDP Gaussian heatmaps and visibility weights from raw keypoint coordinates."""
		heatmap_height = output_size[0] // downscale_factor

Uh oh!

Conversation

Sainava commented Jul 9, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Testing

Before submitting

Who can review?

Uh oh!

guarin left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

HuggingFaceDocBuilderDev commented Jul 10, 2026

Uh oh!

github-actions Bot commented Jul 11, 2026

Uh oh!

github-actions Bot commented Jul 11, 2026

CI recap

Uh oh!

Sainava commented Jul 11, 2026

Uh oh!

guarin left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Sainava commented Jul 9, 2026 •

edited by github-actions Bot

Loading