core/transforms/image_transform.py

# Copyright (c) Meta Platforms, Inc. and affiliates.

import math
from functools import reduce
from logging import getLogger
from typing import Any, Callable, Tuple

import numpy as np
import torch
import torchvision.transforms as tv
from PIL import Image
from torchvision.transforms import functional as F
from torchvision.transforms.functional import InterpolationMode

logger = getLogger()


MEAN = (0.5, 0.5, 0.5)
STD = (0.5, 0.5, 0.5)


def get_image_transform(
    vision_input_type: str = "vanilla",
    image_res: int = 336,
    max_num_tiles: int = 1,
    normalize_img: bool = True,
) -> Tuple[Callable, int]:

    if vision_input_type == "thumb+tile":
        transforms = VariableSizeImageTransform(
            size=image_res,
            max_num_tiles=max_num_tiles,
            normalize_img=normalize_img,
            use_thumbnail="before",
        )
    else:
        transforms = ImageTransform(
            size=image_res,
            normalize_img=normalize_img,
        )

    logger.info(
        f"Initalized transforms with: vision_input_type: '{vision_input_type}' and max_num_tiles: {max_num_tiles}."
    )

    return transforms


class ImageTransform(object):
    """
    Image transform will resize the longer edge to a given size and pad the shorter edge with mean pixel value of the image.
    """

    def __init__(
        self,
        size: int = 336,
        normalize_img: bool = True,
    ) -> None:
        self.size = size
        self._mean = MEAN
        self._std = STD

        logger.info(f"ImageTransform size: {self.size}")

        self.to_tensor = tv.ToTensor()
        self.normalize = (
            tv.Normalize(
                mean=self._mean,
                std=self._std,
                inplace=True,
            )
            if normalize_img
            else lambda x: x
        )

    def __call__(self, image: Image.Image):
        w, h = image.size
        image = F.resize(
            image, (self.size, self.size), interpolation=InterpolationMode.BICUBIC
        )
        image = self.to_tensor(image)
        image = self.normalize(image)

        # Add chunk dim to make it compatible with existing dataloaders
        image = image.view(1, 3, self.size, self.size)
        return image, (w, h)

    def _transform_torch_tensor(self, image: torch.Tensor):
        h, w = image.shape[-2:]  # Image shape (C, H, W) or (N, C, H, W)
        image = F.resize(
            image, size=(self.size, self.size), interpolation=InterpolationMode.BICUBIC
        )
        image = (
            image.to(torch.float32) / 255.0
        )  # Convert to float and scale to [0, 1] range
        image = self.normalize(image)
        return image, (w, h)


class VariableSizeImageTransform(object):
    """
    The variable size image transform will resize the image dynamically
    based on the image aspect ratio and the number of image chunks we allow.

    The algorithm will not upsample low-res images to fit a certain aspect
    ratio, because that leads to a significant degradation in image quality.

    For example, if an input image is of size 300x800, and we want to allow
    a maximum of 16 image chunks, it will find the closest aspect ratio that
    is allowed within 16 image chunks, i.e., 2:5 = 2 horizontal patches and
    5 vertical patches, giving a total of 10 chunks.

    The image will then be resized to products of the base size (default is
    224px because MetaCLIP takes that), so in this case it will  be resized to
    2*224:5*224 = 448:1120, where we maintain the original aspect ratio and
    pad with the mean value for the rest. This approach minimizes the amount
    of padding required for any arbitrary resolution.

    The final output will therefore be of shape (11, 3, 224, 224), where 10
    patches are coming from the resizing and chunking, and the first patch
    is a downsampled version of the image that preserves aspect ratios.
    """

    def __init__(
        self,
        size: int = 336,
        normalize_img: bool = True,
        max_num_tiles: int = 1,
        use_thumbnail: str = "no",
        area_limit: bool = False,
    ) -> None:
        self.size = size
        self._mean = MEAN
        self._std = STD

        logger.info(f"VariableSizeImageTransform size: {self.size}")

        self.to_tensor = tv.ToTensor()
        self.normalize = (
            tv.Normalize(
                mean=self._mean,
                std=self._std,
                inplace=True,
            )
            if normalize_img
            else lambda x: x
        )
        self.area_limit = area_limit
        self.max_num_tiles = max_num_tiles
        self.use_thumbnail = use_thumbnail
        if self.use_thumbnail != "no":
            self.thumbnail_transform = ImageTransform(
                size=self.size,
                normalize_img=normalize_img,
            )

    @staticmethod
    def _factors(n: int):
        """Return all factors of a number."""
        return set(
            reduce(
                list.__add__,
                ([i, n // i] for i in range(1, int(n**0.5) + 1) if n % i == 0),
            )
        )

    def _find_supported_aspect_ratios(self):
        """
        This function computes all the allowed aspect ratios for a fixed
        number of input chunks.

        For example, with `num_tiles=5`, it will return:
        {
            0.2: [(1, 5)],
            5.0: [(5, 1)],
            0.25: [(1, 4)],
            1.0: [(2, 2), (1, 1)],
            4.0: [(4, 1)],
            0.3333333333333333: [(1, 3)],
            3.0: [(3, 1)],
            0.5: [(1, 2)],
            2.0: [(2, 1)]
        }
        """
        asp_dict = {}
        for chunk_size in range(self.max_num_tiles, 0, -1):
            _factors = sorted(VariableSizeImageTransform._factors(chunk_size))
            _asp_ratios = [(x, chunk_size // x) for x in _factors]
            for ratio in _asp_ratios:
                k = ratio[0] / ratio[1]
                if k not in asp_dict:
                    asp_dict[k] = [ratio]
                else:
                    asp_dict[k].append(ratio)
        return asp_dict

    def _find_closest_aspect_ratio(self, img_width: int, img_height: int) -> Tuple:
        """
        Given an image width, height and target number of chunks
        this function will find the closest supported aspect ratio.
        """
        tgt_ar = img_width / img_height
        asp_dict = self._find_supported_aspect_ratios()
        cl_d, cl_p = 1e23, None
        if tgt_ar >= 1:
            cl_p = min(
                [k for k in asp_dict.keys() if k <= tgt_ar],
                key=lambda x: abs(x - tgt_ar),
            )
            v = asp_dict[cl_p]
            # select width
            widths = [(idx, self.size * vv[0]) for idx, vv in enumerate(v)]
            tgt_idx = max(widths, key=lambda x: x[1])[0]
        else:
            cl_p = min(
                [k for k in asp_dict.keys() if k > tgt_ar],
                key=lambda x: abs(1 / x - 1 / tgt_ar),
            )
            v = asp_dict[cl_p]
            # select height
            heights = [(idx, self.size * vv[1]) for idx, vv in enumerate(v)]
            tgt_idx = max(heights, key=lambda x: x[1])[0]
        out = v[tgt_idx]
        return out

    def _resize(
        self, image: Image.Image, target_width: int, target_height: int
    ) -> Image.Image:
        # Resize longer edge to given size.
        w, h = image.size
        scale = w / h

        if scale > 1.0:
            # width > height
            new_w = target_width
            new_h = math.floor(new_w / scale)
        else:
            # height >= width
            new_h = target_height
            new_w = math.floor(new_h * scale)

        image = F.resize(image, (new_h, new_w))
        return image

    def _pad(self, image: Image.Image, new_width: int, new_height: int) -> Image.Image:
        mean_per_channel = tuple(
            np.clip(np.array(image).mean(axis=(0, 1)), 0, 255).astype(np.uint8)
        )
        new_im = Image.new(mode="RGB", size=(new_width, new_height), color=(0, 0, 0))  # type: ignore
        new_im.paste(image)
        return new_im

    def _split(self, image: torch.Tensor, ncw: int, nch: int) -> torch.Tensor:
        # Split image into number of required tiles (width x height)
        num_channels, height, width = image.size()
        image = image.view(num_channels, nch, height // nch, ncw, width // ncw)
        # Permute dimensions to reorder the axes
        image = image.permute(1, 3, 0, 2, 4).contiguous()
        # Reshape into the desired output shape (batch_size * 4, num_channels, width/2, height/2)
        image = image.view(ncw * nch, num_channels, height // nch, width // ncw)
        return image

    def _get_image_height_width(
        self, image_width: int, image_height: int, target_width: int, target_height: int
    ) -> Tuple[int, int]:
        """
        Given image width, height and target width, height for the canvas, return the dimensions of how the image would be resized
        with aspect ratio preservation.
        """
        scale = image_width / image_height

        if scale > 1.0:
            # Width is larger than height

            # Rescaling factor is the minimum of the two scaling factors. Else one side would be outside of the canvas.
            rescaling_factor = min(
                target_width / image_width, target_height / image_height
            )

            # Set new width to target width and height to the rescaled height.
            new_w = rescaling_factor * image_width
            new_h = math.floor(new_w / scale)

        else:
            # Height is larger than width

            # Rescaling factor is the minimum of the two scaling factors. Else one side would be outside of the canvas.
            rescaling_factor = min(
                target_width / image_width, target_height / image_height
            )

            # Set new height to target height and width to the rescaled width.
            new_h = rescaling_factor * image_height
            new_w = math.floor(new_h * scale)

        return new_w, new_h

    def _fit_image_to_canvas(
        self, img_width: int, img_height: int, area_limit: bool
    ) -> Any:
        """
        Given an image width, height and target number of chunks this function will see if the image
        can be fit into any of the canvases that can be build from arranging the tiles in a grid.
        If the image can be fit onto several canvases, it will return the canvas where the shorter edge
        of the image will be largest.

        If area_limit is set to True, the tie-breaking prefers the canvas where area is less than 2x the original area.
        """
        # Initialize the optimal canvas to None. If no canvas is found where image fits, function returns None.
        optimal_canvas = None
        optimal_image_width_height = None

        scale = img_width / img_height

        # Gather all potential supported image resolutions and iterate through them to find best match
        potential_arrangements = [
            item
            for sublist in self._find_supported_aspect_ratios().values()
            for item in sublist
        ]
        for n_w, n_h in potential_arrangements:
            # Compute the canvas size
            canvas_width, canvas_height = n_w * self.size, n_h * self.size

            # Check if image can fit into the canvas without downsampling
            if canvas_width >= img_width and canvas_height >= img_height:
                # If we did not find a good canvas yet, we will use the current one
                if optimal_canvas is None:
                    # Set optimal canvas and determine the actual image height and width in the canvas with aspect ratio preserving resampling
                    optimal_canvas = (n_w, n_h)
                    optimal_image_width_height = self._get_image_height_width(
                        image_width=img_width,
                        image_height=img_height,
                        target_width=n_w * self.size,
                        target_height=n_h * self.size,
                    )
                else:
                    # If we already found an optimal canvas before, we will check if the shorter edge of the image will be larger than the current optimal canvas.
                    # This means we can potentially upsample the image resolution which is beneficial to performance.
                    image_width_height = self._get_image_height_width(
                        image_width=img_width,
                        image_height=img_height,
                        target_width=n_w * self.size,
                        target_height=n_h * self.size,
                    )
                    if area_limit:
                        # Prioritize aspect ratio, and choose best within area limit when tied.
                        curr_scale = image_width_height[0] / image_width_height[1]
                        optim_scale = (
                            optimal_image_width_height[0]
                            / optimal_image_width_height[1]
                        )
                        if abs(scale - curr_scale) < abs(scale - optim_scale):
                            # 1. optimize aspect ratio
                            optimal_canvas = (n_w, n_h)
                            optimal_image_width_height = image_width_height
                        elif abs(scale - curr_scale) == abs(scale - optim_scale):
                            # 2. optimize area
                            if (
                                image_width_height[0] * image_width_height[1]
                                < 2 * img_width * img_height
                            ):
                                # 2.1 area is less than 2x the original area
                                optimal_canvas = (n_w, n_h)
                                optimal_image_width_height = image_width_height
                    else:
                        # NOTE: L3V dynamid tiling. Priortize biggest canvas.
                        if (
                            scale < 1.0
                            and (image_width_height[0] >= optimal_image_width_height[0])
                        ) or (
                            scale >= 1.0
                            and (image_width_height[1] >= optimal_image_width_height[1])
                        ):
                            optimal_canvas = (n_w, n_h)
                            optimal_image_width_height = image_width_height
        return optimal_canvas

    def __call__(self, image: Image.Image) -> Tuple[Any, Any]:
        assert isinstance(image, Image.Image), type(image)
        if self.use_thumbnail != "no":
            thumbnail = self.thumbnail_transform(image)[0]

        w, h = image.size
        # Check if the image can be fit to the canvas without downsampling
        ar = self._fit_image_to_canvas(
            img_width=w, img_height=h, area_limit=self.area_limit
        )
        if ar is None:
            # If we did not find a canvas, we have to find the closest aspect ratio and downsample the image
            ar = self._find_closest_aspect_ratio(img_width=w, img_height=h)

        image = F.resize(
            image,
            (ar[1] * self.size, ar[0] * self.size),  # (h, w)
            interpolation=InterpolationMode.BICUBIC,
        )
        image = self._pad(image, ar[0] * self.size, ar[1] * self.size)
        image = self.to_tensor(image)
        image = self.normalize(image)
        image = self._split(image, ar[0], ar[1])  # type: ignore
        if self.use_thumbnail == "before":
            image = torch.cat((thumbnail, image), dim=0)
        elif self.use_thumbnail == "after":
            image = torch.cat((image, thumbnail), dim=0)
        elif self.use_thumbnail == "both":
            image = torch.cat((thumbnail, image, thumbnail), dim=0)

        return image, ar