codingutf8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 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 image processor class for beit import warnings from typing import any dict list optional tuple union import numpy as np from imageprocessingutils import baseimageprocessor batchfeature getsizedict from imagetransforms import resize tochanneldimensionformat from imageutils import imagenetstandardmean imagenetstandardstd channeldimension imageinput pilimageresampling inferchanneldimensionformat isscaledimage makelistofimages tonumpyarray validimages from utils import tensortype istorchavailable istorchtensor isvisionavailable logging if isvisionavailable import pil if istorchavailable import torch logger logging getloggername class beitimageprocessorbaseimageprocessor r constructs a beit image processor args doresize bool optional defaults to true whether to resize the image s height width dimensions to the specified size can be overridden by the doresize parameter in the preprocess method size dictstr int optional defaults to height 256 width 256 size of the output image after resizing can be overridden by the size parameter in the preprocess method resample pilimageresampling optional defaults to resampling bicubic resampling filter to use if resizing the image can be overridden by the resample parameter in the preprocess method docentercrop bool optional defaults to true whether to center crop the image if the input size is smaller than cropsize along any edge the image is padded with 0 s and then center cropped can be overridden by the docentercrop parameter in the preprocess method cropsize dictstr int optional defaults to height 224 width 224 desired output size when applying centercropping only has an effect if docentercrop is set to true can be overridden by the cropsize parameter in the preprocess method rescalefactor int or float optional defaults to 1255 scale factor to use if rescaling the image can be overridden by the rescalefactor parameter in the preprocess method dorescale bool optional defaults to true whether to rescale the image by the specified scale rescalefactor can be overridden by the dorescale parameter in the preprocess method donormalize bool optional defaults to true whether to normalize the image can be overridden by the donormalize parameter in the preprocess method imagemean float or listfloat optional defaults to imagenetstandardmean the mean to use if normalizing the image this is a float or list of floats of length of the number of channels of the image can be overridden by the imagemean parameter in the preprocess method imagestd float or listfloat optional defaults to imagenetstandardstd the standard deviation to use if normalizing the image this is a float or list of floats of length of the number of channels of the image can be overridden by the imagestd parameter in the preprocess method doreducelabels bool optional defaults to false 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 can be overridden by the doreducelabels parameter in the preprocess method overrides the fromdict method from the base class to make sure reducelabels is updated if image processor is created using fromdict and kwargs e g beitimageprocessor frompretrainedcheckpoint reducelabelstrue resize an image to sizeheight sizewidth args image np ndarray image to resize size dictstr int size of the output image resample pilimageresampling optional defaults to pil image bicubic resampling filter to use when resiizing the image dataformat str or channeldimension optional the channel dimension format of the image if not provided it will be the same as the input image inputdataformat str or channeldimension optional the channel dimension format of the input image if not provided it will be inferred avoid using underflow conversion preprocesses a single image all transformations expect numpy arrays image tonumpyarrayimage if isscaledimageimage and dorescale logger warningonce it looks like you are trying to rescale already rescaled images if the input images have pixel values between 0 and 1 set dorescalefalse to avoid rescaling them again if inputdataformat is none inputdataformat inferchanneldimensionformatimage image self preprocess image doreducelabelsfalse doresizedoresize sizesize resampleresample docentercropdocentercrop cropsizecropsize dorescaledorescale rescalefactorrescalefactor donormalizedonormalize imagemeanimagemean imagestdimagestd inputdataformatinputdataformat if dataformat is not none image tochanneldimensionformatimage dataformat inputchanneldiminputdataformat return image def preprocesssegmentationmap self segmentationmap imageinput doresize bool none size dictstr int none resample pilimageresampling none docentercrop bool none cropsize dictstr int none doreducelabels bool none inputdataformat optionalunionstr channeldimension none all transformations expect numpy arrays add an axis to the segmentation maps for transformations remove extra axis if added overrides the call method of the preprocessor class such that the images and segmentation maps can both be passed in as positional arguments preprocess an image or batch of images args images imageinput image to preprocess expects a single or batch of images with pixel values ranging from 0 to 255 if passing in images with pixel values between 0 and 1 set dorescalefalse doresize bool optional defaults to self doresize whether to resize the image size dictstr int optional defaults to self size size of the image after resizing resample int optional defaults to self resample resampling filter to use if resizing the image this can be one of the enum pilimageresampling only has an effect if doresize is set to true docentercrop bool optional defaults to self docentercrop whether to center crop the image cropsize dictstr int optional defaults to self cropsize size of the image after center crop if one edge the image is smaller than cropsize it will be padded with zeros and then cropped dorescale bool optional defaults to self dorescale whether to rescale the image values between 0 1 rescalefactor float optional defaults to self rescalefactor rescale factor to rescale the image by if dorescale is set to true donormalize bool optional defaults to self donormalize whether to normalize the image imagemean float or listfloat optional defaults to self imagemean image mean imagestd float or listfloat optional defaults to self imagestd image standard deviation doreducelabels bool optional defaults to self doreducelabels 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 returntensors str or tensortype optional the type of tensors to return can be one of unset return a list of np ndarray tensortype tensorflow or tf return a batch of type tf tensor tensortype pytorch or pt return a batch of type torch tensor tensortype numpy or np return a batch of type np ndarray tensortype jax or jax return a batch of type jax numpy ndarray dataformat channeldimension or str optional defaults to channeldimension first the channel dimension format for the output image can be one of channelsfirst or channeldimension first image in numchannels height width format channelslast or channeldimension last image in height width numchannels format unset use the channel dimension format of the input image inputdataformat channeldimension or str optional the channel dimension format for the input image if unset the channel dimension format is inferred from the input image can be one of channelsfirst or channeldimension first image in numchannels height width format channelslast or channeldimension last image in height width numchannels format none or channeldimension none image in height width format converts the output of beitforsemanticsegmentation into semantic segmentation maps only supports pytorch args outputs beitforsemanticsegmentation raw outputs of the model targetsizes listtuple of length batchsize optional list of tuples corresponding to the requested final size height width of each prediction if unset predictions will not be resized returns semanticsegmentation listtorch tensor of length batchsize where each item is a semantic segmentation map of shape height width corresponding to the targetsizes entry if targetsizes is specified each entry of each torch tensor correspond to a semantic class id todo add support for other frameworks resize logits and compute semantic segmentation maps coding utf 8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 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 image processor class for beit constructs a beit image processor args do_resize bool optional defaults to true whether to resize the image s height width dimensions to the specified size can be overridden by the do_resize parameter in the preprocess method size dict str int optional defaults to height 256 width 256 size of the output image after resizing can be overridden by the size parameter in the preprocess method resample pilimageresampling optional defaults to resampling bicubic resampling filter to use if resizing the image can be overridden by the resample parameter in the preprocess method do_center_crop bool optional defaults to true whether to center crop the image if the input size is smaller than crop_size along any edge the image is padded with 0 s and then center cropped can be overridden by the do_center_crop parameter in the preprocess method crop_size dict str int optional defaults to height 224 width 224 desired output size when applying center cropping only has an effect if do_center_crop is set to true can be overridden by the crop_size parameter in the preprocess method rescale_factor int or float optional defaults to 1 255 scale factor to use if rescaling the image can be overridden by the rescale_factor parameter in the preprocess method do_rescale bool optional defaults to true whether to rescale the image by the specified scale rescale_factor can be overridden by the do_rescale parameter in the preprocess method do_normalize bool optional defaults to true whether to normalize the image can be overridden by the do_normalize parameter in the preprocess method image_mean float or list float optional defaults to imagenet_standard_mean the mean to use if normalizing the image this is a float or list of floats of length of the number of channels of the image can be overridden by the image_mean parameter in the preprocess method image_std float or list float optional defaults to imagenet_standard_std the standard deviation to use if normalizing the image this is a float or list of floats of length of the number of channels of the image can be overridden by the image_std parameter in the preprocess method do_reduce_labels bool optional defaults to false 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 can be overridden by the do_reduce_labels parameter in the preprocess method overrides the from_dict method from the base class to make sure reduce_labels is updated if image processor is created using from_dict and kwargs e g beitimageprocessor from_pretrained checkpoint reduce_labels true resize an image to size height size width args image np ndarray image to resize size dict str int size of the output image resample pilimageresampling optional defaults to pil image bicubic resampling filter to use when resiizing the image data_format str or channeldimension optional the channel dimension format of the image if not provided it will be the same as the input image input_data_format str or channeldimension optional the channel dimension format of the input image if not provided it will be inferred avoid using underflow conversion preprocesses a single image all transformations expect numpy arrays preprocesses a single segmentation map all transformations expect numpy arrays add an axis to the segmentation maps for transformations remove extra axis if added overrides the __call__ method of the preprocessor class such that the images and segmentation maps can both be passed in as positional arguments preprocess an image or batch of images args images imageinput image to preprocess expects a single or batch of images with pixel values ranging from 0 to 255 if passing in images with pixel values between 0 and 1 set do_rescale false do_resize bool optional defaults to self do_resize whether to resize the image size dict str int optional defaults to self size size of the image after resizing resample int optional defaults to self resample resampling filter to use if resizing the image this can be one of the enum pilimageresampling only has an effect if do_resize is set to true do_center_crop bool optional defaults to self do_center_crop whether to center crop the image crop_size dict str int optional defaults to self crop_size size of the image after center crop if one edge the image is smaller than crop_size it will be padded with zeros and then cropped do_rescale bool optional defaults to self do_rescale whether to rescale the image values between 0 1 rescale_factor float optional defaults to self rescale_factor rescale factor to rescale the image by if do_rescale is set to true do_normalize bool optional defaults to self do_normalize whether to normalize the image image_mean float or list float optional defaults to self image_mean image mean image_std float or list float optional defaults to self image_std image standard deviation 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 return_tensors str or tensortype optional the type of tensors to return can be one of unset return a list of np ndarray tensortype tensorflow or tf return a batch of type tf tensor tensortype pytorch or pt return a batch of type torch tensor tensortype numpy or np return a batch of type np ndarray tensortype jax or jax return a batch of type jax numpy ndarray data_format channeldimension or str optional defaults to channeldimension first the channel dimension format for the output image can be one of channels_first or channeldimension first image in num_channels height width format channels_last or channeldimension last image in height width num_channels format unset use the channel dimension format of the input image input_data_format channeldimension or str optional the channel dimension format for the input image if unset the channel dimension format is inferred from the input image can be one of channels_first or channeldimension first image in num_channels height width format channels_last or channeldimension last image in height width num_channels format none or channeldimension none image in height width format converts the output of beitforsemanticsegmentation into semantic segmentation maps only supports pytorch args outputs beitforsemanticsegmentation 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 todo add support for other frameworks resize logits and compute semantic segmentation maps
import warnings
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import resize, to_channel_dimension_format
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    infer_channel_dimension_format,
    is_scaled_image,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_torch_available, is_torch_tensor, is_vision_available, logging


if is_vision_available():
    import PIL

if is_torch_available():
    import torch


logger = logging.get_logger(__name__)


class BeitImageProcessor(BaseImageProcessor):
    r

    model_input_names = ["pixel_values"]

    def __init__(
        self,
        do_resize: bool = True,
        size: Dict[str, int] = None,
        resample: PILImageResampling = PILImageResampling.BICUBIC,
        do_center_crop: bool = True,
        crop_size: Dict[str, int] = None,
        rescale_factor: Union[int, float] = 1 / 255,
        do_rescale: bool = True,
        do_normalize: bool = True,
        image_mean: Optional[Union[float, List[float]]] = None,
        image_std: Optional[Union[float, List[float]]] = None,
        do_reduce_labels: bool = False,
        **kwargs,
    ) -> None:
        if "reduce_labels" in kwargs:
            warnings.warn(
                "The `reduce_labels` parameter is deprecated and will be removed in a future version. Please use"
                " `do_reduce_labels` instead.",
                FutureWarning,
            )
            do_reduce_labels = kwargs.pop("reduce_labels")
        super().__init__(**kwargs)
        size = size if size is not None else {"height": 256, "width": 256}
        size = get_size_dict(size)
        crop_size = crop_size if crop_size is not None else {"height": 224, "width": 224}
        crop_size = get_size_dict(crop_size, param_name="crop_size")
        self.do_resize = do_resize
        self.size = size
        self.resample = resample
        self.do_center_crop = do_center_crop
        self.crop_size = crop_size
        self.do_rescale = do_rescale
        self.rescale_factor = rescale_factor
        self.do_normalize = do_normalize
        self.image_mean = image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
        self.image_std = image_std if image_std is not None else IMAGENET_STANDARD_STD
        self.do_reduce_labels = do_reduce_labels

    @classmethod
    def from_dict(cls, image_processor_dict: Dict[str, Any], **kwargs):
        
        image_processor_dict = image_processor_dict.copy()
        if "reduce_labels" in kwargs:
            image_processor_dict["reduce_labels"] = kwargs.pop("reduce_labels")
        return super().from_dict(image_processor_dict, **kwargs)

    def resize(
        self,
        image: np.ndarray,
        size: Dict[str, int],
        resample: PILImageResampling = PILImageResampling.BICUBIC,
        data_format: Optional[Union[str, ChannelDimension]] = None,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
        **kwargs,
    ) -> np.ndarray:
        
        size = get_size_dict(size, default_to_square=True, param_name="size")
        if "height" not in size or "width" not in size:
            raise ValueError(f"The `size` argument must contain `height` and `width` keys. Got {size.keys()}")
        return resize(
            image,
            size=(size["height"], size["width"]),
            resample=resample,
            data_format=data_format,
            input_data_format=input_data_format,
            **kwargs,
        )

    def reduce_label(self, label: ImageInput) -> np.ndarray:
        label = to_numpy_array(label)
        
        label[label == 0] = 255
        label = label - 1
        label[label == 254] = 255
        return label

    def _preprocess(
        self,
        image: ImageInput,
        do_reduce_labels: bool = None,
        do_resize: bool = None,
        size: Dict[str, int] = None,
        resample: PILImageResampling = None,
        do_center_crop: bool = None,
        crop_size: Dict[str, int] = None,
        do_rescale: bool = None,
        rescale_factor: float = None,
        do_normalize: bool = None,
        image_mean: Optional[Union[float, List[float]]] = None,
        image_std: Optional[Union[float, List[float]]] = None,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
    ):
        if do_reduce_labels:
            image = self.reduce_label(image)

        if do_resize:
            image = self.resize(image=image, size=size, resample=resample, input_data_format=input_data_format)

        if do_center_crop:
            image = self.center_crop(image=image, size=crop_size, input_data_format=input_data_format)

        if do_rescale:
            image = self.rescale(image=image, scale=rescale_factor, input_data_format=input_data_format)

        if do_normalize:
            image = self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)

        return image

    def _preprocess_image(
        self,
        image: ImageInput,
        do_resize: bool = None,
        size: Dict[str, int] = None,
        resample: PILImageResampling = None,
        do_center_crop: bool = None,
        crop_size: Dict[str, int] = None,
        do_rescale: bool = None,
        rescale_factor: float = None,
        do_normalize: bool = None,
        image_mean: Optional[Union[float, List[float]]] = None,
        image_std: Optional[Union[float, List[float]]] = None,
        data_format: Optional[Union[str, ChannelDimension]] = None,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
    ) -> np.ndarray:
        
        
        image = to_numpy_array(image)
        if is_scaled_image(image) and do_rescale:
            logger.warning_once(
                "It looks like you are trying to rescale already rescaled images. If the input"
                " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again."
            )
        if input_data_format is None:
            input_data_format = infer_channel_dimension_format(image)
        image = self._preprocess(
            image,
            do_reduce_labels=False,
            do_resize=do_resize,
            size=size,
            resample=resample,
            do_center_crop=do_center_crop,
            crop_size=crop_size,
            do_rescale=do_rescale,
            rescale_factor=rescale_factor,
            do_normalize=do_normalize,
            image_mean=image_mean,
            image_std=image_std,
            input_data_format=input_data_format,
        )
        if data_format is not None:
            image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
        return image

    def _preprocess_segmentation_map(
        self,
        segmentation_map: ImageInput,
        do_resize: bool = None,
        size: Dict[str, int] = None,
        resample: PILImageResampling = None,
        do_center_crop: bool = None,
        crop_size: Dict[str, int] = None,
        do_reduce_labels: bool = None,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
    ):
        
        
        segmentation_map = to_numpy_array(segmentation_map)
        
        if segmentation_map.ndim == 2:
            segmentation_map = segmentation_map[None, ...]
            added_dimension = True
            input_data_format = ChannelDimension.FIRST
        else:
            added_dimension = False
            if input_data_format is None:
                input_data_format = infer_channel_dimension_format(segmentation_map, num_channels=1)
        segmentation_map = self._preprocess(
            image=segmentation_map,
            do_reduce_labels=do_reduce_labels,
            do_resize=do_resize,
            resample=resample,
            size=size,
            do_center_crop=do_center_crop,
            crop_size=crop_size,
            do_normalize=False,
            do_rescale=False,
            input_data_format=ChannelDimension.FIRST,
        )
        
        if added_dimension:
            segmentation_map = np.squeeze(segmentation_map, axis=0)
        segmentation_map = segmentation_map.astype(np.int64)
        return segmentation_map

    def __call__(self, images, segmentation_maps=None, **kwargs):
        
        
        return super().__call__(images, segmentation_maps=segmentation_maps, **kwargs)

    def preprocess(
        self,
        images: ImageInput,
        segmentation_maps: Optional[ImageInput] = None,
        do_resize: bool = None,
        size: Dict[str, int] = None,
        resample: PILImageResampling = None,
        do_center_crop: bool = None,
        crop_size: Dict[str, int] = None,
        do_rescale: bool = None,
        rescale_factor: float = None,
        do_normalize: bool = None,
        image_mean: Optional[Union[float, List[float]]] = None,
        image_std: Optional[Union[float, List[float]]] = None,
        do_reduce_labels: Optional[bool] = None,
        return_tensors: Optional[Union[str, TensorType]] = None,
        data_format: ChannelDimension = ChannelDimension.FIRST,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
        **kwargs,
    ) -> PIL.Image.Image:
        
        do_resize = do_resize if do_resize is not None else self.do_resize
        size = size if size is not None else self.size
        size = get_size_dict(size, default_to_square=True, param_name="size")
        resample = resample if resample is not None else self.resample
        do_center_crop = do_center_crop if do_center_crop is not None else self.do_center_crop
        crop_size = crop_size if crop_size is not None else self.crop_size
        crop_size = get_size_dict(crop_size, default_to_square=True, param_name="crop_size")
        do_rescale = do_rescale if do_rescale is not None else self.do_rescale
        rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
        image_mean = image_mean if image_mean is not None else self.image_mean
        image_std = image_std if image_std is not None else self.image_std
        do_reduce_labels = do_reduce_labels if do_reduce_labels is not None else self.do_reduce_labels

        images = make_list_of_images(images)
        if segmentation_maps is not None:
            segmentation_maps = make_list_of_images(segmentation_maps, expected_ndims=2)

        if not valid_images(images):
            raise ValueError(
                "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
                "torch.Tensor, tf.Tensor or jax.ndarray."
            )

        if segmentation_maps is not None and not valid_images(segmentation_maps):
            raise ValueError(
                "Invalid segmentation map type. Must be of type PIL.Image.Image, numpy.ndarray, "
                "torch.Tensor, tf.Tensor or jax.ndarray."
            )

        if do_resize and size is None or resample is None:
            raise ValueError("Size and resample must be specified if do_resize is True.")

        if do_center_crop and crop_size is None:
            raise ValueError("Crop size must be specified if do_center_crop is True.")

        if do_rescale and rescale_factor is None:
            raise ValueError("Rescale factor must be specified if do_rescale is True.")

        if do_normalize and (image_mean is None or image_std is None):
            raise ValueError("Image mean and std must be specified if do_normalize is True.")

        images = [
            self._preprocess_image(
                image=img,
                do_resize=do_resize,
                do_center_crop=do_center_crop,
                do_rescale=do_rescale,
                do_normalize=do_normalize,
                resample=resample,
                size=size,
                rescale_factor=rescale_factor,
                crop_size=crop_size,
                image_mean=image_mean,
                image_std=image_std,
                data_format=data_format,
                input_data_format=input_data_format,
            )
            for img in images
        ]

        data = {"pixel_values": images}

        if segmentation_maps is not None:
            segmentation_maps = [
                self._preprocess_segmentation_map(
                    segmentation_map=segmentation_map,
                    do_reduce_labels=do_reduce_labels,
                    do_resize=do_resize,
                    resample=resample,
                    size=size,
                    do_center_crop=do_center_crop,
                    crop_size=crop_size,
                )
                for segmentation_map in segmentation_maps
            ]
            data["labels"] = segmentation_maps

        return BatchFeature(data=data, tensor_type=return_tensors)

    def post_process_semantic_segmentation(self, outputs, target_sizes: List[Tuple] = None):
        
        
        logits = outputs.logits

        
        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 is_torch_tensor(target_sizes):
                target_sizes = target_sizes.numpy()

            semantic_segmentation = []

            for idx in range(len(logits)):
                resized_logits = torch.nn.functional.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