fan345/Falcon-OCR

1import io2import math3 4import einops as E5import numpy as np6import requests7import torch8from PIL import Image9from transformers.image_processing_utils import BaseImageProcessor10from transformers.image_transforms import convert_to_rgb, resize11from transformers.image_utils import (12    ImageInput,13    get_image_size,14    infer_channel_dimension_format,15    to_numpy_array,16    valid_images,17    validate_preprocess_arguments,18)19 20IMAGE_MEAN = [0.5, 0.5, 0.5]21IMAGE_STD = [0.5, 0.5, 0.5]22 23 24def load_image(image):25    if image is None:26        return None27    if isinstance(image, Image.Image):28        return image29    if isinstance(image, str):30        if image.startswith(("http://", "https://")):31            response = requests.get(image, timeout=10)32            response.raise_for_status()33            return Image.open(io.BytesIO(response.content))34        if image.endswith(".npy"):35            img_array = io.BytesIO(np.load(image))36            return Image.open(img_array)37        return Image.open(image)38    if isinstance(image, np.bytes_):39        return Image.open(io.BytesIO(image))40    if isinstance(image, np.ndarray):41        return Image.fromarray(image)42    raise TypeError(f"Unknown image format {image}")43 44 45def load_images(images_input, min_dimension: int, max_dimension: int):46    images = []47    if images_input is not None:48        for inp in images_input:49            img = load_image(inp)50            img = resize_image_if_necessary(img, min_dimension, max_dimension)51            images.append(img)52    return images53 54 55def resize_image_if_necessary(56    image,57    shortest_dimension=224,58    longest_dimension=896,59):60    original_width, original_height = image.size61    aspect_ratio = original_width / original_height62 63    if (64        shortest_dimension <= original_width <= longest_dimension65        and shortest_dimension <= original_height <= longest_dimension66    ):67        return image68 69    is_vertical_image = original_width < original_height70    if original_width < shortest_dimension or original_height < shortest_dimension:71        if is_vertical_image:72            new_width = shortest_dimension73            new_height = int(new_width / aspect_ratio)74        else:75            new_height = shortest_dimension76            new_width = int(new_height * aspect_ratio)77    else:78        if is_vertical_image:79            new_width = longest_dimension80            new_height = int(new_width / aspect_ratio)81        else:82            new_height = longest_dimension83            new_width = int(new_height * aspect_ratio)84 85    if new_width > longest_dimension:86        new_width = longest_dimension87        new_height = int(new_width / aspect_ratio)88    if new_height > longest_dimension:89        new_height = longest_dimension90        new_width = int(new_height * aspect_ratio)91 92    resized_image = image.resize((new_width, new_height))93    return resized_image94 95 96def smart_resize(97    image,98    factor: int,99    resample,100    input_data_format,101    min_pixels: int = 56 * 56,102    max_pixels: int = 14 * 14 * 4 * 1280,103):104    height, width = get_image_size(image, channel_dim=input_data_format)105    if height < factor or width < factor:106        raise ValueError(f"{height=} or {width=} must be larger than {factor=}")107    if max(height, width) / min(height, width) > 200:108        raise ValueError(109            f"absolute aspect ratio must be smaller than 200, got {max(height, width) / min(height, width)}"110        )111    h_bar = round(height / factor) * factor112    w_bar = round(width / factor) * factor113    if h_bar * w_bar > max_pixels:114        beta = np.sqrt((height * width) / max_pixels)115        h_bar = math.floor(height / beta / factor) * factor116        w_bar = math.floor(width / beta / factor) * factor117    elif h_bar * w_bar < min_pixels:118        beta = np.sqrt(min_pixels / (height * width))119        h_bar = math.ceil(height * beta / factor) * factor120        w_bar = math.ceil(width * beta / factor) * factor121    image = resize(122        image,123        size=(h_bar, w_bar),124        resample=resample,125        input_data_format=input_data_format,126    )127    return image128 129 130class ImageProcessor(BaseImageProcessor):131    def __init__(132        self,133        patch_size,134        merge_size,135        do_resize: bool = True,136        resample: Image.Resampling = Image.Resampling.BICUBIC,137        do_rescale: bool = True,138        rescale_factor: float = 1 / 255,139        do_normalize: bool = True,140        image_mean: float | list[float] | None = None,141        image_std: float | list[float] | None = None,142        do_convert_rgb: bool = True,143        min_pixels: int = 56 * 56,144        max_pixels: int = 28 * 28 * 1280,145        **kwargs,146    ) -> None:147        super().__init__(**kwargs)148        self.do_resize = do_resize149        self.resample = resample150        self.do_rescale = do_rescale151        self.rescale_factor = rescale_factor152        self.do_normalize = do_normalize153        self.image_mean = image_mean or IMAGE_MEAN154        self.image_std = image_std or IMAGE_STD155        self.min_pixels = min_pixels156        self.max_pixels = max_pixels157        self.patch_size = patch_size158        self.merge_size = merge_size159        self.size = {"min_pixels": min_pixels, "max_pixels": max_pixels}160        self.do_convert_rgb = do_convert_rgb161        validate_preprocess_arguments(162            rescale_factor=self.rescale_factor,163            do_normalize=self.do_normalize,164            image_mean=self.image_mean,165            image_std=self.image_std,166            do_resize=self.do_resize,167            size=self.size,168            resample=self.resample,169        )170 171    def _preprocess(self, image: ImageInput, do_rescale=None, do_normalize=None):172        if self.do_convert_rgb:173            image = convert_to_rgb(image)174        image = to_numpy_array(image)175        input_data_format = infer_channel_dimension_format(image)176        if self.do_resize:177            image = smart_resize(178                image,179                factor=self.patch_size * self.merge_size,180                resample=self.resample,181                input_data_format=input_data_format,182                min_pixels=self.min_pixels,183                max_pixels=self.max_pixels,184            )185        if do_rescale or self.do_rescale:186            image = self.rescale(image, scale=self.rescale_factor, input_data_format=input_data_format)187        if do_normalize or self.do_normalize:188            image = self.normalize(189                image=image, mean=self.image_mean, std=self.image_std,190                input_data_format=input_data_format,191            )192        return image193 194    def preprocess(self, images: list[ImageInput] | None, do_rescale=None, do_normalize=None, **kwargs):195        del kwargs196        if images is None:197            return []198        images = [item for item in images if item is not None]199        if not valid_images(images):200            raise ValueError(201                "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "202                "torch.Tensor, tf.Tensor or jax.ndarray."203            )204        pixel_values = []205        for image in images:206            processed_image = self._preprocess(image, do_rescale, do_normalize)207            processed_image = processed_image[None, ...]208            pixel_values.append(processed_image)209        return pixel_values210 211    def batch_images_with_mask(self, pixel_values, max_image_height, max_image_width):212        if pixel_values is None:213            return None214        pixel_values = [item for item in pixel_values if item is not None and len(item) != 0]215        if len(pixel_values) == 0:216            return None217        pixel_values = [torch.from_numpy(img) for img in pixel_values]218        max_temporal = max(img.shape[0] for img in pixel_values)219 220        def pad_image_and_mask(img):221            time_steps, height, width, channels = img.shape222            if channels != 3:223                raise ValueError(f"Expected 3-channel RGB images, got {channels} channels.")224            padding = (0, 0, 0, max_image_width - width, 0, max_image_height - height, 0, max_temporal - time_steps)225            padded_image = torch.nn.functional.pad(img, padding)226            mask = torch.zeros((max_temporal, max_image_height, max_image_width), dtype=torch.long)227            mask[:time_steps, :height, :width] = 1228            return padded_image, mask229 230        padded_pixel_values, padding_masks = zip(*[pad_image_and_mask(img) for img in pixel_values])231        padded_pixel_values = torch.stack(list(padded_pixel_values))232        padding_masks = torch.stack(list(padding_masks))233        return {"pixel_values": padded_pixel_values, "padding_mask": padding_masks}234 235 236# ---------------------------------------------------------------------------237# Positional encoding helpers238# ---------------------------------------------------------------------------239 240def _compute_image_spatial_positions(241    pixel_mask_THW: torch.Tensor,242    spatial_patch_size: int,243    temporal_patch_size: int = 1,244) -> tuple[torch.Tensor, torch.Tensor]:245    mask_thw = E.reduce(246        pixel_mask_THW,247        "(t tp) (h hp) (w wp) -> t h w",248        reduction="any",249        tp=temporal_patch_size,250        hp=spatial_patch_size,251        wp=spatial_patch_size,252    )253    width = E.reduce(mask_thw.sum(dim=-1).int(), "t h -> ", reduction="max")254    height = E.reduce(mask_thw.sum(dim=-2).int(), "t w -> ", reduction="max")255    xlim = torch.sqrt(width / height)256    ylim = torch.sqrt(height / width)257    xpos = torch.linspace(-xlim, xlim, int(width))258    ypos = torch.linspace(-ylim, ylim, int(height))259    wpos, hpos = torch.meshgrid(xpos, ypos, indexing="xy")260    return hpos.flatten(), wpos.flatten()261 262 263def _get_image_token_masks(tokens, config):264    spatial_mask = tokens == config.img_id265    no_increase_mask = (266        spatial_mask267        | (tokens == config.image_reg_1_token_id)268        | (tokens == config.image_reg_2_token_id)269        | (tokens == config.image_reg_3_token_id)270        | (tokens == config.image_reg_4_token_id)271        | (tokens == config.img_end_id)272    )273    return spatial_mask, no_increase_mask274 275 276def get_pos_thw(277    tokens: torch.Tensor,278    pixel_masks_NTHW: torch.Tensor,279    config,280    spatial_patch_size: int,281    temporal_patch_size: int = 1,282    pad_token_id: int = None,283):284    assert pad_token_id is not None285    assert tokens.ndim == 2286    assert pixel_masks_NTHW.ndim == 4287 288    spatial_img_token_mask_BS, no_increase_idx_img_token_mask_BS = _get_image_token_masks(tokens, config)289 290    hpos_parts, wpos_parts = [], []291    for i in range(pixel_masks_NTHW.shape[0]):292        h, w = _compute_image_spatial_positions(pixel_masks_NTHW[i], spatial_patch_size, temporal_patch_size)293        hpos_parts.append(h)294        wpos_parts.append(w)295 296    hpos_N = torch.cat(hpos_parts) if hpos_parts else torch.empty(0)297    wpos_N = torch.cat(wpos_parts) if wpos_parts else torch.empty(0)298 299    expected_tokens = spatial_img_token_mask_BS.sum().item()300    actual_tokens = hpos_N.numel()301    assert actual_tokens == expected_tokens, (302        f"Mismatch between spatial image tokens ({expected_tokens}) and generated positions ({actual_tokens})."303    )304 305    hpos_BS = torch.full_like(tokens, fill_value=torch.nan, dtype=torch.float, device=tokens.device)306    wpos_BS = torch.full_like(tokens, fill_value=torch.nan, dtype=torch.float, device=tokens.device)307    hpos_BS = hpos_BS.masked_scatter_(spatial_img_token_mask_BS, hpos_N)308    wpos_BS = wpos_BS.masked_scatter_(spatial_img_token_mask_BS, wpos_N)309 310    tpos_BS = torch.ones_like(tokens, dtype=torch.float, device=tokens.device)311    tpos_BS[no_increase_idx_img_token_mask_BS] = 0312    tpos_BS = torch.cumsum(tpos_BS, dim=1) - 1313    tpos_BS[tokens == pad_token_id] = 0314 315    hw_pos_BS2 = torch.stack([hpos_BS, wpos_BS], dim=-1)316    return tpos_BS.long(), hw_pos_BS2317 318 319def calculate_image_tokens(image, patch_size, merge_size):320    height, width = get_image_size(image)321    return int((height * width) / (patch_size * patch_size * merge_size * merge_size))322 323 324def tokenize_inputs(prompt, images, tokenizer, config, patch_size, merge_size, max_length):325    img_reg_ids = [326        config.image_reg_1_token_id,327        config.image_reg_2_token_id,328        config.image_reg_3_token_id,329        config.image_reg_4_token_id,330    ]331 332    if images is not None and len(images) > 0:333        image_token_counts = [calculate_image_tokens(image, patch_size, merge_size) for image in images]334    else:335        image_token_counts = []336 337    image_token = tokenizer.convert_ids_to_tokens(config.img_id)338    prompt_chunks = [tokenizer.encode(chunk) for chunk in prompt.split(image_token)]339 340    def insert_separator(X, sep):341        return [ele for sublist in zip(X, sep) for ele in sublist][:-1]342 343    input_ids = []344    offset = 0345    bos_id = getattr(tokenizer, "bos_token_id", None)346    if len(prompt_chunks) > 0 and len(prompt_chunks[0]) > 0 and bos_id is not None and prompt_chunks[0][0] == bos_id:347        offset = 1348        input_ids.append(prompt_chunks[0][0])349 350    separators = []351    for count in image_token_counts:352        tokens = [config.img_id] * count353        image_block = [config.image_cls_token_id, *img_reg_ids, *tokens, config.img_end_id]354        separators.append(image_block)355 356    if len(separators) != 0 and len(separators) != len(prompt_chunks):357        separators.append(separators[-1])358 359    selected_images = []360    if len(separators) == 0:361        input_ids = prompt_chunks[0]362    else:363        for index, x in enumerate(insert_separator(prompt_chunks, separators)):364            if index % 2 != 0:365                if (len(input_ids) + len(x)) < max_length:366                    input_ids.extend(x)367                    selected_images.append(images[index // 2])368            elif index % 2 == 0:369                input_ids.extend(x[offset:])370 371    input_ids = torch.LongTensor(input_ids)372    return input_ids, selected_images373 374 375def process_batch(376    tokenizer,377    config,378    image_prompt_pairs,379    max_length,380    min_dimension,381    max_dimension,382    patch_size=16,383    merge_size=1,384):385    """386    Process a batch of images with text prompts.387    Uses LEFT PADDING for proper batch generation with causal models.388    """389    all_input_ids = []390    all_selected_images = []391    processor_local = ImageProcessor(patch_size, merge_size)392 393    for img_input, prompt in image_prompt_pairs:394        img = load_image(img_input)395        if img is not None:396            img = resize_image_if_necessary(img, min_dimension, max_dimension)397        images = processor_local.preprocess(images=[img] if img else [])398        input_ids, selected_images = tokenize_inputs(399            prompt, images, tokenizer, config, patch_size, merge_size, max_length,400        )401        all_input_ids.append(input_ids)402        all_selected_images.extend(selected_images)403 404    pad_token_id = tokenizer.convert_tokens_to_ids("<|pad|>")405    padded_input_ids = torch.nn.utils.rnn.pad_sequence(406        all_input_ids, batch_first=True, padding_value=pad_token_id, padding_side="left",407    )408 409    processed = processor_local.batch_images_with_mask(all_selected_images, max_dimension, max_dimension)410    assert processed is not None411 412    pos_t, pos_hw = get_pos_thw(413        padded_input_ids, processed["padding_mask"], config, patch_size, pad_token_id=pad_token_id,414    )415 416    return {417        "tokens": padded_input_ids,418        "pixel_values": processed["pixel_values"],419        "pixel_mask": processed["padding_mask"],420        "pos_t": pos_t,421        "pos_hw": pos_hw,422        "pad_token_id": pad_token_id,423    }424