| |
| |
| |
| """Image processor class for BrainOCR.""" |
|
|
| |
| import math |
|
|
| import numpy as np |
| import torchvision.transforms as transforms |
| from transformers import AutoImageProcessor |
| from transformers.image_processing_utils import BaseImageProcessor, BatchFeature |
| from transformers.image_transforms import ( |
| convert_to_rgb, |
| ) |
| from transformers.image_utils import ( |
| OPENAI_CLIP_MEAN, |
| OPENAI_CLIP_STD, |
| ChannelDimension, |
| ImageInput, |
| PILImageResampling, |
| make_flat_list_of_images, |
| make_list_of_images, |
| valid_images, |
| validate_preprocess_arguments, |
| ) |
| from transformers.utils import TensorType, logging |
| from transformers.video_utils import VideoInput, make_batched_videos |
|
|
| logger = logging.get_logger(__name__) |
|
|
|
|
| def smart_resize( |
| height: int, |
| width: int, |
| factor: int = 16, |
| min_pixels: int = 512 * 512, |
| max_pixels: int = 2048 * 2048, |
| ): |
| """Rescales the image so that the following conditions are met: |
| |
| 1. Both dimensions (height and width) are divisible by 'factor'. |
| |
| 2. The total number of pixels is within the range ['min_pixels', 'max_pixels']. |
| |
| 3. The aspect ratio of the image is maintained as closely as possible. |
| |
| """ |
| if max(height, width) / min(height, width) > 200: |
| raise ValueError( |
| "absolute aspect ratio must be smaller than 200, got " |
| f"{max(height, width) / min(height, width)}" |
| ) |
| h_bar = round(height / factor) * factor |
| w_bar = round(width / factor) * factor |
| if h_bar * w_bar > max_pixels: |
| beta = math.sqrt((height * width) / max_pixels) |
| h_bar = max(factor, math.floor(height / beta / factor) * factor) |
| w_bar = max(factor, math.floor(width / beta / factor) * factor) |
| elif h_bar * w_bar < min_pixels: |
| beta = math.sqrt(min_pixels / (height * width)) |
| h_bar = math.ceil(height * beta / factor) * factor |
| w_bar = math.ceil(width * beta / factor) * factor |
| return h_bar, w_bar |
|
|
|
|
| class BrainOCRImageProcessor(BaseImageProcessor): |
| model_input_names = [ |
| "pixel_values", |
| "image_grid_thw", |
| "pixel_values_videos", |
| "video_grid_thw", |
| ] |
|
|
| def __init__( |
| self, |
| do_resize: bool = True, |
| size: dict[str, int] | None = None, |
| resample: PILImageResampling = PILImageResampling.BICUBIC, |
| do_rescale: bool = True, |
| rescale_factor: int | float = 1 / 255, |
| do_normalize: bool = True, |
| image_mean: float | list[float] | None = None, |
| image_std: float | list[float] | None = None, |
| do_convert_rgb: bool = True, |
| min_pixels: int | None = None, |
| max_pixels: int | None = None, |
| patch_size: int = 16, |
| temporal_patch_size: int = 2, |
| merge_size: int = 2, |
| **kwargs, |
| ) -> None: |
| super().__init__(**kwargs) |
| if size is not None and ( |
| "shortest_edge" not in size or "longest_edge" not in size |
| ): |
| raise ValueError( |
| "size must contain 'shortest_edge' and 'longest_edge' keys." |
| ) |
| else: |
| size = {"shortest_edge": 512 * 512, "longest_edge": 2048 * 2048} |
| if min_pixels is not None: |
| size["shortest_edge"] = min_pixels |
| if max_pixels is not None: |
| size["longest_edge"] = max_pixels |
| self.min_pixels = size["shortest_edge"] |
| self.max_pixels = size["longest_edge"] |
| self.size = size |
|
|
| self.do_resize = do_resize |
| self.resample = resample |
| self.do_rescale = do_rescale |
| self.rescale_factor = rescale_factor |
| 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 |
|
|
| self.patch_size = patch_size |
| self.temporal_patch_size = temporal_patch_size |
| self.merge_size = merge_size |
| self.do_convert_rgb = do_convert_rgb |
|
|
| def _preprocess( |
| self, |
| images: ImageInput | VideoInput, |
| do_resize: bool | None = None, |
| size: dict[str, int] | None = None, |
| resample: PILImageResampling = None, |
| do_rescale: bool | None = None, |
| rescale_factor: float | None = None, |
| do_normalize: bool | None = None, |
| image_mean: float | list[float] | None = None, |
| image_std: float | list[float] | None = None, |
| patch_size: int = 16, |
| temporal_patch_size: int = 2, |
| merge_size: int = 2, |
| do_convert_rgb: bool | None = None, |
| data_format: ChannelDimension | None = ChannelDimension.FIRST, |
| input_data_format: str | ChannelDimension | None = None, |
| ): |
| images = make_list_of_images(images) |
|
|
| if do_convert_rgb: |
| images = [convert_to_rgb(image) for image in images] |
|
|
| width, height = images[0].width, images[0].height |
| resized_width, resized_height = width, height |
| processed_images = [] |
| for image in images: |
| if do_resize: |
| resized_height, resized_width = smart_resize( |
| height=height, |
| width=width, |
| factor=patch_size * merge_size, |
| min_pixels=self.min_pixels, |
| max_pixels=self.max_pixels, |
| ) |
| image = image.resize((resized_width, resized_height)) |
|
|
| if do_normalize: |
| image = transforms.Compose( |
| [ |
| transforms.ToTensor(), |
| transforms.Normalize(self.image_mean, self.image_std), |
| ] |
| )(image) |
| processed_images.append(image) |
|
|
| patches = np.array(processed_images) |
| channel = patches.shape[1] |
| grid_t = patches.shape[0] // temporal_patch_size |
| grid_h, grid_w = resized_height // patch_size, resized_width // patch_size |
| patches = patches.reshape( |
| 1, |
| channel, |
| grid_h // merge_size, |
| merge_size, |
| patch_size, |
| grid_w // merge_size, |
| merge_size, |
| patch_size, |
| ) |
| patches = patches.transpose(0, 2, 3, 5, 6, 1, 4, 7) |
| flatten_patches = patches.reshape( |
| 1 * grid_h * grid_w, channel * patch_size * patch_size |
| ) |
|
|
| return flatten_patches, (grid_t, grid_h, grid_w) |
|
|
| def preprocess( |
| self, |
| images: ImageInput, |
| videos: VideoInput = None, |
| do_resize: bool | None = None, |
| size: dict[str, int] | None = None, |
| min_pixels: int | None = None, |
| max_pixels: int | None = None, |
| resample: PILImageResampling = None, |
| do_rescale: bool | None = None, |
| rescale_factor: float | None = None, |
| do_normalize: bool | None = None, |
| image_mean: float | list[float] | None = None, |
| image_std: float | list[float] | None = None, |
| patch_size: int | None = None, |
| temporal_patch_size: int | None = None, |
| merge_size: int | None = None, |
| do_convert_rgb: bool | None = None, |
| return_tensors: str | TensorType | None = None, |
| data_format: ChannelDimension | None = ChannelDimension.FIRST, |
| input_data_format: str | ChannelDimension | None = None, |
| ): |
| min_pixels = min_pixels if min_pixels is not None else self.min_pixels |
| max_pixels = max_pixels if max_pixels is not None else self.max_pixels |
|
|
| if size is not None: |
| if "shortest_edge" not in size or "longest_edge" not in size: |
| raise ValueError( |
| "size must contain 'shortest_edge' and 'longest_edge' keys." |
| ) |
| min_pixels = size["shortest_edge"] |
| elif min_pixels is not None and max_pixels is not None: |
| size = {"shortest_edge": min_pixels, "longest_edge": max_pixels} |
| else: |
| size = {**self.size} |
|
|
| do_resize = do_resize if do_resize is not None else self.do_resize |
| resample = resample if resample is not None else self.resample |
| do_rescale = do_rescale if do_rescale is not None else self.do_rescale |
| rescale_factor = ( |
| rescale_factor if rescale_factor is not None else self.rescale_factor |
| ) |
| do_normalize = do_normalize if do_normalize is not None else self.do_normalize |
| image_mean = image_mean if image_mean is not None else self.image_mean |
| image_std = image_std if image_std is not None else self.image_std |
| patch_size = patch_size if patch_size is not None else self.patch_size |
| temporal_patch_size = ( |
| temporal_patch_size |
| if temporal_patch_size is not None |
| else self.temporal_patch_size |
| ) |
| merge_size = merge_size if merge_size is not None else self.merge_size |
| do_convert_rgb = ( |
| do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb |
| ) |
|
|
| if images is not None: |
| images = make_flat_list_of_images(images) |
|
|
| if images is not None and not valid_images(images): |
| raise ValueError( |
| "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, " |
| "torch.Tensor, tf.Tensor or jax.ndarray." |
| ) |
|
|
| validate_preprocess_arguments( |
| rescale_factor=rescale_factor, |
| do_normalize=do_normalize, |
| image_mean=image_mean, |
| image_std=image_std, |
| do_resize=do_resize, |
| size=size, |
| resample=resample, |
| ) |
|
|
| data = {} |
| if images is not None: |
| pixel_values, vision_grid_thws = [], [] |
| for image in images: |
| patches, image_grid_thw = self._preprocess( |
| image, |
| do_resize=do_resize, |
| size=size, |
| resample=resample, |
| do_rescale=do_rescale, |
| rescale_factor=rescale_factor, |
| do_normalize=do_normalize, |
| image_mean=image_mean, |
| image_std=image_std, |
| patch_size=patch_size, |
| temporal_patch_size=temporal_patch_size, |
| merge_size=merge_size, |
| data_format=data_format, |
| do_convert_rgb=do_convert_rgb, |
| input_data_format=input_data_format, |
| ) |
| pixel_values.extend(patches) |
| vision_grid_thws.append(image_grid_thw) |
| pixel_values = np.array(pixel_values) |
| vision_grid_thws = np.array(vision_grid_thws) |
| data.update( |
| {"pixel_values": pixel_values, "image_grid_thw": vision_grid_thws} |
| ) |
|
|
| if videos is not None: |
| videos = make_batched_videos(videos) |
| pixel_values_videos, vision_grid_thws_videos = [], [] |
| for images in videos: |
| patches, video_grid_thw = self._preprocess( |
| images, |
| do_resize=do_resize, |
| size=size, |
| resample=resample, |
| do_rescale=do_rescale, |
| rescale_factor=rescale_factor, |
| do_normalize=do_normalize, |
| image_mean=image_mean, |
| image_std=image_std, |
| patch_size=patch_size, |
| temporal_patch_size=temporal_patch_size, |
| merge_size=merge_size, |
| data_format=data_format, |
| do_convert_rgb=do_convert_rgb, |
| input_data_format=input_data_format, |
| ) |
| pixel_values_videos.extend(patches) |
| vision_grid_thws_videos.append(video_grid_thw) |
| data.update( |
| { |
| "pixel_values_videos": np.array(pixel_values_videos), |
| "video_grid_thw": np.array(vision_grid_thws_videos), |
| } |
| ) |
|
|
| return BatchFeature(data=data, tensor_type=return_tensors) |
|
|
| def get_number_of_image_patches(self, height: int, width: int, images_kwargs=None): |
| min_pixels = ( |
| images_kwargs["min_pixels"] |
| if "min_pixels" in images_kwargs |
| else self.size["shortest_edge"] |
| ) |
| max_pixels = ( |
| images_kwargs["max_pixels"] |
| if "max_pixels" in images_kwargs |
| else self.size["longest_edge"] |
| ) |
| patch_size = images_kwargs.get("patch_size", self.patch_size) |
| merge_size = images_kwargs.get("merge_size", self.merge_size) |
|
|
| factor = patch_size * merge_size |
| resized_height, resized_width = smart_resize( |
| height, width, factor, min_pixels=min_pixels, max_pixels=max_pixels |
| ) |
| grid_h, grid_w = resized_height // patch_size, resized_width // patch_size |
| return grid_h * (grid_w + 1) + 2 |
|
|
|
|
| AutoImageProcessor.register("BrainOCRImageProcessor", BrainOCRImageProcessor) |
|
|
