Aluode/PerceptionLabPortable

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 15from collections import defaultdict16from collections.abc import Collection, Iterable17from math import ceil18from typing import Optional, Union19 20import numpy as np21 22from .image_utils import (23    ChannelDimension,24    ImageInput,25    get_channel_dimension_axis,26    get_image_size,27    infer_channel_dimension_format,28)29from .utils import ExplicitEnum, TensorType, is_jax_tensor, is_tf_tensor, is_torch_tensor30from .utils.import_utils import (31    is_flax_available,32    is_tf_available,33    is_torch_available,34    is_vision_available,35    requires_backends,36)37 38 39if is_vision_available():40    import PIL41 42    from .image_utils import PILImageResampling43 44if is_torch_available():45    import torch46 47if is_tf_available():48    import tensorflow as tf49 50if is_flax_available():51    import jax.numpy as jnp52 53 54def to_channel_dimension_format(55    image: np.ndarray,56    channel_dim: Union[ChannelDimension, str],57    input_channel_dim: Optional[Union[ChannelDimension, str]] = None,58) -> np.ndarray:59    """60    Converts `image` to the channel dimension format specified by `channel_dim`. The input61    can have arbitrary number of leading dimensions. Only last three dimension will be permuted62    to format the `image`.63 64    Args:65        image (`numpy.ndarray`):66            The image to have its channel dimension set.67        channel_dim (`ChannelDimension`):68            The channel dimension format to use.69        input_channel_dim (`ChannelDimension`, *optional*):70            The channel dimension format of the input image. If not provided, it will be inferred from the input image.71 72    Returns:73        `np.ndarray`: The image with the channel dimension set to `channel_dim`.74    """75    if not isinstance(image, np.ndarray):76        raise TypeError(f"Input image must be of type np.ndarray, got {type(image)}")77 78    if input_channel_dim is None:79        input_channel_dim = infer_channel_dimension_format(image)80 81    target_channel_dim = ChannelDimension(channel_dim)82    if input_channel_dim == target_channel_dim:83        return image84 85    if target_channel_dim == ChannelDimension.FIRST:86        axes = list(range(image.ndim - 3)) + [image.ndim - 1, image.ndim - 3, image.ndim - 2]87        image = image.transpose(axes)88    elif target_channel_dim == ChannelDimension.LAST:89        axes = list(range(image.ndim - 3)) + [image.ndim - 2, image.ndim - 1, image.ndim - 3]90        image = image.transpose(axes)91    else:92        raise ValueError(f"Unsupported channel dimension format: {channel_dim}")93 94    return image95 96 97def rescale(98    image: np.ndarray,99    scale: float,100    data_format: Optional[ChannelDimension] = None,101    dtype: np.dtype = np.float32,102    input_data_format: Optional[Union[str, ChannelDimension]] = None,103) -> np.ndarray:104    """105    Rescales `image` by `scale`.106 107    Args:108        image (`np.ndarray`):109            The image to rescale.110        scale (`float`):111            The scale to use for rescaling the image.112        data_format (`ChannelDimension`, *optional*):113            The channel dimension format of the image. If not provided, it will be the same as the input image.114        dtype (`np.dtype`, *optional*, defaults to `np.float32`):115            The dtype of the output image. Defaults to `np.float32`. Used for backwards compatibility with feature116            extractors.117        input_data_format (`ChannelDimension`, *optional*):118            The channel dimension format of the input image. If not provided, it will be inferred from the input image.119 120    Returns:121        `np.ndarray`: The rescaled image.122    """123    if not isinstance(image, np.ndarray):124        raise TypeError(f"Input image must be of type np.ndarray, got {type(image)}")125 126    rescaled_image = image.astype(np.float64) * scale  # Numpy type promotion has changed, so always upcast first127    if data_format is not None:128        rescaled_image = to_channel_dimension_format(rescaled_image, data_format, input_data_format)129 130    rescaled_image = rescaled_image.astype(dtype)  # Finally downcast to the desired dtype at the end131 132    return rescaled_image133 134 135def _rescale_for_pil_conversion(image):136    """137    Detects whether or not the image needs to be rescaled before being converted to a PIL image.138 139    The assumption is that if the image is of type `np.float` and all values are between 0 and 1, it needs to be140    rescaled.141    """142    if image.dtype == np.uint8:143        do_rescale = False144    elif np.allclose(image, image.astype(int)):145        if np.all(0 <= image) and np.all(image <= 255):146            do_rescale = False147        else:148            raise ValueError(149                "The image to be converted to a PIL image contains values outside the range [0, 255], "150                f"got [{image.min()}, {image.max()}] which cannot be converted to uint8."151            )152    elif np.all(0 <= image) and np.all(image <= 1):153        do_rescale = True154    else:155        raise ValueError(156            "The image to be converted to a PIL image contains values outside the range [0, 1], "157            f"got [{image.min()}, {image.max()}] which cannot be converted to uint8."158        )159    return do_rescale160 161 162def to_pil_image(163    image: Union[np.ndarray, "PIL.Image.Image", "torch.Tensor", "tf.Tensor", "jnp.ndarray"],164    do_rescale: Optional[bool] = None,165    image_mode: Optional[str] = None,166    input_data_format: Optional[Union[str, ChannelDimension]] = None,167) -> "PIL.Image.Image":168    """169    Converts `image` to a PIL Image. Optionally rescales it and puts the channel dimension back as the last axis if170    needed.171 172    Args:173        image (`PIL.Image.Image` or `numpy.ndarray` or `torch.Tensor` or `tf.Tensor`):174            The image to convert to the `PIL.Image` format.175        do_rescale (`bool`, *optional*):176            Whether or not to apply the scaling factor (to make pixel values integers between 0 and 255). Will default177            to `True` if the image type is a floating type and casting to `int` would result in a loss of precision,178            and `False` otherwise.179        image_mode (`str`, *optional*):180            The mode to use for the PIL image. If unset, will use the default mode for the input image type.181        input_data_format (`ChannelDimension`, *optional*):182            The channel dimension format of the input image. If unset, will use the inferred format from the input.183 184    Returns:185        `PIL.Image.Image`: The converted image.186    """187    requires_backends(to_pil_image, ["vision"])188 189    if isinstance(image, PIL.Image.Image):190        return image191 192    # Convert all tensors to numpy arrays before converting to PIL image193    if is_torch_tensor(image) or is_tf_tensor(image):194        image = image.numpy()195    elif is_jax_tensor(image):196        image = np.array(image)197    elif not isinstance(image, np.ndarray):198        raise ValueError(f"Input image type not supported: {type(image)}")199 200    # If the channel has been moved to first dim, we put it back at the end.201    image = to_channel_dimension_format(image, ChannelDimension.LAST, input_data_format)202 203    # If there is a single channel, we squeeze it, as otherwise PIL can't handle it.204    image = np.squeeze(image, axis=-1) if image.shape[-1] == 1 else image205 206    # PIL.Image can only store uint8 values so we rescale the image to be between 0 and 255 if needed.207    do_rescale = _rescale_for_pil_conversion(image) if do_rescale is None else do_rescale208 209    if do_rescale:210        image = rescale(image, 255)211 212    image = image.astype(np.uint8)213    return PIL.Image.fromarray(image, mode=image_mode)214 215 216def get_size_with_aspect_ratio(image_size, size, max_size=None) -> tuple[int, int]:217    """218    Computes the output image size given the input image size and the desired output size.219 220    Args:221        image_size (`tuple[int, int]`):222            The input image size.223        size (`int`):224            The desired output size.225        max_size (`int`, *optional*):226            The maximum allowed output size.227    """228    height, width = image_size229    raw_size = None230    if max_size is not None:231        min_original_size = float(min((height, width)))232        max_original_size = float(max((height, width)))233        if max_original_size / min_original_size * size > max_size:234            raw_size = max_size * min_original_size / max_original_size235            size = int(round(raw_size))236 237    if (height <= width and height == size) or (width <= height and width == size):238        oh, ow = height, width239    elif width < height:240        ow = size241        if max_size is not None and raw_size is not None:242            oh = int(raw_size * height / width)243        else:244            oh = int(size * height / width)245    else:246        oh = size247        if max_size is not None and raw_size is not None:248            ow = int(raw_size * width / height)249        else:250            ow = int(size * width / height)251 252    return (oh, ow)253 254 255# Logic adapted from torchvision resizing logic: https://github.com/pytorch/vision/blob/511924c1ced4ce0461197e5caa64ce5b9e558aab/torchvision/transforms/functional.py#L366256def get_resize_output_image_size(257    input_image: np.ndarray,258    size: Union[int, tuple[int, int], list[int], tuple[int, ...]],259    default_to_square: bool = True,260    max_size: Optional[int] = None,261    input_data_format: Optional[Union[str, ChannelDimension]] = None,262) -> tuple:263    """264    Find the target (height, width) dimension of the output image after resizing given the input image and the desired265    size.266 267    Args:268        input_image (`np.ndarray`):269            The image to resize.270        size (`int` or `tuple[int, int]` or list[int] or `tuple[int]`):271            The size to use for resizing the image. If `size` is a sequence like (h, w), output size will be matched to272            this.273 274            If `size` is an int and `default_to_square` is `True`, then image will be resized to (size, size). If275            `size` is an int and `default_to_square` is `False`, then smaller edge of the image will be matched to this276            number. i.e, if height > width, then image will be rescaled to (size * height / width, size).277        default_to_square (`bool`, *optional*, defaults to `True`):278            How to convert `size` when it is a single int. If set to `True`, the `size` will be converted to a square279            (`size`,`size`). If set to `False`, will replicate280            [`torchvision.transforms.Resize`](https://pytorch.org/vision/stable/transforms.html#torchvision.transforms.Resize)281            with support for resizing only the smallest edge and providing an optional `max_size`.282        max_size (`int`, *optional*):283            The maximum allowed for the longer edge of the resized image: if the longer edge of the image is greater284            than `max_size` after being resized according to `size`, then the image is resized again so that the longer285            edge is equal to `max_size`. As a result, `size` might be overruled, i.e the smaller edge may be shorter286            than `size`. Only used if `default_to_square` is `False`.287        input_data_format (`ChannelDimension`, *optional*):288            The channel dimension format of the input image. If unset, will use the inferred format from the input.289 290    Returns:291        `tuple`: The target (height, width) dimension of the output image after resizing.292    """293    if isinstance(size, (tuple, list)):294        if len(size) == 2:295            return tuple(size)296        elif len(size) == 1:297            # Perform same logic as if size was an int298            size = size[0]299        else:300            raise ValueError("size must have 1 or 2 elements if it is a list or tuple")301 302    if default_to_square:303        return (size, size)304 305    height, width = get_image_size(input_image, input_data_format)306    short, long = (width, height) if width <= height else (height, width)307    requested_new_short = size308 309    new_short, new_long = requested_new_short, int(requested_new_short * long / short)310 311    if max_size is not None:312        if max_size <= requested_new_short:313            raise ValueError(314                f"max_size = {max_size} must be strictly greater than the requested "315                f"size for the smaller edge size = {size}"316            )317        if new_long > max_size:318            new_short, new_long = int(max_size * new_short / new_long), max_size319 320    return (new_long, new_short) if width <= height else (new_short, new_long)321 322 323def resize(324    image: np.ndarray,325    size: tuple[int, int],326    resample: Optional["PILImageResampling"] = None,327    reducing_gap: Optional[int] = None,328    data_format: Optional[ChannelDimension] = None,329    return_numpy: bool = True,330    input_data_format: Optional[Union[str, ChannelDimension]] = None,331) -> np.ndarray:332    """333    Resizes `image` to `(height, width)` specified by `size` using the PIL library.334 335    Args:336        image (`np.ndarray`):337            The image to resize.338        size (`tuple[int, int]`):339            The size to use for resizing the image.340        resample (`int`, *optional*, defaults to `PILImageResampling.BILINEAR`):341            The filter to user for resampling.342        reducing_gap (`int`, *optional*):343            Apply optimization by resizing the image in two steps. The bigger `reducing_gap`, the closer the result to344            the fair resampling. See corresponding Pillow documentation for more details.345        data_format (`ChannelDimension`, *optional*):346            The channel dimension format of the output image. If unset, will use the inferred format from the input.347        return_numpy (`bool`, *optional*, defaults to `True`):348            Whether or not to return the resized image as a numpy array. If False a `PIL.Image.Image` object is349            returned.350        input_data_format (`ChannelDimension`, *optional*):351            The channel dimension format of the input image. If unset, will use the inferred format from the input.352 353    Returns:354        `np.ndarray`: The resized image.355    """356    requires_backends(resize, ["vision"])357 358    resample = resample if resample is not None else PILImageResampling.BILINEAR359 360    if not len(size) == 2:361        raise ValueError("size must have 2 elements")362 363    # For all transformations, we want to keep the same data format as the input image unless otherwise specified.364    # The resized image from PIL will always have channels last, so find the input format first.365    if input_data_format is None:366        input_data_format = infer_channel_dimension_format(image)367    data_format = input_data_format if data_format is None else data_format368 369    # To maintain backwards compatibility with the resizing done in previous image feature extractors, we use370    # the pillow library to resize the image and then convert back to numpy371    do_rescale = False372    if not isinstance(image, PIL.Image.Image):373        do_rescale = _rescale_for_pil_conversion(image)374        image = to_pil_image(image, do_rescale=do_rescale, input_data_format=input_data_format)375    height, width = size376    # PIL images are in the format (width, height)377    resized_image = image.resize((width, height), resample=resample, reducing_gap=reducing_gap)378 379    if return_numpy:380        resized_image = np.array(resized_image)381        # If the input image channel dimension was of size 1, then it is dropped when converting to a PIL image382        # so we need to add it back if necessary.383        resized_image = np.expand_dims(resized_image, axis=-1) if resized_image.ndim == 2 else resized_image384        # The image is always in channels last format after converting from a PIL image385        resized_image = to_channel_dimension_format(386            resized_image, data_format, input_channel_dim=ChannelDimension.LAST387        )388        # If an image was rescaled to be in the range [0, 255] before converting to a PIL image, then we need to389        # rescale it back to the original range.390        resized_image = rescale(resized_image, 1 / 255) if do_rescale else resized_image391    return resized_image392 393 394def normalize(395    image: np.ndarray,396    mean: Union[float, Collection[float]],397    std: Union[float, Collection[float]],398    data_format: Optional[ChannelDimension] = None,399    input_data_format: Optional[Union[str, ChannelDimension]] = None,400) -> np.ndarray:401    """402    Normalizes `image` using the mean and standard deviation specified by `mean` and `std`.403 404    image = (image - mean) / std405 406    Args:407        image (`np.ndarray`):408            The image to normalize.409        mean (`float` or `Collection[float]`):410            The mean to use for normalization.411        std (`float` or `Collection[float]`):412            The standard deviation to use for normalization.413        data_format (`ChannelDimension`, *optional*):414            The channel dimension format of the output image. If unset, will use the inferred format from the input.415        input_data_format (`ChannelDimension`, *optional*):416            The channel dimension format of the input image. If unset, will use the inferred format from the input.417    """418    if not isinstance(image, np.ndarray):419        raise TypeError("image must be a numpy array")420 421    if input_data_format is None:422        input_data_format = infer_channel_dimension_format(image)423 424    channel_axis = get_channel_dimension_axis(image, input_data_format=input_data_format)425    num_channels = image.shape[channel_axis]426 427    # We cast to float32 to avoid errors that can occur when subtracting uint8 values.428    # We preserve the original dtype if it is a float type to prevent upcasting float16.429    if not np.issubdtype(image.dtype, np.floating):430        image = image.astype(np.float32)431 432    if isinstance(mean, Collection):433        if len(mean) != num_channels:434            raise ValueError(f"mean must have {num_channels} elements if it is an iterable, got {len(mean)}")435    else:436        mean = [mean] * num_channels437    mean = np.array(mean, dtype=image.dtype)438 439    if isinstance(std, Collection):440        if len(std) != num_channels:441            raise ValueError(f"std must have {num_channels} elements if it is an iterable, got {len(std)}")442    else:443        std = [std] * num_channels444    std = np.array(std, dtype=image.dtype)445 446    if input_data_format == ChannelDimension.LAST:447        image = (image - mean) / std448    else:449        image = ((image.T - mean) / std).T450 451    image = to_channel_dimension_format(image, data_format, input_data_format) if data_format is not None else image452    return image453 454 455def center_crop(456    image: np.ndarray,457    size: tuple[int, int],458    data_format: Optional[Union[str, ChannelDimension]] = None,459    input_data_format: Optional[Union[str, ChannelDimension]] = None,460) -> np.ndarray:461    """462    Crops the `image` to the specified `size` using a center crop. Note that if the image is too small to be cropped to463    the size given, it will be padded (so the returned result will always be of size `size`).464 465    Args:466        image (`np.ndarray`):467            The image to crop.468        size (`tuple[int, int]`):469            The target size for the cropped image.470        data_format (`str` or `ChannelDimension`, *optional*):471            The channel dimension format for the output image. Can be one of:472                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.473                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.474            If unset, will use the inferred format of the input image.475        input_data_format (`str` or `ChannelDimension`, *optional*):476            The channel dimension format for the input image. Can be one of:477                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.478                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.479            If unset, will use the inferred format of the input image.480    Returns:481        `np.ndarray`: The cropped image.482    """483    requires_backends(center_crop, ["vision"])484 485    if not isinstance(image, np.ndarray):486        raise TypeError(f"Input image must be of type np.ndarray, got {type(image)}")487 488    if not isinstance(size, Iterable) or len(size) != 2:489        raise ValueError("size must have 2 elements representing the height and width of the output image")490 491    if input_data_format is None:492        input_data_format = infer_channel_dimension_format(image)493    output_data_format = data_format if data_format is not None else input_data_format494 495    # We perform the crop in (C, H, W) format and then convert to the output format496    image = to_channel_dimension_format(image, ChannelDimension.FIRST, input_data_format)497 498    orig_height, orig_width = get_image_size(image, ChannelDimension.FIRST)499    crop_height, crop_width = size500    crop_height, crop_width = int(crop_height), int(crop_width)501 502    # In case size is odd, (image_shape[0] + size[0]) // 2 won't give the proper result.503    top = (orig_height - crop_height) // 2504    bottom = top + crop_height505    # In case size is odd, (image_shape[1] + size[1]) // 2 won't give the proper result.506    left = (orig_width - crop_width) // 2507    right = left + crop_width508 509    # Check if cropped area is within image boundaries510    if top >= 0 and bottom <= orig_height and left >= 0 and right <= orig_width:511        image = image[..., top:bottom, left:right]512        image = to_channel_dimension_format(image, output_data_format, ChannelDimension.FIRST)513        return image514 515    # Otherwise, we may need to pad if the image is too small. Oh joy...516    new_height = max(crop_height, orig_height)517    new_width = max(crop_width, orig_width)518    new_shape = image.shape[:-2] + (new_height, new_width)519    new_image = np.zeros_like(image, shape=new_shape)520 521    # If the image is too small, pad it with zeros522    top_pad = ceil((new_height - orig_height) / 2)523    bottom_pad = top_pad + orig_height524    left_pad = ceil((new_width - orig_width) / 2)525    right_pad = left_pad + orig_width526    new_image[..., top_pad:bottom_pad, left_pad:right_pad] = image527 528    top += top_pad529    bottom += top_pad530    left += left_pad531    right += left_pad532 533    new_image = new_image[..., max(0, top) : min(new_height, bottom), max(0, left) : min(new_width, right)]534    new_image = to_channel_dimension_format(new_image, output_data_format, ChannelDimension.FIRST)535 536    return new_image537 538 539def _center_to_corners_format_torch(bboxes_center: "torch.Tensor") -> "torch.Tensor":540    center_x, center_y, width, height = bboxes_center.unbind(-1)541    bbox_corners = torch.stack(542        # top left x, top left y, bottom right x, bottom right y543        [(center_x - 0.5 * width), (center_y - 0.5 * height), (center_x + 0.5 * width), (center_y + 0.5 * height)],544        dim=-1,545    )546    return bbox_corners547 548 549def _center_to_corners_format_numpy(bboxes_center: np.ndarray) -> np.ndarray:550    center_x, center_y, width, height = bboxes_center.T551    bboxes_corners = np.stack(552        # top left x, top left y, bottom right x, bottom right y553        [center_x - 0.5 * width, center_y - 0.5 * height, center_x + 0.5 * width, center_y + 0.5 * height],554        axis=-1,555    )556    return bboxes_corners557 558 559def _center_to_corners_format_tf(bboxes_center: "tf.Tensor") -> "tf.Tensor":560    center_x, center_y, width, height = tf.unstack(bboxes_center, axis=-1)561    bboxes_corners = tf.stack(562        # top left x, top left y, bottom right x, bottom right y563        [center_x - 0.5 * width, center_y - 0.5 * height, center_x + 0.5 * width, center_y + 0.5 * height],564        axis=-1,565    )566    return bboxes_corners567 568 569# 2 functions below inspired by https://github.com/facebookresearch/detr/blob/master/util/box_ops.py570def center_to_corners_format(bboxes_center: TensorType) -> TensorType:571    """572    Converts bounding boxes from center format to corners format.573 574    center format: contains the coordinate for the center of the box and its width, height dimensions575        (center_x, center_y, width, height)576    corners format: contains the coordinates for the top-left and bottom-right corners of the box577        (top_left_x, top_left_y, bottom_right_x, bottom_right_y)578    """579    # Function is used during model forward pass, so we use the input framework if possible, without580    # converting to numpy581    if is_torch_tensor(bboxes_center):582        return _center_to_corners_format_torch(bboxes_center)583    elif isinstance(bboxes_center, np.ndarray):584        return _center_to_corners_format_numpy(bboxes_center)585    elif is_tf_tensor(bboxes_center):586        return _center_to_corners_format_tf(bboxes_center)587 588    raise ValueError(f"Unsupported input type {type(bboxes_center)}")589 590 591def _corners_to_center_format_torch(bboxes_corners: "torch.Tensor") -> "torch.Tensor":592    top_left_x, top_left_y, bottom_right_x, bottom_right_y = bboxes_corners.unbind(-1)593    b = [594        (top_left_x + bottom_right_x) / 2,  # center x595        (top_left_y + bottom_right_y) / 2,  # center y596        (bottom_right_x - top_left_x),  # width597        (bottom_right_y - top_left_y),  # height598    ]599    return torch.stack(b, dim=-1)600 601 602def _corners_to_center_format_numpy(bboxes_corners: np.ndarray) -> np.ndarray:603    top_left_x, top_left_y, bottom_right_x, bottom_right_y = bboxes_corners.T604    bboxes_center = np.stack(605        [606            (top_left_x + bottom_right_x) / 2,  # center x607            (top_left_y + bottom_right_y) / 2,  # center y608            (bottom_right_x - top_left_x),  # width609            (bottom_right_y - top_left_y),  # height610        ],611        axis=-1,612    )613    return bboxes_center614 615 616def _corners_to_center_format_tf(bboxes_corners: "tf.Tensor") -> "tf.Tensor":617    top_left_x, top_left_y, bottom_right_x, bottom_right_y = tf.unstack(bboxes_corners, axis=-1)618    bboxes_center = tf.stack(619        [620            (top_left_x + bottom_right_x) / 2,  # center x621            (top_left_y + bottom_right_y) / 2,  # center y622            (bottom_right_x - top_left_x),  # width623            (bottom_right_y - top_left_y),  # height624        ],625        axis=-1,626    )627    return bboxes_center628 629 630def corners_to_center_format(bboxes_corners: TensorType) -> TensorType:631    """632    Converts bounding boxes from corners format to center format.633 634    corners format: contains the coordinates for the top-left and bottom-right corners of the box635        (top_left_x, top_left_y, bottom_right_x, bottom_right_y)636    center format: contains the coordinate for the center of the box and its the width, height dimensions637        (center_x, center_y, width, height)638    """639    # Inverse function accepts different input types so implemented here too640    if is_torch_tensor(bboxes_corners):641        return _corners_to_center_format_torch(bboxes_corners)642    elif isinstance(bboxes_corners, np.ndarray):643        return _corners_to_center_format_numpy(bboxes_corners)644    elif is_tf_tensor(bboxes_corners):645        return _corners_to_center_format_tf(bboxes_corners)646 647    raise ValueError(f"Unsupported input type {type(bboxes_corners)}")648 649 650# 2 functions below copied from https://github.com/cocodataset/panopticapi/blob/master/panopticapi/utils.py651# Copyright (c) 2018, Alexander Kirillov652# All rights reserved.653def rgb_to_id(color):654    """655    Converts RGB color to unique ID.656    """657    if isinstance(color, np.ndarray) and len(color.shape) == 3:658        if color.dtype == np.uint8:659            color = color.astype(np.int32)660        return color[:, :, 0] + 256 * color[:, :, 1] + 256 * 256 * color[:, :, 2]661    return int(color[0] + 256 * color[1] + 256 * 256 * color[2])662 663 664def id_to_rgb(id_map):665    """666    Converts unique ID to RGB color.667    """668    if isinstance(id_map, np.ndarray):669        id_map_copy = id_map.copy()670        rgb_shape = tuple(list(id_map.shape) + [3])671        rgb_map = np.zeros(rgb_shape, dtype=np.uint8)672        for i in range(3):673            rgb_map[..., i] = id_map_copy % 256674            id_map_copy //= 256675        return rgb_map676    color = []677    for _ in range(3):678        color.append(id_map % 256)679        id_map //= 256680    return color681 682 683class PaddingMode(ExplicitEnum):684    """685    Enum class for the different padding modes to use when padding images.686    """687 688    CONSTANT = "constant"689    REFLECT = "reflect"690    REPLICATE = "replicate"691    SYMMETRIC = "symmetric"692 693 694def pad(695    image: np.ndarray,696    padding: Union[int, tuple[int, int], Iterable[tuple[int, int]]],697    mode: PaddingMode = PaddingMode.CONSTANT,698    constant_values: Union[float, Iterable[float]] = 0.0,699    data_format: Optional[Union[str, ChannelDimension]] = None,700    input_data_format: Optional[Union[str, ChannelDimension]] = None,701) -> np.ndarray:702    """703    Pads the `image` with the specified (height, width) `padding` and `mode`.704 705    Args:706        image (`np.ndarray`):707            The image to pad.708        padding (`int` or `tuple[int, int]` or `Iterable[tuple[int, int]]`):709            Padding to apply to the edges of the height, width axes. Can be one of three formats:710            - `((before_height, after_height), (before_width, after_width))` unique pad widths for each axis.711            - `((before, after),)` yields same before and after pad for height and width.712            - `(pad,)` or int is a shortcut for before = after = pad width for all axes.713        mode (`PaddingMode`):714            The padding mode to use. Can be one of:715                - `"constant"`: pads with a constant value.716                - `"reflect"`: pads with the reflection of the vector mirrored on the first and last values of the717                  vector along each axis.718                - `"replicate"`: pads with the replication of the last value on the edge of the array along each axis.719                - `"symmetric"`: pads with the reflection of the vector mirrored along the edge of the array.720        constant_values (`float` or `Iterable[float]`, *optional*):721            The value to use for the padding if `mode` is `"constant"`.722        data_format (`str` or `ChannelDimension`, *optional*):723            The channel dimension format for the output image. Can be one of:724                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.725                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.726            If unset, will use same as the input image.727        input_data_format (`str` or `ChannelDimension`, *optional*):728            The channel dimension format for the input image. Can be one of:729                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.730                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.731            If unset, will use the inferred format of the input image.732 733    Returns:734        `np.ndarray`: The padded image.735 736    """737    if input_data_format is None:738        input_data_format = infer_channel_dimension_format(image)739 740    def _expand_for_data_format(values):741        """742        Convert values to be in the format expected by np.pad based on the data format.743        """744        if isinstance(values, (int, float)):745            values = ((values, values), (values, values))746        elif isinstance(values, tuple) and len(values) == 1:747            values = ((values[0], values[0]), (values[0], values[0]))748        elif isinstance(values, tuple) and len(values) == 2 and isinstance(values[0], int):749            values = (values, values)750        elif isinstance(values, tuple) and len(values) == 2 and isinstance(values[0], tuple):751            pass752        else:753            raise ValueError(f"Unsupported format: {values}")754 755        # add 0 for channel dimension756        values = ((0, 0), *values) if input_data_format == ChannelDimension.FIRST else (*values, (0, 0))757 758        # Add additional padding if there's a batch dimension759        values = ((0, 0), *values) if image.ndim == 4 else values760        return values761 762    padding = _expand_for_data_format(padding)763 764    if mode == PaddingMode.CONSTANT:765        constant_values = _expand_for_data_format(constant_values)766        image = np.pad(image, padding, mode="constant", constant_values=constant_values)767    elif mode == PaddingMode.REFLECT:768        image = np.pad(image, padding, mode="reflect")769    elif mode == PaddingMode.REPLICATE:770        image = np.pad(image, padding, mode="edge")771    elif mode == PaddingMode.SYMMETRIC:772        image = np.pad(image, padding, mode="symmetric")773    else:774        raise ValueError(f"Invalid padding mode: {mode}")775 776    image = to_channel_dimension_format(image, data_format, input_data_format) if data_format is not None else image777    return image778 779 780# TODO (Amy): Accept 1/3/4 channel numpy array as input and return np.array as default781def convert_to_rgb(image: ImageInput) -> ImageInput:782    """783    Converts an image to RGB format. Only converts if the image is of type PIL.Image.Image, otherwise returns the image784    as is.785    Args:786        image (Image):787            The image to convert.788    """789    requires_backends(convert_to_rgb, ["vision"])790 791    if not isinstance(image, PIL.Image.Image):792        return image793 794    if image.mode == "RGB":795        return image796 797    image = image.convert("RGB")798    return image799 800 801def flip_channel_order(802    image: np.ndarray,803    data_format: Optional[ChannelDimension] = None,804    input_data_format: Optional[Union[str, ChannelDimension]] = None,805) -> np.ndarray:806    """807    Flips the channel order of the image.808 809    If the image is in RGB format, it will be converted to BGR and vice versa.810 811    Args:812        image (`np.ndarray`):813            The image to flip.814        data_format (`ChannelDimension`, *optional*):815            The channel dimension format for the output image. Can be one of:816                - `ChannelDimension.FIRST`: image in (num_channels, height, width) format.817                - `ChannelDimension.LAST`: image in (height, width, num_channels) format.818            If unset, will use same as the input image.819        input_data_format (`ChannelDimension`, *optional*):820            The channel dimension format for the input image. Can be one of:821                - `ChannelDimension.FIRST`: image in (num_channels, height, width) format.822                - `ChannelDimension.LAST`: image in (height, width, num_channels) format.823            If unset, will use the inferred format of the input image.824    """825    input_data_format = infer_channel_dimension_format(image) if input_data_format is None else input_data_format826 827    if input_data_format == ChannelDimension.LAST:828        image = image[..., ::-1]829    elif input_data_format == ChannelDimension.FIRST:830        image = image[::-1, ...]831    else:832        raise ValueError(f"Unsupported channel dimension: {input_data_format}")833 834    if data_format is not None:835        image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)836    return image837 838 839def _cast_tensor_to_float(x):840    if x.is_floating_point():841        return x842    return x.float()843 844 845def _group_images_by_shape(nested_images, is_nested: bool = False):846    """Helper function to flatten a single level of nested image structures and group by shape."""847    grouped_images = defaultdict(list)848    grouped_images_index = {}849    nested_images = [nested_images] if not is_nested else nested_images850    for i, sublist in enumerate(nested_images):851        for j, image in enumerate(sublist):852            key = (i, j) if is_nested else j853            shape = image.shape[1:]854            grouped_images[shape].append(image)855            grouped_images_index[key] = (shape, len(grouped_images[shape]) - 1)856 857    return grouped_images, grouped_images_index858 859 860def _reconstruct_nested_structure(indices, processed_images):861    """Helper function to reconstruct a single level nested structure."""862    # Find the maximum outer index863    max_outer_idx = max(idx[0] for idx in indices)864 865    # Create the outer list866    result = [None] * (max_outer_idx + 1)867 868    # Group indices by outer index869    nested_indices = defaultdict(list)870    for i, j in indices:871        nested_indices[i].append(j)872 873    for i in range(max_outer_idx + 1):874        if i in nested_indices:875            inner_max_idx = max(nested_indices[i])876            inner_list = [None] * (inner_max_idx + 1)877            for j in range(inner_max_idx + 1):878                if (i, j) in indices:879                    shape, idx = indices[(i, j)]880                    inner_list[j] = processed_images[shape][idx]881            result[i] = inner_list882 883    return result884 885 886def group_images_by_shape(887    images: Union[list["torch.Tensor"], "torch.Tensor"],888    disable_grouping: bool,889    is_nested: bool = False,890) -> tuple[891    dict[tuple[int, int], list["torch.Tensor"]], dict[Union[int, tuple[int, int]], tuple[tuple[int, int], int]]892]:893    """894    Groups images by shape.895    Returns a dictionary with the shape as key and a list of images with that shape as value,896    and a dictionary with the index of the image in the original list as key and the shape and index in the grouped list as value.897 898    The function supports both flat lists of tensors and nested structures.899    The input must be either all flat or all nested, not a mix of both.900 901    Args:902        images (Union[list["torch.Tensor"], "torch.Tensor"]):903            A list of images or a single tensor904        disable_grouping (bool):905            Whether to disable grouping. If None, will be set to True if the images are on CPU, and False otherwise.906            This choice is based on empirical observations, as detailed here: https://github.com/huggingface/transformers/pull/38157907        is_nested (bool, *optional*, defaults to False):908            Whether the images are nested.909 910    Returns:911        tuple[dict[tuple[int, int], list["torch.Tensor"]], dict[Union[int, tuple[int, int]], tuple[tuple[int, int], int]]]:912            - A dictionary with shape as key and list of images with that shape as value913            - A dictionary mapping original indices to (shape, index) tuples914    """915    # If disable grouping is not explicitly provided, we favor disabling it if the images are on CPU, and enabling it otherwise.916    if disable_grouping is None:917        device = images[0][0].device if is_nested else images[0].device918        disable_grouping = device == "cpu"919 920    if disable_grouping:921        if is_nested:922            return {(i, j): images[i][j].unsqueeze(0) for i in range(len(images)) for j in range(len(images[i]))}, {923                (i, j): ((i, j), 0) for i in range(len(images)) for j in range(len(images[i]))924            }925        else:926            return {i: images[i].unsqueeze(0) for i in range(len(images))}, {i: (i, 0) for i in range(len(images))}927 928    # Handle single level nested structure929    grouped_images, grouped_images_index = _group_images_by_shape(images, is_nested)930 931    # Stack images with the same shape932    grouped_images = {shape: torch.stack(images_list, dim=0) for shape, images_list in grouped_images.items()}933 934    return grouped_images, grouped_images_index935 936 937def reorder_images(938    processed_images: dict[tuple[int, int], "torch.Tensor"],939    grouped_images_index: dict[Union[int, tuple[int, int]], tuple[tuple[int, int], int]],940    is_nested: bool = False,941) -> Union[list["torch.Tensor"], "torch.Tensor"]:942    """943    Reconstructs images in the original order, preserving the original structure (nested or not).944    The input structure is either all flat or all nested.945 946    Args:947        processed_images (dict[tuple[int, int], "torch.Tensor"]):948            Dictionary mapping shapes to batched processed images.949        grouped_images_index (dict[Union[int, tuple[int, int]], tuple[tuple[int, int], int]]):950            Dictionary mapping original indices to (shape, index) tuples.951        is_nested (bool, *optional*, defaults to False):952            Whether the images are nested. Cannot be inferred from the input, as some processing functions outputs nested images.953            even with non nested images,e.g functions splitting images into patches. We thus can't deduce is_nested from the input.954 955 956    Returns:957        Union[list["torch.Tensor"], "torch.Tensor"]:958            Images in the original structure.959    """960    if not is_nested:961        return [962            processed_images[grouped_images_index[i][0]][grouped_images_index[i][1]]963            for i in range(len(grouped_images_index))964        ]965 966    return _reconstruct_nested_structure(grouped_images_index, processed_images)967 968 969class NumpyToTensor:970    """971    Convert a numpy array to a PyTorch tensor.972    """973 974    def __call__(self, image: np.ndarray):975        # Same as in PyTorch, we assume incoming numpy images are in HWC format976        # c.f. https://github.com/pytorch/vision/blob/61d97f41bc209e1407dcfbd685d2ee2da9c1cdad/torchvision/transforms/functional.py#L154977        return torch.from_numpy(image.transpose(2, 0, 1)).contiguous()978