File size: 3,203 Bytes
8966f37 bfeb81e 8966f37 badc3e1 8966f37 badc3e1 8966f37 749251c 8966f37 bfeb81e 8966f37 bfeb81e 8966f37 bfeb81e 8966f37 bfeb81e badc3e1 8966f37 | 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 | from typing import Optional, Union
import numpy as np
import torch
from PIL import Image
from torchvision.transforms import v2
from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
class CondUNetImageProcessor(BaseImageProcessor):
model_input_names = ["pixel_values"]
def __init__(
self,
image_size=512,
keep_aspect_ratio=True,
self_normalize=True,
mean=None,
std=None,
**kwargs,
):
super().__init__(**kwargs)
self.image_size = image_size
self.keep_aspect_ratio = keep_aspect_ratio
self.self_normalize = self_normalize
self.mean = mean or [123.675, 116.28, 103.53]
self.std = std or [58.395, 57.12, 57.375]
def preprocess(
self,
images: Union[Image.Image, np.ndarray, torch.Tensor, list],
return_tensors: Optional[Union[str, torch.Tensor]] = None,
**kwargs,
):
if not isinstance(images, (list, tuple)):
images = [images]
pixel_values = [self._preprocess_image(image) for image in images]
return BatchFeature(
data={"pixel_values": torch.stack(pixel_values)},
tensor_type=return_tensors,
)
def _preprocess_image(self, image):
if isinstance(image, Image.Image):
image = np.array(image.convert("RGB"), copy=True)
if isinstance(image, np.ndarray):
image = torch.from_numpy(image)
if image.ndim != 3:
raise ValueError("Expected an HWC or CHW RGB image.")
if image.shape[-1] in (1, 3):
image = image.permute(2, 0, 1)
if image.shape[0] == 1:
image = image.expand(3, -1, -1)
if image.shape[0] != 3:
raise ValueError("Cond-UNet requires one or three input channels.")
height, width = image.shape[-2:]
if self.keep_aspect_ratio:
resize_factor = max(height, width) / self.image_size
new_height = int(height / resize_factor)
new_width = int(width / resize_factor)
new_height += new_height % 2
new_width += new_width % 2
image = v2.functional.resize(image, [new_height, new_width])
pad_left = (self.image_size - new_width) // 2
pad_top = (self.image_size - new_height) // 2
image = v2.functional.pad(image, fill=0, padding=[pad_left, pad_top])
else:
image = v2.functional.resize(image, [self.image_size, self.image_size])
image = image.to(dtype=torch.float32)
if image.max() <= 1:
image = image * 255.0
if self.self_normalize:
mask = (image > 0).any(dim=0)
if mask.any():
valid_pixels = image[:, mask]
mean = valid_pixels.mean()
std = valid_pixels.std()
if std > 1e-8:
return (image - mean) / std
return image - mean
return image.clone()
mean = torch.tensor(self.mean, dtype=image.dtype).view(-1, 1, 1)
std = torch.tensor(self.std, dtype=image.dtype).view(-1, 1, 1)
return (image - mean) / std
|