image_processor.py

from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
from typing import Optional, Union, Tuple, Dict, List, Iterable
from transformers.image_transforms import to_channel_dimension_format, PaddingMode
from transformers.image_utils import ChannelDimension, to_numpy_array, make_list_of_images, get_image_size, infer_channel_dimension_format
from transformers.utils import TensorType
from PIL import Image
import numpy as np
try:
    from torchvision.transforms import InterpolationMode
    BICUBIC = InterpolationMode.BICUBIC
except ImportError:
    BICUBIC = Image.BICUBIC

import torch
from transformers.utils import (
    TensorType,
    is_torch_device,
    is_torch_dtype,
    requires_backends,
)

from torchvision.transforms import Compose, ToTensor, Normalize, ToPILImage, RandomResizedCrop, Resize

try:
    from torchvision.transforms import InterpolationMode
    BICUBIC = InterpolationMode.BICUBIC
except ImportError:
    BICUBIC = Image.BICUBIC

from PIL import Image
import torch
import numpy as np
import os
processor_for_vllm = int(os.getenv("PROCESSOR_FOR_VLLM", 0))

def select_best_resolution(original_size, possible_resolutions):
    """
    Selects the best resolution from a list of possible resolutions based on the original size.

    Args:
        original_size (tuple): The original size of the image in the format (width, height).
        possible_resolutions (list): A list of possible resolutions in the format [(width1, height1), (width2, height2), ...].

    Returns:
        tuple: The best fit resolution in the format (width, height).
    """
    original_width, original_height = original_size
    best_fit = None
    max_effective_resolution = 0
    min_wasted_resolution = float("inf")

    for width, height in possible_resolutions:
        # Calculate the downscaled size to keep the aspect ratio
        scale = min(width / original_width, height / original_height)
        downscaled_width, downscaled_height = int(original_width * scale), int(original_height * scale)

        # Calculate effective and wasted resolutions
        effective_resolution = min(downscaled_width * downscaled_height, original_width * original_height)
        wasted_resolution = (width * height) - effective_resolution

        if effective_resolution > max_effective_resolution or (effective_resolution == max_effective_resolution and wasted_resolution < min_wasted_resolution):
            max_effective_resolution = effective_resolution
            min_wasted_resolution = wasted_resolution
            best_fit = (width, height)

    return best_fit 

def divide_to_patches(image, patch_size):
    """
    Divides an image into patches of a specified size.

    Args:
        image (PIL.Image.Image): The input image.
        patch_size (int): The size of each patch.

    Returns:
        list: A list of PIL.Image.Image objects representing the patches.
    """
    patches = []
    width, height = image.size
    for i in range(0, height, patch_size):
        for j in range(0, width, patch_size):
            box = (j, i, j + patch_size, i + patch_size)
            patch = image.crop(box)
            patches.append(patch)

    return patches

def image_size_to_num_patches(image_size, grid_pinpoints, patch_size):
    if not isinstance(grid_pinpoints, list):
        raise TypeError("grid_pinpoints should be a list of tuples or lists")

    # ! VERY IMPORTANT if image_size is tensor, must convert to into tuple, otherwise it will cause wrong calculate
    if not isinstance(image_size, (list, tuple)):
        if not isinstance(image_size, (torch.Tensor, np.ndarray)):
            raise TypeError(f"image_size invalid type {type(image_size)} with value {image_size}")
        image_size = image_size.tolist()

    best_resolution = select_best_resolution(image_size, grid_pinpoints)
    width, height = best_resolution
    num_patches = 0
    # consider change to ceil(height/patch_size)*ceil(width/patch_size) + 1
    for i in range(0, height, patch_size):
        for j in range(0, width, patch_size):
            num_patches += 1
    # add the base patch
    num_patches += 1
    return num_patches

def get_anyres_image_grid_shape(image_size, grid_pinpoints, patch_size):
    """
    Calculate the shape of the image patch grid after the preprocessing for images of any resolution.

    Args:
        image_size (`tuple`):
            The size of the input image in the format (width, height).
        grid_pinpoints (`List`):
            A list containing possible resolutions. Each item in the list should be a tuple or list
            of the form `(height, width)`.
        patch_size (`int`):
            The size of each image patch.

    Returns:
        tuple: The shape of the image patch grid in the format (width, height).
    """
    if not isinstance(grid_pinpoints, list):
        raise TypeError("grid_pinpoints should be a list of tuples or lists")

    # ! VERY IMPORTANT if image_size is tensor, must convert to into tuple, otherwise it will cause wrong calculate
    if not isinstance(image_size, (list, tuple)):
        if not isinstance(image_size, (torch.Tensor, np.ndarray)):
            raise TypeError(
                f"image_size invalid type: {type(image_size)} not valid, should be either list, tuple, np.ndarray or tensor"
            )
        image_size = image_size.tolist()

    width, height = select_best_resolution(image_size, grid_pinpoints)
    return width // patch_size, height // patch_size


# custom transform
class KeeyRatioResize(object):
    def __init__(self, size):
        self.size = size

    def __call__(self, image):
        return keepratio_resize(image, self.size)

def keepratio_resize(image, size, return_scale=False):
    # Resize the image to keep the ratio
    w, h = image.size
    resized_w, resized_h = size
    if w / h > resized_w / resized_h:
        # resize and pad to the right and left
        new_h = int(resized_w*h/w)
        resized_image = image.resize((resized_w, new_h), Image.BICUBIC)

        image = Image.new('RGB', (resized_w, resized_h), (0, 0, 0))
        pad_h = (resized_h - new_h) // 2
        image.paste(resized_image, (0, pad_h))
        scale = resized_w / w
        #image.paste(resized_image, (0, 0))
    else:
        # resize and pad to the top and bottom
        new_w = int(resized_h*w/h)
        resized_image = image.resize((new_w, resized_h), Image.BICUBIC)
        image = Image.new('RGB', (resized_w, resized_h), (0, 0, 0))
        #image.paste(resized_image, (0, 0))
        pad_w = (resized_w - new_w) // 2
        image.paste(resized_image, (pad_w, 0))
        scale = resized_h / h
    if return_scale:
        return image, scale
    return image

def _convert_image_to_rgb(image):
    return image.convert("RGB")

def _transform(img_h, img_w, image_mean=(0.48145466, 0.4578275, 0.40821073), image_std=(0.26862954, 0.26130258, 0.27577711)):
    return Compose([
        # ToPILImage(),
        #RandomResizedCrop((img_h, img_w), scale=(0.5, 1.0), interpolation=BICUBIC),
        #Resize((img_h, img_w), interpolation=BICUBIC),
        _convert_image_to_rgb,
        ToTensor(),
        Normalize(image_mean, image_std),
    ])


def get_hw_multiple_of(image_size, multiple, max_size=None):
    w, h = image_size
    new_w = w if w % multiple == 0 else w + (multiple - w % multiple)
    new_h = h if h % multiple == 0 else h + (multiple - h % multiple)
    if max_size is not None:
        assert isinstance(max_size, (list, tuple)) and len(max_size) == 2
        max_w, max_h = max_size
        assert max_w % multiple == 0 and max_h % multiple == 0
        if new_w > max_w or new_h > max_h:
            # ratio = min(max_w / new_w, max_h / new_h)
            # new_w = int(new_w * ratio)
            # new_h = int(new_h * ratio)
            new_w = min((new_w * max_w) // new_w, (new_w * max_h) // new_h)
            new_h = min((new_h * max_w) // new_w, (new_h * max_h) // new_h)

            new_w = new_w if new_w % multiple == 0 else new_w + (multiple - new_w % multiple)
            new_h = new_h if new_h % multiple == 0 else new_h + (multiple - new_h % multiple)
        assert new_w % multiple == 0 and new_h % multiple == 0
        assert new_w <= max_w and new_h <= max_h
    return new_w, new_h

def resize_multiple_of(image, multiple, max_size=None):
    """
    Resize the image to the multiple of a number.

    Args:
        image (PIL.Image.Image): The input image.
        multiple (int): The number to which the image should be resized.

    Returns:
        PIL.Image.Image: The resized image.
    """
    width, height = image.size
    new_width, new_height = get_hw_multiple_of((width, height), multiple, max_size)
    return image.resize((new_width, new_height), Image.BICUBIC)



class CustomBatchFeature(BatchFeature):
    def convert_to_tensors(self, tensor_type: Optional[Union[str, TensorType]] = None):
        """
        Convert the inner content to tensors.

        Args:
            tensor_type (`str` or [`~utils.TensorType`], *optional*):
                The type of tensors to use. If `str`, should be one of the values of the enum [`~utils.TensorType`]. If
                `None`, no modification is done.
        """
        if tensor_type is None:
            return self

        is_tensor, as_tensor = self._get_is_as_tensor_fns(tensor_type)

        # Do the tensor conversion in batch
        for key, value in self.items():
            if key == "pixel_values":
                for i, image in enumerate(value):
                    if not is_tensor(image):
                        tensor = as_tensor(image)
                        self[key][i] = tensor
                continue
            try:
                if not is_tensor(value):
                    tensor = as_tensor(value)

                    self[key] = tensor
            except:  # noqa E722
                if key == "overflowing_values":
                    raise ValueError("Unable to create tensor returning overflowing values of different lengths. ")
                raise ValueError(
                    "Unable to create tensor, you should probably activate padding "
                    "with 'padding=True' to have batched tensors with the same length."
                )

        return self

    def to(self, *args, **kwargs) -> "BatchFeature":
        """
        Send all values to device by calling `v.to(*args, **kwargs)` (PyTorch only). This should support casting in
        different `dtypes` and sending the `BatchFeature` to a different `device`.

        Args:
            args (`Tuple`):
                Will be passed to the `to(...)` function of the tensors.
            kwargs (`Dict`, *optional*):
                Will be passed to the `to(...)` function of the tensors.

        Returns:
            [`BatchFeature`]: The same instance after modification.
        """
        requires_backends(self, ["torch"])
        import torch  # noqa

        new_data = {}
        device = kwargs.get("device")
        # Check if the args are a device or a dtype
        if device is None and len(args) > 0:
            # device should be always the first argument
            arg = args[0]
            if is_torch_dtype(arg):
                # The first argument is a dtype
                pass
            elif isinstance(arg, str) or is_torch_device(arg) or isinstance(arg, int):
                device = arg
            else:
                # it's something else
                raise ValueError(f"Attempting to cast a BatchFeature to type {str(arg)}. This is not supported.")
        # We cast only floating point tensors to avoid issues with tokenizers casting `LongTensor` to `FloatTensor`
        for k, v in self.items():
            if k == "pixel_values":
                new_data[k] = [v[i].to(*args, **kwargs) for i in range(len(v))]
                continue
            # check if v is a floating point
            if torch.is_floating_point(v):
                # cast and send to device
                new_data[k] = v.to(*args, **kwargs)
            elif device is not None:
                new_data[k] = v.to(device=device)
            else:
                new_data[k] = v
        self.data = new_data
        return self


def as_tensor(value):
    if isinstance(value, (list, tuple)) and len(value) > 0:
        if isinstance(value[0], np.ndarray):
            value = np.array(value)
        elif (
            isinstance(value[0], (list, tuple))
            and len(value[0]) > 0
            and isinstance(value[0][0], np.ndarray)
        ):
            value = np.array(value)
    if isinstance(value, np.ndarray):
        return torch.from_numpy(value)
    else:
        return torch.tensor(value)

class ImageProcessor(BaseImageProcessor):
    model_input_names = ["pixel_values"]

    def __init__(
        self,
        size: Optional[Union[int, Tuple[int, int], Dict[str, int]]] = None,
        image_mean: Optional[Union[float, List[float]]] = None,
        image_std: Optional[Union[float, List[float]]] = None,
        process_image_mode: Optional[str] = 'resize',
        patch_size: Optional[int] = 14,
        image_grid_pinpoints: List = None,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        self.size = size # (width, height)
        self.image_mean = image_mean
        self.image_std = image_std
        self.process_image_mode = process_image_mode
        image_grid_pinpoints = (
            image_grid_pinpoints
            if image_grid_pinpoints is not None
            else [[336, 672], [672, 336], [672, 672], [1008, 336], [336, 1008]]
        )
        self.image_grid_pinpoints = image_grid_pinpoints
        self.patch_size = patch_size

    def preprocess(self,
                    images,
                    return_tensors: Optional[Union[str, TensorType]] = None,
                    data_format: Optional[ChannelDimension] = ChannelDimension.FIRST,
                    input_data_format: Optional[Union[str, ChannelDimension]] = None,
                    **kwargs,
                    ):
        if self.process_image_mode == 'resize':
            return self.resize_preprocess(images, return_tensors, data_format, input_data_format, **kwargs)
        elif self.process_image_mode == 'anyres':
            if processor_for_vllm == 1:
                return self.anyres_for_vllm_preprocess(images, return_tensors, data_format, input_data_format, **kwargs)
            return self.anyres_preprocess(images, return_tensors, data_format, input_data_format, **kwargs)
        elif self.process_image_mode == 'keepratio_resize':
            return self.keepratio_resize_preprocess(images, return_tensors, data_format, input_data_format, **kwargs)
        elif self.process_image_mode == 'dynamic_res':
            return self.dynamic_res_preprocess(images, return_tensors, data_format, input_data_format, **kwargs)
        else:
            raise ValueError(f"Invalid process_image_mode: {self.process_image_mode}")
    
    def resize_preprocess(self, images, return_tensors: Optional[Union[str, TensorType]] = None, data_format: Optional[ChannelDimension] = ChannelDimension.FIRST, input_data_format: Optional[Union[str, ChannelDimension]] = None, **kwargs):
        images = make_list_of_images(images)
        all_images = []
        for image in images:
            resized_image = image.resize(self.size, Image.BICUBIC)
            transform_img = _transform(self.size[1], self.size[0], self.image_mean, self.image_std)(resized_image)
            all_images.append(to_numpy_array(transform_img))

        images = [
            to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
            for image in all_images
        ]

        data = {"pixel_values": images}
        return CustomBatchFeature(data=data, tensor_type=return_tensors)

    def keepratio_resize_preprocess(self, images, return_tensors: Optional[Union[str, TensorType]] = None, data_format: Optional[ChannelDimension] = ChannelDimension.FIRST, input_data_format: Optional[Union[str, ChannelDimension]] = None, **kwargs):
        images = make_list_of_images(images)
        all_images = []
        for image in images:
            resized_image = keepratio_resize(image, self.size)
            transform_img = _transform(self.size[1], self.size[0], self.image_mean, self.image_std)(resized_image)
            all_images.append(to_numpy_array(transform_img))

        images = [
            to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
            for image in all_images
        ]

        data = {"pixel_values": images}
        return CustomBatchFeature(data=data, tensor_type=return_tensors)

    def dynamic_res_preprocess(self, images, return_tensors: Optional[Union[str, TensorType]] = None, data_format: Optional[ChannelDimension] = ChannelDimension.FIRST, input_data_format: Optional[Union[str, ChannelDimension]] = None, **kwargs):
        images = make_list_of_images(images)
        all_images = []
        image_sizes = []
        for image in images:
            ori_w, ori_h = image.size
            image_sizes.append([ori_h, ori_w])
            resized_image = resize_multiple_of(image, self.patch_size, max_size=self.size)
            resized_w, resized_h = resized_image.size
            transform_img = _transform(resized_h, resized_w, self.image_mean, self.image_std)(resized_image)
            all_images.append(to_numpy_array(transform_img))

        images = [
            as_tensor(to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format))
            for image in all_images
        ]

        # data = {"pixel_values": images, "image_sizes": as_tensor(image_sizes)}
        # return data
        data = {"pixel_values": images, "image_sizes": image_sizes}
        #return BatchFeature(data=data, data_format=data_format, tensor_type=return_tensors)
        
        return CustomBatchFeature(data=data, tensor_type=return_tensors)

    def get_image_patches(
        self,
        data: Image,
        image_grid_pinpoints,
    ):
        if not isinstance(image_grid_pinpoints, list):
            raise TypeError("grid_pinpoints must be a list of possible resolutions.")


        best_resolution = select_best_resolution(data.size, image_grid_pinpoints)

        resized_data, scale = keepratio_resize(data, best_resolution, return_scale=True)
        resized_data = divide_to_patches(resized_data, self.size[0])
        ori_data = data.resize(self.size, Image.BICUBIC)
        data = [ori_data] + resized_data
        return data
    
    def pad(
        self,
        image: np.ndarray,
        padding: Union[int, Tuple[int, int], Iterable[Tuple[int, int]]],
        mode: PaddingMode = PaddingMode.CONSTANT,
        constant_values: Union[float, Iterable[float]] = 0.0,
        data_format: Optional[Union[str, ChannelDimension]] = None,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
    ) -> np.ndarray:
        """
        Pads the `image` with the specified `padding` and `mode`. Padding can be in the (`height`, `width`)
        dimension of in the (`num_patches`) dimension. In the second case an iterable if tuples is expected
        as input.

        Args:
            image (`np.ndarray`):
                The image to pad.
            padding (`int` or `Tuple[int, int]` or `Iterable[Tuple[int, int]]`):
                Padding to apply to the edges of the height, width axes. Can be one of three formats:
                - `((before_height, after_height), (before_width, after_width))` unique pad widths for each axis.
                - `((before, after),)` yields same before and after pad for height and width.
                - `(pad,)` or int is a shortcut for before = after = pad width for all axes.
            mode (`PaddingMode`):
                The padding mode to use. Can be one of:
                    - `"constant"`: pads with a constant value.
                    - `"reflect"`: pads with the reflection of the vector mirrored on the first and last values of the
                    vector along each axis.
                    - `"replicate"`: pads with the replication of the last value on the edge of the array along each axis.
                    - `"symmetric"`: pads with the reflection of the vector mirrored along the edge of the array.
            constant_values (`float` or `Iterable[float]`, *optional*):
                The value to use for the padding if `mode` is `"constant"`.
            data_format (`str` or `ChannelDimension`, *optional*):
                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.
                If unset, will use same as the input image.
            input_data_format (`str` or `ChannelDimension`, *optional*):
                The channel dimension format for 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.
                If unset, will use the inferred format of the input image.

        Returns:
            `np.ndarray`: The padded image.

        """

        # call the general `pad` if padding on `height/width`, otherwise it's the `num_patched` dim
        if isinstance(padding, int) or len(padding) != 4:
            return pad(image, padding, mode, constant_values, data_format, input_data_format)

        if input_data_format is None:
            input_data_format = infer_channel_dimension_format(image)
        if mode == PaddingMode.CONSTANT:
            image = np.pad(image, padding, mode="constant", constant_values=constant_values)
        elif mode == PaddingMode.REFLECT:
            image = np.pad(image, padding, mode="reflect")
        elif mode == PaddingMode.REPLICATE:
            image = np.pad(image, padding, mode="edge")
        elif mode == PaddingMode.SYMMETRIC:
            image = np.pad(image, padding, mode="symmetric")
        else:
            raise ValueError(f"Invalid padding mode: {mode}")
        image = (
            to_channel_dimension_format(image, data_format, input_data_format) if data_format is not None else image
        )
        return image

    def _pad_for_batching(
        self,
        pixel_values: List[np.ndarray],
        data_format: Optional[Union[str, ChannelDimension]] = None,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
    ):
        """
        Pads images on the `num_of_patches` dimension with zeros to form a batch of same number of patches.

        Args:
            pixel_values (`List[np.ndarray]`):
                An array of pixel values of each images of shape (`batch_size`, `num_patches`, `image_in_3D`)
            data_format (`str` or `ChannelDimension`, *optional*):
                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.
                If unset, will use same as the input image.
            input_data_format (`str` or `ChannelDimension`, *optional*):
                The channel dimension format for 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.
                If unset, will use the inferred format of the input image.

        Returns:
            List[`np.ndarray`]: The padded images.
        """
        max_patch = max(len(x) for x in pixel_values)
        pixel_values = [
            self.pad(
                image,
                padding=((0, max_patch - image.shape[0]), (0, 0), (0, 0), (0, 0)),
                data_format=data_format,
                input_data_format=input_data_format,
            )
            for image in pixel_values
        ]

        return pixel_values

    def anyres_for_vllm_preprocess(self, images, return_tensors: Optional[Union[str, TensorType]] = None, data_format: Optional[ChannelDimension] = ChannelDimension.FIRST, input_data_format: Optional[Union[str, ChannelDimension]] = None, do_pad: Optional[bool] = None, **kwargs):
        
        images = make_list_of_images(images)
        new_images = []
        image_sizes = []

        for image in images:
            ori_w, ori_h = image.size
            image_sizes.append([ori_h, ori_w])
            image_patches = self.get_image_patches(
                image,
                self.image_grid_pinpoints
            )
            all_images = []
            for image in image_patches:
                transform_img = _transform(self.size[0], self.size[1], self.image_mean, self.image_std)(image)
                img_array = to_numpy_array(transform_img)
                img_array = to_channel_dimension_format(img_array, data_format, input_channel_dim=input_data_format)
                all_images.append(img_array)
                #new_images.append(img_array)
            pixel_values = np.array(all_images)
            new_images.append(pixel_values)
        

        new_images = self._pad_for_batching(new_images)

        data = {"pixel_values": new_images, "image_sizes": image_sizes}
        return BatchFeature(data=data, tensor_type=return_tensors)

    
    def anyres_preprocess(self, images, return_tensors: Optional[Union[str, TensorType]] = None, data_format: Optional[ChannelDimension] = ChannelDimension.FIRST, input_data_format: Optional[Union[str, ChannelDimension]] = None, do_pad: Optional[bool] = None, **kwargs):
        
        images = make_list_of_images(images)
        new_images = []
        image_sizes = []

        for image in images:
            ori_w, ori_h = image.size
            image_sizes.append([ori_h, ori_w])
            image_patches = self.get_image_patches(
                image,
                self.image_grid_pinpoints
            )
            #all_images = []
            for image in image_patches:
                transform_img = _transform(self.size[0], self.size[1], self.image_mean, self.image_std)(image)
                img_array = to_numpy_array(transform_img)
                img_array = to_channel_dimension_format(img_array, data_format, input_channel_dim=input_data_format)
                #all_images.append(img_array)
                new_images.append(img_array)
            #pixel_values = np.array(all_images)
            #new_images.append(pixel_values)
        
        # if do_pad:
        #     new_images = self._pad_for_batching(new_images)

        data = {"pixel_values": new_images, "image_sizes": image_sizes}
        return CustomBatchFeature(data=data, tensor_type=return_tensors)