Aluode/PerceptionLabPortable

1# coding=utf-82# Copyright 2022 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 MobileViT."""16 17from typing import Optional, Union18 19import numpy as np20 21from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict22from ...image_transforms import flip_channel_order, get_resize_output_image_size, resize, to_channel_dimension_format23from ...image_utils import (24    ChannelDimension,25    ImageInput,26    PILImageResampling,27    infer_channel_dimension_format,28    is_scaled_image,29    make_flat_list_of_images,30    to_numpy_array,31    valid_images,32    validate_preprocess_arguments,33)34from ...utils import (35    TensorType,36    filter_out_non_signature_kwargs,37    is_torch_available,38    is_torch_tensor,39    is_vision_available,40    logging,41)42from ...utils.import_utils import requires43 44 45if is_vision_available():46    import PIL47 48if is_torch_available():49    import torch50 51 52logger = logging.get_logger(__name__)53 54 55@requires(backends=("vision",))56class MobileViTImageProcessor(BaseImageProcessor):57    r"""58    Constructs a MobileViT image processor.59 60    Args:61        do_resize (`bool`, *optional*, defaults to `True`):62            Whether to resize the image's (height, width) dimensions to the specified `size`. Can be overridden by the63            `do_resize` parameter in the `preprocess` method.64        size (`dict[str, int]` *optional*, defaults to `{"shortest_edge": 224}`):65            Controls the size of the output image after resizing. Can be overridden by the `size` parameter in the66            `preprocess` method.67        resample (`PILImageResampling`, *optional*, defaults to `Resampling.BILINEAR`):68            Defines the resampling filter to use if resizing the image. Can be overridden by the `resample` parameter69            in the `preprocess` method.70        do_rescale (`bool`, *optional*, defaults to `True`):71            Whether to rescale the image by the specified scale `rescale_factor`. Can be overridden by the `do_rescale`72            parameter in the `preprocess` method.73        rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):74            Scale factor to use if rescaling the image. Can be overridden by the `rescale_factor` parameter in the75            `preprocess` method.76        do_center_crop (`bool`, *optional*, defaults to `True`):77            Whether to crop the input at the center. If the input size is smaller than `crop_size` along any edge, the78            image is padded with 0's and then center cropped. Can be overridden by the `do_center_crop` parameter in79            the `preprocess` method.80        crop_size (`dict[str, int]`, *optional*, defaults to `{"height": 256, "width": 256}`):81            Desired output size `(size["height"], size["width"])` when applying center-cropping. Can be overridden by82            the `crop_size` parameter in the `preprocess` method.83        do_flip_channel_order (`bool`, *optional*, defaults to `True`):84            Whether to flip the color channels from RGB to BGR. Can be overridden by the `do_flip_channel_order`85            parameter in the `preprocess` method.86        do_reduce_labels (`bool`, *optional*, defaults to `False`):87            Whether or not to reduce all label values of segmentation maps by 1. Usually used for datasets where 0 is88            used for background, and background itself is not included in all classes of a dataset (e.g. ADE20k). The89            background label will be replaced by 255. Can be overridden by the `do_reduce_labels` parameter in the90            `preprocess` method.91    """92 93    model_input_names = ["pixel_values"]94 95    def __init__(96        self,97        do_resize: bool = True,98        size: Optional[dict[str, int]] = None,99        resample: PILImageResampling = PILImageResampling.BILINEAR,100        do_rescale: bool = True,101        rescale_factor: Union[int, float] = 1 / 255,102        do_center_crop: bool = True,103        crop_size: Optional[dict[str, int]] = None,104        do_flip_channel_order: bool = True,105        do_reduce_labels: bool = False,106        **kwargs,107    ) -> None:108        super().__init__(**kwargs)109        size = size if size is not None else {"shortest_edge": 224}110        size = get_size_dict(size, default_to_square=False)111        crop_size = crop_size if crop_size is not None else {"height": 256, "width": 256}112        crop_size = get_size_dict(crop_size, param_name="crop_size")113 114        self.do_resize = do_resize115        self.size = size116        self.resample = resample117        self.do_rescale = do_rescale118        self.rescale_factor = rescale_factor119        self.do_center_crop = do_center_crop120        self.crop_size = crop_size121        self.do_flip_channel_order = do_flip_channel_order122        self.do_reduce_labels = do_reduce_labels123 124    # Copied from transformers.models.mobilenet_v1.image_processing_mobilenet_v1.MobileNetV1ImageProcessor.resize with PILImageResampling.BICUBIC->PILImageResampling.BILINEAR125    def resize(126        self,127        image: np.ndarray,128        size: dict[str, int],129        resample: PILImageResampling = PILImageResampling.BILINEAR,130        data_format: Optional[Union[str, ChannelDimension]] = None,131        input_data_format: Optional[Union[str, ChannelDimension]] = None,132        **kwargs,133    ) -> np.ndarray:134        """135        Resize an image. The shortest edge of the image is resized to size["shortest_edge"], with the longest edge136        resized to keep the input aspect ratio.137 138        Args:139            image (`np.ndarray`):140                Image to resize.141            size (`dict[str, int]`):142                Size of the output image.143            resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BILINEAR`):144                Resampling filter to use when resiizing the image.145            data_format (`str` or `ChannelDimension`, *optional*):146                The channel dimension format of the image. If not provided, it will be the same as the input image.147            input_data_format (`ChannelDimension` or `str`, *optional*):148                The channel dimension format of the input image. If not provided, it will be inferred.149        """150        default_to_square = True151        if "shortest_edge" in size:152            size = size["shortest_edge"]153            default_to_square = False154        elif "height" in size and "width" in size:155            size = (size["height"], size["width"])156        else:157            raise ValueError("Size must contain either 'shortest_edge' or 'height' and 'width'.")158 159        output_size = get_resize_output_image_size(160            image,161            size=size,162            default_to_square=default_to_square,163            input_data_format=input_data_format,164        )165        return resize(166            image,167            size=output_size,168            resample=resample,169            data_format=data_format,170            input_data_format=input_data_format,171            **kwargs,172        )173 174    def flip_channel_order(175        self,176        image: np.ndarray,177        data_format: Optional[Union[str, ChannelDimension]] = None,178        input_data_format: Optional[Union[str, ChannelDimension]] = None,179    ) -> np.ndarray:180        """181        Flip the color channels from RGB to BGR or vice versa.182 183        Args:184            image (`np.ndarray`):185                The image, represented as a numpy array.186            data_format (`ChannelDimension` or `str`, *optional*):187                The channel dimension format of the image. If not provided, it will be the same as the input image.188            input_data_format (`ChannelDimension` or `str`, *optional*):189                The channel dimension format of the input image. If not provided, it will be inferred.190        """191        return flip_channel_order(image, data_format=data_format, input_data_format=input_data_format)192 193    # Copied from transformers.models.beit.image_processing_beit.BeitImageProcessor.reduce_label194    def reduce_label(self, label: ImageInput) -> np.ndarray:195        label = to_numpy_array(label)196        # Avoid using underflow conversion197        label[label == 0] = 255198        label = label - 1199        label[label == 254] = 255200        return label201 202    def __call__(self, images, segmentation_maps=None, **kwargs):203        """204        Preprocesses a batch of images and optionally segmentation maps.205 206        Overrides the `__call__` method of the `Preprocessor` class so that both images and segmentation maps can be207        passed in as positional arguments.208        """209        return super().__call__(images, segmentation_maps=segmentation_maps, **kwargs)210 211    def _preprocess(212        self,213        image: ImageInput,214        do_reduce_labels: bool,215        do_resize: bool,216        do_rescale: bool,217        do_center_crop: bool,218        do_flip_channel_order: bool,219        size: Optional[dict[str, int]] = None,220        resample: Optional[PILImageResampling] = None,221        rescale_factor: Optional[float] = None,222        crop_size: Optional[dict[str, int]] = None,223        input_data_format: Optional[Union[str, ChannelDimension]] = None,224    ):225        if do_reduce_labels:226            image = self.reduce_label(image)227 228        if do_resize:229            image = self.resize(image=image, size=size, resample=resample, input_data_format=input_data_format)230 231        if do_rescale:232            image = self.rescale(image=image, scale=rescale_factor, input_data_format=input_data_format)233 234        if do_center_crop:235            image = self.center_crop(image=image, size=crop_size, input_data_format=input_data_format)236 237        if do_flip_channel_order:238            image = self.flip_channel_order(image, input_data_format=input_data_format)239 240        return image241 242    def _preprocess_image(243        self,244        image: ImageInput,245        do_resize: Optional[bool] = None,246        size: Optional[dict[str, int]] = None,247        resample: Optional[PILImageResampling] = None,248        do_rescale: Optional[bool] = None,249        rescale_factor: Optional[float] = None,250        do_center_crop: Optional[bool] = None,251        crop_size: Optional[dict[str, int]] = None,252        do_flip_channel_order: Optional[bool] = None,253        data_format: Optional[Union[str, ChannelDimension]] = None,254        input_data_format: Optional[Union[str, ChannelDimension]] = None,255    ) -> np.ndarray:256        """Preprocesses a single image."""257        # All transformations expect numpy arrays.258        image = to_numpy_array(image)259        if do_rescale and is_scaled_image(image):260            logger.warning_once(261                "It looks like you are trying to rescale already rescaled images. If the input"262                " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again."263            )264        if input_data_format is None:265            input_data_format = infer_channel_dimension_format(image)266 267        image = self._preprocess(268            image=image,269            do_reduce_labels=False,270            do_resize=do_resize,271            size=size,272            resample=resample,273            do_rescale=do_rescale,274            rescale_factor=rescale_factor,275            do_center_crop=do_center_crop,276            crop_size=crop_size,277            do_flip_channel_order=do_flip_channel_order,278            input_data_format=input_data_format,279        )280 281        image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)282 283        return image284 285    def _preprocess_mask(286        self,287        segmentation_map: ImageInput,288        do_reduce_labels: Optional[bool] = None,289        do_resize: Optional[bool] = None,290        size: Optional[dict[str, int]] = None,291        do_center_crop: Optional[bool] = None,292        crop_size: Optional[dict[str, int]] = None,293        input_data_format: Optional[Union[str, ChannelDimension]] = None,294    ) -> np.ndarray:295        """Preprocesses a single mask."""296        segmentation_map = to_numpy_array(segmentation_map)297        # Add channel dimension if missing - needed for certain transformations298        if segmentation_map.ndim == 2:299            added_channel_dim = True300            segmentation_map = segmentation_map[None, ...]301            input_data_format = ChannelDimension.FIRST302        else:303            added_channel_dim = False304            if input_data_format is None:305                input_data_format = infer_channel_dimension_format(segmentation_map, num_channels=1)306 307        segmentation_map = self._preprocess(308            image=segmentation_map,309            do_reduce_labels=do_reduce_labels,310            do_resize=do_resize,311            size=size,312            resample=PILImageResampling.NEAREST,313            do_rescale=False,314            do_center_crop=do_center_crop,315            crop_size=crop_size,316            do_flip_channel_order=False,317            input_data_format=input_data_format,318        )319        # Remove extra channel dimension if added for processing320        if added_channel_dim:321            segmentation_map = segmentation_map.squeeze(0)322        segmentation_map = segmentation_map.astype(np.int64)323        return segmentation_map324 325    @filter_out_non_signature_kwargs()326    def preprocess(327        self,328        images: ImageInput,329        segmentation_maps: Optional[ImageInput] = None,330        do_resize: Optional[bool] = None,331        size: Optional[dict[str, int]] = None,332        resample: Optional[PILImageResampling] = None,333        do_rescale: Optional[bool] = None,334        rescale_factor: Optional[float] = None,335        do_center_crop: Optional[bool] = None,336        crop_size: Optional[dict[str, int]] = None,337        do_flip_channel_order: Optional[bool] = None,338        do_reduce_labels: Optional[bool] = None,339        return_tensors: Optional[Union[str, TensorType]] = None,340        data_format: ChannelDimension = ChannelDimension.FIRST,341        input_data_format: Optional[Union[str, ChannelDimension]] = None,342    ) -> PIL.Image.Image:343        """344        Preprocess an image or batch of images.345 346        Args:347            images (`ImageInput`):348                Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If349                passing in images with pixel values between 0 and 1, set `do_rescale=False`.350            segmentation_maps (`ImageInput`, *optional*):351                Segmentation map to preprocess.352            do_resize (`bool`, *optional*, defaults to `self.do_resize`):353                Whether to resize the image.354            size (`dict[str, int]`, *optional*, defaults to `self.size`):355                Size of the image after resizing.356            resample (`int`, *optional*, defaults to `self.resample`):357                Resampling filter to use if resizing the image. This can be one of the enum `PILImageResampling`, Only358                has an effect if `do_resize` is set to `True`.359            do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):360                Whether to rescale the image by rescale factor.361            rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):362                Rescale factor to rescale the image by if `do_rescale` is set to `True`.363            do_center_crop (`bool`, *optional*, defaults to `self.do_center_crop`):364                Whether to center crop the image.365            crop_size (`dict[str, int]`, *optional*, defaults to `self.crop_size`):366                Size of the center crop if `do_center_crop` is set to `True`.367            do_flip_channel_order (`bool`, *optional*, defaults to `self.do_flip_channel_order`):368                Whether to flip the channel order of the image.369            do_reduce_labels (`bool`, *optional*, defaults to `self.do_reduce_labels`):370                Whether or not to reduce all label values of segmentation maps by 1. Usually used for datasets where 0371                is used for background, and background itself is not included in all classes of a dataset (e.g.372                ADE20k). The background label will be replaced by 255.373            return_tensors (`str` or `TensorType`, *optional*):374                The type of tensors to return. Can be one of:375                    - Unset: Return a list of `np.ndarray`.376                    - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.377                    - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.378                    - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.379                    - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.380            data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):381                The channel dimension format for the output image. Can be one of:382                    - `ChannelDimension.FIRST`: image in (num_channels, height, width) format.383                    - `ChannelDimension.LAST`: image in (height, width, num_channels) format.384            input_data_format (`ChannelDimension` or `str`, *optional*):385                The channel dimension format for the input image. If unset, the channel dimension format is inferred386                from the input image. Can be one of:387                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.388                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.389                - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.390        """391        do_resize = do_resize if do_resize is not None else self.do_resize392        resample = resample if resample is not None else self.resample393        do_rescale = do_rescale if do_rescale is not None else self.do_rescale394        rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor395        do_center_crop = do_center_crop if do_center_crop is not None else self.do_center_crop396        do_flip_channel_order = (397            do_flip_channel_order if do_flip_channel_order is not None else self.do_flip_channel_order398        )399 400        size = size if size is not None else self.size401        size = get_size_dict(size, default_to_square=False)402        crop_size = crop_size if crop_size is not None else self.crop_size403        crop_size = get_size_dict(crop_size, param_name="crop_size")404 405        do_reduce_labels = do_reduce_labels if do_reduce_labels is not None else self.do_reduce_labels406 407        images = make_flat_list_of_images(images)408 409        if segmentation_maps is not None:410            segmentation_maps = make_flat_list_of_images(segmentation_maps, expected_ndims=2)411 412        images = make_flat_list_of_images(images)413 414        if not valid_images(images):415            raise ValueError(416                "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "417                "torch.Tensor, tf.Tensor or jax.ndarray."418            )419 420        if segmentation_maps is not None and not valid_images(segmentation_maps):421            raise ValueError(422                "Invalid segmentation map type. Must be of type PIL.Image.Image, numpy.ndarray, "423                "torch.Tensor, tf.Tensor or jax.ndarray."424            )425 426        validate_preprocess_arguments(427            do_rescale=do_rescale,428            rescale_factor=rescale_factor,429            do_center_crop=do_center_crop,430            crop_size=crop_size,431            do_resize=do_resize,432            size=size,433            resample=resample,434        )435 436        images = [437            self._preprocess_image(438                image=img,439                do_resize=do_resize,440                size=size,441                resample=resample,442                do_rescale=do_rescale,443                rescale_factor=rescale_factor,444                do_center_crop=do_center_crop,445                crop_size=crop_size,446                do_flip_channel_order=do_flip_channel_order,447                data_format=data_format,448                input_data_format=input_data_format,449            )450            for img in images451        ]452 453        data = {"pixel_values": images}454 455        if segmentation_maps is not None:456            segmentation_maps = [457                self._preprocess_mask(458                    segmentation_map=segmentation_map,459                    do_reduce_labels=do_reduce_labels,460                    do_resize=do_resize,461                    size=size,462                    do_center_crop=do_center_crop,463                    crop_size=crop_size,464                    input_data_format=input_data_format,465                )466                for segmentation_map in segmentation_maps467            ]468 469            data["labels"] = segmentation_maps470 471        return BatchFeature(data=data, tensor_type=return_tensors)472 473    # Copied from transformers.models.beit.image_processing_beit.BeitImageProcessor.post_process_semantic_segmentation with Beit->MobileViT474    def post_process_semantic_segmentation(self, outputs, target_sizes: Optional[list[tuple]] = None):475        """476        Converts the output of [`MobileViTForSemanticSegmentation`] into semantic segmentation maps. Only supports PyTorch.477 478        Args:479            outputs ([`MobileViTForSemanticSegmentation`]):480                Raw outputs of the model.481            target_sizes (`list[Tuple]` of length `batch_size`, *optional*):482                List of tuples corresponding to the requested final size (height, width) of each prediction. If unset,483                predictions will not be resized.484 485        Returns:486            semantic_segmentation: `list[torch.Tensor]` of length `batch_size`, where each item is a semantic487            segmentation map of shape (height, width) corresponding to the target_sizes entry (if `target_sizes` is488            specified). Each entry of each `torch.Tensor` correspond to a semantic class id.489        """490        # TODO: add support for other frameworks491        logits = outputs.logits492 493        # Resize logits and compute semantic segmentation maps494        if target_sizes is not None:495            if len(logits) != len(target_sizes):496                raise ValueError(497                    "Make sure that you pass in as many target sizes as the batch dimension of the logits"498                )499 500            if is_torch_tensor(target_sizes):501                target_sizes = target_sizes.numpy()502 503            semantic_segmentation = []504 505            for idx in range(len(logits)):506                resized_logits = torch.nn.functional.interpolate(507                    logits[idx].unsqueeze(dim=0), size=target_sizes[idx], mode="bilinear", align_corners=False508                )509                semantic_map = resized_logits[0].argmax(dim=0)510                semantic_segmentation.append(semantic_map)511        else:512            semantic_segmentation = logits.argmax(dim=1)513            semantic_segmentation = [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])]514 515        return semantic_segmentation516 517 518__all__ = ["MobileViTImageProcessor"]519