|
from transformers.image_processing_utils import ImageProcessingMixin, BatchFeature |
|
|
|
from torchvision.transforms import transforms as tf |
|
import torchvision.transforms.functional as F |
|
from PIL import Image |
|
import torch |
|
|
|
|
|
class CondViTProcessor(ImageProcessingMixin): |
|
def __init__( |
|
self, |
|
bkg_color=255, |
|
input_resolution=224, |
|
image_mean=(0.48145466, 0.4578275, 0.40821073), |
|
image_std=(0.26862954, 0.26130258, 0.27577711), |
|
**kwargs, |
|
): |
|
super().__init__(**kwargs) |
|
|
|
self.bkg_color = bkg_color |
|
self.input_resolution = input_resolution |
|
self.image_mean = image_mean |
|
self.image_std = image_std |
|
|
|
def square_pad(self, image): |
|
max_wh = max(image.size) |
|
p_left, p_top = [(max_wh - s) // 2 for s in image.size] |
|
p_right, p_bottom = [ |
|
max_wh - (s + pad) for s, pad in zip(image.size, [p_left, p_top]) |
|
] |
|
padding = (p_left, p_top, p_right, p_bottom) |
|
return F.pad(image, padding, self.bkg_color, "constant") |
|
|
|
def process_img(self, image): |
|
img = self.square_pad(image) |
|
img = F.resize(img, self.input_resolution) |
|
img = F.to_tensor(img) |
|
img = F.normalize(img, self.image_mean, self.image_std) |
|
return img |
|
|
|
def __call__(self, images, texts=None): |
|
""" |
|
Parameters |
|
---------- |
|
images : Union[Image.Image, List[Image.Image]] |
|
Image or list of images to process |
|
texts : Union[str, List[str]] |
|
Text or list of texts to process. Pass through, no operation is performed. |
|
|
|
Returns |
|
------- |
|
BatchFeature |
|
pixel_values : torch.Tensor |
|
Processed image tensor (B C H W) |
|
texts : Union[str, List[str]] |
|
""" |
|
|
|
data = {} |
|
if isinstance(images, Image.Image): |
|
data["pixel_values"] = self.process_img(images) |
|
else: |
|
data["pixel_values"] = torch.stack( |
|
[self.process_img(img) for img in images] |
|
) |
|
|
|
if texts is not None: |
|
data["texts"] = texts |
|
|
|
return BatchFeature(data=data) |
|
|
