|
|
| """Image processor for SARMAE models (self-contained for trust_remote_code)."""
|
|
|
| from typing import Optional, Union
|
|
|
| import numpy as np
|
|
|
| from transformers.image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
|
| from transformers.image_transforms import resize, to_channel_dimension_format
|
| from transformers.image_utils import (
|
| ChannelDimension,
|
| ImageInput,
|
| PILImageResampling,
|
| infer_channel_dimension_format,
|
| to_numpy_array,
|
| valid_images,
|
| validate_preprocess_arguments,
|
| )
|
| from transformers.utils import TensorType, filter_out_non_signature_kwargs, logging
|
|
|
|
|
| logger = logging.get_logger(__name__)
|
|
|
|
|
| def _repeat_grayscale_channels(image: np.ndarray, target_channels: int, input_data_format: ChannelDimension) -> np.ndarray:
|
| if input_data_format == ChannelDimension.FIRST:
|
| num_channels = image.shape[0]
|
| if num_channels == target_channels:
|
| return image
|
| if num_channels == 1:
|
| return np.repeat(image, target_channels, axis=0)
|
| return image[:target_channels]
|
| num_channels = image.shape[-1]
|
| if num_channels == target_channels:
|
| return image
|
| if num_channels == 1:
|
| return np.repeat(image, target_channels, axis=-1)
|
| return image[..., :target_channels]
|
|
|
|
|
| def _prepare_image_batch(images: ImageInput) -> list:
|
| if isinstance(images, np.ndarray):
|
| images = [images]
|
| elif not isinstance(images, (list, tuple)):
|
| images = [images]
|
|
|
| prepared = []
|
| for image in images:
|
| array = to_numpy_array(image)
|
| if array.ndim == 2:
|
| array = np.expand_dims(array, axis=-1)
|
| prepared.append(array)
|
| return prepared
|
|
|
|
|
| class SarmaeImageProcessor(BaseImageProcessor):
|
| model_input_names = ["pixel_values"]
|
|
|
| def __init__(
|
| self,
|
| do_resize: bool = True,
|
| size: Optional[dict[str, int]] = None,
|
| resample: PILImageResampling = PILImageResampling.BILINEAR,
|
| do_rescale: bool = True,
|
| rescale_factor: float = 1 / 255.0,
|
| do_normalize: bool = True,
|
| image_mean: Optional[Union[float, list[float]]] = None,
|
| image_std: Optional[Union[float, list[float]]] = None,
|
| do_convert_rgb: bool = False,
|
| repeat_grayscale_channels: bool = True,
|
| **kwargs,
|
| ):
|
| super().__init__(**kwargs)
|
| size = size if size is not None else {"height": 224, "width": 224}
|
| self.do_resize = do_resize
|
| self.size = size
|
| self.resample = resample
|
| self.do_rescale = do_rescale
|
| self.rescale_factor = rescale_factor
|
| self.do_normalize = do_normalize
|
| self.image_mean = image_mean
|
| self.image_std = image_std
|
| self.do_convert_rgb = do_convert_rgb
|
| self.repeat_grayscale_channels = repeat_grayscale_channels
|
|
|
| @filter_out_non_signature_kwargs()
|
| def preprocess(
|
| self,
|
| images: ImageInput,
|
| do_resize: Optional[bool] = None,
|
| size: Optional[dict[str, int]] = None,
|
| resample: Optional[PILImageResampling] = None,
|
| do_rescale: Optional[bool] = None,
|
| rescale_factor: Optional[float] = None,
|
| do_normalize: Optional[bool] = None,
|
| image_mean: Optional[Union[float, list[float]]] = None,
|
| image_std: Optional[Union[float, list[float]]] = None,
|
| return_tensors: Optional[Union[str, TensorType]] = None,
|
| data_format: Union[str, ChannelDimension] = ChannelDimension.FIRST,
|
| input_data_format: Optional[Union[str, ChannelDimension]] = None,
|
| do_convert_rgb: Optional[bool] = None,
|
| repeat_grayscale_channels: Optional[bool] = None,
|
| ):
|
| do_resize = do_resize if do_resize is not None else self.do_resize
|
| size = get_size_dict(size if size is not None else self.size, default_to_square=True)
|
| resample = resample if resample is not None else self.resample
|
| do_rescale = do_rescale if do_rescale is not None else self.do_rescale
|
| rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
|
| do_normalize = do_normalize if do_normalize is not None else self.do_normalize
|
| image_mean = image_mean if image_mean is not None else self.image_mean
|
| image_std = image_std if image_std is not None else self.image_std
|
| do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
|
| repeat_grayscale_channels = (
|
| repeat_grayscale_channels if repeat_grayscale_channels is not None else self.repeat_grayscale_channels
|
| )
|
|
|
| if do_normalize and (image_mean is None or image_std is None):
|
| raise ValueError("Normalization requires `image_mean` and `image_std` with one value per channel.")
|
|
|
| images = _prepare_image_batch(images)
|
| if not valid_images(images):
|
| raise ValueError("Invalid image type. Must be PIL, numpy, or torch tensor.")
|
|
|
| validate_preprocess_arguments(
|
| do_rescale=do_rescale,
|
| rescale_factor=rescale_factor,
|
| do_normalize=do_normalize,
|
| image_mean=image_mean,
|
| image_std=image_std,
|
| do_resize=do_resize,
|
| size=size,
|
| resample=resample,
|
| )
|
|
|
| processed_images = []
|
| for image in images:
|
| image = to_numpy_array(image)
|
| if do_convert_rgb:
|
| image = self._convert_image_to_rgb(image)
|
| if input_data_format is None:
|
| try:
|
| input_data_format = infer_channel_dimension_format(image)
|
| except ValueError:
|
| input_data_format = ChannelDimension.LAST
|
| if repeat_grayscale_channels:
|
| image = _repeat_grayscale_channels(image, target_channels=3, input_data_format=input_data_format)
|
| if do_resize:
|
| image = resize(
|
| image,
|
| size=(size["height"], size["width"]),
|
| resample=resample,
|
| input_data_format=input_data_format,
|
| )
|
| if do_rescale:
|
| image = image * rescale_factor
|
| if do_normalize:
|
| image = self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)
|
| image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
|
| processed_images.append(image)
|
|
|
| return BatchFeature(data={"pixel_values": processed_images}, tensor_type=return_tensors)
|
|
|
|
|
| __all__ = ["SarmaeImageProcessor"]
|
|
|
