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"""Fast Image processor class for ConvNeXT."""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    group_images_by_shape,27    reorder_images,28)29from ...image_transforms import get_resize_output_image_size30from ...image_utils import (31    IMAGENET_STANDARD_MEAN,32    IMAGENET_STANDARD_STD,33    ChannelDimension,34    ImageInput,35    PILImageResampling,36)37from ...processing_utils import Unpack38from ...utils import (39    TensorType,40    auto_docstring,41)42 43 44class ConvNextFastImageProcessorKwargs(DefaultFastImageProcessorKwargs):45    """46    crop_pct (`float`, *optional*):47        Percentage of the image to crop. Only has an effect if size < 384. Can be48        overridden by `crop_pct` in the`preprocess` method.49    """50 51    crop_pct: Optional[float]52 53 54@auto_docstring55class ConvNextImageProcessorFast(BaseImageProcessorFast):56    resample = PILImageResampling.BILINEAR57    image_mean = IMAGENET_STANDARD_MEAN58    image_std = IMAGENET_STANDARD_STD59    size = {"shortest_edge": 384}60    default_to_square = False61    do_resize = True62    do_rescale = True63    do_normalize = True64    crop_pct = 224 / 25665    valid_kwargs = ConvNextFastImageProcessorKwargs66 67    def __init__(self, **kwargs: Unpack[ConvNextFastImageProcessorKwargs]):68        super().__init__(**kwargs)69 70    @auto_docstring71    def preprocess(self, images: ImageInput, **kwargs: Unpack[ConvNextFastImageProcessorKwargs]) -> BatchFeature:72        return super().preprocess(images, **kwargs)73 74    def resize(75        self,76        image: "torch.Tensor",77        size: dict[str, int],78        crop_pct: float,79        interpolation: PILImageResampling = PILImageResampling.BICUBIC,80        **kwargs,81    ) -> "torch.Tensor":82        """83        Resize an image.84 85        Args:86            image (`torch.Tensor`):87                Image to resize.88            size (`dict[str, int]`):89                Dictionary of the form `{"shortest_edge": int}`, specifying the size of the output image. If90                `size["shortest_edge"]` >= 384 image is resized to `(size["shortest_edge"], size["shortest_edge"])`.91                Otherwise, the smaller edge of the image will be matched to `int(size["shortest_edge"] / crop_pct)`,92                after which the image is cropped to `(size["shortest_edge"], size["shortest_edge"])`.93            crop_pct (`float`):94                Percentage of the image to crop. Only has an effect if size < 384.95            resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BICUBIC`):96                Resampling filter to use when resizing the image.97 98        Returns:99            `torch.Tensor`: Resized image.100        """101        if not size.shortest_edge:102            raise ValueError(f"Size dictionary must contain 'shortest_edge' key. Got {size.keys()}")103        shortest_edge = size["shortest_edge"]104 105        if shortest_edge < 384:106            # maintain same ratio, resizing shortest edge to shortest_edge/crop_pct107            resize_shortest_edge = int(shortest_edge / crop_pct)108            resize_size = get_resize_output_image_size(109                image, size=resize_shortest_edge, default_to_square=False, input_data_format=ChannelDimension.FIRST110            )111            image = F.resize(112                image,113                resize_size,114                interpolation=interpolation,115                **kwargs,116            )117            # then crop to (shortest_edge, shortest_edge)118            return F.center_crop(119                image,120                (shortest_edge, shortest_edge),121                **kwargs,122            )123        else:124            # warping (no cropping) when evaluated at 384 or larger125            return F.resize(126                image,127                (shortest_edge, shortest_edge),128                interpolation=interpolation,129                **kwargs,130            )131 132    def _preprocess(133        self,134        images: list["torch.Tensor"],135        do_resize: bool,136        size: dict[str, int],137        crop_pct: float,138        interpolation: Optional["F.InterpolationMode"],139        do_center_crop: bool,140        crop_size: int,141        do_rescale: bool,142        rescale_factor: float,143        do_normalize: bool,144        image_mean: Optional[Union[float, list[float]]],145        image_std: Optional[Union[float, list[float]]],146        disable_grouping: Optional[bool],147        return_tensors: Optional[Union[str, TensorType]],148        **kwargs,149    ) -> BatchFeature:150        # Group images by size for batched resizing151        grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping)152        resized_images_grouped = {}153        for shape, stacked_images in grouped_images.items():154            if do_resize:155                stacked_images = self.resize(156                    image=stacked_images, size=size, crop_pct=crop_pct, interpolation=interpolation157                )158            resized_images_grouped[shape] = stacked_images159        resized_images = reorder_images(resized_images_grouped, grouped_images_index)160 161        # Group images by size for further processing162        # Needed in case do_resize is False, or resize returns images with different sizes163        grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping)164        processed_images_grouped = {}165        for shape, stacked_images in grouped_images.items():166            if do_center_crop:167                stacked_images = self.center_crop(stacked_images, crop_size)168            # Fused rescale and normalize169            stacked_images = self.rescale_and_normalize(170                stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std171            )172            processed_images_grouped[shape] = stacked_images173 174        processed_images = reorder_images(processed_images_grouped, grouped_images_index)175        processed_images = torch.stack(processed_images, dim=0) if return_tensors else processed_images176 177        return BatchFeature(data={"pixel_values": processed_images}, tensor_type=return_tensors)178 179 180__all__ = ["ConvNextImageProcessorFast"]181