image_processing_ax_ve.py

import math
from typing import Dict, List, Optional, Union

import numpy as np

from transformers.image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    get_image_size,
    infer_channel_dimension_format,
    is_scaled_image,
    make_flat_list_of_images,
    make_list_of_images,
    load_image,
    valid_images,
    to_numpy_array,
)
from transformers.image_processing_utils import BatchFeature, BaseImageProcessor
from transformers.image_transforms import (
    convert_to_rgb,
    resize,
    to_channel_dimension_format,    
)
from transformers.utils import TensorType, logging
from transformers.models.auto import AutoImageProcessor


logger = logging.get_logger(__name__)


def smart_resize(
    height: int, width: int, factor: int = 32, min_pixels: int = 16 * 16 * 4 * 64, max_pixels: int = 16 * 16 * 4 * 2048
):
    if height < factor or width < factor:
        raise ValueError(f"height:{height} or width:{width} must be larger than factor:{factor}")

    h_bar = round(height / factor) * factor
    w_bar = round(width / factor) * factor
    if h_bar * w_bar > max_pixels:
        beta = math.sqrt((height * width) / max_pixels)
        h_bar = math.floor(height / beta / factor) * factor
        w_bar = math.floor(width / beta / factor) * factor
    elif h_bar * w_bar < min_pixels:
        beta = math.sqrt(min_pixels / (height * width))
        h_bar = math.ceil(height * beta / factor) * factor
        w_bar = math.ceil(width * beta / factor) * factor
    
    assert h_bar % factor == 0, f"h_bar:{h_bar} must be divisible by factor:{factor}"
    assert w_bar % factor == 0, f"w_bar:{w_bar} must be divisible by factor:{factor}"
    return h_bar, w_bar


def resize_small_images(image, min_size=32):
    # `image` is a numpy array here (see to_numpy_array in _preprocess);
    # get_image_size returns (height, width).
    height, width = get_image_size(image)
    factor = 1.
    if height < min_size:
        factor = min_size / height
    if width < min_size:
        factor = max(factor, min_size / width)

    if factor > 1.:
        new_height = round(height * factor + 0.5)
        new_width = round(width * factor + 0.5)
        # transformers' numpy resize takes size=(height, width); matches the
        # resize() call in _preprocess and avoids PIL's (width, height) ordering.
        image = resize(image, size=(new_height, new_width))
    return image


class AXVEImageProcessor(BaseImageProcessor):

    model_input_names = ["pixel_values_images", "image_grid_hw"]

    def __init__(
        self,
        do_resize: bool = True,
        size: Dict[str, int] = None,
        resample: PILImageResampling = PILImageResampling.BICUBIC,
        do_rescale: bool = True,
        rescale_factor: Union[int, float] = 1 / 255,
        do_normalize: bool = True,
        image_mean: Optional[Union[float, List[float]]] = None,
        image_std: Optional[Union[float, List[float]]] = None,
        do_convert_rgb: bool = True,
        min_pixels: int = 16 * 16 * 4 * 64,
        max_pixels: int = 16777216,
        patch_size: int = 16,
        merge_size: int = 2,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)

        if size is not None:
            print("size will be ignored. use min_pixels and max_pixels instead.")
        
        if min_pixels is None or max_pixels is None:
            raise ValueError("min_pixels and max_pixels must be provided")
            
        self.min_pixels = min_pixels
        self.max_pixels = max_pixels
        self.size = size

        self.do_resize = do_resize
        self.resample = resample
        self.do_rescale = do_rescale
        self.rescale_factor = rescale_factor
        self.do_normalize = do_normalize
        self.image_mean = image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
        self.image_std = image_std if image_std is not None else IMAGENET_STANDARD_STD

        self.patch_size = patch_size
        self.merge_size = merge_size
        self.do_convert_rgb = do_convert_rgb

    def get_size_of_processed_image(self, image): 
        factor = self.patch_size * self.merge_size
        image = resize_small_images(image, min_size=factor)
        input_data_format = infer_channel_dimension_format(image)
        height, width = get_image_size(image, channel_dim=input_data_format)
        resized_height, resized_width = smart_resize(
            height,
            width,
            factor=factor,
            min_pixels=self.min_pixels,
            max_pixels=self.max_pixels,
        )
        return resized_height, resized_width

    def _preprocess(
        self,
        images: ImageInput,
        do_resize: Optional[bool] = None,
        size: Optional[dict[str, int]] = None,
        resample: PILImageResampling = None,
        do_rescale: Optional[bool] = None,
        rescale_factor: Optional[float] = None,
        do_normalize: Optional[bool] = None,
        image_mean: Optional[Union[float, list[float]]] = None,
        image_std: Optional[Union[float, list[float]]] = None,
        do_convert_rgb: Optional[bool] = None,
        patch_size: Optional[int] = None,
        merge_size: Optional[int] = None,
        data_format: Optional[ChannelDimension] = ChannelDimension.FIRST,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
    ):
        images = make_list_of_images(images)

        if do_convert_rgb:
            images = [convert_to_rgb(image) for image in images]

        # All transformations expect numpy arrays.
        images = [to_numpy_array(image) for image in images]

        if do_rescale and is_scaled_image(images[0]):
            logger.warning_once(
                "It looks like you are trying to rescale already rescaled images. If the input"
                " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again."
            )
        if input_data_format is None:
            # We assume that all images have the same channel dimension format.
            input_data_format = infer_channel_dimension_format(images[0])
      
        patches_per_turn = []
        hw_grids_per_turn = []
        for image in images:
            if do_resize:
                image = resize_small_images(image, min_size=patch_size * merge_size)
                height, width = get_image_size(image, channel_dim=input_data_format)
                resized_height, resized_width = smart_resize(
                    height,
                    width,
                    factor=patch_size * merge_size,
                    min_pixels=self.min_pixels,
                    max_pixels=self.max_pixels,
                )
                image = resize(
                    image, size=(resized_height, resized_width), resample=resample, input_data_format=input_data_format
                )

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

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

            image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
            
            channel = image.shape[0]
            grid_h, grid_w = resized_height // self.patch_size, resized_width // self.patch_size
            patches = image.reshape(
                1,
                channel,
                grid_h // self.merge_size,
                self.merge_size,
                self.patch_size,
                grid_w // self.merge_size,
                self.merge_size,
                self.patch_size,
            )
            patches = patches.transpose(0, 2, 5, 3, 6, 1, 4, 7)
            patches = patches.reshape(-1, channel, patch_size, patch_size)
            patches = patches.reshape(
                grid_h * grid_w, channel * self.patch_size * self.patch_size
            )     

            hw_grid = (grid_h, grid_w)

            assert patches.shape[0] == hw_grid[0] * hw_grid[1], f"patches.shape[0]: {patches.shape[0]} != hw_grid[0] * hw_grid[1]: {hw_grid[0] * hw_grid[1]}"

            patches_per_turn.append(patches)
            hw_grids_per_turn.append(hw_grid)
        
        return np.concatenate(patches_per_turn, axis=0), hw_grids_per_turn

    def preprocess(
        self,
        images: ImageInput,
        do_resize: Optional[bool] = None,
        size: Optional[dict[str, int]] = None,
        min_pixels: Optional[int] = None,
        max_pixels: Optional[int] = None,
        resample: Optional[PILImageResampling] = None,
        do_rescale: Optional[bool] = None,
        rescale_factor: Optional[float] = None,
        do_normalize: Optional[bool] = None,
        image_mean: Optional[Union[float, list[float]]] = None,
        image_std: Optional[Union[float, list[float]]] = None,
        do_convert_rgb: Optional[bool] = None,
        patch_size: Optional[int] = None,
        merge_size: Optional[int] = None,
        return_tensors: Optional[Union[str, TensorType]] = None,
        data_format: Optional[ChannelDimension] = ChannelDimension.FIRST,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
    ):
        min_pixels = min_pixels if min_pixels is not None else self.min_pixels
        max_pixels = max_pixels if max_pixels is not None else self.max_pixels

        if size is not None:
            raise ValueError("size is not supported. use min_pixels and max_pixels instead")

        do_resize = do_resize if do_resize is not None else self.do_resize

        resample = resample if resample is not None else self.resample
        do_rescale = do_rescale if do_rescale is not None else self.do_rescale
        rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
        image_mean = image_mean if image_mean is not None else self.image_mean
        image_std = image_std if image_std is not None else self.image_std
        patch_size = patch_size if patch_size is not None else self.patch_size
        merge_size = merge_size if merge_size is not None else self.merge_size
        do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb

        if images is not None:
            images = self.fetch_images(images)
            images = make_flat_list_of_images(images)

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

        data = {}
        if images is not None:
            pixel_values, grid_hws = [], []            
            for image in images:
                patches, grid_hw = self._preprocess(
                    image,
                    do_resize=do_resize,
                    size=size,
                    resample=resample,
                    do_rescale=do_rescale,
                    rescale_factor=rescale_factor,
                    do_normalize=do_normalize,
                    image_mean=image_mean,
                    image_std=image_std,
                    patch_size=patch_size,
                    merge_size=merge_size,
                    data_format=data_format,
                    do_convert_rgb=do_convert_rgb,
                    input_data_format=input_data_format,
                )
                pixel_values.extend(patches)
                grid_hws += grid_hw
            pixel_values = np.array(pixel_values)
            grid_hws = np.array(grid_hws)
            data.update({"pixel_values_images": pixel_values, "image_grid_hw": grid_hws})

        return BatchFeature(data=data, tensor_type=return_tensors)

    def get_number_of_image_patches(self, height: int, width: int, images_kwargs=None):

        images_kwargs = images_kwargs or {}
        min_pixels = images_kwargs.get("min_pixels", self.min_pixels)
        max_pixels = images_kwargs.get("max_pixels", self.max_pixels)
        patch_size = images_kwargs.get("patch_size", self.patch_size)
        merge_size = images_kwargs.get("merge_size", self.merge_size)

        factor = patch_size * merge_size
        resized_height, resized_width = smart_resize(
            height, width, factor, min_pixels=min_pixels, max_pixels=max_pixels
        )
        grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
        return grid_h * grid_w    

    def fetch_images(self, image_url_or_urls: Union[str, list[str], list[list[str]]]):

        if isinstance(image_url_or_urls, list):
            return [self.fetch_images(x) for x in image_url_or_urls]
        elif isinstance(image_url_or_urls, str):
            return load_image(image_url_or_urls)
        elif valid_images(image_url_or_urls):
            return image_url_or_urls
        else:
            raise TypeError(f"only a single or a list of entries is supported but got type={type(image_url_or_urls)}")