| """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: |
| |
| 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"] |
|
|
