Add pose estimation keypoint preprocessing to Sapiens2ImageProcessor

@@ -380,7 +380,13 @@ def _validate_preprocess_kwargs(         )      @auto_docstring-    def preprocess(self, images: ImageInput, *args, **kwargs: Unpack[ImagesKwargs]) -> BatchFeature:+    def preprocess(+        self,+        images: ImageInput,+        *args,+        image_like_kwargs: dict[str, Any] | None = None,+        **kwargs: Unpack[ImagesKwargs],+    ) -> BatchFeature:         """         Preprocess an image or a batch of images.         """@@ -397,7 +403,9 @@ def preprocess(self, images: ImageInput, *args, **kwargs: Unpack[ImagesKwargs])         # Validate kwargs         self._validate_preprocess_kwargs(**kwargs) -        return self._preprocess_image_like_inputs(images, *args, **kwargs)+        image_like_kwargs = {} if image_like_kwargs is None else image_like_kwargs++        return self._preprocess_image_like_inputs(images, *args, **image_like_kwargs, **kwargs)      def to_dict(self) -> dict[str, Any]:         processor_dict = super().to_dict()

@@ -22,8 +22,7 @@ import torch.nn.functional as F from torchvision.transforms.v2 import functional as tvF -from transformers.image_processing_backends import TorchvisionBackend-+from ...image_processing_backends import TorchvisionBackend from ...image_processing_outputs import SemanticSegmentationPostProcessorOutput from ...image_processing_utils import BatchFeature from ...image_transforms import group_images_by_shape, reorder_images@@ -42,15 +41,22 @@   class Sapiens2ImageProcessorKwargs(ImagesKwargs, total=False):-    r"""+    """     do_reduce_labels (`bool`, *optional*, defaults to `self.do_reduce_labels`):         Whether or not to reduce all label values of segmentation maps by 1. Usually used for datasets where 0         is used for background, and background itself is not included in all classes of a dataset (e.g.         ADE20k). The background label will be replaced by 255.+    keypoint_heatmap_downscale_factor (`int`, *optional*, defaults to 4):+        The downscale factor for the target heatmap size relative to the model input size.+    keypoint_heatmap_sigma (`float`, *optional*, defaults to 6.0):+        The standard deviation (sigma) for the 2D Gaussian distributions used to generate the heatmaps.     """      do_reduce_labels: bool +    keypoint_heatmap_downscale_factor: int+    keypoint_heatmap_sigma: float+  def box_xywh_to_xyxy(x):     x, y, w, h = x.unbind(-1)@@ -292,6 +298,111 @@ def post_dark_unbiased_data_processing(     return keypoints - torch.cat([offset_x, offset_y], dim=-1)  +def generate_udp_gaussian_heatmaps(+    boxes: list[list[list[float]]],+    keypoints: list[list[list[list[float]]]],+    output_size: tuple[int, int],+    downscale_factor: int,+    sigma: float,+    device: Union[str, "torch.device"] | None = None,+) -> tuple[list[torch.Tensor], list[torch.Tensor]]:+    """Generates UDP Gaussian heatmaps and visibility weights from raw keypoint coordinates.++    Args:+        boxes (`list[list[list[float]]]`):+            List of bounding boxes for each image in COCO format `(top_left_x, top_left_y, width, height)`.+        keypoints (`list[list[list[list[float]]]]`):+            List of keypoints for each person in each image. Expected format is COCO-style `[x, y, visibility]`.+        output_size (`tuple[int, int]`):+            The target size `(height, width)` of the cropped images.+        downscale_factor (`int`, *optional*, defaults to 4):+            The downscale factor for the target heatmap size relative to the output size.+        sigma (`float`, *optional*, defaults to 6.0):+            The standard deviation (sigma) for the 2D Gaussian distributions.+        device (`str` or `torch.device`, *optional*):+            The device to put the resulting tensors on.++    Returns:+        tuple:+        - heatmaps_list (list[torch.Tensor]): The generated heatmaps. Each tensor has shape+          `(num_persons, num_keypoints, heatmap_height, heatmap_width)`.+        - weights_list (list[torch.Tensor]): The target weights. Each tensor has shape+          `(num_persons, num_keypoints)`.+    """+    heatmap_height = output_size[0] // downscale_factor+    heatmap_width = output_size[1] // downscale_factor++    heatmaps_list = []+    weights_list = []++    grid_y, grid_x = torch.meshgrid(+        torch.arange(heatmap_height, dtype=torch.float32, device=device),+        torch.arange(heatmap_width, dtype=torch.float32, device=device),+        indexing="ij",+    )+    heatmap_size = torch.tensor([heatmap_width - 1, heatmap_height - 1], dtype=torch.float32, device=device)+    radius = sigma * 3++    for image_boxes, image_keypoints in zip(boxes, keypoints):+        boxes_tensor = box_xywh_to_cxcywh(torch.tensor(image_boxes, dtype=torch.float32, device=device))++        centers, scales = boxes_to_crop_params(boxes_tensor, output_size=output_size)++        for person_idx in range(len(image_boxes)):+            person_keypoints = image_keypoints[person_idx]++            if len(person_keypoints) == 0:+                heatmaps_list.append(+                    torch.zeros((0, heatmap_height, heatmap_width), dtype=torch.float32, device=device)+                )+                weights_list.append(torch.zeros((0,), dtype=torch.float32, device=device))+                continue++            person_keypoints_tensor = torch.tensor(person_keypoints, dtype=torch.float32, device=device)+            raw_coords = person_keypoints_tensor[:, :2]++            center = centers[person_idx]+            scale = scales[person_idx]++            heatmap_coords = ((raw_coords - center) / scale + 0.5) * heatmap_size++            x_coords = heatmap_coords[:, 0].view(-1, 1, 1)+            y_coords = heatmap_coords[:, 1].view(-1, 1, 1)++            distance_squared = (grid_x.unsqueeze(0) - x_coords) ** 2 + (grid_y.unsqueeze(0) - y_coords) ** 2+            person_heatmaps = torch.exp(-distance_squared / (2 * (sigma**2)))++            if person_keypoints_tensor.shape[1] > 2:+                visibilities = person_keypoints_tensor[:, 2]+                mask = (visibilities > 0).float()+            else:+                mask = torch.ones(person_keypoints_tensor.shape[0], dtype=torch.float32, device=device)++            # 3-sigma bounds check matching original implementation+            mu = (heatmap_coords + 0.5).to(torch.int64)+            left = mu[:, 0] - int(radius)+            top = mu[:, 1] - int(radius)+            right = mu[:, 0] + int(radius) + 1+            bottom = mu[:, 1] + int(radius) + 1++            out_of_bounds = (left >= heatmap_width) | (top >= heatmap_height) | (right < 0) | (bottom < 0)+            valid_mask = mask * (~out_of_bounds).float()++            spatial_mask = (+                (grid_x.unsqueeze(0) >= left.view(-1, 1, 1))+                & (grid_x.unsqueeze(0) < right.view(-1, 1, 1))+                & (grid_y.unsqueeze(0) >= top.view(-1, 1, 1))+                & (grid_y.unsqueeze(0) < bottom.view(-1, 1, 1))+            )++            person_heatmaps = person_heatmaps * valid_mask.view(-1, 1, 1) * spatial_mask.float()++            heatmaps_list.append(person_heatmaps)+            weights_list.append(valid_mask)++    return heatmaps_list, weights_list++ @auto_docstring class Sapiens2ImageProcessor(TorchvisionBackend):     """PIL backend for Sapiens2 with reduce_label support."""@@ -310,6 +421,8 @@ class Sapiens2ImageProcessor(TorchvisionBackend):     do_normalize = True     do_reduce_labels = False     do_pad = False  # Set to True for normal, albedo, and pointmap estimation+    keypoint_heatmap_downscale_factor = 4+    keypoint_heatmap_sigma = 6.0      def __init__(self, **kwargs: Unpack[Sapiens2ImageProcessorKwargs]):         super().__init__(**kwargs)@@ -320,6 +433,7 @@ def preprocess(         images: ImageInput,         segmentation_maps: ImageInput | None = None,         boxes: list[list[list[float]]] | None = None,+        keypoints: list[list[list[list[float]]]] | None = None,         **kwargs: Unpack[Sapiens2ImageProcessorKwargs],     ) -> BatchFeature:         r"""@@ -330,8 +444,17 @@ def preprocess(             representing the bounding box coordinates in COCO format             (top_left_x, top_left_y, width, height). When provided, each person crop is             affine-warped to the model input size instead of resizing the full image.+        keypoints (`list[list[list[list[float]]]]`, *optional*):+            List of keypoints for each person in each image. Expected format is COCO-style `[x, y, visibility]`.+            The `x` and `y` values are expected to be absolute image pixel coordinates.+            Format is `[images -> persons -> keypoints -> [x, y, visibility]]`. Used to generate+            ground-truth heatmaps and visibility weights for pose estimation fine-tuning.         """-        return super().preprocess(images, segmentation_maps, boxes, **kwargs)+        return super().preprocess(+            images,+            image_like_kwargs={"segmentation_maps": segmentation_maps, "boxes": boxes, "keypoints": keypoints},+            **kwargs,+        )      def _preprocess_image_like_inputs(         self,@@ -341,18 +464,26 @@ def _preprocess_image_like_inputs(         do_convert_rgb: bool,         input_data_format: ChannelDimension,         return_tensors: str | TensorType | None,-        device: Union[str, "torch.device"] | None = None,+        device: Union[str, "torch.device"] | None,+        keypoints: list[list[list[list[float]]]] | None,+        keypoint_heatmap_downscale_factor: int | None = None,+        keypoint_heatmap_sigma: float | None = None,         **kwargs,     ) -> BatchFeature:         """Handle extra inputs beyond images."""-        kwargs["boxes"] = boxes  # modular trick+        if segmentation_maps is not None and keypoints is not None:+            raise ValueError(+                "Cannot process both `segmentation_maps` and `keypoints` in the same forward pass. "+                "Please provide only one depending on the task you want to perform."+            )+         images = self._prepare_image_like_inputs(             images=images, do_convert_rgb=do_convert_rgb, input_data_format=input_data_format, device=device         )         images_kwargs = kwargs.copy()         images_kwargs["do_reduce_labels"] = False         data = {}-        data["pixel_values"] = self._preprocess(images, **images_kwargs)+        data["pixel_values"] = self._preprocess(images, **images_kwargs, boxes=boxes)          # Prepare segmentation maps if provided         if segmentation_maps is not None:@@ -367,7 +498,7 @@ def _preprocess_image_like_inputs(             segmentation_maps_kwargs = kwargs.copy()             segmentation_maps_kwargs.update({"do_normalize": False, "do_rescale": False})             processed_segmentation_maps = self._preprocess(-                images=processed_segmentation_maps, **segmentation_maps_kwargs+                images=processed_segmentation_maps, **segmentation_maps_kwargs, boxes=boxes             )              # Convert to int64 and squeeze channel dimension@@ -377,6 +508,34 @@ def _preprocess_image_like_inputs(             ]             data["labels"] = processed_segmentation_maps +        # Prepare pose estimation keypoints if provided+        if keypoints is not None:+            if boxes is None:+                raise ValueError("Bounding `boxes` must be provided when passing `keypoints` for pose estimation.")+            if keypoint_heatmap_downscale_factor is None:+                raise ValueError(+                    "`keypoint_heatmap_downscale_factor` must be provided when passing `keypoints` for pose estimation."+                )+            if keypoint_heatmap_sigma is None:+                raise ValueError(+                    "`keypoint_heatmap_sigma` must be provided when passing `keypoints` for pose estimation."+                )++            # Extract dynamic size override if it exists, otherwise fall back to default+            target_size = kwargs.get("size", self.size)++            heatmaps_list, weights_list = generate_udp_gaussian_heatmaps(+                boxes=boxes,+                keypoints=keypoints,+                output_size=(target_size["height"], target_size["width"]),+                downscale_factor=keypoint_heatmap_downscale_factor,+                sigma=keypoint_heatmap_sigma,+                device=device,+            )++            data["labels"] = heatmaps_list+            data["label_weights"] = weights_list+         return BatchFeature(data=data, tensor_type=return_tensors)      def reduce_label(self, labels: list["torch.Tensor"]) -> list["torch.Tensor"]:

@@ -20,12 +20,10 @@ from torch import nn from torchvision.transforms.v2 import functional as tvF -from transformers.image_processing_backends import TorchvisionBackend-from transformers.models.dinov3_vit.modeling_dinov3_vit import DINOv3ViTBackboneOutput- from ... import initialization as init from ...activations import ACT2FN from ...configuration_utils import PreTrainedConfig+from ...image_processing_backends import TorchvisionBackend from ...image_processing_utils import BatchFeature from ...image_transforms import group_images_by_shape, reorder_images from ...image_utils import (@@ -48,6 +46,7 @@ from ..dinov3_vit.modeling_dinov3_vit import (     DINOv3ViTAttention,     DINOv3ViTBackbone,+    DINOv3ViTBackboneOutput,     DINOv3ViTEmbeddings,     DINOv3ViTEncoder,     DINOv3ViTLayer,@@ -407,8 +406,125 @@ def post_dark_unbiased_data_processing(     return keypoints - torch.cat([offset_x, offset_y], dim=-1)  +def generate_udp_gaussian_heatmaps(+    boxes: list[list[list[float]]],+    keypoints: list[list[list[list[float]]]],+    output_size: tuple[int, int],+    downscale_factor: int,+    sigma: float,+    device: Union[str, "torch.device"] | None = None,+) -> tuple[list[torch.Tensor], list[torch.Tensor]]:+    """Generates UDP Gaussian heatmaps and visibility weights from raw keypoint coordinates.++    Args:+        boxes (`list[list[list[float]]]`):+            List of bounding boxes for each image in COCO format `(top_left_x, top_left_y, width, height)`.+        keypoints (`list[list[list[list[float]]]]`):+            List of keypoints for each person in each image. Expected format is COCO-style `[x, y, visibility]`.+        output_size (`tuple[int, int]`):+            The target size `(height, width)` of the cropped images.+        downscale_factor (`int`, *optional*, defaults to 4):+            The downscale factor for the target heatmap size relative to the output size.+        sigma (`float`, *optional*, defaults to 6.0):+            The standard deviation (sigma) for the 2D Gaussian distributions.+        device (`str` or `torch.device`, *optional*):+            The device to put the resulting tensors on.++    Returns:+        tuple:+        - heatmaps_list (list[torch.Tensor]): The generated heatmaps. Each tensor has shape+          `(num_persons, num_keypoints, heatmap_height, heatmap_width)`.+        - weights_list (list[torch.Tensor]): The target weights. Each tensor has shape+          `(num_persons, num_keypoints)`.+    """+    heatmap_height = output_size[0] // downscale_factor+    heatmap_width = output_size[1] // downscale_factor++    heatmaps_list = []+    weights_list = []++    grid_y, grid_x = torch.meshgrid(+        torch.arange(heatmap_height, dtype=torch.float32, device=device),+        torch.arange(heatmap_width, dtype=torch.float32, device=device),+        indexing="ij",+    )+    heatmap_size = torch.tensor([heatmap_width - 1, heatmap_height - 1], dtype=torch.float32, device=device)+    radius = sigma * 3++    for image_boxes, image_keypoints in zip(boxes, keypoints):+        boxes_tensor = box_xywh_to_cxcywh(torch.tensor(image_boxes, dtype=torch.float32, device=device))++        centers, scales = boxes_to_crop_params(boxes_tensor, output_size=output_size)++        for person_idx in range(len(image_boxes)):+            person_keypoints = image_keypoints[person_idx]++            if len(person_keypoints) == 0:+                heatmaps_list.append(+                    torch.zeros((0, heatmap_height, heatmap_width), dtype=torch.float32, device=device)+                )+                weights_list.append(torch.zeros((0,), dtype=torch.float32, device=device))+                continue++            person_keypoints_tensor = torch.tensor(person_keypoints, dtype=torch.float32, device=device)+            raw_coords = person_keypoints_tensor[:, :2]++            center = centers[person_idx]+            scale = scales[person_idx]++            heatmap_coords = ((raw_coords - center) / scale + 0.5) * heatmap_size++            x_coords = heatmap_coords[:, 0].view(-1, 1, 1)+            y_coords = heatmap_coords[:, 1].view(-1, 1, 1)++            distance_squared = (grid_x.unsqueeze(0) - x_coords) ** 2 + (grid_y.unsqueeze(0) - y_coords) ** 2+            person_heatmaps = torch.exp(-distance_squared / (2 * (sigma**2)))++            if person_keypoints_tensor.shape[1] > 2:+                visibilities = person_keypoints_tensor[:, 2]+                mask = (visibilities > 0).float()+            else:+                mask = torch.ones(person_keypoints_tensor.shape[0], dtype=torch.float32, device=device)++            # 3-sigma bounds check matching original implementation+            mu = (heatmap_coords + 0.5).to(torch.int64)+            left = mu[:, 0] - int(radius)+            top = mu[:, 1] - int(radius)+            right = mu[:, 0] + int(radius) + 1+            bottom = mu[:, 1] + int(radius) + 1++            out_of_bounds = (left >= heatmap_width) | (top >= heatmap_height) | (right < 0) | (bottom < 0)+            valid_mask = mask * (~out_of_bounds).float()++            spatial_mask = (+                (grid_x.unsqueeze(0) >= left.view(-1, 1, 1))+                & (grid_x.unsqueeze(0) < right.view(-1, 1, 1))+                & (grid_y.unsqueeze(0) >= top.view(-1, 1, 1))+                & (grid_y.unsqueeze(0) < bottom.view(-1, 1, 1))+            )++            person_heatmaps = person_heatmaps * valid_mask.view(-1, 1, 1) * spatial_mask.float()++            heatmaps_list.append(person_heatmaps)+            weights_list.append(valid_mask)++    return heatmaps_list, weights_list++ class Sapiens2ImageProcessorKwargs(BeitImageProcessorKwargs, total=False):-    pass+    """+    do_reduce_labels (`bool`, *optional*, defaults to `self.do_reduce_labels`):+        Whether or not to reduce all label values of segmentation maps by 1. Usually used for datasets where 0+        is used for background, and background itself is not included in all classes of a dataset (e.g.+        ADE20k). The background label will be replaced by 255.+    keypoint_heatmap_downscale_factor (`int`, *optional*, defaults to 4):+        The downscale factor for the target heatmap size relative to the model input size.+    keypoint_heatmap_sigma (`float`, *optional*, defaults to 6.0):+        The standard deviation (sigma) for the 2D Gaussian distributions used to generate the heatmaps.+    """++    keypoint_heatmap_downscale_factor: int+    keypoint_heatmap_sigma: float   class Sapiens2ImageProcessor(BeitImageProcessor):@@ -419,6 +535,8 @@ class Sapiens2ImageProcessor(BeitImageProcessor):     image_std = IMAGENET_DEFAULT_STD     size = {"height": 1024, "width": 768}     do_pad = False  # Set to True for normal, albedo, and pointmap estimation+    keypoint_heatmap_downscale_factor = 4+    keypoint_heatmap_sigma = 6.0      def __init__(self, **kwargs: Unpack[Sapiens2ImageProcessorKwargs]):         super().__init__(**kwargs)@@ -429,6 +547,7 @@ def preprocess(         images: ImageInput,         segmentation_maps: ImageInput | None = None,         boxes: list[list[list[float]]] | None = None,+        keypoints: list[list[list[list[float]]]] | None = None,         **kwargs: Unpack[Sapiens2ImageProcessorKwargs],     ) -> BatchFeature:         r"""@@ -439,8 +558,18 @@ def preprocess(             representing the bounding box coordinates in COCO format             (top_left_x, top_left_y, width, height). When provided, each person crop is             affine-warped to the model input size instead of resizing the full image.+        keypoints (`list[list[list[list[float]]]]`, *optional*):+            List of keypoints for each person in each image. Expected format is COCO-style `[x, y, visibility]`.+            The `x` and `y` values are expected to be absolute image pixel coordinates.+            Format is `[images -> persons -> keypoints -> [x, y, visibility]]`. Used to generate+            ground-truth heatmaps and visibility weights for pose estimation fine-tuning.         """-        return TorchvisionBackend.preprocess(images, segmentation_maps, boxes, **kwargs)+        return TorchvisionBackend.preprocess(+            self,+            images,+            image_like_kwargs={"segmentation_maps": segmentation_maps, "boxes": boxes, "keypoints": keypoints},+            **kwargs,+        )      def _preprocess_image_like_inputs(         self,@@ -450,21 +579,79 @@ def _preprocess_image_like_inputs(         do_convert_rgb: bool,         input_data_format: ChannelDimension,         return_tensors: str | TensorType | None,-        device: Union[str, "torch.device"] | None = None,+        device: Union[str, "torch.device"] | None,+        keypoints: list[list[list[list[float]]]] | None,+        keypoint_heatmap_downscale_factor: int | None = None,+        keypoint_heatmap_sigma: float | None = None,         **kwargs,     ) -> BatchFeature:         """Handle extra inputs beyond images."""-        kwargs["boxes"] = boxes  # modular trick-        return super()._preprocess_image_like_inputs(-            self,-            images=images,-            segmentation_maps=segmentation_maps,-            do_convert_rgb=do_convert_rgb,-            input_data_format=input_data_format,-            return_tensors=return_tensors,-            device=device,-            **kwargs,+        if segmentation_maps is not None and keypoints is not None:+            raise ValueError(+                "Cannot process both `segmentation_maps` and `keypoints` in the same forward pass. "+                "Please provide only one depending on the task you want to perform."+            )++        images = self._prepare_image_like_inputs(+            images=images, do_convert_rgb=do_convert_rgb, input_data_format=input_data_format, device=device         )+        images_kwargs = kwargs.copy()+        images_kwargs["do_reduce_labels"] = False+        data = {}+        data["pixel_values"] = self._preprocess(images, **images_kwargs, boxes=boxes)++        # Prepare segmentation maps if provided+        if segmentation_maps is not None:+            processed_segmentation_maps = self._prepare_image_like_inputs(+                images=segmentation_maps,+                expected_ndims=2,+                do_convert_rgb=False,+                input_data_format=ChannelDimension.FIRST,+            )++            # Process segmentation maps with do_normalize=False and do_rescale=False+            segmentation_maps_kwargs = kwargs.copy()+            segmentation_maps_kwargs.update({"do_normalize": False, "do_rescale": False})+            processed_segmentation_maps = self._preprocess(+                images=processed_segmentation_maps, **segmentation_maps_kwargs, boxes=boxes+            )++            # Convert to int64 and squeeze channel dimension+            processed_segmentation_maps = [+                processed_segmentation_map.squeeze(0).to(torch.int64)+                for processed_segmentation_map in processed_segmentation_maps+            ]+            data["labels"] = processed_segmentation_maps++        # Prepare pose estimation keypoints if provided+        if keypoints is not None:+            if boxes is None:+                raise ValueError("Bounding `boxes` must be provided when passing `keypoints` for pose estimation.")+            if keypoint_heatmap_downscale_factor is None:+                raise ValueError(+                    "`keypoint_heatmap_downscale_factor` must be provided when passing `keypoints` for pose estimation."+                )+            if keypoint_heatmap_sigma is None:+                raise ValueError(+                    "`keypoint_heatmap_sigma` must be provided when passing `keypoints` for pose estimation."+                )++            # Extract dynamic size override if it exists, otherwise fall back to default+            target_size = kwargs.get("size", self.size)++            heatmaps_list, weights_list = generate_udp_gaussian_heatmaps(+                boxes=boxes,+                keypoints=keypoints,+                output_size=(target_size["height"], target_size["width"]),+                downscale_factor=keypoint_heatmap_downscale_factor,+                sigma=keypoint_heatmap_sigma,+                device=device,+            )++            data["labels"] = heatmaps_list+            data["label_weights"] = weights_list++        return BatchFeature(data=data, tensor_type=return_tensors)      def _preprocess(         self,

@@ -13,6 +13,9 @@ # limitations under the License.  import unittest+from itertools import product++import numpy as np  from transformers.testing_utils import require_torch, require_vision from transformers.utils import is_torch_available@@ -28,6 +31,11 @@     import torch      from transformers import Sapiens2ImageProcessor+    from transformers.models.sapiens2.image_processing_sapiens2 import (+        box_xywh_to_cxcywh,+        boxes_to_crop_params,+        generate_udp_gaussian_heatmaps,+    )     from transformers.models.sapiens2.modeling_sapiens2 import (         Sapiens2ImageMattingOutput,         Sapiens2NormalEstimatorOutput,@@ -243,3 +251,112 @@ def test_post_process_image_matting(self):         # mismatched batch size raises ValueError         with self.assertRaises(ValueError):             image_processor.post_process_image_matting(outputs, target_sizes=[(100, 100)])++    def test_pose_estimation_keypoint_preprocessing(self):+        image_inputs = self.image_processor_tester.prepare_image_inputs(equal_resolution=False, numpify=True)+        image = image_inputs[0]++        for image_processing_class in self.image_processing_classes.values():+            image_processor = image_processing_class()++            boxes = [[[50.0, 50.0, 200.0, 400.0]]]+            keypoints = [[[[60.0, 70.0, 1.0], [80.0, 90.0, 0.0]]]]++            inputs = image_processor(images=image, boxes=boxes, keypoints=keypoints, return_tensors="pt")++            self.assertEqual(inputs["pixel_values"].shape, (1, 3, 1024, 768))+            self.assertEqual(inputs["labels"].shape, (1, 2, 256, 192))+            self.assertEqual(inputs["label_weights"].shape, (1, 2))++            self.assertEqual(inputs["label_weights"][0, 1].max().item(), 0.0)++            with self.assertRaises(ValueError):+                image_processor(images=image, keypoints=keypoints, return_tensors="pt")++    def test_generate_udp_gaussian_heatmaps_parity(self):+        # 1. Original Meta Implementation for exact parity testing+        def original_generate_udp_gaussian_heatmaps(heatmap_size, keypoints, keypoints_visible, sigma):+            N, K, _ = keypoints.shape+            W, H = heatmap_size+            heatmaps = np.zeros((K, H, W), dtype=np.float32)+            keypoint_weights = keypoints_visible.copy()+            radius = sigma * 3+            gaussian_size = 2 * radius + 1+            x = np.arange(0, gaussian_size, 1, dtype=np.float32)+            y = x[:, None]++            for n, k in product(range(N), range(K)):+                if keypoints_visible[n, k] < 0.5:+                    continue+                mu = (keypoints[n, k] + 0.5).astype(np.int64)+                left, top = (mu - radius).astype(np.int64)+                right, bottom = (mu + radius + 1).astype(np.int64)++                if left >= W or top >= H or right < 0 or bottom < 0:+                    keypoint_weights[n, k] = 0+                    continue++                mu_ac = keypoints[n, k]+                x0 = y0 = gaussian_size // 2+                x0 += mu_ac[0] - mu[0]+                y0 += mu_ac[1] - mu[1]+                gaussian = np.exp(-((x - x0) ** 2 + (y - y0) ** 2) / (2 * sigma**2))++                g_x1, g_x2 = max(0, -left), min(W, right) - left+                g_y1, g_y2 = max(0, -top), min(H, bottom) - top+                h_x1, h_x2 = max(0, left), min(W, right)+                h_y1, h_y2 = max(0, top), min(H, bottom)++                heatmap_region = heatmaps[k, h_y1:h_y2, h_x1:h_x2]+                gaussian_regsion = gaussian[g_y1:g_y2, g_x1:g_x2]+                _ = np.maximum(heatmap_region, gaussian_regsion, out=heatmap_region)++            return heatmaps, keypoint_weights++        # 2. Setup dummy inputs+        output_size = (1024, 768)+        downscale_factor = 4+        sigma = 6.0++        heatmap_height = output_size[0] // downscale_factor+        heatmap_width = output_size[1] // downscale_factor+        heatmap_size_array = np.array([heatmap_width - 1, heatmap_height - 1], dtype=np.float32)++        # 1 box, 2 keypoints (one in bounds, one completely out of bounds to test the mask)+        boxes = [[[50.0, 50.0, 200.0, 400.0]]]+        keypoints = [[[[100.0, 150.0, 1.0], [3000.0, 4000.0, 1.0]]]]++        # 3. Run our PyTorch implementation+        pt_heatmaps, pt_weights = generate_udp_gaussian_heatmaps(+            boxes=boxes,+            keypoints=keypoints,+            output_size=output_size,+            downscale_factor=downscale_factor,+            sigma=sigma,+            device="cpu",+        )+        pt_heatmaps = pt_heatmaps[0].numpy()+        pt_weights = pt_weights[0].numpy()++        # 4. Prepare inputs for the original NumPy function+        boxes_tensor = box_xywh_to_cxcywh(torch.tensor(boxes[0], dtype=torch.float32))+        centers, scales = boxes_to_crop_params(boxes_tensor, output_size=output_size)++        raw_coords = np.array(keypoints[0][0])[:, :2]+        visibilities = np.array(keypoints[0][0])[:, 2]++        center = centers[0].numpy()+        scale = scales[0].numpy()+        heatmap_coords = ((raw_coords - center) / scale + 0.5) * heatmap_size_array++        # 5. Run the original Meta NumPy implementation+        np_heatmaps, np_weights = original_generate_udp_gaussian_heatmaps(+            heatmap_size=(heatmap_width, heatmap_height),+            keypoints=np.expand_dims(heatmap_coords, axis=0),+            keypoints_visible=np.expand_dims(visibilities, axis=0),+            sigma=sigma,+        )++        # 6. Assert strict parity+        np.testing.assert_allclose(pt_heatmaps, np_heatmaps, atol=1e-5)+        np.testing.assert_allclose(pt_weights, np_weights[0], atol=1e-5)