Aluode/PerceptionLabPortable
0
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 Vilt."""16 17from collections.abc import Iterable18from typing import Any, Optional, Union19 20import numpy as np21 22from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict23from ...image_transforms import PaddingMode, pad, resize, to_channel_dimension_format24from ...image_utils import (25 IMAGENET_STANDARD_MEAN,26 IMAGENET_STANDARD_STD,27 ChannelDimension,28 ImageInput,29 PILImageResampling,30 get_image_size,31 infer_channel_dimension_format,32 is_scaled_image,33 make_flat_list_of_images,34 to_numpy_array,35 valid_images,36 validate_preprocess_arguments,37)38from ...utils import TensorType, filter_out_non_signature_kwargs, is_vision_available, logging39from ...utils.import_utils import requires40 41 42if is_vision_available():43 import PIL44 45 46logger = logging.get_logger(__name__)47 48 49def max_across_indices(values: Iterable[Any]) -> list[Any]:50 """51 Return the maximum value across all indices of an iterable of values.52 """53 return [max(values_i) for values_i in zip(*values)]54 55 56def make_pixel_mask(57 image: np.ndarray, output_size: tuple[int, int], input_data_format: Optional[Union[str, ChannelDimension]] = None58) -> np.ndarray:59 """60 Make a pixel mask for the image, where 1 indicates a valid pixel and 0 indicates padding.61 62 Args:63 image (`np.ndarray`):64 Image to make the pixel mask for.65 output_size (`tuple[int, int]`):66 Output size of the mask.67 """68 input_height, input_width = get_image_size(image, channel_dim=input_data_format)69 mask = np.zeros(output_size, dtype=np.int64)70 mask[:input_height, :input_width] = 171 return mask72 73 74def get_max_height_width(75 images: list[np.ndarray], input_data_format: Optional[Union[str, ChannelDimension]] = None76) -> list[int]:77 """78 Get the maximum height and width across all images in a batch.79 """80 if input_data_format is None:81 input_data_format = infer_channel_dimension_format(images[0])82 83 if input_data_format == ChannelDimension.FIRST:84 _, max_height, max_width = max_across_indices([img.shape for img in images])85 elif input_data_format == ChannelDimension.LAST:86 max_height, max_width, _ = max_across_indices([img.shape for img in images])87 else:88 raise ValueError(f"Invalid channel dimension format: {input_data_format}")89 return (max_height, max_width)90 91 92def get_resize_output_image_size(93 input_image: np.ndarray,94 shorter: int = 800,95 longer: int = 1333,96 size_divisor: int = 32,97 input_data_format: Optional[Union[str, ChannelDimension]] = None,98) -> tuple[int, int]:99 input_height, input_width = get_image_size(input_image, input_data_format)100 min_size, max_size = shorter, longer101 102 scale = min_size / min(input_height, input_width)103 104 if input_height < input_width:105 new_height = min_size106 new_width = scale * input_width107 else:108 new_height = scale * input_height109 new_width = min_size110 111 if max(new_height, new_width) > max_size:112 scale = max_size / max(new_height, new_width)113 new_height = scale * new_height114 new_width = scale * new_width115 116 new_height, new_width = int(new_height + 0.5), int(new_width + 0.5)117 new_height = new_height // size_divisor * size_divisor118 new_width = new_width // size_divisor * size_divisor119 120 return new_height, new_width121 122 123@requires(backends=("vision",))124class ViltImageProcessor(BaseImageProcessor):125 r"""126 Constructs a ViLT image processor.127 128 Args:129 do_resize (`bool`, *optional*, defaults to `True`):130 Whether to resize the image's (height, width) dimensions to the specified `size`. Can be overridden by the131 `do_resize` parameter in the `preprocess` method.132 size (`dict[str, int]` *optional*, defaults to `{"shortest_edge": 384}`):133 Resize the shorter side of the input to `size["shortest_edge"]`. The longer side will be limited to under134 `int((1333 / 800) * size["shortest_edge"])` while preserving the aspect ratio. Only has an effect if135 `do_resize` is set to `True`. Can be overridden by the `size` parameter in the `preprocess` method.136 size_divisor (`int`, *optional*, defaults to 32):137 The size by which to make sure both the height and width can be divided. Only has an effect if `do_resize`138 is set to `True`. Can be overridden by the `size_divisor` parameter in the `preprocess` method.139 resample (`PILImageResampling`, *optional*, defaults to `Resampling.BICUBIC`):140 Resampling filter to use if resizing the image. Only has an effect if `do_resize` is set to `True`. Can be141 overridden by the `resample` parameter in the `preprocess` method.142 do_rescale (`bool`, *optional*, defaults to `True`):143 Wwhether to rescale the image by the specified scale `rescale_factor`. Can be overridden by the144 `do_rescale` parameter in the `preprocess` method.145 rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):146 Scale factor to use if rescaling the image. Only has an effect if `do_rescale` is set to `True`. Can be147 overridden by the `rescale_factor` parameter in the `preprocess` method.148 do_normalize (`bool`, *optional*, defaults to `True`):149 Whether to normalize the image. Can be overridden by the `do_normalize` parameter in the `preprocess`150 method. Can be overridden by the `do_normalize` parameter in the `preprocess` method.151 image_mean (`float` or `list[float]`, *optional*, defaults to `IMAGENET_STANDARD_MEAN`):152 Mean to use if normalizing the image. This is a float or list of floats the length of the number of153 channels in the image. Can be overridden by the `image_mean` parameter in the `preprocess` method. Can be154 overridden by the `image_mean` parameter in the `preprocess` method.155 image_std (`float` or `list[float]`, *optional*, defaults to `IMAGENET_STANDARD_STD`):156 Standard deviation to use if normalizing the image. This is a float or list of floats the length of the157 number of channels in the image. Can be overridden by the `image_std` parameter in the `preprocess` method.158 Can be overridden by the `image_std` parameter in the `preprocess` method.159 do_pad (`bool`, *optional*, defaults to `True`):160 Whether to pad the image to the `(max_height, max_width)` of the images in the batch. Can be overridden by161 the `do_pad` parameter in the `preprocess` method.162 """163 164 model_input_names = ["pixel_values"]165 166 def __init__(167 self,168 do_resize: bool = True,169 size: Optional[dict[str, int]] = None,170 size_divisor: int = 32,171 resample: PILImageResampling = PILImageResampling.BICUBIC,172 do_rescale: bool = True,173 rescale_factor: Union[int, float] = 1 / 255,174 do_normalize: bool = True,175 image_mean: Optional[Union[float, list[float]]] = None,176 image_std: Optional[Union[float, list[float]]] = None,177 do_pad: bool = True,178 **kwargs,179 ) -> None:180 if "pad_and_return_pixel_mask" in kwargs:181 do_pad = kwargs.pop("pad_and_return_pixel_mask")182 183 super().__init__(**kwargs)184 size = size if size is not None else {"shortest_edge": 384}185 size = get_size_dict(size, default_to_square=False)186 187 self.do_resize = do_resize188 self.size = size189 self.size_divisor = size_divisor190 self.resample = resample191 self.do_rescale = do_rescale192 self.rescale_factor = rescale_factor193 self.do_normalize = do_normalize194 self.image_mean = image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN195 self.image_std = image_std if image_std is not None else IMAGENET_STANDARD_STD196 self.do_pad = do_pad197 198 @classmethod199 def from_dict(cls, image_processor_dict: dict[str, Any], **kwargs):200 """201 Overrides the `from_dict` method from the base class to make sure `pad_and_return_pixel_mask` is updated if image processor202 is created using from_dict and kwargs e.g. `ViltImageProcessor.from_pretrained(checkpoint,203 pad_and_return_pixel_mask=False)`204 """205 image_processor_dict = image_processor_dict.copy()206 if "pad_and_return_pixel_mask" in kwargs:207 image_processor_dict["pad_and_return_pixel_mask"] = kwargs.pop("pad_and_return_pixel_mask")208 return super().from_dict(image_processor_dict, **kwargs)209 210 def resize(211 self,212 image: np.ndarray,213 size: dict[str, int],214 size_divisor: int = 32,215 resample: PILImageResampling = PILImageResampling.BICUBIC,216 data_format: Optional[Union[str, ChannelDimension]] = None,217 input_data_format: Optional[Union[str, ChannelDimension]] = None,218 **kwargs,219 ) -> np.ndarray:220 """221 Resize an image.222 223 Resizes the shorter side of the image to `size["shortest_edge"]` while preserving the aspect ratio. If the224 longer side is larger than the max size `(int(`size["shortest_edge"]` * 1333 / 800))`, the longer side is then225 resized to the max size while preserving the aspect ratio.226 227 Args:228 image (`np.ndarray`):229 Image to resize.230 size (`dict[str, int]`):231 Controls the size of the output image. Should be of the form `{"shortest_edge": int}`.232 size_divisor (`int`, *optional*, defaults to 32):233 The image is resized to a size that is a multiple of this value.234 resample (`PILImageResampling` filter, *optional*, defaults to `PILImageResampling.BICUBIC`):235 Resampling filter to use when resiizing the image.236 data_format (`str` or `ChannelDimension`, *optional*):237 The channel dimension format of the image. If not provided, it will be the same as the input image.238 input_data_format (`str` or `ChannelDimension`, *optional*):239 The channel dimension format of the input image. If not provided, it will be inferred.240 """241 size = get_size_dict(size, default_to_square=False)242 if "shortest_edge" not in size:243 raise ValueError(f"The `size` dictionary must contain the key `shortest_edge`. Got {size.keys()}")244 shorter = size["shortest_edge"]245 longer = int(1333 / 800 * shorter)246 output_size = get_resize_output_image_size(247 image, shorter=shorter, longer=longer, size_divisor=size_divisor, input_data_format=input_data_format248 )249 return resize(250 image,251 size=output_size,252 resample=resample,253 data_format=data_format,254 input_data_format=input_data_format,255 **kwargs,256 )257 258 def _pad_image(259 self,260 image: np.ndarray,261 output_size: tuple[int, int],262 constant_values: Union[float, Iterable[float]] = 0,263 data_format: Optional[ChannelDimension] = None,264 input_data_format: Optional[Union[str, ChannelDimension]] = None,265 ) -> np.ndarray:266 """267 Pad an image with zeros to the given size.268 """269 input_height, input_width = get_image_size(image, channel_dim=input_data_format)270 output_height, output_width = output_size271 272 pad_bottom = output_height - input_height273 pad_right = output_width - input_width274 padding = ((0, pad_bottom), (0, pad_right))275 padded_image = pad(276 image,277 padding,278 mode=PaddingMode.CONSTANT,279 constant_values=constant_values,280 data_format=data_format,281 input_data_format=input_data_format,282 )283 return padded_image284 285 def pad(286 self,287 images: list[np.ndarray],288 constant_values: Union[float, Iterable[float]] = 0,289 return_pixel_mask: bool = True,290 return_tensors: Optional[Union[str, TensorType]] = None,291 data_format: Optional[ChannelDimension] = None,292 input_data_format: Optional[Union[str, ChannelDimension]] = None,293 ) -> BatchFeature:294 """295 Pads a batch of images to the bottom and right of the image with zeros to the size of largest height and width296 in the batch and optionally returns their corresponding pixel mask.297 298 Args:299 image (`np.ndarray`):300 Image to pad.301 constant_values (`float` or `Iterable[float]`, *optional*):302 The value to use for the padding if `mode` is `"constant"`.303 return_pixel_mask (`bool`, *optional*, defaults to `True`):304 Whether to return a pixel mask.305 return_tensors (`str` or `TensorType`, *optional*):306 The type of tensors to return. Can be one of:307 - Unset: Return a list of `np.ndarray`.308 - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.309 - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.310 - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.311 - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.312 data_format (`str` or `ChannelDimension`, *optional*):313 The channel dimension format of the image. If not provided, it will be the same as the input image.314 input_data_format (`ChannelDimension` or `str`, *optional*):315 The channel dimension format of the input image. If not provided, it will be inferred.316 """317 pad_size = get_max_height_width(images, input_data_format=input_data_format)318 319 padded_images = [320 self._pad_image(321 image,322 pad_size,323 constant_values=constant_values,324 data_format=data_format,325 input_data_format=input_data_format,326 )327 for image in images328 ]329 data = {"pixel_values": padded_images}330 331 if return_pixel_mask:332 masks = [333 make_pixel_mask(image=image, output_size=pad_size, input_data_format=input_data_format)334 for image in images335 ]336 data["pixel_mask"] = masks337 338 return BatchFeature(data=data, tensor_type=return_tensors)339 340 @filter_out_non_signature_kwargs()341 def preprocess(342 self,343 images: ImageInput,344 do_resize: Optional[bool] = None,345 size: Optional[dict[str, int]] = None,346 size_divisor: Optional[int] = None,347 resample: Optional[PILImageResampling] = None,348 do_rescale: Optional[bool] = None,349 rescale_factor: Optional[float] = None,350 do_normalize: Optional[bool] = None,351 image_mean: Optional[Union[float, list[float]]] = None,352 image_std: Optional[Union[float, list[float]]] = None,353 do_pad: Optional[bool] = None,354 return_tensors: Optional[Union[str, TensorType]] = None,355 data_format: ChannelDimension = ChannelDimension.FIRST,356 input_data_format: Optional[Union[str, ChannelDimension]] = None,357 ) -> PIL.Image.Image:358 """359 Preprocess an image or batch of images.360 361 Args:362 images (`ImageInput`):363 Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If364 passing in images with pixel values between 0 and 1, set `do_rescale=False`.365 do_resize (`bool`, *optional*, defaults to `self.do_resize`):366 Whether to resize the image.367 size (`dict[str, int]`, *optional*, defaults to `self.size`):368 Controls the size of the image after `resize`. The shortest edge of the image is resized to369 `size["shortest_edge"]` whilst preserving the aspect ratio. If the longest edge of this resized image370 is > `int(size["shortest_edge"] * (1333 / 800))`, then the image is resized again to make the longest371 edge equal to `int(size["shortest_edge"] * (1333 / 800))`.372 size_divisor (`int`, *optional*, defaults to `self.size_divisor`):373 The image is resized to a size that is a multiple of this value.374 resample (`PILImageResampling`, *optional*, defaults to `self.resample`):375 Resampling filter to use if resizing the image. Only has an effect if `do_resize` is set to `True`.376 do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):377 Whether to rescale the image values between [0 - 1].378 rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):379 Rescale factor to rescale the image by if `do_rescale` is set to `True`.380 do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):381 Whether to normalize the image.382 image_mean (`float` or `list[float]`, *optional*, defaults to `self.image_mean`):383 Image mean to normalize the image by if `do_normalize` is set to `True`.384 image_std (`float` or `list[float]`, *optional*, defaults to `self.image_std`):385 Image standard deviation to normalize the image by if `do_normalize` is set to `True`.386 do_pad (`bool`, *optional*, defaults to `self.do_pad`):387 Whether to pad the image to the (max_height, max_width) in the batch. If `True`, a pixel mask is also388 created and returned.389 return_tensors (`str` or `TensorType`, *optional*):390 The type of tensors to return. Can be one of:391 - Unset: Return a list of `np.ndarray`.392 - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.393 - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.394 - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.395 - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.396 data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):397 The channel dimension format for the output image. Can be one of:398 - `ChannelDimension.FIRST`: image in (num_channels, height, width) format.399 - `ChannelDimension.LAST`: image in (height, width, num_channels) format.400 input_data_format (`ChannelDimension` or `str`, *optional*):401 The channel dimension format for the input image. If unset, the channel dimension format is inferred402 from the input image. Can be one of:403 - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.404 - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.405 - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.406 """407 do_resize = do_resize if do_resize is not None else self.do_resize408 size_divisor = size_divisor if size_divisor is not None else self.size_divisor409 resample = resample if resample is not None else self.resample410 do_rescale = do_rescale if do_rescale is not None else self.do_rescale411 rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor412 do_normalize = do_normalize if do_normalize is not None else self.do_normalize413 image_mean = image_mean if image_mean is not None else self.image_mean414 image_std = image_std if image_std is not None else self.image_std415 do_pad = do_pad if do_pad is not None else self.do_pad416 417 size = size if size is not None else self.size418 size = get_size_dict(size, default_to_square=False)419 420 images = make_flat_list_of_images(images)421 422 if not valid_images(images):423 raise ValueError(424 "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "425 "torch.Tensor, tf.Tensor or jax.ndarray."426 )427 428 # Here the pad() method does not require any additional argument as it takes the maximum of (height, width).429 # Hence, it does not need to be passed to a validate_preprocess_arguments() method.430 validate_preprocess_arguments(431 do_rescale=do_rescale,432 rescale_factor=rescale_factor,433 do_normalize=do_normalize,434 image_mean=image_mean,435 image_std=image_std,436 do_resize=do_resize,437 size=size,438 resample=resample,439 )440 441 # All transformations expect numpy arrays.442 images = [to_numpy_array(image) for image in images]443 444 if do_rescale and is_scaled_image(images[0]):445 logger.warning_once(446 "It looks like you are trying to rescale already rescaled images. If the input"447 " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again."448 )449 450 if input_data_format is None:451 # We assume that all images have the same channel dimension format.452 input_data_format = infer_channel_dimension_format(images[0])453 454 if do_resize:455 images = [456 self.resize(457 image=image,458 size=size,459 size_divisor=size_divisor,460 resample=resample,461 input_data_format=input_data_format,462 )463 for image in images464 ]465 466 if do_rescale:467 images = [468 self.rescale(image=image, scale=rescale_factor, input_data_format=input_data_format)469 for image in images470 ]471 472 if do_normalize:473 images = [474 self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)475 for image in images476 ]477 478 images = [479 to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format) for image in images480 ]481 482 if do_pad:483 encoded_outputs = self.pad(484 images, return_pixel_mask=True, return_tensors=return_tensors, input_data_format=data_format485 )486 else:487 encoded_outputs = BatchFeature(data={"pixel_values": images}, tensor_type=return_tensors)488 489 return encoded_outputs490 491 492__all__ = ["ViltImageProcessor"]493