rhymes-ai/Aria-sequential_mlp

1# Copyright 2024 Rhymes AI. All rights reserved.2#3# Licensed to the Apache Software Foundation (ASF) under one4# or more contributor license agreements.  See the NOTICE file5# distributed with this work for additional information6# regarding copyright ownership.  The ASF licenses this file7# to you under the Apache License, Version 2.0 (the8# "License"); you may not use this file except in compliance9# with the License.  You may obtain a copy of the License at10#11#   http://www.apache.org/licenses/LICENSE-2.012#13# Unless required by applicable law or agreed to in writing,14# software distributed under the License is distributed on an15# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY16# KIND, either express or implied.  See the License for the17# specific language governing permissions and limitations18# under the License.19 20from typing import List, Optional, Union21 22import numpy as np23import torch24from PIL import Image, ImageOps25from torchvision import transforms26from transformers import BaseImageProcessor, BatchFeature, TensorType27 28 29def _select_best_resolution(30    img_width: int, img_height: int, target_ratios: List[List[int]], patch_size: int31):32    """33    Selects the best resolution from a list of possible resolutions based on the original size.34 35    Args:36        img_width: the original widths of images.37        img_height: the original heights of images.38        target_ratios (2d numpy array): dimension size (M,2)39        patch_size (int): image patch size40 41    Returns:42        tuple: The best fit resolution in the format (width, height).43    """44 45    aspect_ratio = img_width / img_height46    best_ratio_diff = float("inf")47    best_ratio_w, best_ratio_h = 1, 148    area = np.int32(img_height) * np.int32(img_height)49    for ratio in target_ratios:50        target_aspect_ratio = ratio[0] / ratio[1]51        ratio_diff = abs(aspect_ratio - target_aspect_ratio)52        if ratio_diff < best_ratio_diff:53            best_ratio_diff = ratio_diff54            best_ratio_w, best_ratio_h = ratio[0], ratio[1]55        elif (56            ratio_diff == best_ratio_diff57            and area > 0.5 * patch_size * patch_size * ratio[0] * ratio[1]58        ):59            best_ratio_w, best_ratio_h = ratio[0], ratio[1]60 61    return best_ratio_w, best_ratio_h62 63 64def _split_image(65    image: Image.Image,66    split_image: bool,67    split_ratio: List[List[int]],68    patch_size: int,69) -> List[Image.Image]:70    """71    Split image into multiple patches72 73    Args:74        image (PIL.Image): Input image.75        split_image (bool): Whether to split the image into patches.76        split_ratio (2d numpy array): dimension size (M,2)77        patch_size (int): image patch size78 79    Returns:80        List[PIL.Image]: List of splitted images.81    """82    if split_image:83        ratio_width, ratio_height = _select_best_resolution(84            image.width, image.height, split_ratio, patch_size85        )86        resize_width = patch_size * ratio_width87        resize_height = patch_size * ratio_height88        blocks = ratio_width * ratio_height89        resized_img = image.resize((resize_width, resize_height))90        processed_images = []91        for i in range(blocks):92            box = (93                (i % (resize_width // patch_size)) * patch_size,94                (i // (resize_width // patch_size)) * patch_size,95                ((i % (resize_width // patch_size)) + 1) * patch_size,96                ((i // (resize_width // patch_size)) + 1) * patch_size,97            )98            # split the image99            split_img = resized_img.crop(box)100            processed_images.append(split_img)101        assert len(processed_images) == blocks102        if len(processed_images) != 1:103            processed_images.insert(0, image)104        return processed_images105    else:106        return [image]107 108 109def keep_ratio_resize_and_pixel_mask(110    img: Image.Image, max_size, min_size=336, padding_value=0111):112    """113    Resize an image while maintaining aspect ratio and create a pixel mask.114 115    Args:116        img (PIL.Image): Input image.117        max_size (int): Maximum size for the larger dimension of the image.118        min_size (int, optional): Minimum size for the smaller dimension. Defaults to 336.119        padding_value (int, optional): Value used for padding. Defaults to 0.120 121    Returns:122        tuple: A tuple containing:123            - PIL.Image: Resized and padded image.124            - torch.Tensor: Boolean pixel mask. This mask is a 2D tensor of shape (max_size, max_size) where:125                - True (1) values indicate pixels that belong to the original resized image.126                - False (0) values indicate pixels that are part of the padding.127              The mask helps distinguish between actual image content and padded areas in subsequent processing steps.128    """129    img = img.convert("RGB")130    # rescale the given image, keep the aspect ratio131    scale = max_size / max(img.size)132 133    w, h = img.size134    if w >= h:135        new_size = (max_size, max(int(h * scale), min_size))  # w, h136    else:137        new_size = (max(int(w * scale), min_size), max_size)  # w, h138 139    img_resized = img.resize(new_size, resample=Image.Resampling.BICUBIC)140 141    # padding the right/bottom142    padding_right, padding_bottom = max_size - new_size[0], max_size - new_size[1]143    img_padded = ImageOps.expand(144        img_resized, (0, 0, padding_right, padding_bottom), fill=padding_value145    )146 147    # Create a pixel mask148    pixel_mask = torch.zeros(max_size, max_size)149    pixel_mask[: new_size[1], : new_size[0]] = 1150    pixel_mask = pixel_mask.bool()151    return img_padded, pixel_mask152 153 154class AriaVisionProcessor(BaseImageProcessor):155    """156    A vision processor for the Aria model that handles image preprocessing.157    """158 159    def __init__(160        self,161        max_image_size=980,162        min_image_size=336,163        image_mean=[0.5, 0.5, 0.5],164        image_std=[0.5, 0.5, 0.5],165        **kwargs,166    ):167        """168        Initialize the AriaVisionProcessor.169 170        Args:171            max_image_size (int, optional): Maximum image size. Defaults to 980.172            min_image_size (int, optional): Minimum image size. Defaults to 336.173            mean (list, optional): Mean values for normalization. Defaults to [0.5, 0.5, 0.5].174            std (list, optional): Standard deviation values for normalization. Defaults to [0.5, 0.5, 0.5].175        """176        super().__init__(**kwargs)177 178        self.max_image_size = max_image_size179        self.min_image_size = min_image_size180        self.image_mean = image_mean181        self.image_std = image_std182        self.auto_map = {183            "AutoProcessor": "processing_aria.AriaProcessor",184            "AutoImageProcessor": "vision_processor.AriaVisionProcessor",185        }186 187        # we make the transform a property so that it is lazily initialized,188        # this could avoid the error "TypeError: Object of type Normalize is not JSON serializable"189        # when we used save_pretrained or from_pretrained.190        self._transform = None191        self._set_processor_class("AriaProcessor")192 193    @property194    def transform(self):195        if self._transform is None:196            # Recreate the transform when accessed197            self._transform = transforms.Compose(198                [199                    transforms.ToTensor(),200                    transforms.Normalize(self.image_mean, self.image_std),201                ]202            )203        return self._transform204 205    def __call__(206        self,207        images: Union[Image.Image, List[Image.Image]],208        max_image_size: Optional[int] = 980,209        min_image_size: Optional[int] = 336,210        return_tensors: Optional[Union[str, TensorType]] = "pt",211        split_image: Optional[bool] = False,212        split_ratio: Optional[List[List[int]]] = [213            [1, 2],214            [1, 3],215            [1, 4],216            [1, 5],217            [1, 6],218            [1, 7],219            [1, 8],220            [2, 4],221            [2, 3],222            [2, 2],223            [2, 1],224            [3, 1],225            [3, 2],226            [4, 1],227            [4, 2],228            [5, 1],229            [6, 1],230            [7, 1],231            [8, 1],232        ],233    ):234        """235        Process a list of images.236 237        Args:238            images (list): List of PIL.Image objects.239            max_image_size (int, optional): Override the default max image size. Defaults to None.240            return_tensors (str or TensorType, optional): The type of tensor to return. Defaults to "pt".241            split_image (bool, optional): Whether to split the image. Defaults to False.242            split_ratio (list, optional): The ratio for splitting the image. Defaults to a list of common split ratios.243        Returns:244            BatchFeature: A BatchFeature object containing:245                - 'pixel_values': Tensor of processed image pixel values.246                - 'pixel_mask': Boolean pixel mask. This mask is a 2D tensor of shape (max_size, max_size) where:247                    - True (1) values indicate pixels that belong to the original resized image.248                    - False (0) values indicate pixels that are part of the padding.249                  The mask helps distinguish between actual image content and padded areas in subsequent processing steps.250                - 'num_crops': Tensor of the number of crops for each image.251        """252        max_size = self.max_image_size if max_image_size is None else max_image_size253        min_size = self.min_image_size if min_image_size is None else min_image_size254 255        if max_size not in [490, 980]:256            raise ValueError("max_image_size must be either 490 or 980")257 258        if isinstance(images, Image.Image):259            images = [images]260 261        pixel_values = []262        pixel_masks = []263        num_crops = []264 265        for image in images:266            crop_images = _split_image(image, split_image, split_ratio, max_size)267            num_crops.append(torch.tensor(len(crop_images)))268            for crop_image in crop_images:269                img_padded, pixel_mask = keep_ratio_resize_and_pixel_mask(270                    crop_image, max_size, min_size271                )272                img_padded = self.transform(img_padded)273                pixel_values.append(img_padded)274                pixel_masks.append(pixel_mask)275 276        return BatchFeature(277            data={278                "pixel_values": torch.stack(pixel_values),279                "pixel_mask": torch.stack(pixel_masks),280                "num_crops": torch.stack(num_crops),281            },282            tensor_type=return_tensors,283        )284 285    def preprocess(286        self,287        images,288        max_image_size=None,289        min_image_size=None,290        return_tensors: Optional[Union[str, TensorType]] = None,291        split_image: Optional[bool] = False,292        split_ratio: Optional[List[List[int]]] = [293            [1, 2],294            [1, 3],295            [1, 4],296            [1, 5],297            [1, 6],298            [1, 7],299            [1, 8],300            [2, 4],301            [2, 3],302            [2, 2],303            [2, 1],304            [3, 1],305            [3, 2],306            [4, 1],307            [4, 2],308            [5, 1],309            [6, 1],310            [7, 1],311            [8, 1],312        ],313    ):314        return self.__call__(315            images,316            max_image_size=max_image_size,317            min_image_size=min_image_size,318            return_tensors=return_tensors,319            split_image=split_image,320            split_ratio=split_ratio,321        )322