| from transformers.models.clip.image_processing_clip import CLIPImageProcessor |
| from transformers.image_utils import ( |
| OPENAI_CLIP_MEAN, |
| OPENAI_CLIP_STD, |
| ChannelDimension, |
| ImageInput, |
| PILImageResampling, |
| infer_channel_dimension_format, |
| is_scaled_image, |
| make_list_of_images, |
| to_numpy_array, |
| valid_images, |
| validate_kwargs, |
| validate_preprocess_arguments, |
| ) |
| from transformers.utils import TensorType, is_vision_available, logging |
|
|
| from typing import Dict, List, Optional, Union |
| from math import ceil |
| from torchvision.transforms import Resize |
|
|
| def get_resize_output_image_size_long( |
| image_size, PATCH_SIZE=32, MAX_RESOLUTION = 1024, MIN_RESOLUTION = 448, |
| ) -> tuple: |
| l1, l2 = image_size |
| short, long = (l2, l1) if l2 <= l1 else (l1, l2) |
|
|
| |
| requested_new_long = min( |
| [ |
| ceil(long / PATCH_SIZE) * PATCH_SIZE, |
| MAX_RESOLUTION, |
| ] |
| ) |
|
|
| requested_new_long = max(requested_new_long, MIN_RESOLUTION) |
|
|
| new_long, new_short = requested_new_long, int(requested_new_long * short / long) |
| |
| new_short = ceil(new_short / PATCH_SIZE) * PATCH_SIZE |
| return (new_long, new_short) if l2 <= l1 else (new_short, new_long) |
|
|
|
|
| class SoloCLIPImageProcessor(CLIPImageProcessor): |
|
|
| def __init__( |
| self, |
| do_resize: bool = True, |
| size: Dict[str, int] = None, |
| resample: PILImageResampling = PILImageResampling.BICUBIC, |
| do_center_crop: bool = True, |
| crop_size: Dict[str, int] = None, |
| do_rescale: bool = True, |
| rescale_factor: Union[int, float] = 1 / 255, |
| do_normalize: bool = True, |
| image_mean: Optional[Union[float, List[float]]] = None, |
| image_std: Optional[Union[float, List[float]]] = None, |
| do_convert_rgb: bool = True, |
| PATCH_SIZE=32, |
| MAX_RESOLUTION=1024, |
| MIN_RESOLUTION=448, |
| **kwargs, |
| ) -> None: |
| super(SoloCLIPImageProcessor, self).__init__( |
| do_resize=do_resize, |
| size=size, |
| resample=resample, |
| do_center_crop=do_center_crop, |
| crop_size=crop_size, |
| do_rescale=do_rescale, |
| rescale_factor=rescale_factor, |
| do_normalize=do_normalize, |
| image_mean=image_mean, |
| image_std=image_std, |
| do_convert_rgb=do_convert_rgb, |
| **kwargs, |
| ) |
| self.PATCH_SIZE = PATCH_SIZE |
| self.MAX_RESOLUTION = MAX_RESOLUTION |
| self.MIN_RESOLUTION = MIN_RESOLUTION |
|
|
| def preprocess( |
| self, |
| images: ImageInput, |
| do_resize: bool = None, |
| size: Dict[str, int] = None, |
| resample: PILImageResampling = None, |
| do_center_crop: bool = None, |
| crop_size: int = None, |
| do_rescale: bool = None, |
| rescale_factor: float = None, |
| do_normalize: bool = None, |
| image_mean: Optional[Union[float, List[float]]] = None, |
| image_std: Optional[Union[float, List[float]]] = None, |
| do_convert_rgb: bool = None, |
| return_tensors: Optional[Union[str, TensorType]] = None, |
| data_format: Optional[ChannelDimension] = ChannelDimension.FIRST, |
| input_data_format: Optional[Union[str, ChannelDimension]] = None, |
| **kwargs, |
| ): |
| return_dict = super(SoloCLIPImageProcessor, self).preprocess( |
| images=images, |
| do_resize=do_resize, |
| size=size, |
| resample=resample, |
| do_center_crop=do_center_crop, |
| crop_size=crop_size, |
| do_rescale=do_rescale, |
| rescale_factor=rescale_factor, |
| do_normalize=do_normalize, |
| image_mean=image_mean, |
| image_std=image_std, |
| do_convert_rgb=do_convert_rgb, |
| return_tensors=return_tensors, |
| data_format=data_format, |
| input_data_format=input_data_format, |
| **kwargs, |
| ) |
| pixel_values = return_dict['pixel_values'][0] |
| _, height, width = pixel_values.size() |
| height, width = get_resize_output_image_size_long( |
| (height, width), |
| PATCH_SIZE=self.PATCH_SIZE, |
| MAX_RESOLUTION=self.MAX_RESOLUTION, |
| MIN_RESOLUTION=self.MIN_RESOLUTION, |
| ) |
| pixel_values = Resize(size=(height, width))(pixel_values) |
| return_dict['pixel_values'] = pixel_values.unsqueeze(0) |
| return return_dict |
