Aluode/PerceptionLabPortable

1# coding=utf-82# Copyright 2024 The HuggingFace Inc. team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""Image processor class for VitPose."""16 17import itertools18import math19from typing import TYPE_CHECKING, Optional, Union20 21import numpy as np22 23from ...image_processing_utils import BaseImageProcessor, BatchFeature24from ...image_transforms import to_channel_dimension_format25from ...image_utils import (26    IMAGENET_DEFAULT_MEAN,27    IMAGENET_DEFAULT_STD,28    ChannelDimension,29    ImageInput,30    infer_channel_dimension_format,31    is_scaled_image,32    make_flat_list_of_images,33    to_numpy_array,34    valid_images,35)36from ...utils import TensorType, is_scipy_available, is_torch_available, is_vision_available, logging37 38 39if is_torch_available():40    import torch41 42if is_vision_available():43    import PIL44 45if is_scipy_available():46    from scipy.linalg import inv47    from scipy.ndimage import affine_transform, gaussian_filter48 49if TYPE_CHECKING:50    from .modeling_vitpose import VitPoseEstimatorOutput51 52logger = logging.get_logger(__name__)53 54 55# inspired by https://github.com/ViTAE-Transformer/ViTPose/blob/d5216452796c90c6bc29f5c5ec0bdba94366768a/mmpose/datasets/datasets/base/kpt_2d_sview_rgb_img_top_down_dataset.py#L13256def box_to_center_and_scale(57    box: Union[tuple, list, np.ndarray],58    image_width: int,59    image_height: int,60    normalize_factor: float = 200.0,61    padding_factor: float = 1.25,62):63    """64    Encodes a bounding box in COCO format into (center, scale).65 66    Args:67        box (`Tuple`, `List`, or `np.ndarray`):68            Bounding box in COCO format (top_left_x, top_left_y, width, height).69        image_width (`int`):70            Image width.71        image_height (`int`):72            Image height.73        normalize_factor (`float`):74            Width and height scale factor.75        padding_factor (`float`):76            Bounding box padding factor.77 78    Returns:79        tuple: A tuple containing center and scale.80 81        - `np.ndarray` [float32](2,): Center of the bbox (x, y).82        - `np.ndarray` [float32](2,): Scale of the bbox width & height.83    """84 85    top_left_x, top_left_y, width, height = box[:4]86    aspect_ratio = image_width / image_height87    center = np.array([top_left_x + width * 0.5, top_left_y + height * 0.5], dtype=np.float32)88 89    if width > aspect_ratio * height:90        height = width * 1.0 / aspect_ratio91    elif width < aspect_ratio * height:92        width = height * aspect_ratio93 94    scale = np.array([width / normalize_factor, height / normalize_factor], dtype=np.float32)95    scale = scale * padding_factor96 97    return center, scale98 99 100def coco_to_pascal_voc(bboxes: np.ndarray) -> np.ndarray:101    """102    Converts bounding boxes from the COCO format to the Pascal VOC format.103 104    In other words, converts from (top_left_x, top_left_y, width, height) format105    to (top_left_x, top_left_y, bottom_right_x, bottom_right_y).106 107    Args:108        bboxes (`np.ndarray` of shape `(batch_size, 4)):109            Bounding boxes in COCO format.110 111    Returns:112        `np.ndarray` of shape `(batch_size, 4) in Pascal VOC format.113    """114    bboxes[:, 2] = bboxes[:, 2] + bboxes[:, 0] - 1115    bboxes[:, 3] = bboxes[:, 3] + bboxes[:, 1] - 1116 117    return bboxes118 119 120def get_keypoint_predictions(heatmaps: np.ndarray) -> tuple[np.ndarray, np.ndarray]:121    """Get keypoint predictions from score maps.122 123    Args:124        heatmaps (`np.ndarray` of shape `(batch_size, num_keypoints, height, width)`):125            Model predicted heatmaps.126 127    Returns:128        tuple: A tuple containing aggregated results.129 130        - coords (`np.ndarray` of shape `(batch_size, num_keypoints, 2)`):131            Predicted keypoint location.132        - scores (`np.ndarray` of shape `(batch_size, num_keypoints, 1)`):133            Scores (confidence) of the keypoints.134    """135    if not isinstance(heatmaps, np.ndarray):136        raise TypeError("Heatmaps should be np.ndarray")137    if heatmaps.ndim != 4:138        raise ValueError("Heatmaps should be 4-dimensional")139 140    batch_size, num_keypoints, _, width = heatmaps.shape141    heatmaps_reshaped = heatmaps.reshape((batch_size, num_keypoints, -1))142    idx = np.argmax(heatmaps_reshaped, 2).reshape((batch_size, num_keypoints, 1))143    scores = np.amax(heatmaps_reshaped, 2).reshape((batch_size, num_keypoints, 1))144 145    preds = np.tile(idx, (1, 1, 2)).astype(np.float32)146    preds[:, :, 0] = preds[:, :, 0] % width147    preds[:, :, 1] = preds[:, :, 1] // width148 149    preds = np.where(np.tile(scores, (1, 1, 2)) > 0.0, preds, -1)150    return preds, scores151 152 153def post_dark_unbiased_data_processing(coords: np.ndarray, batch_heatmaps: np.ndarray, kernel: int = 3) -> np.ndarray:154    """DARK post-pocessing. Implemented by unbiased_data_processing.155 156    Paper references:157    - Huang et al. The Devil is in the Details: Delving into Unbiased Data Processing for Human Pose Estimation (CVPR 2020).158    - Zhang et al. Distribution-Aware Coordinate Representation for Human Pose Estimation (CVPR 2020).159 160    Args:161        coords (`np.ndarray` of shape `(num_persons, num_keypoints, 2)`):162            Initial coordinates of human pose.163        batch_heatmaps (`np.ndarray` of shape `(batch_size, num_keypoints, height, width)`):164            Batched heatmaps as predicted by the model.165            A batch_size of 1 is used for the bottom up paradigm where all persons share the same heatmap.166            A batch_size of `num_persons` is used for the top down paradigm where each person has its own heatmaps.167        kernel (`int`, *optional*, defaults to 3):168            Gaussian kernel size (K) for modulation.169 170    Returns:171        `np.ndarray` of shape `(num_persons, num_keypoints, 2)` ):172            Refined coordinates.173    """174    batch_size, num_keypoints, height, width = batch_heatmaps.shape175    num_coords = coords.shape[0]176    if not (batch_size == 1 or batch_size == num_coords):177        raise ValueError("The batch size of heatmaps should be 1 or equal to the batch size of coordinates.")178    radius = int((kernel - 1) // 2)179    batch_heatmaps = np.array(180        [181            [gaussian_filter(heatmap, sigma=0.8, radius=(radius, radius), axes=(0, 1)) for heatmap in heatmaps]182            for heatmaps in batch_heatmaps183        ]184    )185    batch_heatmaps = np.clip(batch_heatmaps, 0.001, 50)186    batch_heatmaps = np.log(batch_heatmaps)187 188    batch_heatmaps_pad = np.pad(batch_heatmaps, ((0, 0), (0, 0), (1, 1), (1, 1)), mode="edge").flatten()189 190    # calculate indices for coordinates191    index = coords[..., 0] + 1 + (coords[..., 1] + 1) * (width + 2)192    index += (width + 2) * (height + 2) * np.arange(0, batch_size * num_keypoints).reshape(-1, num_keypoints)193    index = index.astype(int).reshape(-1, 1)194    i_ = batch_heatmaps_pad[index]195    ix1 = batch_heatmaps_pad[index + 1]196    iy1 = batch_heatmaps_pad[index + width + 2]197    ix1y1 = batch_heatmaps_pad[index + width + 3]198    ix1_y1_ = batch_heatmaps_pad[index - width - 3]199    ix1_ = batch_heatmaps_pad[index - 1]200    iy1_ = batch_heatmaps_pad[index - 2 - width]201 202    # calculate refined coordinates using Newton's method203    dx = 0.5 * (ix1 - ix1_)204    dy = 0.5 * (iy1 - iy1_)205    derivative = np.concatenate([dx, dy], axis=1)206    derivative = derivative.reshape(num_coords, num_keypoints, 2, 1)207    dxx = ix1 - 2 * i_ + ix1_208    dyy = iy1 - 2 * i_ + iy1_209    dxy = 0.5 * (ix1y1 - ix1 - iy1 + i_ + i_ - ix1_ - iy1_ + ix1_y1_)210    hessian = np.concatenate([dxx, dxy, dxy, dyy], axis=1)211    hessian = hessian.reshape(num_coords, num_keypoints, 2, 2)212    hessian = np.linalg.inv(hessian + np.finfo(np.float32).eps * np.eye(2))213    coords -= np.einsum("ijmn,ijnk->ijmk", hessian, derivative).squeeze()214    return coords215 216 217def transform_preds(coords: np.ndarray, center: np.ndarray, scale: np.ndarray, output_size: np.ndarray) -> np.ndarray:218    """Get final keypoint predictions from heatmaps and apply scaling and219    translation to map them back to the image.220 221    Note:222        num_keypoints: K223 224    Args:225        coords (`np.ndarray` of shape `(num_keypoints, ndims)`):226 227            * If ndims=2, corrds are predicted keypoint location.228            * If ndims=4, corrds are composed of (x, y, scores, tags)229            * If ndims=5, corrds are composed of (x, y, scores, tags,230              flipped_tags)231 232        center (`np.ndarray` of shape `(2,)`):233            Center of the bounding box (x, y).234        scale (`np.ndarray` of shape `(2,)`):235            Scale of the bounding box wrt original image of width and height.236        output_size (`np.ndarray` of shape `(2,)`):237            Size of the destination heatmaps in (height, width) format.238 239    Returns:240        np.ndarray: Predicted coordinates in the images.241    """242    if coords.shape[1] not in (2, 4, 5):243        raise ValueError("Coordinates need to have either 2, 4 or 5 dimensions.")244    if len(center) != 2:245        raise ValueError("Center needs to have 2 elements, one for x and one for y.")246    if len(scale) != 2:247        raise ValueError("Scale needs to consist of a width and height")248    if len(output_size) != 2:249        raise ValueError("Output size needs to consist of a height and width")250 251    # Recover the scale which is normalized by a factor of 200.252    scale = scale * 200.0253 254    # We use unbiased data processing255    scale_y = scale[1] / (output_size[0] - 1.0)256    scale_x = scale[0] / (output_size[1] - 1.0)257 258    target_coords = np.ones_like(coords)259    target_coords[:, 0] = coords[:, 0] * scale_x + center[0] - scale[0] * 0.5260    target_coords[:, 1] = coords[:, 1] * scale_y + center[1] - scale[1] * 0.5261 262    return target_coords263 264 265def get_warp_matrix(theta: float, size_input: np.ndarray, size_dst: np.ndarray, size_target: np.ndarray):266    """267    Calculate the transformation matrix under the constraint of unbiased. Paper ref: Huang et al. The Devil is in the268    Details: Delving into Unbiased Data Processing for Human Pose Estimation (CVPR 2020).269 270    Source: https://github.com/open-mmlab/mmpose/blob/master/mmpose/core/post_processing/post_transforms.py271 272    Args:273        theta (`float`):274            Rotation angle in degrees.275        size_input (`np.ndarray`):276            Size of input image [width, height].277        size_dst (`np.ndarray`):278            Size of output image [width, height].279        size_target (`np.ndarray`):280            Size of ROI in input plane [w, h].281 282    Returns:283        `np.ndarray`: A matrix for transformation.284    """285    theta = np.deg2rad(theta)286    matrix = np.zeros((2, 3), dtype=np.float32)287    scale_x = size_dst[0] / size_target[0]288    scale_y = size_dst[1] / size_target[1]289    matrix[0, 0] = math.cos(theta) * scale_x290    matrix[0, 1] = -math.sin(theta) * scale_x291    matrix[0, 2] = scale_x * (292        -0.5 * size_input[0] * math.cos(theta) + 0.5 * size_input[1] * math.sin(theta) + 0.5 * size_target[0]293    )294    matrix[1, 0] = math.sin(theta) * scale_y295    matrix[1, 1] = math.cos(theta) * scale_y296    matrix[1, 2] = scale_y * (297        -0.5 * size_input[0] * math.sin(theta) - 0.5 * size_input[1] * math.cos(theta) + 0.5 * size_target[1]298    )299    return matrix300 301 302def scipy_warp_affine(src, M, size):303    """304    This function implements cv2.warpAffine function using affine_transform in scipy. See https://docs.scipy.org/doc/scipy/reference/generated/scipy.ndimage.affine_transform.html and https://docs.opencv.org/4.x/d4/d61/tutorial_warp_affine.html for more details.305 306    Note: the original implementation of cv2.warpAffine uses cv2.INTER_LINEAR.307    """308    channels = [src[..., i] for i in range(src.shape[-1])]309 310    # Convert to a 3x3 matrix used by SciPy311    M_scipy = np.vstack([M, [0, 0, 1]])312    # If you have a matrix for the ‘push’ transformation, use its inverse (numpy.linalg.inv) in this function.313    M_inv = inv(M_scipy)314    M_inv[0, 0], M_inv[0, 1], M_inv[1, 0], M_inv[1, 1], M_inv[0, 2], M_inv[1, 2] = (315        M_inv[1, 1],316        M_inv[1, 0],317        M_inv[0, 1],318        M_inv[0, 0],319        M_inv[1, 2],320        M_inv[0, 2],321    )322 323    new_src = [affine_transform(channel, M_inv, output_shape=size, order=1) for channel in channels]324    new_src = np.stack(new_src, axis=-1)325    return new_src326 327 328class VitPoseImageProcessor(BaseImageProcessor):329    r"""330    Constructs a VitPose image processor.331 332    Args:333        do_affine_transform (`bool`, *optional*, defaults to `True`):334            Whether to apply an affine transformation to the input images.335        size (`dict[str, int]` *optional*, defaults to `{"height": 256, "width": 192}`):336            Resolution of the image after `affine_transform` is applied. Only has an effect if `do_affine_transform` is set to `True`. Can337            be overridden by `size` in the `preprocess` method.338        do_rescale (`bool`, *optional*, defaults to `True`):339            Whether or not to apply the scaling factor (to make pixel values floats between 0. and 1.).340        rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):341            Scale factor to use if rescaling the image. Can be overridden by `rescale_factor` in the `preprocess`342            method.343        do_normalize (`bool`, *optional*, defaults to `True`):344            Whether or not to normalize the input with mean and standard deviation.345        image_mean (`list[int]`, defaults to `[0.485, 0.456, 0.406]`, *optional*):346            The sequence of means for each channel, to be used when normalizing images.347        image_std (`list[int]`, defaults to `[0.229, 0.224, 0.225]`, *optional*):348            The sequence of standard deviations for each channel, to be used when normalizing images.349    """350 351    model_input_names = ["pixel_values"]352 353    def __init__(354        self,355        do_affine_transform: bool = True,356        size: Optional[dict[str, int]] = None,357        do_rescale: bool = True,358        rescale_factor: Union[int, float] = 1 / 255,359        do_normalize: bool = True,360        image_mean: Optional[Union[float, list[float]]] = None,361        image_std: Optional[Union[float, list[float]]] = None,362        **kwargs,363    ):364        super().__init__(**kwargs)365        self.do_affine_transform = do_affine_transform366        self.size = size if size is not None else {"height": 256, "width": 192}367        self.do_rescale = do_rescale368        self.rescale_factor = rescale_factor369        self.do_normalize = do_normalize370        self.image_mean = image_mean if image_mean is not None else IMAGENET_DEFAULT_MEAN371        self.image_std = image_std if image_std is not None else IMAGENET_DEFAULT_STD372        self.normalize_factor = 200.0373 374    def affine_transform(375        self,376        image: np.ndarray,377        center: tuple[float],378        scale: tuple[float],379        rotation: float,380        size: dict[str, int],381        data_format: Optional[ChannelDimension] = None,382        input_data_format: Optional[Union[str, ChannelDimension]] = None,383    ) -> np.ndarray:384        """385        Apply an affine transformation to an image.386 387        Args:388            image (`np.ndarray`):389                Image to transform.390            center (`tuple[float]`):391                Center of the bounding box (x, y).392            scale (`tuple[float]`):393                Scale of the bounding box with respect to height/width.394            rotation (`float`):395                Rotation angle in degrees.396            size (`dict[str, int]`):397                Size of the destination image.398            data_format (`ChannelDimension`, *optional*, defaults to `ChannelDimension.FIRST`):399                The channel dimension format of the output image.400            input_data_format (`str` or `ChannelDimension`, *optional*):401                The channel dimension format of the input image.402        """403 404        data_format = input_data_format if data_format is None else data_format405 406        size = (size["width"], size["height"])407 408        # one uses a pixel standard deviation of 200 pixels409        transformation = get_warp_matrix(rotation, center * 2.0, np.array(size) - 1.0, scale * 200.0)410 411        # input image requires channels last format412        image = (413            image414            if input_data_format == ChannelDimension.LAST415            else to_channel_dimension_format(image, ChannelDimension.LAST, input_data_format)416        )417        image = scipy_warp_affine(src=image, M=transformation, size=(size[1], size[0]))418 419        image = to_channel_dimension_format(image, data_format, ChannelDimension.LAST)420 421        return image422 423    def preprocess(424        self,425        images: ImageInput,426        boxes: Union[list[list[float]], np.ndarray],427        do_affine_transform: Optional[bool] = None,428        size: Optional[dict[str, int]] = None,429        do_rescale: Optional[bool] = None,430        rescale_factor: Optional[float] = None,431        do_normalize: Optional[bool] = None,432        image_mean: Optional[Union[float, list[float]]] = None,433        image_std: Optional[Union[float, list[float]]] = None,434        return_tensors: Optional[Union[str, TensorType]] = None,435        data_format: Union[str, ChannelDimension] = ChannelDimension.FIRST,436        input_data_format: Optional[Union[str, ChannelDimension]] = None,437    ) -> PIL.Image.Image:438        """439        Preprocess an image or batch of images.440 441        Args:442            images (`ImageInput`):443                Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If444                passing in images with pixel values between 0 and 1, set `do_rescale=False`.445 446            boxes (`list[list[list[float]]]` or `np.ndarray`):447                List or array of bounding boxes for each image. Each box should be a list of 4 floats representing the bounding448                box coordinates in COCO format (top_left_x, top_left_y, width, height).449 450            do_affine_transform (`bool`, *optional*, defaults to `self.do_affine_transform`):451                Whether to apply an affine transformation to the input images.452            size (`dict[str, int]` *optional*, defaults to `self.size`):453                Dictionary in the format `{"height": h, "width": w}` specifying the size of the output image after454                resizing.455            do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):456                Whether to rescale the image values between [0 - 1].457            rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):458                Rescale factor to rescale the image by if `do_rescale` is set to `True`.459            do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):460                Whether to normalize the image.461            image_mean (`float` or `list[float]`, *optional*, defaults to `self.image_mean`):462                Image mean to use if `do_normalize` is set to `True`.463            image_std (`float` or `list[float]`, *optional*, defaults to `self.image_std`):464                Image standard deviation to use if `do_normalize` is set to `True`.465            return_tensors (`str` or [`~utils.TensorType`], *optional*, defaults to `'np'`):466                If set, will return tensors of a particular framework. Acceptable values are:467 468                - `'tf'`: Return TensorFlow `tf.constant` objects.469                - `'pt'`: Return PyTorch `torch.Tensor` objects.470                - `'np'`: Return NumPy `np.ndarray` objects.471                - `'jax'`: Return JAX `jnp.ndarray` objects.472 473        Returns:474            [`BatchFeature`]: A [`BatchFeature`] with the following fields:475 476            - **pixel_values** -- Pixel values to be fed to a model, of shape (batch_size, num_channels, height,477              width).478        """479        do_affine_transform = do_affine_transform if do_affine_transform is not None else self.do_affine_transform480        size = size if size is not None else self.size481        do_rescale = do_rescale if do_rescale is not None else self.do_rescale482        rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor483        do_normalize = do_normalize if do_normalize is not None else self.do_normalize484        image_mean = image_mean if image_mean is not None else self.image_mean485        image_std = image_std if image_std is not None else self.image_std486 487        images = make_flat_list_of_images(images)488 489        if not valid_images(images):490            raise ValueError(491                "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "492                "torch.Tensor, tf.Tensor or jax.ndarray."493            )494 495        if isinstance(boxes, list) and len(images) != len(boxes):496            raise ValueError(f"Batch of images and boxes mismatch : {len(images)} != {len(boxes)}")497        elif isinstance(boxes, np.ndarray) and len(images) != boxes.shape[0]:498            raise ValueError(f"Batch of images and boxes mismatch : {len(images)} != {boxes.shape[0]}")499 500        # All transformations expect numpy arrays.501        images = [to_numpy_array(image) for image in images]502 503        if is_scaled_image(images[0]) and do_rescale:504            logger.warning_once(505                "It looks like you are trying to rescale already rescaled images. If the input"506                " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again."507            )508 509        if input_data_format is None:510            # We assume that all images have the same channel dimension format.511            input_data_format = infer_channel_dimension_format(images[0])512 513        # transformations (affine transformation + rescaling + normalization)514        if self.do_affine_transform:515            new_images = []516            for image, image_boxes in zip(images, boxes):517                for box in image_boxes:518                    center, scale = box_to_center_and_scale(519                        box,520                        image_width=size["width"],521                        image_height=size["height"],522                        normalize_factor=self.normalize_factor,523                    )524                    transformed_image = self.affine_transform(525                        image, center, scale, rotation=0, size=size, input_data_format=input_data_format526                    )527                    new_images.append(transformed_image)528            images = new_images529 530        # For batch processing, the number of boxes must be consistent across all images in the batch.531        # When using a list input, the number of boxes can vary dynamically per image.532        # The image processor creates pixel_values of shape (batch_size*num_persons, num_channels, height, width)533 534        all_images = []535        for image in images:536            if do_rescale:537                image = self.rescale(image=image, scale=rescale_factor, input_data_format=input_data_format)538 539            if do_normalize:540                image = self.normalize(541                    image=image, mean=image_mean, std=image_std, input_data_format=input_data_format542                )543 544            all_images.append(image)545        images = [546            to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)547            for image in all_images548        ]549 550        data = {"pixel_values": images}551        encoded_inputs = BatchFeature(data=data, tensor_type=return_tensors)552 553        return encoded_inputs554 555    def keypoints_from_heatmaps(556        self,557        heatmaps: np.ndarray,558        center: np.ndarray,559        scale: np.ndarray,560        kernel: int = 11,561    ):562        """563        Get final keypoint predictions from heatmaps and transform them back to564        the image.565 566        Args:567            heatmaps (`np.ndarray` of shape `(batch_size, num_keypoints, height, width])`):568                Model predicted heatmaps.569            center (`np.ndarray` of shape `(batch_size, 2)`):570                Center of the bounding box (x, y).571            scale (`np.ndarray` of shape `(batch_size, 2)`):572                Scale of the bounding box wrt original images of width and height.573            kernel (int, *optional*, defaults to 11):574                Gaussian kernel size (K) for modulation, which should match the heatmap gaussian sigma when training.575                K=17 for sigma=3 and k=11 for sigma=2.576 577        Returns:578            tuple: A tuple containing keypoint predictions and scores.579 580            - preds (`np.ndarray` of shape `(batch_size, num_keypoints, 2)`):581                Predicted keypoint location in images.582            - scores (`np.ndarray` of shape `(batch_size, num_keypoints, 1)`):583                Scores (confidence) of the keypoints.584        """585        batch_size, _, height, width = heatmaps.shape586 587        coords, scores = get_keypoint_predictions(heatmaps)588 589        preds = post_dark_unbiased_data_processing(coords, heatmaps, kernel=kernel)590 591        # Transform back to the image592        for i in range(batch_size):593            preds[i] = transform_preds(preds[i], center=center[i], scale=scale[i], output_size=[height, width])594 595        return preds, scores596 597    def post_process_pose_estimation(598        self,599        outputs: "VitPoseEstimatorOutput",600        boxes: Union[list[list[list[float]]], np.ndarray],601        kernel_size: int = 11,602        threshold: Optional[float] = None,603        target_sizes: Union[TensorType, list[tuple]] = None,604    ):605        """606        Transform the heatmaps into keypoint predictions and transform them back to the image.607 608        Args:609            outputs (`VitPoseEstimatorOutput`):610                VitPoseForPoseEstimation model outputs.611            boxes (`list[list[list[float]]]` or `np.ndarray`):612                List or array of bounding boxes for each image. Each box should be a list of 4 floats representing the bounding613                box coordinates in COCO format (top_left_x, top_left_y, width, height).614            kernel_size (`int`, *optional*, defaults to 11):615                Gaussian kernel size (K) for modulation.616            threshold (`float`, *optional*, defaults to None):617                Score threshold to keep object detection predictions.618            target_sizes (`torch.Tensor` or `list[tuple[int, int]]`, *optional*):619                Tensor of shape `(batch_size, 2)` or list of tuples (`tuple[int, int]`) containing the target size620                `(height, width)` of each image in the batch. If unset, predictions will be resize with the default value.621        Returns:622            `list[list[Dict]]`: A list of dictionaries, each dictionary containing the keypoints and boxes for an image623            in the batch as predicted by the model.624        """625 626        # First compute centers and scales for each bounding box627        batch_size, num_keypoints, _, _ = outputs.heatmaps.shape628 629        if target_sizes is not None:630            if batch_size != len(target_sizes):631                raise ValueError(632                    "Make sure that you pass in as many target sizes as the batch dimension of the logits"633                )634 635        centers = np.zeros((batch_size, 2), dtype=np.float32)636        scales = np.zeros((batch_size, 2), dtype=np.float32)637        flattened_boxes = list(itertools.chain(*boxes))638        for i in range(batch_size):639            if target_sizes is not None:640                image_width, image_height = target_sizes[i][0], target_sizes[i][1]641                scale_factor = np.array([image_width, image_height, image_width, image_height])642                flattened_boxes[i] = flattened_boxes[i] * scale_factor643            width, height = self.size["width"], self.size["height"]644            center, scale = box_to_center_and_scale(flattened_boxes[i], image_width=width, image_height=height)645            centers[i, :] = center646            scales[i, :] = scale647 648        preds, scores = self.keypoints_from_heatmaps(649            outputs.heatmaps.cpu().numpy(), centers, scales, kernel=kernel_size650        )651 652        all_boxes = np.zeros((batch_size, 4), dtype=np.float32)653        all_boxes[:, 0:2] = centers[:, 0:2]654        all_boxes[:, 2:4] = scales[:, 0:2]655 656        poses = torch.tensor(preds)657        scores = torch.tensor(scores)658        labels = torch.arange(0, num_keypoints)659        bboxes_xyxy = torch.tensor(coco_to_pascal_voc(all_boxes))660 661        results: list[list[dict[str, torch.Tensor]]] = []662 663        pose_bbox_pairs = zip(poses, scores, bboxes_xyxy)664 665        for image_bboxes in boxes:666            image_results: list[dict[str, torch.Tensor]] = []667            for _ in image_bboxes:668                # Unpack the next pose and bbox_xyxy from the iterator669                pose, score, bbox_xyxy = next(pose_bbox_pairs)670                score = score.squeeze()671                keypoints_labels = labels672                if threshold is not None:673                    keep = score > threshold674                    pose = pose[keep]675                    score = score[keep]676                    keypoints_labels = keypoints_labels[keep]677                pose_result = {"keypoints": pose, "scores": score, "labels": keypoints_labels, "bbox": bbox_xyxy}678                image_results.append(pose_result)679            results.append(image_results)680 681        return results682 683 684__all__ = ["VitPoseImageProcessor"]685