allenai/MolmoPoint-GUI-8B

1"""Image processor class for Molmo2"""2from typing import Optional, Union3import numpy as np4import einops5import torch6import torchvision.transforms7 8from transformers.image_utils import (9    IMAGENET_STANDARD_MEAN,10    IMAGENET_STANDARD_STD,11    ImageInput,12    PILImageResampling,13    make_flat_list_of_images,14    valid_images,15    to_numpy_array,16)17from transformers.image_transforms import convert_to_rgb18from transformers.processing_utils import ImagesKwargs19from transformers.image_processing_utils import BaseImageProcessor, get_size_dict20from transformers.utils import logging21from transformers.feature_extraction_utils import BatchFeature22from transformers.utils import TensorType, logging23 24 25logger = logging.get_logger(__name__)26 27 28def normalize_image(29    image: np.ndarray,30    image_mean: list[float],31    image_std: list[float],32) -> np.ndarray:33    image -= np.array(image_mean, dtype=np.float32)[None, None, :]34    image /= np.array(image_std, dtype=np.float32)[None, None, :]35    return image36 37 38def resize_image(39    image: np.ndarray,40    desired_output_size: list[int],41    resample: PILImageResampling,42) -> np.ndarray:43    image = torch.permute(torch.from_numpy(image), [2, 0, 1])44    dtype = image.dtype45    if torch.is_floating_point(image):46        in_min = 0.047        in_max = 1.048        resized = torchvision.transforms.Resize(49            desired_output_size,50            resample,51            antialias=False,52        )(image)53        resized = torch.clip(resized, 0.0, 1.0).to(dtype)54    else:55        assert image.dtype == torch.uint8, "SigLIP expects float images or uint8 images, but got {}".format(image.dtype)56        in_min = 0.057        in_max = 255.058        resized = torchvision.transforms.Resize(59            desired_output_size,60            resample,61            antialias=False,62        )(image)63        resized = torch.clip(resized, 0, 255).to(dtype)64 65    resized = resized.to(torch.float32)66    resized = (resized - in_min) / (in_max - in_min)67 68    resized = torch.permute(resized, [1, 2, 0]).numpy()69 70    return resized71 72 73def select_tiling(h, w, patch_size, max_num_crops):74    """Divide in image of size [w, h] in up to max_num_patches of size patch_size"""75    original_size = np.stack([h, w])  # [1, 2]76    original_res = h * w77    tilings = []78    for i in range(1, max_num_crops + 1):79        for j in range(1, max_num_crops + 1):80            if i*j <= max_num_crops:81                tilings.append((i, j))82    # sort so argmin and argmax favour smaller tilings in the event of a tie83    tilings.sort(key=lambda x: (x[0]*x[1], x[0]))84    candidate_tilings = np.array(tilings, dtype=np.int32)  # [n_resolutions, 2]85    candidate_resolutions = candidate_tilings * patch_size  # [n_resolutions, 2]86 87    # How much we would need to scale the image to fit exactly in each tiling88    original_size = np.stack([h, w], dtype=np.float32)  # [1, 2]89 90    # The original size can be zero in rare cases if the image is smaller than the margin91    # In those cases letting the scale become infinite means the tiling is based on the92    # other side, or falls back to the smallest tiling93    with np.errstate(divide='ignore'):94        required_scale_d = candidate_resolutions.astype(np.float32) / original_size,95    required_scale = np.min(required_scale_d, axis=-1, keepdims=True)  # [n_resolutions, 1]96    if np.all(required_scale < 1):97        # We are forced to downscale, so try to minimize the amount of downscaling98        ix = np.argmax(required_scale)99    else:100        # Pick the resolution that required the least upscaling so that it most closely fits the image101        required_scale = np.where(required_scale < 1.0, 10e9, required_scale)102        ix = np.argmin(required_scale)103    return candidate_tilings[ix]104 105 106def build_resized_image(107    image: np.ndarray,108    base_image_input_size: list[int],109    resample: PILImageResampling,110    image_mean: list[float],111    image_std: list[float],112    image_patch_size: int,113) -> tuple[np.ndarray, np.ndarray]:114    resized = resize_image(115        image, base_image_input_size, resample,116    )117    resized = normalize_image(resized, image_mean, image_std)118    if len(resized.shape) == 3:119        resized = np.expand_dims(resized, 0)120    crop_patch_w = base_image_input_size[1] // image_patch_size121    crop_patch_h = base_image_input_size[0] // image_patch_size122    resize_idx = np.arange(crop_patch_w*crop_patch_h).reshape([crop_patch_h, crop_patch_w])123    return resized, resize_idx124 125 126def build_overlapping_crops(127    image: np.ndarray,128    max_crops: int,129    overlap_margins: list[int],130    base_image_input_size: list[int],131    resample: PILImageResampling,132    image_mean: list[float],133    image_std: list[float],134    image_patch_size: int,135) -> tuple[np.ndarray, np.ndarray]:136    """Decompose an image into a set of overlapping crops137 138    :return crop_arr: [n_crops, h, w, 3] The crops139    :return patch_idx: [overlap_patch_h, overlap_patch_w] For each patch in the resized image140                        the crops were extracted from, what patch in `crop_arr` it corresponds to141    """142    original_image_h, original_image_w = image.shape[:2]143    crop_size = base_image_input_size[0]144    assert base_image_input_size[0] == base_image_input_size[1]145 146    left_margin, right_margin = overlap_margins147    total_margin_pixels = image_patch_size * (right_margin + left_margin)  # pixels removed per dim148    crop_patches = base_image_input_size[0] // image_patch_size  # patches per crop dim149    crop_window_patches = crop_patches - (right_margin + left_margin)  # usable patches150    crop_window_size = crop_window_patches * image_patch_size151    crop_patch_w = base_image_input_size[1] // image_patch_size152    crop_patch_h = base_image_input_size[0] // image_patch_size153    original_image_h, original_image_w = image.shape[:2]154    crop_size = base_image_input_size[0]155 156    # Decide how to tile the image, to account for the overlap margins we compute the tiling157    # as if we had an image without the margins and were using a crop size without the margins158    tiling = select_tiling(159        original_image_h - total_margin_pixels,160        original_image_w - total_margin_pixels,161        crop_window_size,162        max_crops,163    )164 165    src = resize_image(166        image,167        [tiling[0]*crop_window_size+total_margin_pixels, tiling[1]*crop_window_size+total_margin_pixels],168        resample,169    )170    src = normalize_image(src, image_mean, image_std)171 172    # Now we have to split the image into crops, and track what patches came from173    # where in `patch_idx_arr`174    n_crops = tiling[0] * tiling[1]175    crop_arr = np.zeros([n_crops, crop_size, crop_size, 3], dtype=src.dtype)176    patch_idx_arr = np.zeros([n_crops, crop_patch_h, crop_patch_w], dtype=np.int32)177    on_crop = 0178    for i in range(tiling[0]):179        # Slide over `src` by `crop_window_size` steps, but extract crops of size `crops_size`180        # which results in overlapping crop windows181        y0 = i*crop_window_size182        for j in range(tiling[1]):183            x0 = j*crop_window_size184            crop_arr[on_crop] = src[y0:y0+crop_size, x0:x0+crop_size]185            patch_idx = np.arange(crop_patch_w*crop_patch_h).reshape(crop_patch_h, crop_patch_w)186            patch_idx += on_crop * crop_patch_h * crop_patch_w187 188            # Mask out idx that are in the overlap region189            if i != 0:190                patch_idx[:left_margin, :] = -1191            if j != 0:192                patch_idx[:, :left_margin] = -1193            if i != tiling[0]-1:194                patch_idx[-right_margin:, :] = -1195            if j != tiling[1]-1:196                patch_idx[:, -right_margin:] = -1197            patch_idx_arr[on_crop] = patch_idx198            on_crop += 1199 200    # `patch_idx_arr` is ordered crop-by-crop, here we transpose `patch_idx_arr`201    # so it is ordered left-to-right order202    patch_idx_arr = np.reshape(203        patch_idx_arr,204        [tiling[0], tiling[1], crop_patch_h, crop_patch_w]205    )206    patch_idx_arr = np.transpose(patch_idx_arr, [0, 2, 1, 3])207    patch_idx_arr = np.reshape(patch_idx_arr, [-1])208 209    # Now get the parts not in the overlap region, so it should map each patch in `src`210    # to the correct patch it should come from in `crop_arr`211    patch_idx_arr = patch_idx_arr[patch_idx_arr >= 0].reshape(212        src.shape[0]//image_patch_size,213        src.shape[1]//image_patch_size,214    )215    return crop_arr, patch_idx_arr216 217 218def batch_pixels_to_patches(array: np.ndarray, patch_size: int) -> np.ndarray:219    """Reshape images of [n_images, h, w, 3] -> [n_images, n_patches, pixels_per_patch]"""220    if len(array.shape) == 3:221        n_crops, h, w = array.shape222        h_patches = h//patch_size223        w_patches = w//patch_size224        array = np.reshape(array, [n_crops, h_patches, patch_size, w_patches, patch_size])225        array = np.transpose(array, [0, 1, 3, 2, 4])226        array = np.reshape(array, [n_crops, h_patches*w_patches, patch_size*patch_size])227        return array228    else:229        n_crops, h, w, c = array.shape230        h_patches = h//patch_size231        w_patches = w//patch_size232        array = np.reshape(array, [n_crops, h_patches, patch_size, w_patches, patch_size, c])233        array = np.transpose(array, [0, 1, 3, 2, 4, 5])234        array = np.reshape(array, [n_crops, h_patches*w_patches, patch_size*patch_size*c])235        return array236 237 238def arange_for_pooling(239    idx_arr: np.ndarray,240    pool_h: int,241    pool_w: int,242) -> np.ndarray:243    h_pad = pool_h * ((idx_arr.shape[0] + pool_h - 1) // pool_h) - idx_arr.shape[0]244    w_pad = pool_w * ((idx_arr.shape[1] + pool_w - 1) // pool_w) - idx_arr.shape[1]245    idx_arr = np.pad(idx_arr, [[h_pad//2, (h_pad+1)//2], [w_pad//2, (w_pad+1)//2]],246                     mode='constant',constant_values=-1)247    return einops.rearrange(248        idx_arr, "(h dh) (w dw) -> h w (dh dw)", dh=pool_h, dw=pool_w)249 250 251def image_to_patches_and_grids(252    image: np.ndarray,253    max_crops: int,254    overlap_margins: list[int],255    base_image_input_size: list[int],256    resample: PILImageResampling,257    image_mean: list[float],258    image_std: list[float],259    image_patch_size: int,260    image_pooling_w: int,261    image_pooling_h: int,262) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:263    """264    :return image_grids, the shape of each (low-res, high-res) image after pooling265    :return crops, the image crops to processes with the ViT266    :return pooled_patch_idx, for each patch_id tokens in `image_tokens`, the indices of the267                                patches in `crops` to pool for that token, masked with -1268    :rturn patch_idx_arr, map patch coordiantes to patch ids269    """270    if isinstance(base_image_input_size, int):271        base_image_input_size = (base_image_input_size, base_image_input_size)272    273    base_image_input_d = image_patch_size274    pooling_w = image_pooling_w275    pooling_h = image_pooling_h276    crop_patch_w = base_image_input_size[1] // base_image_input_d277    crop_patch_h = base_image_input_size[0] // base_image_input_d278 279    crop_arr, patch_idx_arr = build_overlapping_crops(280        image,281        max_crops,282        overlap_margins,283        base_image_input_size,284        resample,285        image_mean,286        image_std,287        image_patch_size,288    )289    pooling_idx = arange_for_pooling(patch_idx_arr, pooling_h, pooling_w)290    h, w = pooling_idx.shape[:2]291    pooling_idx = pooling_idx.reshape([-1, pooling_h*pooling_w])292    293    # Finally do the same for the global image294    resized, resize_idx = build_resized_image(295        image,296        base_image_input_size,297        resample,298        image_mean,299        image_std,300        image_patch_size,301    )302    patch_idx_arr += crop_patch_h*crop_patch_w303    crop_arr = np.concatenate([resized, crop_arr], 0)304 305    resize_idx = arange_for_pooling(resize_idx, pooling_h, pooling_w)306    resized_h, resized_w = resize_idx.shape[:2]307    resize_idx = resize_idx.reshape([-1, pooling_h*pooling_w])308 309    # Global image goes first, so the order of patches in previous crops gets increased310    pooling_idx = np.where(311        pooling_idx >= 0,312        pooling_idx + crop_patch_h*crop_patch_w,313        -1314    )315    pooling_idx = np.concatenate([resize_idx, pooling_idx])316    image_grid = [np.array([resized_h, resized_w, h, w])]317 318    return (319        np.stack(image_grid, 0),320        batch_pixels_to_patches(crop_arr, image_patch_size),321        pooling_idx,322        patch_idx_arr323    )324 325 326class Molmo2ImagesKwargs(ImagesKwargs, total=False):327    max_crops: Optional[int]328    overlap_margins: Optional[list[int]]329    patch_size: Optional[int]330    pooling_size: Optional[list[int]]331 332 333class Molmo2ImageProcessor(BaseImageProcessor):334    r"""335    Constructs a Molmo2 image processor that preprocesses images for the model.336 337    Args:338        size (`dict[str, int]` *optional*, defaults to `{"height": 378, "width": 378}`):339            Size of the image after resizing.340        resample (`PILImageResampling`, *optional*, defaults to `Resampling.BILINEAR`):341            Resampling filter to use when resizing the image.342        image_mean (`float` or `list[float]`, *optional*, defaults to `[0.5, 0.5, 0.5]`):343            Mean to use if normalizing the image. This is a float or list of floats for each channel in the image.344        image_std (`float` or `list[float]`, *optional*, defaults to `[0.5, 0.5, 0.5]`):345            Standard deviation to use if normalizing the image. This is a float or list of floats for each channel in the image.346        do_convert_rgb (`bool`, *optional*, defaults to `True`):347            Whether to convert the image to RGB.348        max_crops (`int`, *optional*, defaults to `8`):349            Maximum number of crops to use per image.350        overlap_margins (`list[int]`, *optional*, defaults to `[4, 4]`):351            Overlap margins to use.352        patch_size (`int`, *optional*, defaults to 14):353            The spatial patch size of the vision encoder.354        pooling_size (`list[int]`, *optional*, defaults to `[2, 2]`):355            The pooling size of the vision adapter.356    """357 358    model_input_names = ["pixel_values", "image_token_pooling", "image_grids", "image_num_crops"]359 360    def __init__(361        self,362        size: Optional[dict[str, int]] = None,363        resample: PILImageResampling = PILImageResampling.BILINEAR,364        image_mean: Optional[Union[float, list[float]]] = None,365        image_std: Optional[Union[float, list[float]]] = None,366        do_convert_rgb: bool = True,367        max_crops: int = 8,368        overlap_margins: list[int] = [4, 4],369        patch_size: int = 14,370        pooling_size: list[int] = [2, 2],371        **kwargs,372    ) -> None:373        super().__init__(**kwargs)374        size = size if size is not None else {"height": 378, "width": 378}375        size = get_size_dict(size, default_to_square=True)376        self.size = size377 378        self.resample = resample379        self.image_mean = image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN380        self.image_std = image_std if image_std is not None else IMAGENET_STANDARD_STD381        self.do_convert_rgb = do_convert_rgb382 383        self.max_crops = max_crops384        self.overlap_margins = overlap_margins385        self.patch_size = patch_size386        self.pooling_size = pooling_size387    388    def preprocess(389        self,390        images: ImageInput,391        size: Optional[dict[str, int]] = None,392        resample: Optional[PILImageResampling] = None,393        image_mean: Optional[Union[float, list[float]]] = None,394        image_std: Optional[Union[float, list[float]]] = None,395        do_convert_rgb: Optional[bool] = None,396        max_crops: Optional[int] = None,397        overlap_margins: Optional[list[int]] = None,398        patch_size: Optional[int] = None,399        pooling_size: Optional[list[int]] = None,400        return_tensors: Optional[Union[str, TensorType]] = None,401        return_pointing_metadata: bool = False,402        **kwargs,403    ) -> BatchFeature:404        """405        Args:406            images (`ImageInput`):407                Image to preprocess.408            size (`dict[str, int]`, *optional*, defaults to `self.size`):409                Size of the image after resizing.410            resample (`PILImageResampling`, *optional*, defaults to `self.resample`):411                Resampling filter to use when resizing the image. This can be one of the enum `PILImageResampling`. Only412                has an effect if `do_resize` is set to `True`.413            image_mean (`float` or `list[float]`, *optional*, defaults to `self.image_mean`):414                Image mean to use for normalization. Only has an effect if `do_normalize` is set to `True`.415            image_std (`float` or `list[float]`, *optional*, defaults to `self.image_std`):416                Image standard deviation to use for normalization. Only has an effect if `do_normalize` is set to417                `True`.418            do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`):419                Whether to convert the image to RGB.420            max_crops (`int`, *optional*, defaults to `self.max_crops`):421                Maximum number of crops to use per image.422            overlap_margins (`list[int]`, *optional*, defaults to `self.overlap_margins`):423                Overlap margins to use.424            patch_size (`int`, *optional*, defaults to `self.patch_size`):425                The spatial patch size of the vision encoder.426            pooling_size (`list[int]`, *optional*, defaults to `self.pooling_size`):427                The pooling size of the vision adapter.428            return_tensors (`str` or `TensorType`, *optional*):429                The type of tensors to return. Can be one of:430                - Unset: Return a list of `np.ndarray`.431                - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.432                - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.433                - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.434                - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.435            return_patch_mappings (bool, optional):436                Whether to return patch mappings used for decoding MolmoPoint points437 438        Returns:439            A `BatchFeature` containing the following keys:440                - `pixel_values`: The preprocessed images.441                - `image_token_pooling`: The indices of the patches in `crops` to pool for each token in `image_tokens`.442                - `image_grids`: The image grids.443                - `image_num_crops`: The number of crops for each image.444        """445        if size is not None:446            if "height" not in size or "width" not in size:447                raise ValueError("size must contain 'height' and 'width' keys.")448        else:449            size = {**self.size}450        451        base_image_input_size = [size["height"], size["width"]]452        453        resample = resample or self.resample454        image_mean = image_mean or self.image_mean455        image_std = image_std or self.image_std456        do_convert_rgb = do_convert_rgb or self.do_convert_rgb457 458        max_crops = max_crops or self.max_crops459        overlap_margins = overlap_margins or self.overlap_margins460        patch_size = patch_size or self.patch_size461        pooling_size = pooling_size or self.pooling_size462 463        image_pooling_h, image_pooling_w = pooling_size464 465        if images is not None:466            images = self.fetch_images(images)467            images = make_flat_list_of_images(images)468        469        if images is not None and not valid_images(images):470            raise ValueError(471                "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "472                "torch.Tensor, tf.Tensor or jax.ndarray."473            )474 475        if do_convert_rgb:476            images = [convert_to_rgb(image) for image in images]477 478        # All transformations expect numpy arrays.479        images = [to_numpy_array(image) for image in images]480 481        data = {}482        patch_mappings = []483        absolute_token_pooling = []484        offset = 0485        if images is not None:486            batch_grids = []487            batch_crops = []488            batch_pooled_patches_idx = []489            batch_num_crops = []490 491            for image in images:492                image_grid, crops, pooled_idx, patch_mapping = image_to_patches_and_grids(493                    image,494                    max_crops,495                    overlap_margins,496                    base_image_input_size,497                    resample,498                    image_mean,499                    image_std,500                    patch_size,501                    image_pooling_w,502                    image_pooling_h,503                )504                batch_grids.append(image_grid)505                batch_crops.append(crops)506                batch_pooled_patches_idx.append(pooled_idx)507                batch_num_crops.append(crops.shape[0])508                if return_pointing_metadata:509                    absolute_token_pooling.append(510                        np.where(pooled_idx >= 0, pooled_idx + offset, -1))511                    patch_mappings.append(patch_mapping + offset)512                    n_patches = np.prod(crops.shape[:2])513                    offset += n_patches514            515            pixel_values = np.concatenate(batch_crops, 0)516            image_token_pooling = np.concatenate(batch_pooled_patches_idx, 0)517            image_grids = np.concatenate(batch_grids, 0)518            image_num_crops = np.array(batch_num_crops)519 520            data.update(521                pixel_values=pixel_values,522                image_token_pooling=image_token_pooling,523                image_grids=image_grids,524                image_num_crops=image_num_crops,525            )526 527        data = BatchFeature(data, tensor_type=return_tensors)528        if return_pointing_metadata:529            data["image_token_pooling_np"] = np.concatenate(absolute_token_pooling, 0) if len(images) else None530            data["subpatch_mapping"] = patch_mappings531            data["image_sizes"] = [x.shape[:2][::-1] for x in images]532        return data533 534 535Molmo2ImageProcessor.register_for_auto_class()