#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
#           This file was automatically generated from src/transformers/models/sapiens2/modular_sapiens2.py.
#               Do NOT edit this file manually as any edits will be overwritten by the generation of
#             the file from the modular. If any change should be done, please apply the change to the
#                          modular_sapiens2.py file directly. One of our CI enforces this.
#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
# Copyright 2026 Meta Platforms, Inc. and the HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Sapiens2 License. You may obtain a copy of the License at
#
#     https://github.com/facebookresearch/sapiens2/blob/main/LICENSE.md
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Union

import torch
import torch.nn.functional as F
from torchvision.transforms.v2 import functional as tvF

from transformers.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 (
    IMAGENET_DEFAULT_MEAN,
    IMAGENET_DEFAULT_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    SizeDict,
    get_image_size_for_max_height_width,
    make_list_of_images,
)
from ...processing_utils import ImagesKwargs, Unpack
from ...utils import TensorType, auto_docstring, is_torch_available
from .modeling_sapiens2 import (
    Sapiens2ImageMattingOutput,
    Sapiens2NormalEstimatorOutput,
    Sapiens2PointmapEstimatorOutput,
    Sapiens2PoseEstimatorOutput,
)


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.
    """

    do_reduce_labels: bool


def box_xywh_to_xyxy(x):
    x, y, w, h = x.unbind(-1)
    b = [(x), (y), (x + w), (y + h)]
    return torch.stack(b, dim=-1)


def box_xywh_to_cxcywh(x):
    x, y, w, h = x.unbind(-1)
    b = [(x + 0.5 * w), (y + 0.5 * h), (w), (h)]
    return torch.stack(b, dim=-1)


def boxes_to_crop_params(
    boxes: torch.Tensor,
    output_size: tuple[int, int],
    padding: float = 1.25,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute crop center and scale from bounding boxes, applying padding and aspect ratio correction.

    Accepts either a single box `(4,)` or multiple boxes `(num_boxes, 4)` and returns center/scale with a matching
    leading dimension.

    Args:
        boxes (`torch.Tensor` of shape `(4,)` or `(num_boxes, 4)`): Bounding box in
            (center-x, center-y, width, height) format, with values in absolute pixel coordinates.
        output_size (`tuple[int, int]`): Target output size as `(height, width)`, used to compute
            the aspect ratio for scale correction.
        padding (`float`, *optional*, defaults to `1.25`): Multiplicative factor applied to the
            bounding box dimensions, adding context around the region of interest.

    Returns:
        `tuple[torch.Tensor, torch.Tensor]`: A pair `(center, scale)` where `center` has shape
        `(..., 2)` with (x, y) in input-image pixel coordinates, and `scale` has shape `(..., 2)`
        with (width, height) in input-image pixels representing the dimensions of the padded,
        aspect-ratio-corrected crop window.
    """
    center_x, center_y, width, height = boxes.unbind(-1)
    center = torch.stack([center_x, center_y], dim=-1)
    scaled_width = width * padding
    scaled_height = height * padding
    output_height, output_width = output_size
    aspect_ratio = output_width / output_height
    scale = torch.where(
        (scaled_width > scaled_height * aspect_ratio)[..., None],
        torch.stack([scaled_width, scaled_width / aspect_ratio], dim=-1),
        torch.stack([scaled_height * aspect_ratio, scaled_height], dim=-1),
    )
    return center, scale


def crop_and_resize(
    image: torch.Tensor,
    boxes: torch.Tensor,
    output_size: tuple[int, int],
    padding: float = 1.25,
) -> torch.Tensor:
    """Crops and resizes bounding box regions from the input image to the target output size.

    Applies padding and aspect ratio correction to each crop before resizing.
    Uses bilinear interpolation for downscaling and bicubic for upscaling.

    This implementation is equivalent to the cv2 affine warp with rotation=0 used in the original
    Sapiens2 codebase. Rotation is always zero because we don't support rotated bounding boxes.

    Args:
        image (`torch.Tensor`): Input image tensor of shape `(C, H, W)` in float32.
        boxes (`torch.Tensor`): Bounding boxes in (center-x, center-y, width, height) format,
            shape `(num_boxes, 4)`, with values in absolute pixel coordinates.
        output_size (`tuple[int, int]`): Target output size as `(height, width)`.
        padding (`float`, *optional*, defaults to `1.25`): Multiplicative factor applied to the
            bounding box dimensions before cropping, adding context around the region of interest.

    Returns:
        `torch.Tensor`: Cropped and resized images of shape `(num_boxes, C, output_height, output_width)`.
    """
    output_height, output_width = output_size
    num_channels, input_height, input_width = image.shape
    center, scale = boxes_to_crop_params(boxes, output_size=output_size, padding=padding)
    center_x, center_y = center.unbind(-1)
    boxes_width, boxes_height = scale.unbind(-1)

    scale_x = (output_width - 1) / boxes_width  # (num_boxes,)
    scale_y = (output_height - 1) / boxes_height  # (num_boxes,)
    is_bilinear = torch.minimum(scale_x, scale_y) < 1.0  # (num_boxes,)

    grid_y, grid_x = torch.meshgrid(
        torch.arange(output_height, dtype=torch.float32, device=image.device),
        torch.arange(output_width, dtype=torch.float32, device=image.device),
        indexing="ij",
    )
    in_x = grid_x / scale_x[:, None, None] + center_x[:, None, None] - 0.5 * boxes_width[:, None, None]
    in_y = grid_y / scale_y[:, None, None] + center_y[:, None, None] - 0.5 * boxes_height[:, None, None]
    # (num_boxes, output_height, output_width, 2)
    grids = torch.stack([2.0 * in_x / (input_width - 1) - 1.0, 2.0 * in_y / (input_height - 1) - 1.0], dim=-1)

    num_boxes = boxes.shape[0]
    output = torch.empty(num_boxes, num_channels, output_height, output_width, device=image.device, dtype=image.dtype)

    # Apply grid sampling separately for upscaling and downscaling to use the appropriate interpolation mode
    image_4d = image.unsqueeze(0)
    for mask, mode in [(is_bilinear, "bilinear"), (~is_bilinear, "bicubic")]:
        if mask.any():
            output[mask] = F.grid_sample(
                image_4d.expand(mask.sum(), -1, -1, -1),
                grids[mask],
                mode=mode,
                padding_mode="zeros",
                align_corners=True,
            )

    return output


def gaussian_blur_preserve_max(heatmaps: torch.Tensor, kernel: int = 11) -> torch.Tensor:
    """Gaussian blur per-keypoint heatmap, preserving the original max value.

    Matches cv2.GaussianBlur with sigma=0 which means that the sigma is automatically
    computed from the kernel size.

    Args:
        heatmaps: Shape `(K, height, width)`.
        kernel: Odd integer kernel size for the Gaussian blur. Must be greater than 1.

    Returns:
        `torch.Tensor`: Blurred heatmaps of the same shape as the input.
    """
    if kernel % 2 == 0 or kernel <= 1:
        raise ValueError("Kernel size must be an odd integer greater than 1.")
    sigma = 0.3 * ((kernel - 1) * 0.5 - 1) + 0.8
    border = (kernel - 1) // 2
    origin_maxes = heatmaps.amax(dim=(1, 2))  # (K,)

    # Padding required to prevent border effect from gaussian blur. Torchvision uses reflect padding internally.
    padded = F.pad(heatmaps, (border, border, border, border), mode="constant", value=0.0)
    blurred = tvF.gaussian_blur(padded, kernel_size=[kernel, kernel], sigma=[sigma, sigma])
    result = blurred[:, border:-border, border:-border]

    result_maxes = result.amax(dim=(1, 2))  # (K,)
    safe_maxes = torch.where(result_maxes > 0, result_maxes, torch.ones_like(result_maxes))
    scale = torch.where(result_maxes > 0, origin_maxes / safe_maxes, torch.ones_like(origin_maxes))
    return result * scale[:, None, None]


def get_keypoint_predictions(heatmaps: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Predict keypoint locations and confidence scores from heatmaps.

    Args:
        heatmaps: Shape `(num_persons, num_keypoints, height, width)`.

    Returns:
        locations: `(num_persons, num_keypoints, 2)` x/y in heatmap pixel coordinates.
        scores: `(num_persons, num_keypoints)` per-keypoint confidence.
    """
    num_persons, num_keypoints, _, heatmap_width = heatmaps.shape
    device = heatmaps.device
    heatmap_flat = heatmaps.reshape(num_persons, num_keypoints, -1)
    scores = heatmap_flat.amax(dim=-1)
    flat_index = heatmap_flat.argmax(dim=-1)
    locations_x = (flat_index % heatmap_width).float()
    locations_y = (flat_index // heatmap_width).float()
    locations = torch.where(
        scores.unsqueeze(-1) > 0.0,
        torch.stack([locations_x, locations_y], dim=-1),
        torch.full((num_persons, num_keypoints, 2), -1.0, device=device),
    )
    return locations, scores


def post_dark_unbiased_data_processing(
    keypoints: torch.Tensor, heatmaps: torch.Tensor, blur_kernel_size: int = 11
) -> torch.Tensor:
    """Sub-pixel refinement via Hessian on log-heatmaps (UDP Dark Pose).

    Args:
        keypoints: Shape `(num_persons, num_keypoints, 2)` x/y in heatmap pixel coordinates.
        heatmaps: Shape `(num_persons, num_keypoints, height, width)`.

    Returns:
        `(num_persons, num_keypoints, 2)` refined keypoint locations.
    """
    num_persons, num_keypoints, heatmap_height, heatmap_width = heatmaps.shape
    device = heatmaps.device

    heatmaps = gaussian_blur_preserve_max(
        heatmaps.reshape(num_persons * num_keypoints, heatmap_height, heatmap_width), blur_kernel_size
    ).reshape(num_persons, num_keypoints, heatmap_height, heatmap_width)
    heatmaps = heatmaps.clamp(1e-3, 50.0).log()  # Clamp values based on original Sapiens2 implementation

    heatmaps_padded = F.pad(heatmaps, (1, 1, 1, 1), mode="replicate")
    heatmaps_flattened = heatmaps_padded.flatten()

    padded_height = heatmap_height + 2
    padded_width = heatmap_width + 2
    keypoint_stride = padded_height * padded_width
    person_stride = num_keypoints * keypoint_stride

    index = keypoints[:, :, 0].long() + 1 + (keypoints[:, :, 1].long() + 1) * padded_width
    index = index + keypoint_stride * torch.arange(num_keypoints, device=device, dtype=torch.long)[None, :]
    index = index + person_stride * torch.arange(num_persons, device=device, dtype=torch.long)[:, None]
    index = index.unsqueeze(-1)  # (num_persons, num_keypoints, 1)

    position_to_index_offset = {
        (0, 0): 0,
        (0, 1): 1,
        (0, -1): -1,
        (1, 0): padded_width,
        (-1, 0): -padded_width,
        (1, 1): padded_width + 1,
        (-1, -1): -(padded_width + 1),
    }
    # Dict mapping from (dx, dy) offsets to the corresponding values in the heatmap
    heatmap_values = {
        (dx, dy): heatmaps_flattened[index + offset] for (dx, dy), offset in position_to_index_offset.items()
    }

    gradient_x = 0.5 * (heatmap_values[0, 1] - heatmap_values[0, -1])
    gradient_y = 0.5 * (heatmap_values[1, 0] - heatmap_values[-1, 0])

    hessian_xx = heatmap_values[0, 1] - 2 * heatmap_values[0, 0] + heatmap_values[0, -1]
    hessian_yy = heatmap_values[1, 0] - 2 * heatmap_values[0, 0] + heatmap_values[-1, 0]
    hessian_xy = 0.5 * (
        heatmap_values[1, 1]
        - heatmap_values[0, 1]
        - heatmap_values[1, 0]
        + heatmap_values[0, 0]
        + heatmap_values[0, 0]
        - heatmap_values[0, -1]
        - heatmap_values[-1, 0]
        + heatmap_values[-1, -1]
    )

    eps = torch.finfo(hessian_xx.dtype).eps
    hessian_xx = hessian_xx + eps
    hessian_yy = hessian_yy + eps
    determinant = hessian_xx * hessian_yy - hessian_xy * hessian_xy
    offset_x = (hessian_yy * gradient_x - hessian_xy * gradient_y) / determinant
    offset_y = (-hessian_xy * gradient_x + hessian_xx * gradient_y) / determinant
    return keypoints - torch.cat([offset_x, offset_y], dim=-1)


@auto_docstring
class Sapiens2ImageProcessor(TorchvisionBackend):
    """PIL backend for Sapiens2 with reduce_label support."""

    valid_kwargs = Sapiens2ImageProcessorKwargs
    # Note: original Sapiens2 uses cv2.INTER_AREA for downsampling and cv2.INTER_CUBIC for upsampling
    resample = PILImageResampling.BILINEAR
    image_mean = IMAGENET_DEFAULT_MEAN
    image_std = IMAGENET_DEFAULT_STD
    size = {"height": 1024, "width": 768}
    default_to_square = True
    crop_size = {"height": 224, "width": 224}
    do_resize = True
    do_center_crop = False
    do_rescale = True
    do_normalize = True
    do_reduce_labels = False
    do_pad = False  # Set to True for normal, albedo, and pointmap estimation

    def __init__(self, **kwargs: Unpack[Sapiens2ImageProcessorKwargs]):
        super().__init__(**kwargs)

    @auto_docstring
    def preprocess(
        self,
        images: ImageInput,
        segmentation_maps: ImageInput | None = None,
        boxes: list[list[list[float]]] | None = None,
        **kwargs: Unpack[Sapiens2ImageProcessorKwargs],
    ) -> BatchFeature:
        r"""
        segmentation_maps (`ImageInput`, *optional*):
            The segmentation maps to preprocess.
        boxes (`list[list[list[float]]]` or `np.ndarray`, *optional*):
            List or array of bounding boxes for each image. Each box should be a list of 4 floats
            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.
        """
        return super().preprocess(images, segmentation_maps, boxes, **kwargs)

    def _preprocess_image_like_inputs(
        self,
        images: ImageInput,
        segmentation_maps: ImageInput | None,
        boxes: list[list[list[float]]] | None,
        do_convert_rgb: bool,
        input_data_format: ChannelDimension,
        return_tensors: str | TensorType | None,
        device: Union[str, "torch.device"] | None = None,
        **kwargs,
    ) -> BatchFeature:
        """Handle extra inputs beyond images."""
        kwargs["boxes"] = boxes  # modular trick
        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)

        # 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
            )

            # 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

        return BatchFeature(data=data, tensor_type=return_tensors)

    def reduce_label(self, labels: list["torch.Tensor"]) -> list["torch.Tensor"]:
        """Reduce label values by 1, replacing 0 with 255."""
        for idx in range(len(labels)):
            label = labels[idx]
            label = torch.where(label == 0, torch.tensor(255, dtype=label.dtype, device=label.device), label)
            label = label - 1
            label = torch.where(label == 254, torch.tensor(255, dtype=label.dtype, device=label.device), label)
            labels[idx] = label
        return labels

    def _preprocess(
        self,
        images: list[torch.Tensor],
        do_resize: bool,
        size: SizeDict,
        resample: "PILImageResampling | tvF.InterpolationMode | int | None",
        do_center_crop: bool,
        crop_size: SizeDict,
        do_rescale: bool,
        rescale_factor: float,
        do_normalize: bool,
        image_mean: float | list[float] | None,
        image_std: float | list[float] | None,
        disable_grouping: bool | None,
        do_reduce_labels: bool = False,
        do_pad: bool = False,
        boxes: list[list[list[float]]] | None = None,
        **kwargs,
    ) -> list[torch.Tensor]:
        """Custom preprocessing for Sapiens2."""
        if boxes is not None:
            output_size = (size["height"], size["width"])
            crops = []
            for image, image_boxes in zip(images, boxes):
                image = tvF.to_dtype_image(image, dtype=torch.float32, scale=False)
                boxes_tensor = box_xywh_to_cxcywh(torch.tensor(image_boxes, dtype=torch.float32, device=image.device))
                crops.extend(crop_and_resize(image, boxes=boxes_tensor, output_size=output_size))
            images = crops
            do_resize = False  # crop_and_resize already produces the target size

        if do_reduce_labels:
            images = self.reduce_label(images)

        grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping)
        resized_images_grouped = {}
        for shape, stacked_images in grouped_images.items():
            if do_resize:
                if do_pad:
                    # Resize while preserving aspect ratio. Then add symmetric padding on all sides to reach the target size.
                    aspect_ratio_size = SizeDict(max_height=size["height"], max_width=size["width"])
                    stacked_images = self.resize(stacked_images, aspect_ratio_size, resample)
                    stacked_images = self.center_crop(stacked_images, size)
                else:
                    stacked_images = self.resize(stacked_images, size, resample)
            resized_images_grouped[shape] = stacked_images
        resized_images = reorder_images(resized_images_grouped, grouped_images_index)

        grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping)
        processed_images_grouped = {}
        for shape, stacked_images in grouped_images.items():
            if do_center_crop:
                stacked_images = self.center_crop(stacked_images, crop_size)
            stacked_images = self.rescale_and_normalize(
                stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std
            )
            processed_images_grouped[shape] = stacked_images

        return reorder_images(processed_images_grouped, grouped_images_index)

    def post_process_semantic_segmentation(self, outputs, target_sizes: list[tuple] | None = None):
        """
        Converts the output of [`Sapiens2ForSemanticSegmentation`] into semantic segmentation maps.

        Args:
            outputs ([`Sapiens2ForSemanticSegmentation`]):
                Raw outputs of the model.
            target_sizes (`list[Tuple]` of length `batch_size`, *optional*):
                List of tuples corresponding to the requested final size (height, width) of each prediction. If unset,
                predictions will not be resized.

        Returns:
            semantic_segmentation: `list[torch.Tensor]` of length `batch_size`, where each item is a semantic
            segmentation map of shape (height, width) corresponding to the target_sizes entry (if `target_sizes` is
            specified). Each entry of each `torch.Tensor` correspond to a semantic class id.
        """
        if not is_torch_available():
            raise ImportError("PyTorch is required for post_process_semantic_segmentation")

        logits = outputs.logits

        # Resize logits and compute semantic segmentation maps
        if target_sizes is not None:
            if len(logits) != len(target_sizes):
                raise ValueError(
                    "Make sure that you pass in as many target sizes as the batch dimension of the logits"
                )

            if isinstance(target_sizes, torch.Tensor):
                target_sizes = target_sizes.numpy()

            semantic_segmentation = []

            for idx in range(len(logits)):
                resized_logits = F.interpolate(
                    logits[idx].unsqueeze(dim=0), size=target_sizes[idx], mode="bilinear", align_corners=False
                )
                semantic_map = resized_logits[0].argmax(dim=0)
                semantic_segmentation.append(semantic_map)
        else:
            semantic_segmentation = logits.argmax(dim=1)
            semantic_segmentation = [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])]

        return semantic_segmentation

    def post_process_pose_estimation(
        self,
        outputs: Sapiens2PoseEstimatorOutput,
        boxes: list[list[list[float]]],
        outputs_flipped: Sapiens2PoseEstimatorOutput | None = None,
        kernel_size: int = 11,
        threshold: float | None = None,
        source_sizes: TensorType | list[tuple[int, int]] | None = None,
        target_sizes: TensorType | list[tuple[int, int]] | None = None,
    ) -> list[list[dict[str, torch.Tensor]]]:
        """
        Converts the output of [`Sapiens2ForPoseEstimation`] into keypoint predictions in image space.

        Args:
            outputs (`Sapiens2PoseEstimatorOutput`):
                Raw outputs of the model. `outputs.heatmaps` must have shape
                `(N_total, num_keypoints, heatmap_height, heatmap_width)` where
                `N_total = sum(len(b) for b in boxes)`.
            boxes (`list[list[list[float]]]` or `np.ndarray`):
                List or array of bounding boxes for each image in absolute pixel coordinates. Each box
                should be a list of 4 floats representing the bounding box coordinates in COCO format
                (top_left_x, top_left_y, width, height). Must match the `boxes` argument passed to
                `preprocess`.
            outputs_flipped (`Sapiens2PoseEstimatorOutput`, *optional*):
                Outputs from running the model on horizontally flipped inputs. When provided, heatmaps
                are averaged with `outputs` before keypoint extraction to improve accuracy:
                `avg_heatmaps = (outputs.heatmaps + outputs_flipped.heatmaps) / 2`.
            kernel_size (`int`, *optional*, defaults to 11):
                Kernel size for the Gaussian blur used in UDP Dark Pose refinement.
            threshold (`float`, *optional*):
                Score threshold. Keypoints with scores at or below this value are
                filtered out from the result dictionaries.
            source_sizes (`torch.Tensor` or `list[tuple[int, int]]` of length `batch_size`, *optional*):
                Original `(height, width)` of each image in pixels. Required when `target_sizes` is
                provided, as the source coordinate space for scaling keypoints and bounding boxes.
            target_sizes (`torch.Tensor` or `list[tuple[int, int]]` of length `batch_size`, *optional*):
                Desired output `(height, width)` coordinate space for each image. When provided
                alongside `source_sizes`, keypoint coordinates and bounding boxes are scaled from
                source to target space.

        Returns:
            `list[list[dict]]`: Outer list is over images, inner list is over persons.
            Each dict contains:
            - `keypoints` (`torch.FloatTensor` of shape `(num_keypoints, 2)`): absolut x/y coordinates in
              the source image space, or in target space if `target_sizes` is provided.
            - `scores` (`torch.FloatTensor` of shape `(num_keypoints,)`): per-keypoint confidence.
            - `labels` (`torch.LongTensor` of shape `(num_keypoints,)`): keypoint indices.
            - `bbox` (`torch.FloatTensor` of shape `(4,)`): bounding box in absolute (x_min, y_min, x_max, y_max)
               format, in the same coordinate space as `keypoints`.
        """
        if isinstance(source_sizes, torch.Tensor):
            source_sizes = source_sizes.tolist()
        if isinstance(target_sizes, torch.Tensor):
            target_sizes = target_sizes.tolist()

        num_images = len(boxes)

        if target_sizes is not None and source_sizes is None:
            raise ValueError("`source_sizes` must be provided when `target_sizes` is specified.")
        if source_sizes is not None and num_images != len(source_sizes):
            raise ValueError("Make sure that you pass in as many source sizes as the number of images.")
        if target_sizes is not None and num_images != len(target_sizes):
            raise ValueError("Make sure that you pass in as many target sizes as the number of images.")

        heatmaps = outputs.heatmaps  # (num_total_persons, num_keypoints, heatmap_height, heatmap_width)
        if outputs_flipped is not None:
            heatmaps = (heatmaps + outputs_flipped.heatmaps) / 2

        device = heatmaps.device
        num_total_persons, num_keypoints, heatmap_height, heatmap_width = heatmaps.shape

        if num_total_persons == 0:
            return [[] for _ in boxes]

        # (num_total_persons, 4)
        boxes_tensor = torch.tensor(
            [box for image_boxes in boxes for box in image_boxes], dtype=torch.float32, device=device
        )

        heatmaps = heatmaps.float()  # For consistency with original numpy/cv2 implementation which uses float32.

        # (num_total_persons, num_keypoints, 2), (num_total_persons, num_keypoints)
        all_keypoints, all_scores = get_keypoint_predictions(heatmaps)
        all_keypoints = post_dark_unbiased_data_processing(
            keypoints=all_keypoints, heatmaps=heatmaps, blur_kernel_size=kernel_size
        )

        # Remap coordinates from heatmap space to original image space
        centers, scales = boxes_to_crop_params(
            box_xywh_to_cxcywh(boxes_tensor), output_size=(self.size["height"], self.size["width"])
        )
        heatmap_size = torch.tensor([heatmap_width - 1, heatmap_height - 1], dtype=torch.float32, device=device)
        all_keypoints = (
            all_keypoints / heatmap_size * scales[:, None, :] + centers[:, None, :] - 0.5 * scales[:, None, :]
        )
        all_boxes = box_xywh_to_xyxy(boxes_tensor)  # (num_total_persons, 4)

        if source_sizes is not None and target_sizes is not None:
            # (num_images, 2)
            per_image_scale = torch.tensor(
                [
                    [target_width / source_width, target_height / source_height]
                    for (source_height, source_width), (target_height, target_width) in zip(source_sizes, target_sizes)
                ],
                dtype=torch.float32,
                device=device,
            )
            # (num_total_persons, 2)
            per_person_scale = torch.cat(
                [
                    per_image_scale[image_index].unsqueeze(0).expand(len(boxes[image_index]), 2)
                    for image_index in range(num_images)
                ]
            )
            all_keypoints = all_keypoints * per_person_scale[:, None, :]
            all_boxes = all_boxes * per_person_scale[:, [0, 1, 0, 1]]

        person_results = []
        for person_index in range(num_total_persons):
            keypoints = all_keypoints[person_index]
            scores = all_scores[person_index]
            labels = torch.arange(num_keypoints, device=device)

            if threshold is not None:
                keep = scores > threshold
                keypoints = keypoints[keep]
                scores = scores[keep]
                labels = labels[keep]

            person_results.append(
                {"keypoints": keypoints, "scores": scores, "labels": labels, "bbox": all_boxes[person_index]}
            )

        # Reassemble into list[list[dict]] grouped by image
        result = []
        person_offset = 0
        for image_boxes in boxes:
            num_persons_in_image = len(image_boxes)
            result.append(person_results[person_offset : person_offset + num_persons_in_image])
            person_offset += num_persons_in_image
        return result

    def post_process_normal_estimation(
        self,
        outputs: Sapiens2NormalEstimatorOutput,
        source_sizes: TensorType | list[tuple[int, int]] | None = None,
        target_sizes: TensorType | list[tuple[int, int]] | None = None,
        do_remove_padding: bool | None = None,
    ) -> list[dict[str, torch.Tensor]]:
        """
        Converts the output of [`Sapiens2ForNormalEstimation`] into L2-normalized surface normal maps.

        Args:
            outputs (`Sapiens2NormalEstimatorOutput`):
                Raw outputs of the model.
            source_sizes (`torch.Tensor` or `list[tuple[int, int]]` of length `batch_size`, *optional*):
                Original `(height, width)` of each image before preprocessing. When provided,
                the padding added during preprocessing is removed and predictions are resized back
                to the original image size (unless `target_sizes` overrides the final size).
            target_sizes (`torch.Tensor` or `list[tuple[int, int]]` of length `batch_size`, *optional*):
                Requested final `(height, width)` for each prediction. When provided, used as the
                resize target instead of `source_sizes`. Resized with bilinear interpolation after
                L2 normalization.
            do_remove_padding (`bool`, *optional*):
                Whether to crop away the zero-padding added during preprocessing before resizing.
                Defaults to `True` when `source_sizes` is provided, `False` otherwise.

        Returns:
            `list[dict[str, torch.Tensor]]` of length `batch_size`. Each dict has a `"normals"` key
            mapping to a tensor of shape `(3, height, width)` with L2-normalized unit vectors in
            `[-1, 1]` per channel (XYZ surface normals).
        """
        normals = F.normalize(outputs.normals, p=2, dim=1, eps=1e-8)
        results = self._post_process_maps(
            maps=normals, source_sizes=source_sizes, target_sizes=target_sizes, do_remove_padding=do_remove_padding
        )
        return [{"normals": result} for result in results]

    def post_process_pointmap_estimation(
        self,
        outputs: Sapiens2PointmapEstimatorOutput,
        source_sizes: TensorType | list[tuple[int, int]] | None = None,
        target_sizes: TensorType | list[tuple[int, int]] | None = None,
        do_remove_padding: bool | None = None,
    ) -> list[dict[str, torch.Tensor]]:
        """
        Converts the output of [`Sapiens2ForPointmapEstimation`] into pointmap tensors in image space.

        Args:
            outputs (`Sapiens2PointmapEstimatorOutput`):
                Raw outputs of the model.
            source_sizes (`torch.Tensor` or `list[tuple[int, int]]` of length `batch_size`, *optional*):
                Original `(height, width)` of each image before preprocessing. When provided,
                the padding added during preprocessing is removed and predictions are resized back
                to the original image size (unless `target_sizes` overrides the final size).
            target_sizes (`torch.Tensor` or `list[tuple[int, int]]` of length `batch_size`, *optional*):
                Requested final `(height, width)` for each prediction. Overrides `source_sizes`
                as the resize target.
            do_remove_padding (`bool`, *optional*):
                Whether to crop away the zero-padding added during preprocessing before resizing.
                Defaults to `True` when `source_sizes` is provided, `False` otherwise.

        Returns:
            `list[dict[str, torch.Tensor]]` of length `batch_size`. Each dict has a `"pointmap"` key
            mapping to a tensor of shape `(3, height, width)` with per-pixel 3D XYZ coordinates in
            canonical camera space, optionally divided by `outputs.scales` to convert to metric coordinates.
        """
        pointmaps = outputs.pointmaps
        if outputs.scales is not None:
            pointmaps = pointmaps / outputs.scales[:, :, None, None]
        results = self._post_process_maps(
            maps=pointmaps, source_sizes=source_sizes, target_sizes=target_sizes, do_remove_padding=do_remove_padding
        )
        return [{"pointmap": result} for result in results]

    def post_process_image_matting(
        self,
        outputs: Sapiens2ImageMattingOutput,
        target_sizes: TensorType | list[tuple[int, int]] | None = None,
        backgrounds: ImageInput | None = None,
    ) -> list[dict[str, torch.Tensor]]:
        """
        Converts the output of [`Sapiens2ForImageMatting`] into alpha mattes and foreground maps.

        Args:
            outputs (`Sapiens2ImageMattingOutput`):
                Raw outputs of the model.
            target_sizes (`torch.Tensor` or `list[tuple[int, int]]` of length `batch_size`, *optional*):
                Requested final `(height, width)` for each prediction. Resized with bilinear
                interpolation. If unset, predictions are returned at the model output resolution.
            backgrounds (`ImageInput`, *optional*):
                Background image(s) to composite over. Can be a single image (applied to every item
                in the batch) or a list of images, one per batch item. Accepts PIL images, numpy
                arrays, or torch tensors of any dtype; integer types (e.g. uint8) are scaled to
                `[0, 1]` automatically. When provided, each result dict gains a `"composite"` key
                with the composited image as a uint8 tensor in `[0, 255]`.

        Returns:
            `list[dict]` of length `batch_size`. Each dict has:
            - `"alpha"` (`torch.Tensor` of shape `(1, height, width)`): alpha values in `[0, 1]`.
            - `"foreground"` (`torch.Tensor` of shape `(3, height, width)`): pre-multiplied RGB in `[0, 1]`.
            - `"composite"` (`torch.Tensor` of shape `(3, height, width)` or `None`): foreground composited
              over `backgrounds` as a uint8 tensor in `[0, 255]`; `None` when `backgrounds` is not provided.
        """
        if isinstance(target_sizes, torch.Tensor):
            target_sizes = target_sizes.tolist()

        batch_size = outputs.foregrounds.shape[0]
        device = outputs.foregrounds.device
        dtype = outputs.foregrounds.dtype

        if target_sizes is not None:
            if batch_size != len(target_sizes):
                raise ValueError(
                    "Make sure that you pass in as many target sizes as the batch dimension of the matting output"
                )
        all_target_sizes_equal = target_sizes is None or all(
            tuple(size) == tuple(target_sizes[0]) for size in target_sizes
        )

        background_tensors = []
        if backgrounds is not None:
            background_list = make_list_of_images(backgrounds)
            if len(background_list) != 1 and len(background_list) != batch_size:
                raise ValueError(
                    "Make sure that you pass in as many backgrounds as the batch dimension of the matting output"
                )
            background_tensors = [
                tvF.to_dtype_image(tvF.to_image(background_image), dtype=dtype, scale=True).to(device)
                for background_image in background_list
            ]
        all_background_sizes_equal = not background_tensors or all(
            background.shape[-2:] == background_tensors[0].shape[-2:] for background in background_tensors
        )

        matting = torch.cat([outputs.foregrounds, outputs.alphas], dim=1)  # (batch_size, 4, height, width)

        if target_sizes is not None and all_target_sizes_equal:
            target_size = tuple(target_sizes[0])
            matting = F.interpolate(
                matting,
                size=target_size,
                mode="bilinear",
                align_corners=False,
                antialias=False,
            )
            matting = matting.clamp(0.0, 1.0)

        result = []
        if all_target_sizes_equal and all_background_sizes_equal:
            # Fast path
            foregrounds = matting[:, :3]
            alphas = matting[:, 3:]
            composites = [None] * batch_size
            if background_tensors:
                background = torch.stack(background_tensors)
                if background.shape[-2:] != matting.shape[-2:]:
                    background = F.interpolate(
                        background,
                        size=matting.shape[-2:],
                        mode="bilinear",
                        align_corners=False,
                        antialias=False,
                    )
                composites = (foregrounds + (1 - alphas) * background).clamp(0.0, 1.0)
                composites = tvF.to_dtype_image(composites, dtype=torch.uint8, scale=True)

            for foreground, alpha, composite in zip(foregrounds, alphas, composites):
                result.append(
                    {
                        "foreground": foreground,
                        "alpha": alpha,
                        "composite": composite,
                    }
                )

        else:
            # Slow path
            for index in range(len(matting)):
                matting_item = matting[index]

                if target_sizes and not all_target_sizes_equal:
                    matting_item = F.interpolate(
                        matting_item.unsqueeze(0),
                        size=target_sizes[index],
                        mode="bilinear",
                        align_corners=False,
                        antialias=False,
                    )[0]
                    matting_item = matting_item.clamp(0.0, 1.0)

                foreground = matting_item[:3]
                alpha = matting_item[3:]
                composite = None

                if background_tensors:
                    background = background_tensors[0] if len(background_tensors) == 1 else background_tensors[index]
                    if background.shape[-2:] != matting_item.shape[-2:]:
                        background = F.interpolate(
                            background.unsqueeze(0),
                            size=matting_item.shape[-2:],
                            mode="bilinear",
                            align_corners=False,
                            antialias=False,
                        )[0]
                    composite = (foreground + (1 - alpha) * background).clamp(0.0, 1.0)
                    composite = tvF.to_dtype_image(composite, dtype=torch.uint8, scale=True)

                result.append({"foreground": foreground, "alpha": alpha, "composite": composite})

        return result

    def _post_process_maps(
        self,
        maps: torch.Tensor,
        source_sizes: TensorType | list[tuple[int, int]] | None,
        target_sizes: TensorType | list[tuple[int, int]] | None,
        do_remove_padding: bool | None,
    ) -> list[torch.Tensor]:
        if isinstance(source_sizes, torch.Tensor):
            source_sizes = source_sizes.tolist()
        if isinstance(target_sizes, torch.Tensor):
            target_sizes = target_sizes.tolist()
        if do_remove_padding is None:
            do_remove_padding = source_sizes is not None
        if do_remove_padding and source_sizes is None:
            raise ValueError("`source_sizes` must be provided when `do_remove_padding=True`.")

        if source_sizes is not None and len(maps) != len(source_sizes):
            raise ValueError("Make sure that you pass in as many source sizes as the batch dimension of the outputs")
        if target_sizes is not None and len(maps) != len(target_sizes):
            raise ValueError("Make sure that you pass in as many target sizes as the batch dimension of the outputs")

        model_height = self.size["height"]
        model_width = self.size["width"]

        crops = []
        if do_remove_padding:
            for original_height, original_width in source_sizes:
                new_height, new_width = get_image_size_for_max_height_width(
                    (original_height, original_width), model_height, model_width
                )
                pad_top = (model_height - new_height) // 2 if new_height < model_height else 0
                pad_left = (model_width - new_width) // 2 if new_width < model_width else 0
                crops.append(
                    (
                        pad_top,
                        pad_left,
                        pad_top + min(new_height, model_height),
                        pad_left + min(new_width, model_width),
                    )
                )
        all_crops_equal = not crops or all(crop == crops[0] for crop in crops)

        final_sizes = []
        if target_sizes is not None:
            final_sizes = [tuple(size) for size in target_sizes]
        elif source_sizes is not None:
            final_sizes = [tuple(size) for size in source_sizes]
        all_final_sizes_equal = not final_sizes or all(size == final_sizes[0] for size in final_sizes)

        result = []
        if all_crops_equal and all_final_sizes_equal:
            # Fast path
            if do_remove_padding:
                top, left, bottom, right = crops[0]
                maps = maps[:, :, top:bottom, left:right]

            if final_sizes:
                maps = F.interpolate(
                    maps,
                    size=final_sizes[0],
                    mode="bilinear",
                    align_corners=False,
                    antialias=False,
                )

            result = list(maps)
        else:
            # Slow path
            for index in range(len(maps)):
                map_item = maps[index]

                if do_remove_padding:
                    top, left, bottom, right = crops[index]
                    map_item = map_item[:, top:bottom, left:right]

                if final_sizes:
                    map_item = F.interpolate(
                        map_item.unsqueeze(0),
                        size=final_sizes[index],
                        mode="bilinear",
                        align_corners=False,
                        antialias=False,
                    )[0]

                result.append(map_item)

        return result


__all__ = ["Sapiens2ImageProcessor"]
