File size: 2,789 Bytes
ba56501 29f03f1 ba56501 29f03f1 ba56501 |
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 90 91 92 93 94 |
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),
categories=[
"Bags",
"Feet",
"Hands",
"Head",
"Lower Body",
"Neck",
"Outwear",
"Upper Body",
"Waist",
"Whole Body",
],
**kwargs,
):
super().__init__(**kwargs)
self.bkg_color = bkg_color
self.input_resolution = input_resolution
self.image_mean = image_mean
self.image_std = image_std
self.categories = categories
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 process_cat(self, cat):
if cat is not None:
cat = torch.tensor(self.categories.index(cat), dtype=int)
return cat
def __call__(self, images, categories=None):
"""
Parameters
----------
images : Union[Image.Image, List[Image.Image]]
Image or list of images to process
categories : Optional[Union[str, List[str]]]
Category or list of categories to process
Returns
-------
BatchFeature
pixel_values : torch.Tensor
Processed image tensor (B C H W)
category_indices : torch.Tensor
Categories indices (B)
"""
use_cats = categories is not None
# Single Image + Single category
if isinstance(images, Image.Image):
images = [images]
if use_cats:
categories = [categories]
data = {}
data["pixel_values"] = torch.stack([self.process_img(img) for img in images])
if use_cats:
data["category_indices"] = torch.stack(
[self.process_cat(c) for c in categories]
)
return BatchFeature(data=data)
|