| import numpy as np |
| import torch |
| from PIL import Image |
| from transformers.image_processing_utils import BaseImageProcessor |
| from transformers.utils import logging |
|
|
| logger = logging.get_logger(__name__) |
|
|
|
|
| class VQModelImageProcessor(BaseImageProcessor): |
| def __init__( |
| self, |
| size: int = 256, |
| convert_rgb: bool = False, |
| resample: Image.Resampling = Image.Resampling.LANCZOS, |
| **kwargs: dict, |
| ) -> None: |
| self.size = size |
| self.convert_rgb = convert_rgb |
| self.resample = resample |
|
|
| def __call__(self, image: Image.Image) -> dict: |
| return self.preprocess(image) |
|
|
| def preprocess(self, image: Image.Image) -> dict: |
| width, height = image.size |
| size = (self.size, self.size) |
| image = image.resize(size, resample=self.resample) |
| image = image.convert("RGBA") |
|
|
| if self.convert_rgb: |
| |
| image_new = Image.new("RGB", image.size, (255, 255, 255)) |
| image_new.paste(image, mask=image.split()[3]) |
| image = image_new |
|
|
| return { |
| "image": self.to_tensor(image), |
| "width": width, |
| "height": height, |
| } |
|
|
| def to_tensor(self, image: Image.Image) -> torch.Tensor: |
| x = np.array(image) / 127.5 - 1.0 |
| x = x.transpose(2, 0, 1).astype(np.float32) |
| return torch.as_tensor(x) |
|
|
| def postprocess( |
| self, |
| x: torch.Tensor, |
| width: int | None = None, |
| height: int | None = None, |
| ) -> Image.Image: |
| x_np = x.detach().cpu().numpy() |
| x_np = x_np.transpose(1, 2, 0) |
| x_np = (x_np + 1.0) * 127.5 |
| x_np = np.clip(x_np, 0, 255).astype(np.uint8) |
| image = Image.fromarray(x_np) |
|
|
| |
| width = width or self.size |
| height = height or self.size |
| image = image.resize((width, height), resample=self.resample) |
|
|
| return image |
|
|
