File size: 3,572 Bytes
c81207b | 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 95 96 | """Image processor for MolParser Mobile."""
from __future__ import annotations
from typing import List, Sequence, Union
import cv2
import numpy as np
import torch
from PIL import Image
from transformers import BaseImageProcessor
from transformers.feature_extraction_utils import BatchFeature
ImageInput = Union[str, Image.Image, np.ndarray, torch.Tensor]
class MolParserImageProcessor(BaseImageProcessor):
model_input_names = ["pixel_values"]
def __init__(
self,
image_size: int = 224,
do_resize: bool = True,
do_normalize: bool = True,
image_mean: Sequence[float] = (0.485, 0.456, 0.406),
image_std: Sequence[float] = (0.229, 0.224, 0.225),
**kwargs,
):
super().__init__(**kwargs)
self.image_size = int(image_size)
self.do_resize = bool(do_resize)
self.do_normalize = bool(do_normalize)
self.image_mean = list(image_mean)
self.image_std = list(image_std)
@property
def size(self):
return {"height": self.image_size, "width": self.image_size}
def _to_pil(self, image: ImageInput) -> Image.Image:
if isinstance(image, Image.Image):
return image.convert("RGB")
if isinstance(image, str):
return Image.open(image).convert("RGB")
if isinstance(image, torch.Tensor):
tensor = image.detach().cpu()
if tensor.ndim == 3 and tensor.shape[0] in {1, 3, 4}:
tensor = tensor.permute(1, 2, 0)
array = tensor.numpy()
else:
array = np.asarray(image)
if array.dtype != np.uint8:
if array.max() <= 1.0:
array = array * 255.0
array = np.clip(np.rint(array), 0, 255).astype(np.uint8)
if array.ndim == 2:
return Image.fromarray(array, mode="L").convert("RGB")
if array.shape[-1] == 4:
return Image.fromarray(array, mode="RGBA").convert("RGB")
return Image.fromarray(array).convert("RGB")
def _preprocess_one(self, image: ImageInput) -> np.ndarray:
pil_image = self._to_pil(image)
array = np.asarray(pil_image).astype(np.uint8)
if self.do_resize:
# Match deploy/transform.py: albumentations.Resize defaults to OpenCV INTER_LINEAR.
array = cv2.resize(array, (self.image_size, self.image_size), interpolation=cv2.INTER_LINEAR)
array = array.astype(np.float32) / 255.0
if self.do_normalize:
mean = np.asarray(self.image_mean, dtype=np.float32).reshape(1, 1, 3)
std = np.asarray(self.image_std, dtype=np.float32).reshape(1, 1, 3)
array = (array - mean) / std
return np.transpose(array, (2, 0, 1))
def preprocess(
self,
images: Union[ImageInput, Sequence[ImageInput]],
return_tensors: str | None = None,
**kwargs,
) -> BatchFeature:
if not isinstance(images, (list, tuple)):
images = [images]
pixel_values: List[np.ndarray] = [self._preprocess_one(image) for image in images]
data = {"pixel_values": np.stack(pixel_values, axis=0)}
encoded = BatchFeature(data=data)
if return_tensors is not None:
encoded = encoded.convert_to_tensors(return_tensors)
return encoded
def __call__(self, images: Union[ImageInput, Sequence[ImageInput]], return_tensors: str | None = None, **kwargs):
return self.preprocess(images=images, return_tensors=return_tensors, **kwargs)
__all__ = ["MolParserImageProcessor"]
|