# Copyright 2024 SatMAE++ Authors and The HuggingFace Inc. team. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. """Image processor for SatMAE++ models.""" 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, make_flat_list_of_images, 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 _resize_multispectral(image: np.ndarray, size: dict[str, int], input_data_format: ChannelDimension) -> np.ndarray: target_height, target_width = size["height"], size["width"] if input_data_format == ChannelDimension.FIRST: image = np.transpose(image, (1, 2, 0)) height, width, _ = image.shape if height == target_height and width == target_width: resized = image else: try: import cv2 except ImportError as exc: raise ImportError( "Multispectral resize requires OpenCV (`opencv-python`) when input has more than 4 channels." ) from exc resized = cv2.resize(image, (target_width, target_height), interpolation=cv2.INTER_LINEAR) if input_data_format == ChannelDimension.FIRST: return np.transpose(resized, (2, 0, 1)) return resized def _reorder_channels(image: np.ndarray, channel_order: str, input_data_format: ChannelDimension) -> np.ndarray: if channel_order != "bgr": return image if input_data_format == ChannelDimension.FIRST: if image.shape[0] < 3: return image return image[[2, 1, 0], ...] if image.shape[-1] < 3: return image return image[..., [2, 1, 0]] class SatMAEppImageProcessor(BaseImageProcessor): """ Image processor for SatMAE++ satellite encoders. FMoW-RGB checkpoints were trained with BGR channel order. Set `channel_order="bgr"` (default for RGB models) to swap the first three channels from RGB to BGR before inference. """ model_input_names = ["pixel_values"] def __init__( self, do_resize: bool = False, size: Optional[dict[str, int]] = None, resample: PILImageResampling = PILImageResampling.BILINEAR, do_rescale: bool = False, rescale_factor: float = 1.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, channel_order: str = "rgb", **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.channel_order = channel_order @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, channel_order: Optional[str] = None, ): do_resize = do_resize if do_resize is not None else self.do_resize size = size if size is not None else self.size size = get_size_dict(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 channel_order = channel_order if channel_order is not None else self.channel_order 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 = make_flat_list_of_images(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 image = _reorder_channels(image, channel_order=channel_order, input_data_format=input_data_format) if do_resize: num_channels = image.shape[0] if input_data_format == ChannelDimension.FIRST else image.shape[-1] if num_channels > 4: image = _resize_multispectral(image, size=size, input_data_format=input_data_format) else: 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) data = {"pixel_values": processed_images} return BatchFeature(data=data, tensor_type=return_tensors) __all__ = ["SatMAEppImageProcessor"]