File size: 2,669 Bytes
c14f952 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
# coding=utf-8
#
# Code mainly copied from:
# https://github.com/huggingface/transformers/blob/main/src/transformers/models/clip/image_processing_clip.py
# and adjusted for Jina CLIP
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})
|