| from networkx import to_numpy_array
|
| import numpy as np
|
| import torch
|
| from PIL import Image, ImageOps
|
| import math
|
| from functools import partial, reduce
|
| from transformers.image_transforms import (
|
| convert_to_rgb,
|
| center_crop,
|
| normalize,
|
| rescale,
|
| resize,
|
| to_channel_dimension_format,
|
| )
|
|
|
| from transformers.image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
|
| from transformers.image_utils import ImageInput, ChannelDimension, PILImageResampling, to_numpy_array
|
|
|
| class UMMImageProcessor(BaseImageProcessor):
|
| model_input_names = ["pixel_values", "grid_hws"]
|
| def __init__(
|
| self,
|
| image_mean=(0.5, 0.5, 0.5),
|
| image_std=(0.5, 0.5, 0.5),
|
| size=(256, 256),
|
| crop_size = None,
|
| resample=PILImageResampling.BICUBIC,
|
| rescale_factor=1 / 255,
|
| data_format=ChannelDimension.FIRST,
|
| scale_resolution=256,
|
| patch_size=16,
|
| **kwargs,
|
| ):
|
| super().__init__(**kwargs)
|
| crop_size = crop_size if crop_size is not None else {"height": 256, "width": 256}
|
| crop_size = get_size_dict(crop_size, default_to_square=True, param_name="crop_size")
|
| self.image_mean = image_mean
|
| self.image_std = image_std
|
| self.size = size
|
| self.resample = resample
|
| self.rescale_factor = rescale_factor
|
| self.data_format = data_format
|
| self.crop_size = crop_size
|
| self.scale_resolution = scale_resolution
|
| self.patch_size = patch_size
|
|
|
| def preprocess(self, image, max_resolution=None, return_tensors = 'pt', und=True, **kwargs) -> BatchFeature:
|
| if max_resolution is not None:
|
| scale_resolution = max_resolution
|
| else:
|
| scale_resolution = self.scale_resolution
|
| if image is not None:
|
| pixel_values, grid_hws = [], []
|
| if und:
|
| image = self._preprocess_und(image, scale_resolution)
|
| else:
|
| image = self._preprocess_gen(image, scale_resolution)
|
| if not torch.is_tensor(image):
|
| image = torch.tensor(image)
|
| _,H,W = image.shape
|
| grid_h = int(H // self.patch_size)
|
| grid_w = int(W // self.patch_size)
|
| grid_hw = (grid_h, grid_w)
|
| pixel_values = torch.stack([image], dim=0)
|
| grid_hws = torch.tensor([grid_hw])
|
| data = {
|
| "pixel_values": pixel_values,
|
| "grid_hws": grid_hws
|
| }
|
| return BatchFeature(data=data, tensor_type=return_tensors)
|
|
|
| def _preprocess_gen(self, source_image, scale_resolution):
|
| w, h = source_image.size
|
| scale = scale_resolution / min(h, w)
|
| new_h = int(round(h * scale))
|
| new_w = int(round(w * scale))
|
| source_image = source_image.resize((new_w, new_h), Image.Resampling.BICUBIC)
|
| source_image = [source_image]
|
| transforms = [
|
| convert_to_rgb,
|
| to_numpy_array,
|
| ]
|
| transforms.append(partial(center_crop, size=(scale_resolution, scale_resolution)))
|
| transforms.append(partial(rescale, scale=self.rescale_factor, data_format=self.data_format))
|
| transforms.append(partial(to_channel_dimension_format, channel_dim=self.data_format, input_channel_dim=self.data_format))
|
| image = reduce(lambda x, f: [*map(f, x)], transforms, source_image)
|
| return image[0] if len(image) == 1 else image
|
|
|
| def _preprocess_und(self, source_image, scale_resolution):
|
| w, h = source_image.size
|
| scale = min(scale_resolution / h, scale_resolution / w)
|
| new_h = int(round(h * scale))
|
| new_w = int(round(w * scale))
|
| resized_image = source_image.resize((new_w, new_h), Image.Resampling.BICUBIC)
|
|
|
| pad_w = scale_resolution - new_w
|
| pad_h = scale_resolution - new_h
|
|
|
| left = pad_w // 2
|
| right = pad_w - left
|
| top = pad_h // 2
|
| bottom = pad_h - top
|
|
|
| new_image = ImageOps.expand(resized_image, border=(left, top, right, bottom), fill=(0,0,0))
|
|
|
| source_image = [new_image]
|
| transforms = [
|
| convert_to_rgb,
|
| to_numpy_array
|
| ]
|
| transforms.append(partial(rescale, scale=self.rescale_factor, data_format=self.data_format))
|
| transforms.append(partial(to_channel_dimension_format, channel_dim=self.data_format, input_channel_dim=self.data_format))
|
| image = reduce(lambda x, f: [*map(f, x)], transforms, source_image)
|
| return image[0] if len(image) == 1 else image
|
|
|
| __all__ = ["UMMImageProcessor"]
|
|
|
|
|
