mlx-community/PaddleOCR-VL-4bit

1# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15"""Image processor class for PaddleOCR-VL."""16 17import math18from typing import Dict, List, Optional, Union19 20import numpy as np21import torch22from transformers.image_processing_utils import BaseImageProcessor, BatchFeature23from torchvision.transforms import functional as TF24from transformers.image_transforms import (25    convert_to_rgb,26    resize,27    to_channel_dimension_format,28)29from transformers.image_utils import (30    OPENAI_CLIP_MEAN,31    OPENAI_CLIP_STD,32    ChannelDimension,33    PILImageResampling,34    get_image_size,35    infer_channel_dimension_format,36    is_scaled_image,37    is_valid_image,38    make_list_of_images,39    to_numpy_array,40    valid_images,41    validate_preprocess_arguments,42)43from transformers.utils import TensorType, is_vision_available, logging44 45 46logger = logging.get_logger(__name__)47 48 49if is_vision_available():50    from PIL import Image51 52ImageInput = Union[53    "PIL.Image.Image",54    np.ndarray,55    "torch.Tensor",56    List["PIL.Image.Image"],57    List[np.ndarray],58    List["torch.Tensor"],59]  # noqa60 61 62VideoInput = Union[63    List["PIL.Image.Image"],64    "np.ndarray",65    "torch.Tensor",66    List["np.ndarray"],67    List["torch.Tensor"],68    List[List["PIL.Image.Image"]],69    List[List["np.ndarrray"]],70    List[List["torch.Tensor"]],71]  # noqa72 73 74def make_batched_images(images) -> List[List[ImageInput]]:75    """76    Accepts images in list or nested list format, and makes a list of images for preprocessing.77 78    Args:79        images (`Union[List[List[ImageInput]], List[ImageInput], ImageInput]`):80            The input image.81 82    Returns:83        list: A list of images.84    """85    if (86        isinstance(images, (list, tuple))87        and isinstance(images[0], (list, tuple))88        and is_valid_image(images[0][0])89    ):90        return [img for img_list in images for img in img_list]91 92    elif isinstance(images, (list, tuple)) and is_valid_image(images[0]):93        return images94 95    elif is_valid_image(images):96        return [images]97 98    raise ValueError(f"Could not make batched images from {images}")99 100 101def adjust_size(size, patch_size):102    num_patches = size // patch_size103    if num_patches % 2 != 0:  # 如果是奇数，减1104        num_patches -= 1105    return num_patches * patch_size106 107 108def make_batched_videos(videos) -> List[VideoInput]:109    if (110        isinstance(videos, (list, tuple))111        and isinstance(videos[0], (list, tuple))112        and is_valid_image(videos[0][0])113    ):114        return videos115 116    elif isinstance(videos, (list, tuple)) and is_valid_image(videos[0]):117        if isinstance(videos[0], Image.Image):118            return [videos]119        elif len(videos[0].shape) == 4:120            return [list(video) for video in videos]121 122    elif is_valid_image(videos) and len(videos.shape) == 4:123        return [list(videos)]124 125    raise ValueError(f"Could not make batched video from {videos}")126 127 128def smart_resize(129    height: int,130    width: int,131    factor: int = 28,132    min_pixels: int = 28 * 28 * 130,133    max_pixels: int = 28 * 28 * 1280,134):135    """Rescales the image so that the following conditions are met:136 137    1. Both dimensions (height and width) are divisible by 'factor'.138 139    2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].140 141    3. The aspect ratio of the image is maintained as closely as possible.142 143    """144    # if height < factor or width < factor:145    #    raise ValueError(f"height:{height} or width:{width} must be larger than factor:{factor}")146    # if int(height < factor//4) + int(width < factor//4):147    #     raise ValueError(f"height:{height} or width:{width} must be larger than factor:{factor//4}")148 149    if height < factor:150        print(f"smart_resize: height={height} < factor={factor}, reset height=factor")151        width = round((width * factor) / height)152        height = factor153 154    if width < factor:155        print(f"smart_resize: width={width} < factor={factor}, reset width=factor")156        height = round((height * factor) / width)157        width = factor158 159    if max(height, width) / min(height, width) > 200:160        raise ValueError(161            f"absolute aspect ratio must be smaller than 200, got {max(height, width) / min(height, width)}"162        )163    h_bar = round(height / factor) * factor164    w_bar = round(width / factor) * factor165    if h_bar * w_bar > max_pixels:166        beta = math.sqrt((height * width) / max_pixels)167        h_bar = math.floor(height / beta / factor) * factor168        w_bar = math.floor(width / beta / factor) * factor169    elif h_bar * w_bar < min_pixels:170        beta = math.sqrt(min_pixels / (height * width))171        h_bar = math.ceil(height * beta / factor) * factor172        w_bar = math.ceil(width * beta / factor) * factor173    return h_bar, w_bar174 175 176class PaddleOCRVLImageProcessor(BaseImageProcessor):177    r"""178    Constructs a Siglip image processor that dynamically resizes images based on the original images.179 180    Args:181        do_resize (`bool`, *optional*, defaults to `True`):182            Whether to resize the image's (height, width) dimensions.183        resample (`PILImageResampling`, *optional*, defaults to `Resampling.BICUBIC`):184            Resampling filter to use when resizing the image.185        do_rescale (`bool`, *optional*, defaults to `True`):186            Whether to rescale the image by the specified scale `rescale_factor`.187        rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):188            Scale factor to use if rescaling the image.189        do_normalize (`bool`, *optional*, defaults to `True`):190            Whether to normalize the image.191        image_mean (`float` or `List[float]`, *optional*, defaults to `[0.48145466, 0.4578275, 0.40821073]`):192            Mean to use if normalizing the image. This is a float or list of floats for each channel in the image.193        image_std (`float` or `List[float]`, *optional*, defaults to `[0.26862954, 0.26130258, 0.27577711]`):194            Standard deviation to use if normalizing the image. This is a float or list of floats for each channel in the image.195        do_convert_rgb (`bool`, *optional*, defaults to `True`):196            Whether to convert the image to RGB.197        min_pixels (`int`, *optional*, defaults to `28 * 28 * 130`):198            The min pixels of the image to resize the image.199        max_pixels (`int`, *optional*, defaults to `28 * 28 * 1670`):200            The max pixels of the image to resize the image.201        patch_size (`int`, *optional*, defaults to 14):202            The spacial patch size of the vision encoder.203        temporal_patch_size (`int`, *optional*, defaults to 2):204            The temporal patch size of the vision encoder.205        merge_size (`int`, *optional*, defaults to 2):206            The merge size of the vision encoder to llm encoder.207    """208 209    model_input_names = [210        "pixel_values",211        "image_grid_thw",212        "pixel_values_videos",213        "video_grid_thw",214    ]215 216    def __init__(217        self,218        do_resize: bool = True,219        resample: PILImageResampling = PILImageResampling.BICUBIC,220        do_rescale: bool = True,221        rescale_factor: Union[int, float] = 1 / 255,222        do_normalize: bool = True,223        image_mean: Optional[Union[float, List[float]]] = None,224        image_std: Optional[Union[float, List[float]]] = None,225        do_convert_rgb: bool = True,226        min_pixels: int = 28 * 28 * 130,227        max_pixels: int = 28 * 28 * 1280,228        patch_size: int = 14,229        temporal_patch_size: int = 1,230        merge_size: int = 2,231        **kwargs,232    ) -> None:233        super().__init__(**kwargs)234        self.do_resize = do_resize235        self.resample = resample236        self.do_rescale = do_rescale237        self.rescale_factor = rescale_factor238        self.do_normalize = do_normalize239        self.image_mean = image_mean if image_mean is not None else OPENAI_CLIP_MEAN240        self.image_std = image_std if image_std is not None else OPENAI_CLIP_STD241        self.min_pixels = min_pixels242        self.max_pixels = max_pixels243        self.patch_size = patch_size244        self.temporal_patch_size = temporal_patch_size245        self.merge_size = merge_size246        self.size = {"min_pixels": min_pixels, "max_pixels": max_pixels}  # not used247        self.do_convert_rgb = do_convert_rgb248 249    def mvit_rescale(self, image: Image.Image, merge_size: int = 2) -> Image.Image:250        try:251            w, h = image.size252        except:253            raise ValueError(str((type(image), image)))254        patch_size = self.patch_size255 256        if (w // patch_size) * (h // patch_size) > self.in_token_limit:257            scale = math.sqrt(258                self.in_token_limit / ((w // patch_size) * (h // patch_size))259            )260            new_w, new_h = int(w * scale), int(h * scale)261 262            image = image.resize((new_w, new_h), Image.Resampling.BICUBIC)263        if self.pad_input:264            new_w, new_h = image.size265            pad_size_h = merge_size * patch_size266            pad_size_w = merge_size * patch_size267 268            pad_h = (pad_size_h - new_h % pad_size_h) % pad_size_h269            pad_w = (pad_size_w - new_w % pad_size_w) % pad_size_w270 271            image = TF.pad(image, (0, 0, pad_w, pad_h))272        else:273            new_w, new_h = image.size274            new_w = new_w - new_w % patch_size275            new_h = new_h - new_h % patch_size276 277            new_w = adjust_size(new_w, patch_size)278            new_h = adjust_size(new_h, patch_size)279 280            image = TF.center_crop(image, (new_h, new_w))281 282        w, h = image.size283        if w // patch_size >= 512 or h // patch_size >= 512:284            new_h = min(patch_size * 510, h)285            new_w = min(patch_size * 510, w)286            image = TF.center_crop(image, (new_h, new_w))287            # raise ValueError("Exceed pos emb")288        return image289 290    def _preprocess(291        self,292        images: Union[ImageInput, VideoInput],293        do_resize: bool = None,294        resample: PILImageResampling = None,295        do_rescale: bool = None,296        rescale_factor: float = None,297        do_normalize: bool = None,298        image_mean: Optional[Union[float, List[float]]] = None,299        image_std: Optional[Union[float, List[float]]] = None,300        do_convert_rgb: bool = None,301        data_format: Optional[ChannelDimension] = ChannelDimension.FIRST,302        input_data_format: Optional[Union[str, ChannelDimension]] = None,303    ):304        """305        Preprocess an image or batch of images. Copy of the `preprocess` method from `CLIPImageProcessor`.306 307        Args:308            images (`ImageInput`):309                Image or batch of images to preprocess. Expects pixel values ranging from 0 to 255. If pixel values range from 0 to 1, set `do_rescale=False`.310            vision_info (`List[Dict]`, *optional*):311                Optional list of dictionaries containing additional information about vision inputs.312            do_resize (`bool`, *optional*, defaults to `self.do_resize`):313                Whether to resize the image.314            resample (`PILImageResampling`, *optional*, defaults to `self.resample`):315                Resampling filter to use if resizing the image. This can be one of the `PILImageResampling` enums.316            do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):317                Whether to rescale the image.318            rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):319                Scale factor to use if rescaling the image.320            do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):321                Whether to normalize the image.322            image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`):323                Mean to use if normalizing the image. Can be a float or a list of floats corresponding to the number of channels in the image.324            image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):325                Standard deviation to use if normalizing the image. Can be a float or a list of floats corresponding to the number of channels in the image.326            do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`):327                Whether to convert the image to RGB.328            data_format (`ChannelDimension`, *optional*, defaults to `ChannelDimension.FIRST`):329                The channel dimension format for the output image. Can be one of:330                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.331                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.332                - Unset: Use the channel dimension format of the input image.333            input_data_format (`ChannelDimension` or `str`, *optional*):334                The channel dimension format for the input image. Can be one of:335                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.336                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.337                - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.   - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.338        """339        images = make_list_of_images(images)340 341        if do_convert_rgb:342            images = [convert_to_rgb(image) for image in images]343 344        # All transformations expect numpy arrays.345        images = [to_numpy_array(image) for image in images]346 347        if is_scaled_image(images[0]) and do_rescale:348            logger.warning_once(349                "It looks like you are trying to rescale already rescaled images. If the input"350                " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again."351            )352        if input_data_format is None:353            # We assume that all images have the same channel dimension format.354            input_data_format = infer_channel_dimension_format(images[0])355 356        height, width = get_image_size(images[0], channel_dim=input_data_format)357        resized_height, resized_width = height, width358        processed_images = []359 360        for image in images:361            if do_resize:362                resized_height, resized_width = smart_resize(363                    height,364                    width,365                    factor=self.patch_size * self.merge_size,366                    min_pixels=self.min_pixels,367                    max_pixels=self.max_pixels,368                )369                image = resize(370                    image,371                    size=(resized_height, resized_width),372                    resample=resample,373                    input_data_format=input_data_format,374                )375 376            if do_rescale:377                image = self.rescale(378                    image, scale=rescale_factor, input_data_format=input_data_format379                )380 381            if do_normalize:382                image = self.normalize(383                    image=image,384                    mean=image_mean,385                    std=image_std,386                    input_data_format=input_data_format,387                )388            image = to_channel_dimension_format(389                image, data_format, input_channel_dim=input_data_format390            )391            processed_images.append(image)392 393        patches = np.array(processed_images)394        if data_format == ChannelDimension.LAST:395            patches = patches.transpose(0, 3, 1, 2)396        if patches.shape[0] == 1:397            patches = np.tile(patches, (self.temporal_patch_size, 1, 1, 1))398        init_patches = patches399        channel = patches.shape[1]400        grid_t = patches.shape[0] // self.temporal_patch_size401        grid_h, grid_w = (402            resized_height // self.patch_size,403            resized_width // self.patch_size,404        )405        patches = patches.reshape(406            grid_t,407            self.temporal_patch_size,408            channel,409            grid_h,410            self.patch_size,411            grid_w,412            self.patch_size,413        )414        patches = patches.transpose(0, 3, 5, 2, 1, 4, 6)415        assert self.temporal_patch_size == 1416        flatten_patches = patches.reshape(417            grid_t * grid_h * grid_w, channel, self.patch_size, self.patch_size418        )419        return flatten_patches, (grid_t, grid_h, grid_w)420 421    def preprocess(422        self,423        images: ImageInput,424        videos: VideoInput = None,425        do_resize: bool = None,426        size: Dict[str, int] = None,427        resample: PILImageResampling = None,428        do_rescale: bool = None,429        rescale_factor: float = None,430        do_normalize: bool = None,431        image_mean: Optional[Union[float, List[float]]] = None,432        image_std: Optional[Union[float, List[float]]] = None,433        do_convert_rgb: bool = None,434        return_tensors: Optional[Union[str, TensorType]] = None,435        data_format: Optional[ChannelDimension] = ChannelDimension.FIRST,436        input_data_format: Optional[Union[str, ChannelDimension]] = None,437    ):438        """439        Args:440            images (`ImageInput`):441                Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If442                passing in images with pixel values between 0 and 1, set `do_rescale=False`.443            videos (`VideoInput`):444                Video to preprocess. Expects a single or batch of videos with pixel values ranging from 0 to 255. If445                passing in videos with pixel values between 0 and 1, set `do_rescale=False`.446            do_resize (`bool`, *optional*, defaults to `self.do_resize`):447                Whether to resize the image.448            size (`Dict[str, int]`, *optional*, defaults to `self.size`):449                Size of the image after resizing. Shortest edge of the image is resized to size["shortest_edge"], with450                the longest edge resized to keep the input aspect ratio.451            resample (`int`, *optional*, defaults to `self.resample`):452                Resampling filter to use if resizing the image. This can be one of the enum `PILImageResampling`. Only453                has an effect if `do_resize` is set to `True`.454            do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):455                Whether to rescale the image.456            rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):457                Rescale factor to rescale the image by if `do_rescale` is set to `True`.458            do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):459                Whether to normalize the image.460            image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`):461                Image mean to use for normalization. Only has an effect if `do_normalize` is set to `True`.462            image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):463                Image standard deviation to use for normalization. Only has an effect if `do_normalize` is set to464                `True`.465            do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`):466                Whether to convert the image to RGB.467            return_tensors (`str` or `TensorType`, *optional*):468                The type of tensors to return. Can be one of:469                - Unset: Return a list of `np.ndarray`.470                - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.471                - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.472                - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.473                - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.474            data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):475                The channel dimension format for the output image. Can be one of:476                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.477                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.478                - Unset: Use the channel dimension format of the input image.479            input_data_format (`ChannelDimension` or `str`, *optional*):480                The channel dimension format for the input image. If unset, the channel dimension format is inferred481                from the input image. Can be one of:482                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.483                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.484                - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.485 486        """487        do_resize = do_resize if do_resize is not None else self.do_resize488        size = size if size is not None else self.size489        resample = resample if resample is not None else self.resample490        do_rescale = do_rescale if do_rescale is not None else self.do_rescale491        rescale_factor = (492            rescale_factor if rescale_factor is not None else self.rescale_factor493        )494        do_normalize = do_normalize if do_normalize is not None else self.do_normalize495        image_mean = image_mean if image_mean is not None else self.image_mean496        image_std = image_std if image_std is not None else self.image_std497        do_convert_rgb = (498            do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb499        )500 501        if images is not None:502            images = make_batched_images(images)503        if videos is not None:504            videos = make_batched_videos(videos)505 506        if images is not None and not valid_images(images):507            raise ValueError(508                "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "509                "torch.Tensor, tf.Tensor or jax.ndarray."510            )511 512        validate_preprocess_arguments(513            rescale_factor=rescale_factor,514            do_normalize=do_normalize,515            image_mean=image_mean,516            image_std=image_std,517            do_resize=do_resize,518            size=size,519            resample=resample,520        )521 522        if images is not None:523            pixel_values, vision_grid_thws = [], []524            for image in images:525                patches, image_grid_thw = self._preprocess(526                    image,527                    do_resize=do_resize,528                    resample=resample,529                    do_rescale=do_rescale,530                    rescale_factor=rescale_factor,531                    do_normalize=do_normalize,532                    image_mean=image_mean,533                    image_std=image_std,534                    data_format=data_format,535                    do_convert_rgb=do_convert_rgb,536                    input_data_format=input_data_format,537                )538                pixel_values.extend(patches)539                vision_grid_thws.append(image_grid_thw)540            pixel_values = np.array(pixel_values)541            vision_grid_thws = np.array(vision_grid_thws)542            data = {"pixel_values": pixel_values, "image_grid_thw": vision_grid_thws}543 544        if videos is not None:545            pixel_values, vision_grid_thws = [], []546            for images in videos:547                patches, video_grid_thw = self._preprocess(548                    images,549                    do_resize=do_resize,550                    resample=resample,551                    do_rescale=do_rescale,552                    rescale_factor=rescale_factor,553                    do_normalize=do_normalize,554                    image_mean=image_mean,555                    image_std=image_std,556                    data_format=data_format,557                    do_convert_rgb=do_convert_rgb,558                    input_data_format=input_data_format,559                )560                pixel_values.extend(patches)561                vision_grid_thws.append(video_grid_thw)562            pixel_values = np.array(pixel_values)563            vision_grid_thws = np.array(vision_grid_thws)564            data = {565                "pixel_values_videos": pixel_values,566                "video_grid_thw": vision_grid_thws,567            }568 569        return BatchFeature(data=data, tensor_type=return_tensors)570