Aluode/PerceptionLabPortable

1# coding=utf-82# Copyright 2025 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"""Fast Image processor class for Vilt."""16 17from typing import Optional, Union18 19import torch20from torchvision.transforms.v2 import functional as F21 22from ...image_processing_utils import BatchFeature23from ...image_processing_utils_fast import (24    BaseImageProcessorFast,25    DefaultFastImageProcessorKwargs,26    get_max_height_width,27    group_images_by_shape,28    reorder_images,29)30from ...image_utils import IMAGENET_STANDARD_MEAN, IMAGENET_STANDARD_STD, PILImageResampling, SizeDict31from ...utils import (32    TensorType,33    auto_docstring,34)35 36 37# Set maximum size based on the typical aspect ratio of the COCO dataset38MAX_LONGER_EDGE = 133339MAX_SHORTER_EDGE = 80040 41 42class ViltFastImageProcessorKwargs(DefaultFastImageProcessorKwargs):43    """44    Args:45        size_divisor (`int`, *optional*, defaults to 32):46            The size to make the height and width divisible by.47        rescale_factor (`float`, *optional*, defaults to 1/255):48            The factor to rescale the image by.49    """50 51    size_divisor: Optional[int]52    rescale_factor: Optional[float]53 54 55@auto_docstring56class ViltImageProcessorFast(BaseImageProcessorFast):57    resample = PILImageResampling.BICUBIC58    image_mean = IMAGENET_STANDARD_MEAN59    image_std = IMAGENET_STANDARD_STD60    size = {"shortest_edge": 384}61    do_resize = True62    do_rescale = True63    do_normalize = True64    size_divisor = 3265    do_pad = True66    default_to_square = False67    model_input_names = ["pixel_values", "pixel_mask"]68    valid_kwargs = ViltFastImageProcessorKwargs69 70    def _preprocess(71        self,72        images: list["torch.Tensor"],73        do_resize: bool,74        size: SizeDict,75        interpolation: Optional["F.InterpolationMode"],76        size_divisor: Optional[int],77        do_pad: bool,78        do_rescale: bool,79        rescale_factor: float,80        do_normalize: bool,81        image_mean: Optional[Union[float, list[float]]],82        image_std: Optional[Union[float, list[float]]],83        disable_grouping: Optional[bool],84        return_tensors: Optional[Union[str, TensorType]],85        **kwargs,86    ) -> BatchFeature:87        """88        Preprocess an image or batch of images.89 90        This method overrides the base class method to include padding and pixel mask generation.91        """92        # Group images by size for batched resizing93        grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping)94        resized_images_grouped = {}95 96        for shape, stacked_images in grouped_images.items():97            if do_resize:98                stacked_images = self.resize(stacked_images, size, interpolation, size_divisor)99            resized_images_grouped[shape] = stacked_images100        resized_images = reorder_images(resized_images_grouped, grouped_images_index)101 102        # Group images by size for further processing103        grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping)104        processed_images_grouped = {}105 106        for shape, stacked_images in grouped_images.items():107            # Fused rescale and normalize108            stacked_images = self.rescale_and_normalize(109                stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std110            )111            processed_images_grouped[shape] = stacked_images112 113        processed_images = reorder_images(processed_images_grouped, grouped_images_index)114 115        # Handle padding if required116        data = {}117        if do_pad:118            pixel_values, pixel_mask = self._pad_batch(119                processed_images, return_tensors, disable_grouping=disable_grouping120            )121            data = {"pixel_values": pixel_values, "pixel_mask": pixel_mask}122        else:123            # If no padding, just return the processed images124            if return_tensors == "pt":125                processed_images = torch.stack(processed_images)126            data = {"pixel_values": processed_images}127 128        return BatchFeature(data=data, tensor_type=return_tensors)129 130    def resize(131        self,132        images: "torch.Tensor",133        size: SizeDict,134        interpolation: Optional["F.InterpolationMode"] = None,135        size_divisor: Optional[int] = None,136    ) -> "torch.Tensor":137        """138        Resize an image or batch of images to specified size.139 140        Args:141            images (`torch.Tensor`): Image or batch of images to resize.142            size (`dict[str, int]`): Size dictionary with shortest_edge key.143            interpolation (`F.InterpolationMode`, *optional*): Interpolation method to use.144            size_divisor (`int`, *optional*): Value to ensure height/width are divisible by.145 146        Returns:147            `torch.Tensor`: Resized image or batch of images.148        """149        if interpolation is None:150            interpolation = self.resample151 152        # Resize with aspect ratio preservation153        shorter = size.shortest_edge154        longer = int(MAX_LONGER_EDGE / MAX_SHORTER_EDGE * shorter)155 156        heights = images.shape[-2]157        widths = images.shape[-1]158 159        # Determine the new dimensions160        if heights < widths:161            new_heights = shorter162            new_widths = widths * (shorter / heights)163        else:164            new_heights = heights * (shorter / widths)165            new_widths = shorter166 167        # Check if the longer side exceeds max size168        if max(new_heights, new_widths) > longer:169            scale = longer / max(new_heights, new_widths)170            new_heights = new_heights * scale171            new_widths = new_widths * scale172 173        new_heights = int(new_heights + 0.5)174        new_widths = int(new_widths + 0.5)175 176        # Make dimensions divisible by size_divisor177        if size_divisor is not None:178            new_heights = new_heights // size_divisor * size_divisor179            new_widths = new_widths // size_divisor * size_divisor180 181        # Resize the image182        return F.resize(images, [new_heights, new_widths], interpolation=interpolation)183 184    def _pad_batch(185        self,186        images: list["torch.Tensor"],187        return_tensors: Optional[Union[str, TensorType]],188        disable_grouping: Optional[bool],189    ) -> tuple:190        """191        Pad a batch of images to the same size based on the maximum dimensions.192 193        Args:194            images (`list[torch.Tensor]`): List of images to pad.195            return_tensors (`str` or `TensorType`, *optional*): The type of tensors to return.196 197        Returns:198            `tuple`: Tuple containing padded images and pixel masks.199        """200        # Calculate global maximum dimensions across all images201        max_size = get_max_height_width(images)202 203        # Group images by shape before padding204        grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping)205        processed_images = {}206        processed_masks = {}207 208        for shape, stacked_images in grouped_images.items():209            # Create mask template for efficient masking210            if return_tensors == "pt" and len(stacked_images) > 0:211                device = stacked_images.device212                mask_template = torch.zeros(max_size, dtype=torch.int64, device=device)213 214            original_size = stacked_images.shape[-2:]215            needs_padding = original_size[0] != max_size[0] or original_size[1] != max_size[1]216 217            if needs_padding:218                padding_bottom = max_size[0] - original_size[0]219                padding_right = max_size[1] - original_size[1]220                padding = [0, 0, padding_right, padding_bottom]221 222                padded_images = F.pad(stacked_images, padding, fill=0)223                pixel_mask = mask_template.clone()224                pixel_mask[: original_size[0], : original_size[1]].fill_(1)225                pixel_masks = pixel_mask.unsqueeze(0).repeat(stacked_images.shape[0], 1, 1)226            else:227                padded_images = stacked_images228                pixel_masks = torch.ones(229                    (stacked_images.shape[0], max_size[0], max_size[1]),230                    dtype=torch.int64,231                    device=stacked_images.device,232                )233 234            # Store processed group235            processed_images[shape] = padded_images236            processed_masks[shape] = pixel_masks237 238        # Reorder images back to original order239        padded_images = reorder_images(processed_images, grouped_images_index)240        pixel_masks = reorder_images(processed_masks, grouped_images_index)241 242        # Stack if tensors are requested for final result243        if return_tensors == "pt" and padded_images:244            padded_images = torch.stack(padded_images)245            pixel_masks = torch.stack(pixel_masks)246 247        return padded_images, pixel_masks248 249 250__all__ = ["ViltImageProcessorFast"]251