mirror of
https://github.com/allenai/olmocr.git
synced 2025-10-11 16:22:29 +00:00
546 lines
22 KiB
Python
546 lines
22 KiB
Python
"""Image processor class for Molmo"""
|
|
from typing import List, Optional, Union, Mapping
|
|
|
|
import numpy as np
|
|
import einops
|
|
import torch
|
|
import torchvision.transforms
|
|
from torchvision.transforms import InterpolationMode
|
|
from torchvision.transforms.functional import convert_image_dtype
|
|
|
|
from transformers.image_utils import (
|
|
OPENAI_CLIP_MEAN,
|
|
OPENAI_CLIP_STD,
|
|
ImageInput,
|
|
is_valid_image,
|
|
)
|
|
from transformers.processing_utils import ImagesKwargs
|
|
from transformers.image_processing_utils import BaseImageProcessor
|
|
from transformers.utils import logging
|
|
|
|
|
|
logger = logging.get_logger(__name__)
|
|
|
|
|
|
def pad_to_bounding_box(
|
|
image, offset_height, offset_width, target_height,
|
|
target_width, value=0
|
|
):
|
|
height, width = image.shape[:2]
|
|
after_padding_width = target_width - offset_width - width
|
|
after_padding_height = target_height - offset_height - height
|
|
return np.pad(image, [
|
|
[offset_height, after_padding_height],
|
|
[offset_width, after_padding_width],
|
|
[0, 0]
|
|
], constant_values=value)
|
|
|
|
|
|
def normalize_image(image, offset, scale):
|
|
image -= np.array(offset, dtype=np.float32)[None, None, :]
|
|
image /= np.array(scale, dtype=np.float32)[None, None, :]
|
|
return image
|
|
|
|
|
|
def resize_and_pad(
|
|
image,
|
|
desired_output_size,
|
|
resize_method="torch-bilinear",
|
|
pad_value=0,
|
|
normalize=True,
|
|
image_mean=OPENAI_CLIP_MEAN,
|
|
image_std=OPENAI_CLIP_STD,
|
|
):
|
|
desired_height, desired_width = desired_output_size
|
|
height, width = image.shape[:2]
|
|
|
|
# Cast into float32 since the training code did this in float32 and it (very rarely) effects
|
|
# the results after rounding.
|
|
image_scale_y = np.array(desired_height, np.float32) / np.array(height, np.float32)
|
|
image_scale_x = np.array(desired_width, np.float32) / np.array(width, np.float32)
|
|
image_scale = min(image_scale_x, image_scale_y)
|
|
scaled_height = int(np.array(height, np.float32) * image_scale)
|
|
scaled_width = int(np.array(width, np.float32) * image_scale)
|
|
|
|
if resize_method == "tensorflow":
|
|
# This how the original training code did resizing, it can produce slightly different
|
|
# results then using torch resize so we keep it just in case
|
|
import tensorflow as tf
|
|
image = tf.image.convert_image_dtype(tf.constant(image), dtype=tf.float32)
|
|
image = tf.image.resize(
|
|
image,
|
|
[scaled_height, scaled_width],
|
|
method=tf.image.ResizeMethod.BILINEAR,
|
|
antialias=True,
|
|
)
|
|
image = tf.clip_by_value(image, 0.0, 1.0)
|
|
image = image.numpy()
|
|
elif resize_method == "torch-bilinear":
|
|
image = torch.permute(torch.from_numpy(image), [2, 0, 1])
|
|
image = convert_image_dtype(image) # resize in float32 to match the training code
|
|
image = torchvision.transforms.Resize(
|
|
[scaled_height, scaled_width], InterpolationMode.BILINEAR, antialias=True
|
|
)(image)
|
|
image = torch.clip(image, 0.0, 1.0)
|
|
image = torch.permute(image, [1, 2, 0]).numpy()
|
|
else:
|
|
raise NotImplementedError(resize_method)
|
|
|
|
top_pad = (desired_height - scaled_height) // 2
|
|
left_pad = (desired_width - scaled_width) // 2
|
|
padding = [
|
|
[top_pad, desired_height - scaled_height - top_pad],
|
|
[left_pad, desired_width - scaled_width - left_pad],
|
|
[0, 0]
|
|
]
|
|
image_mask = np.pad(np.ones_like(image[:, :, 0], dtype=bool), padding[:2])
|
|
image = np.pad(image, padding, constant_values=pad_value)
|
|
if normalize:
|
|
image = normalize_image(image, offset=image_mean, scale=image_std)
|
|
return image, image_mask
|
|
|
|
|
|
def select_tiling(h, w, patch_size, max_num_patches):
|
|
"""Decide how best to divide in image of size [w, h] in up to max_num_patches of size patch_size"""
|
|
original_size = np.stack([h, w]) # [1, 2]
|
|
original_res = h * w
|
|
tilings = []
|
|
for i in range(1, max_num_patches+1):
|
|
for j in range(1, max_num_patches+1):
|
|
if i*j <= max_num_patches:
|
|
tilings.append((i, j))
|
|
# sort so argmin and argmax favour smaller tilings in the event of a tie
|
|
tilings.sort(key=lambda x: (x[0]*x[1], x[0]))
|
|
candidate_tilings = np.array(tilings, dtype=np.int32) # [n_resolutions, 2]
|
|
candidate_resolutions = candidate_tilings * patch_size # [n_resolutions, 2]
|
|
|
|
# How much we would need to scale the image to fit exactly in each tiling
|
|
original_size = np.stack([h, w], dtype=np.float32) # [1, 2]
|
|
required_scale_d = candidate_resolutions.astype(np.float32) / original_size
|
|
required_scale = np.min(required_scale_d, axis=-1, keepdims=True) # [n_resolutions, 1]
|
|
if np.all(required_scale < 1):
|
|
# We are forced to downscale, so try to minimize the amount of downscaling
|
|
ix = np.argmax(required_scale)
|
|
else:
|
|
# Pick the resolution that required the least upscaling so that it most closely fits the image
|
|
required_scale = np.where(required_scale < 1.0, 10e9, required_scale)
|
|
ix = np.argmin(required_scale)
|
|
return candidate_tilings[ix]
|
|
|
|
|
|
class MolmoImagesKwargs(ImagesKwargs, total=False):
|
|
max_crops: Optional[int]
|
|
overlap_margins: Optional[List[int]]
|
|
base_image_input_size: Optional[List[int]]
|
|
image_token_length_w: Optional[int]
|
|
image_token_length_h: Optional[int]
|
|
image_patch_size: Optional[int]
|
|
image_padding_mask: Optional[bool]
|
|
|
|
|
|
class MolmoImageProcessor(BaseImageProcessor):
|
|
"""Preprocess images and multi-model inputs"""
|
|
|
|
def __init__(
|
|
self,
|
|
max_crops: int = 12,
|
|
overlap_margins: List[int] = (4, 4),
|
|
base_image_input_size: List[int] = (336, 336),
|
|
image_token_length_w: int = 12,
|
|
image_token_length_h: int = 12,
|
|
image_patch_size: int = 14,
|
|
image_padding_mask: bool = True,
|
|
do_normalize: bool = True,
|
|
image_mean: Optional[Union[float, List[float]]] = None,
|
|
image_std: Optional[Union[float, List[float]]] = None,
|
|
**kwargs,
|
|
):
|
|
super().__init__(**kwargs)
|
|
self.max_crops = max_crops
|
|
self.overlap_margins = overlap_margins
|
|
self.base_image_input_size = base_image_input_size
|
|
self.image_token_length_w = image_token_length_w
|
|
self.image_token_length_h = image_token_length_h
|
|
self.image_patch_size = image_patch_size
|
|
self.image_padding_mask = image_padding_mask
|
|
self.do_normalize = do_normalize
|
|
self.image_mean = image_mean if image_mean is not None else OPENAI_CLIP_MEAN
|
|
self.image_std = image_std if image_std is not None else OPENAI_CLIP_STD
|
|
|
|
def image_to_patches_and_tokens(
|
|
self,
|
|
image: ImageInput,
|
|
image_patch_token_id: int,
|
|
image_col_token_id: int,
|
|
image_start_token_id: int,
|
|
image_end_token_id: int,
|
|
max_crops: Optional[int] = None,
|
|
overlap_margins: Optional[List[int]] = None,
|
|
base_image_input_size: Optional[Union[int, List[int]]] = None,
|
|
image_token_length_w: Optional[int] = None,
|
|
image_token_length_h: Optional[int] = None,
|
|
image_patch_size: Optional[int] = None,
|
|
):
|
|
if isinstance(base_image_input_size, int):
|
|
base_image_input_size = (base_image_input_size, base_image_input_size)
|
|
|
|
base_image_input_d = image_patch_size
|
|
tokens_per_image = image_token_length_w * image_token_length_h
|
|
image_base_patch_w = base_image_input_size[1] // base_image_input_d
|
|
image_base_patch_h = base_image_input_size[0] // base_image_input_d
|
|
|
|
original_image_h, original_image_w = image.shape[:2]
|
|
crop_size = base_image_input_size[0]
|
|
|
|
# Discard this many patches from the (left/top, right/bottom) of crops
|
|
left_margin, right_margin = overlap_margins
|
|
# left_margin, right_margin = 2, 2
|
|
assert left_margin % 2 == 0 # Required for compatibility with 2x2 pooling
|
|
total_margin_pixels = base_image_input_d*(right_margin + left_margin) # pixels removed per dim
|
|
crop_patches = base_image_input_size[0] // base_image_input_d # patches per crop dim
|
|
crop_window_patches = crop_patches - (right_margin + left_margin) # usable patches
|
|
crop_window_size = crop_window_patches * base_image_input_d
|
|
tiling = select_tiling(
|
|
original_image_h - total_margin_pixels,
|
|
original_image_w - total_margin_pixels,
|
|
crop_window_size,
|
|
max_crops
|
|
)
|
|
src, img_mask = resize_and_pad(
|
|
image,
|
|
[tiling[0]*crop_window_size+total_margin_pixels, tiling[1]*crop_window_size+total_margin_pixels]
|
|
)
|
|
|
|
# Now we have to split the image into crops, while keeping track of how each patch in the
|
|
# each crop should be ordered in the global image, this require a lot of tricky booking
|
|
n_crops = tiling[0] * tiling[1]
|
|
patches_arr = []
|
|
mask_arr = []
|
|
patch_ordering_arr = []
|
|
|
|
# We assume 2x2 pooling, but can allow padding the right/bottom with extra
|
|
# patches if the number of patches per side is not even
|
|
assert (crop_patches+1)//2 == image_token_length_h
|
|
assert (crop_patches+1)//2 == image_token_length_w
|
|
on = 0
|
|
on_patch = 0
|
|
for i in range(tiling[0]):
|
|
y0 = i*crop_window_size
|
|
if i == 0:
|
|
crop_y0 = 0
|
|
else:
|
|
crop_y0 = left_margin // 2
|
|
|
|
crop_h = image_base_patch_h - (right_margin + left_margin)
|
|
if i == 0:
|
|
crop_h += left_margin
|
|
if i == (tiling[0]-1):
|
|
crop_h += right_margin
|
|
for j in range(tiling[1]):
|
|
x0 = j*crop_window_size
|
|
if j == 0:
|
|
crop_x0 = 0
|
|
else:
|
|
crop_x0 = left_margin // 2
|
|
|
|
crop_w = image_base_patch_w - (right_margin + left_margin)
|
|
if j == 0:
|
|
crop_w += left_margin
|
|
if j == (tiling[1]-1):
|
|
crop_w += right_margin
|
|
|
|
pooled_w = (crop_w + 1) // 2
|
|
pooled_h = (crop_h + 1) // 2
|
|
patch_ordering_arr.append(
|
|
pad_to_bounding_box(
|
|
np.reshape(np.arange(on, on+pooled_h*pooled_w, dtype=np.int32), (pooled_h, pooled_w, 1)),
|
|
crop_y0, crop_x0, image_token_length_h, image_token_length_w, value=-1
|
|
)[:, :, 0]
|
|
)
|
|
patches_arr.append(src[y0:y0+crop_size, x0:x0+crop_size])
|
|
mask_arr.append(img_mask[y0:y0+crop_size, x0:x0+crop_size])
|
|
|
|
on += pooled_h*pooled_w
|
|
on_patch += 1
|
|
patches = np.stack(patches_arr)
|
|
patch_ordering = np.stack(patch_ordering_arr)
|
|
img_mask = np.stack(mask_arr)
|
|
|
|
# Switch to [n_crops, n_patches, pixels_per_patch] format
|
|
image_layout_impatch_w, image_layout_impatch_h = tiling[0], tiling[1]
|
|
patches = einops.rearrange(
|
|
patches, 'p (h dh) (w dw) c -> p (h w) (dh dw c)',
|
|
dh=base_image_input_d,
|
|
dw=base_image_input_d,
|
|
h=image_base_patch_h,
|
|
w=image_base_patch_w
|
|
)
|
|
img_mask = einops.rearrange(
|
|
img_mask, 'p (h dh) (w dw) -> p (h w) (dh dw)',
|
|
dh=base_image_input_d,
|
|
dw=base_image_input_d,
|
|
h=image_base_patch_h,
|
|
w=image_base_patch_w
|
|
)
|
|
|
|
img_mask = img_mask.astype(np.float32).mean(axis=-1)
|
|
patch_ordering = np.reshape(patch_ordering, [-1])
|
|
valid = patch_ordering >= 0
|
|
|
|
# Transpose order, to get left-to-right order instead of crop-by-crop order
|
|
patch_ordering_rh = np.reshape(
|
|
patch_ordering,
|
|
[tiling[0], tiling[1], image_token_length_h, image_token_length_w]
|
|
)
|
|
patch_ordering_rh = np.transpose(patch_ordering_rh, [0, 2, 1, 3])
|
|
patch_ordering_rh = np.reshape(patch_ordering_rh, [-1])
|
|
|
|
# The transpose will screw up which patches are masked, project the
|
|
# new order into sparse structure of `patch_ordering` to fix this
|
|
patch_ordering[valid] = patch_ordering_rh[patch_ordering_rh >= 0]
|
|
|
|
# Now build the output tokens
|
|
h = tiling[0] * crop_window_patches + (right_margin+left_margin)
|
|
w = tiling[1] * crop_window_patches + (right_margin+left_margin)
|
|
per_row = np.full(
|
|
((w+1)//2,),
|
|
image_patch_token_id,
|
|
)
|
|
per_row = np.concatenate([per_row, [image_col_token_id]], 0)
|
|
|
|
joint = np.tile(per_row, [(h+1)//2])
|
|
joint = [
|
|
[image_start_token_id],
|
|
joint,
|
|
[image_end_token_id]
|
|
]
|
|
|
|
# Finally do the same for the global image
|
|
resized, _ = resize_and_pad(image, base_image_input_size)
|
|
resized = einops.rearrange(
|
|
resized, '(h dh) (w dw) c -> (h w) (dh dw c)',
|
|
dh=base_image_input_d,
|
|
dw=base_image_input_d,
|
|
h=image_base_patch_h,
|
|
w=image_base_patch_w
|
|
)
|
|
patches = np.concatenate([np.expand_dims(resized, 0), patches], 0)
|
|
|
|
# Global image goes first, so the order of patches in previous crops gets increased
|
|
patch_ordering = np.where(
|
|
patch_ordering >= 0,
|
|
patch_ordering + tokens_per_image,
|
|
-1
|
|
)
|
|
patch_ordering = np.concatenate([np.arange(0, tokens_per_image), patch_ordering], 0)
|
|
per_row = np.full(
|
|
(image_token_length_w,),
|
|
image_patch_token_id,
|
|
)
|
|
per_row = np.concatenate([per_row, [image_col_token_id]], 0)
|
|
extra_tokens = np.tile(per_row, [image_token_length_h])
|
|
joint = [
|
|
[image_start_token_id],
|
|
extra_tokens,
|
|
[image_end_token_id],
|
|
] + joint
|
|
|
|
joint = np.concatenate(joint, 0)
|
|
img_mask = np.pad(img_mask, [[0, 1], [0, 0]], constant_values=-1)
|
|
return patches, joint, patch_ordering, img_mask
|
|
|
|
def build_image_input_idx(
|
|
self,
|
|
image_tokens: np.ndarray,
|
|
patch_order: np.ndarray,
|
|
image_patch_token_id: int,
|
|
no_image: Optional[bool] = None,
|
|
image_token_length_w: Optional[int] = None,
|
|
image_token_length_h: Optional[int] = None,
|
|
):
|
|
"""Converts `patch_order` into a mapping of token_id -> patch_id"""
|
|
|
|
tokens_per_image = image_token_length_w * image_token_length_h
|
|
if no_image is not None and no_image:
|
|
return np.zeros((0, tokens_per_image), np.int32)
|
|
|
|
# Indices to insert the patches
|
|
image_input_idx = image_tokens == image_patch_token_id
|
|
image_input_idx = np.nonzero(image_input_idx)[0].astype(np.int32)
|
|
|
|
if patch_order is not None:
|
|
n_tokens = image_input_idx.shape[0]
|
|
patch_order = np.reshape(patch_order, [-1])
|
|
n_patches = patch_order.shape[0]
|
|
|
|
valid = patch_order >= 0
|
|
n_valid_patches = valid.sum()
|
|
assert len(image_input_idx) == n_valid_patches
|
|
|
|
sorted_patch_ixs = np.zeros([n_tokens], np.int32)
|
|
sorted_patch_ixs[patch_order[valid]] = np.arange(n_valid_patches, dtype=np.int32)
|
|
|
|
# Project the inverted mapping into same sparse structure
|
|
sorted_patch_ixs_ex = np.full(np.shape(patch_order), -1)
|
|
sorted_patch_ixs_ex[valid] = sorted_patch_ixs
|
|
|
|
# Do the gather and then re-masked outputs that were masked in `sorted_patch_ixs`
|
|
valid = (sorted_patch_ixs_ex >= 0).astype(np.int32)
|
|
image_input_idx = image_input_idx[sorted_patch_ixs_ex*valid]
|
|
image_input_idx = image_input_idx*valid - 100*(1 - valid)
|
|
image_input_idx = np.reshape(image_input_idx, [-1, tokens_per_image])
|
|
return image_input_idx
|
|
|
|
def preprocess(
|
|
self,
|
|
image: np.ndarray,
|
|
image_patch_token_id: int,
|
|
image_col_token_id: int,
|
|
image_start_token_id: int,
|
|
image_end_token_id: int,
|
|
max_crops: Optional[int] = None,
|
|
overlap_margins: Optional[List[int]] = None,
|
|
base_image_input_size: Optional[Union[int, List[int]]] = None,
|
|
image_token_length_w: Optional[int] = None,
|
|
image_token_length_h: Optional[int] = None,
|
|
image_patch_size: Optional[int] = None,
|
|
**kwargs,
|
|
):
|
|
"""Preprocesses an image
|
|
|
|
Returns:
|
|
crops: (n_crops, n_patches, patch_dim) individual crops, `n_crops` might
|
|
change between images but the other dimension are fixed
|
|
tokens: (n_tokens,) int32 tokens, pad tokens indicate where to insert the
|
|
patch features, might include other special tokens as well
|
|
image_idx: (n_crops, n_patches) index in `tokens` to put the patch features from the
|
|
crops after pooling, negative values indicates patches features to exclude
|
|
padding_mask: (n_crops, n_patches) what percent of each crop is padding, can be None
|
|
if the image mask is not being used.
|
|
"""
|
|
|
|
max_crops = max_crops or self.max_crops
|
|
overlap_margins = overlap_margins or self.overlap_margins
|
|
base_image_input_size = base_image_input_size or self.base_image_input_size
|
|
image_token_length_w = image_token_length_w or self.image_token_length_w
|
|
image_token_length_h = image_token_length_h or self.image_token_length_h
|
|
image_patch_size = image_patch_size or self.image_patch_size
|
|
|
|
crops, image_tokens, patch_ordering, img_mask = self.image_to_patches_and_tokens(
|
|
image,
|
|
image_patch_token_id,
|
|
image_col_token_id,
|
|
image_start_token_id,
|
|
image_end_token_id,
|
|
max_crops,
|
|
overlap_margins,
|
|
base_image_input_size,
|
|
image_token_length_w,
|
|
image_token_length_h,
|
|
image_patch_size,
|
|
)
|
|
patch_idx = self.build_image_input_idx(
|
|
image_tokens,
|
|
patch_ordering,
|
|
image_patch_token_id,
|
|
image_token_length_w=image_token_length_w,
|
|
image_token_length_h=image_token_length_h,
|
|
)
|
|
return crops, image_tokens, patch_idx, img_mask
|
|
|
|
def multimodal_preprocess(
|
|
self,
|
|
images: np.ndarray,
|
|
tokens: List[int],
|
|
image_idx: np.ndarray,
|
|
sequence_length: int,
|
|
image_patch_token_id: int,
|
|
image_col_token_id: int,
|
|
image_start_token_id: int,
|
|
image_end_token_id: int,
|
|
**kwargs,
|
|
):
|
|
"""Merge images and text tokens into multi-modal features for the model
|
|
|
|
:param images: images to use as input
|
|
:param tokens: input text tokens
|
|
:param image_idx: where to insert the images into `tokens`
|
|
:params image_patch_token_id: id to use of tokens that will contain image features
|
|
:params image_col_token_id: token id for image column special tokens
|
|
:params image_start_token_id: token id for image start special tokens
|
|
:params image_end_token_id: token id for image end special tokens
|
|
:params kwargs: override preprocessor default args
|
|
"""
|
|
max_total_crops = kwargs.get("max_crops") or self.max_crops
|
|
image_token_length_w = kwargs.get("image_token_length_w") or self.image_token_length_w
|
|
image_token_length_h = kwargs.get("image_token_length_h") or self.image_token_length_h
|
|
image_patch_size = kwargs.get("image_patch_size") or self.image_patch_size
|
|
base_image_input_size = kwargs.get("base_image_input_size") or self.base_image_input_size
|
|
image_num_patch = (
|
|
base_image_input_size[0] // image_patch_size,
|
|
base_image_input_size[1] // image_patch_size,
|
|
)
|
|
image_padding_mask = kwargs.get("image_padding_mask") or self.image_padding_mask
|
|
|
|
tokens_per_image = image_token_length_w * image_token_length_h
|
|
n_pixels = image_patch_size * image_patch_size * 3
|
|
n_patches = image_num_patch[0] * image_num_patch[1]
|
|
|
|
if images is None:
|
|
return {
|
|
"input_ids": tokens,
|
|
}
|
|
else:
|
|
n = len(images)
|
|
all_crops = []
|
|
all_image_idx = []
|
|
out_tokens = []
|
|
all_crop_masks = []
|
|
|
|
for ix in range(n):
|
|
token_ix = image_idx[ix]
|
|
crops, image_tokens, patch_idx, img_mask = self.preprocess(
|
|
images[ix],
|
|
image_patch_token_id,
|
|
image_col_token_id,
|
|
image_start_token_id,
|
|
image_end_token_id,
|
|
**kwargs,
|
|
)
|
|
|
|
if token_ix == -1: # -1 is an image inserted at the very start
|
|
start = 0
|
|
token_ix = 0
|
|
end = 0
|
|
else:
|
|
start = 0 if ix == 0 else image_idx[ix-1] + 1
|
|
end = token_ix + 1
|
|
|
|
all_image_idx.append(patch_idx + token_ix)
|
|
all_crops.append(crops)
|
|
out_tokens.append(tokens[start:token_ix])
|
|
out_tokens.append(image_tokens)
|
|
if ix == (n - 1):
|
|
out_tokens.append(tokens[end:])
|
|
if image_padding_mask:
|
|
all_crop_masks.append(img_mask)
|
|
|
|
input_ids = np.concatenate(out_tokens, 0)
|
|
images = np.concatenate(all_crops, 0)
|
|
image_input_idx = np.concatenate(all_image_idx, 0)
|
|
if image_padding_mask:
|
|
image_masks = np.concatenate(all_crop_masks, 0)
|
|
else:
|
|
image_masks = None
|
|
|
|
out = {
|
|
"input_ids": input_ids,
|
|
"images": images,
|
|
"image_input_idx": image_input_idx
|
|
}
|
|
if image_masks is not None:
|
|
out["image_masks"] = image_masks
|
|
return out
|
|
|
|
|
|
MolmoImageProcessor.register_for_auto_class() |