Aluode/PerceptionLabPortable
0
1# Copyright 2022 The HuggingFace Inc. team.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15import math16from collections.abc import Iterable17from typing import Optional, Union18 19import numpy as np20 21from .image_processing_base import BatchFeature, ImageProcessingMixin22from .image_transforms import center_crop, normalize, rescale23from .image_utils import ChannelDimension, get_image_size24from .utils import logging25from .utils.import_utils import requires26 27 28logger = logging.get_logger(__name__)29 30 31INIT_SERVICE_KWARGS = [32 "processor_class",33 "image_processor_type",34]35 36 37@requires(backends=("vision",))38class BaseImageProcessor(ImageProcessingMixin):39 def __init__(self, **kwargs):40 super().__init__(**kwargs)41 42 @property43 def is_fast(self) -> bool:44 """45 `bool`: Whether or not this image processor is a fast processor (backed by PyTorch and TorchVision).46 """47 return False48 49 def __call__(self, images, **kwargs) -> BatchFeature:50 """Preprocess an image or a batch of images."""51 return self.preprocess(images, **kwargs)52 53 def preprocess(self, images, **kwargs) -> BatchFeature:54 raise NotImplementedError("Each image processor must implement its own preprocess method")55 56 def rescale(57 self,58 image: np.ndarray,59 scale: float,60 data_format: Optional[Union[str, ChannelDimension]] = None,61 input_data_format: Optional[Union[str, ChannelDimension]] = None,62 **kwargs,63 ) -> np.ndarray:64 """65 Rescale an image by a scale factor. image = image * scale.66 67 Args:68 image (`np.ndarray`):69 Image to rescale.70 scale (`float`):71 The scaling factor to rescale pixel values by.72 data_format (`str` or `ChannelDimension`, *optional*):73 The channel dimension format for the output image. If unset, the channel dimension format of the input74 image is used. Can be one of:75 - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.76 - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.77 input_data_format (`ChannelDimension` or `str`, *optional*):78 The channel dimension format for the input image. If unset, the channel dimension format is inferred79 from the input image. Can be one of:80 - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.81 - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.82 83 Returns:84 `np.ndarray`: The rescaled image.85 """86 return rescale(image, scale=scale, data_format=data_format, input_data_format=input_data_format, **kwargs)87 88 def normalize(89 self,90 image: np.ndarray,91 mean: Union[float, Iterable[float]],92 std: Union[float, Iterable[float]],93 data_format: Optional[Union[str, ChannelDimension]] = None,94 input_data_format: Optional[Union[str, ChannelDimension]] = None,95 **kwargs,96 ) -> np.ndarray:97 """98 Normalize an image. image = (image - image_mean) / image_std.99 100 Args:101 image (`np.ndarray`):102 Image to normalize.103 mean (`float` or `Iterable[float]`):104 Image mean to use for normalization.105 std (`float` or `Iterable[float]`):106 Image standard deviation to use for normalization.107 data_format (`str` or `ChannelDimension`, *optional*):108 The channel dimension format for the output image. If unset, the channel dimension format of the input109 image is used. Can be one of:110 - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.111 - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.112 input_data_format (`ChannelDimension` or `str`, *optional*):113 The channel dimension format for the input image. If unset, the channel dimension format is inferred114 from the input image. Can be one of:115 - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.116 - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.117 118 Returns:119 `np.ndarray`: The normalized image.120 """121 return normalize(122 image, mean=mean, std=std, data_format=data_format, input_data_format=input_data_format, **kwargs123 )124 125 def center_crop(126 self,127 image: np.ndarray,128 size: dict[str, int],129 data_format: Optional[Union[str, ChannelDimension]] = None,130 input_data_format: Optional[Union[str, ChannelDimension]] = None,131 **kwargs,132 ) -> np.ndarray:133 """134 Center crop an image to `(size["height"], size["width"])`. If the input size is smaller than `crop_size` along135 any edge, the image is padded with 0's and then center cropped.136 137 Args:138 image (`np.ndarray`):139 Image to center crop.140 size (`dict[str, int]`):141 Size of the output image.142 data_format (`str` or `ChannelDimension`, *optional*):143 The channel dimension format for the output image. If unset, the channel dimension format of the input144 image is used. Can be one of:145 - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.146 - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.147 input_data_format (`ChannelDimension` or `str`, *optional*):148 The channel dimension format for the input image. If unset, the channel dimension format is inferred149 from the input image. Can be one of:150 - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.151 - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.152 """153 size = get_size_dict(size)154 if "height" not in size or "width" not in size:155 raise ValueError(f"The size dictionary must have keys 'height' and 'width'. Got {size.keys()}")156 return center_crop(157 image,158 size=(size["height"], size["width"]),159 data_format=data_format,160 input_data_format=input_data_format,161 **kwargs,162 )163 164 def to_dict(self):165 encoder_dict = super().to_dict()166 encoder_dict.pop("_valid_processor_keys", None)167 return encoder_dict168 169 170VALID_SIZE_DICT_KEYS = (171 {"height", "width"},172 {"shortest_edge"},173 {"shortest_edge", "longest_edge"},174 {"longest_edge"},175 {"max_height", "max_width"},176)177 178 179def is_valid_size_dict(size_dict):180 if not isinstance(size_dict, dict):181 return False182 183 size_dict_keys = set(size_dict.keys())184 for allowed_keys in VALID_SIZE_DICT_KEYS:185 if size_dict_keys == allowed_keys:186 return True187 return False188 189 190def convert_to_size_dict(191 size, max_size: Optional[int] = None, default_to_square: bool = True, height_width_order: bool = True192):193 # By default, if size is an int we assume it represents a tuple of (size, size).194 if isinstance(size, int) and default_to_square:195 if max_size is not None:196 raise ValueError("Cannot specify both size as an int, with default_to_square=True and max_size")197 return {"height": size, "width": size}198 # In other configs, if size is an int and default_to_square is False, size represents the length of199 # the shortest edge after resizing.200 elif isinstance(size, int) and not default_to_square:201 size_dict = {"shortest_edge": size}202 if max_size is not None:203 size_dict["longest_edge"] = max_size204 return size_dict205 # Otherwise, if size is a tuple it's either (height, width) or (width, height)206 elif isinstance(size, (tuple, list)) and height_width_order:207 return {"height": size[0], "width": size[1]}208 elif isinstance(size, (tuple, list)) and not height_width_order:209 return {"height": size[1], "width": size[0]}210 elif size is None and max_size is not None:211 if default_to_square:212 raise ValueError("Cannot specify both default_to_square=True and max_size")213 return {"longest_edge": max_size}214 215 raise ValueError(f"Could not convert size input to size dict: {size}")216 217 218def get_size_dict(219 size: Optional[Union[int, Iterable[int], dict[str, int]]] = None,220 max_size: Optional[int] = None,221 height_width_order: bool = True,222 default_to_square: bool = True,223 param_name="size",224) -> dict:225 """226 Converts the old size parameter in the config into the new dict expected in the config. This is to ensure backwards227 compatibility with the old image processor configs and removes ambiguity over whether the tuple is in (height,228 width) or (width, height) format.229 230 - If `size` is tuple, it is converted to `{"height": size[0], "width": size[1]}` or `{"height": size[1], "width":231 size[0]}` if `height_width_order` is `False`.232 - If `size` is an int, and `default_to_square` is `True`, it is converted to `{"height": size, "width": size}`.233 - If `size` is an int and `default_to_square` is False, it is converted to `{"shortest_edge": size}`. If `max_size`234 is set, it is added to the dict as `{"longest_edge": max_size}`.235 236 Args:237 size (`Union[int, Iterable[int], dict[str, int]]`, *optional*):238 The `size` parameter to be cast into a size dictionary.239 max_size (`Optional[int]`, *optional*):240 The `max_size` parameter to be cast into a size dictionary.241 height_width_order (`bool`, *optional*, defaults to `True`):242 If `size` is a tuple, whether it's in (height, width) or (width, height) order.243 default_to_square (`bool`, *optional*, defaults to `True`):244 If `size` is an int, whether to default to a square image or not.245 """246 if not isinstance(size, dict):247 size_dict = convert_to_size_dict(size, max_size, default_to_square, height_width_order)248 logger.info(249 f"{param_name} should be a dictionary on of the following set of keys: {VALID_SIZE_DICT_KEYS}, got {size}."250 f" Converted to {size_dict}.",251 )252 else:253 size_dict = size254 255 if not is_valid_size_dict(size_dict):256 raise ValueError(257 f"{param_name} must have one of the following set of keys: {VALID_SIZE_DICT_KEYS}, got {size_dict.keys()}"258 )259 return size_dict260 261 262def select_best_resolution(original_size: tuple, possible_resolutions: list) -> tuple:263 """264 Selects the best resolution from a list of possible resolutions based on the original size.265 266 This is done by calculating the effective and wasted resolution for each possible resolution.267 268 The best fit resolution is the one that maximizes the effective resolution and minimizes the wasted resolution.269 270 Args:271 original_size (tuple):272 The original size of the image in the format (height, width).273 possible_resolutions (list):274 A list of possible resolutions in the format [(height1, width1), (height2, width2), ...].275 276 Returns:277 tuple: The best fit resolution in the format (height, width).278 """279 original_height, original_width = original_size280 best_fit = None281 max_effective_resolution = 0282 min_wasted_resolution = float("inf")283 284 for height, width in possible_resolutions:285 scale = min(width / original_width, height / original_height)286 downscaled_width, downscaled_height = int(original_width * scale), int(original_height * scale)287 effective_resolution = min(downscaled_width * downscaled_height, original_width * original_height)288 wasted_resolution = (width * height) - effective_resolution289 290 if effective_resolution > max_effective_resolution or (291 effective_resolution == max_effective_resolution and wasted_resolution < min_wasted_resolution292 ):293 max_effective_resolution = effective_resolution294 min_wasted_resolution = wasted_resolution295 best_fit = (height, width)296 297 return best_fit298 299 300def get_patch_output_size(image, target_resolution, input_data_format):301 """302 Given an image and a target resolution, calculate the output size of the image after cropping to the target303 """304 original_height, original_width = get_image_size(image, channel_dim=input_data_format)305 target_height, target_width = target_resolution306 307 scale_w = target_width / original_width308 scale_h = target_height / original_height309 310 if scale_w < scale_h:311 new_width = target_width312 new_height = min(math.ceil(original_height * scale_w), target_height)313 else:314 new_height = target_height315 new_width = min(math.ceil(original_width * scale_h), target_width)316 317 return new_height, new_width318