| | |
| | |
| | |
| | |
| | |
| |
|
| | from typing import Tuple, Union |
| |
|
| | import torch |
| | from transformers.image_processing_utils import BaseImageProcessor, BatchFeature |
| | from transformers.image_utils import ImageInput, make_list_of_images |
| | from transformers.models.clip import CLIPProcessor |
| |
|
| | from .transform import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD, image_transform |
| |
|
| | """ Jina CLIP processor implementation """ |
| |
|
| |
|
| | class JinaCLIPProcessor(CLIPProcessor): |
| | image_processor_class = 'AutoImageProcessor' |
| | tokenizer_class = 'AutoTokenizer' |
| |
|
| |
|
| | """ Jina CLIP image processor implementation """ |
| |
|
| |
|
| | class JinaCLIPImageProcessor(BaseImageProcessor): |
| | model_input_names = ['pixel_values'] |
| | _valid_processor_keys = [ |
| | 'size', |
| | 'mean', |
| | 'std', |
| | 'resize_mode', |
| | 'interpolation', |
| | 'fill_color', |
| | ] |
| |
|
| | def __init__( |
| | self, |
| | size: Union[int, Tuple[int, int]] = 224, |
| | mean: Union[float, Tuple[float]] = OPENAI_DATASET_MEAN, |
| | std: Union[float, Tuple[float]] = OPENAI_DATASET_STD, |
| | resize_mode: str = 'shortest', |
| | interpolation: str = 'bicubic', |
| | fill_color: int = 0, |
| | **kwargs, |
| | ) -> None: |
| | super().__init__(**kwargs) |
| | self.size = size |
| | self.mean = mean |
| | self.std = std |
| | self.resize_mode = resize_mode |
| | self.interpolation = interpolation |
| | self.fill_color = fill_color |
| | self.transform = self._build_transform() |
| |
|
| | def _build_transform(self): |
| | return image_transform( |
| | image_size=self.size, |
| | is_train=False, |
| | mean=self.mean, |
| | std=self.std, |
| | resize_mode=self.resize_mode, |
| | interpolation=self.interpolation, |
| | fill_color=self.fill_color, |
| | aug_cfg=None, |
| | ) |
| |
|
| | def to_dict(self): |
| | output = super().to_dict() |
| | output.pop('transform') |
| | return output |
| |
|
| | def preprocess(self, images: ImageInput, **kwargs) -> BatchFeature: |
| | _transform_needs_rebuild = False |
| | for k, v in kwargs.items(): |
| | if k in self._valid_processor_keys: |
| | if v != getattr(self, k): |
| | setattr(self, k, v) |
| | _transform_needs_rebuild = True |
| |
|
| | if _transform_needs_rebuild: |
| | self.transform = self._build_transform() |
| |
|
| | images = make_list_of_images(images) |
| | out = torch.stack([self.transform(image) for image in images], dim=0) |
| | return BatchFeature(data={'pixel_values': out}) |
| |
|
