# Copyright (c) IBM Corp. 2024. All rights reserved. # Copyright 2024 Prithvi-EO-2.0 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 Prithvi-EO-2.0 models.""" from __future__ import annotations from typing import Optional, Union import numpy as np from transformers.image_processing_utils import BaseImageProcessor, BatchFeature from transformers.image_transforms import to_channel_dimension_format from transformers.image_utils import ( ChannelDimension, ImageInput, 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__) NO_DATA = -9999 NO_DATA_FLOAT = 0.0001 def _as_temporal_frames(images) -> list: if isinstance(images, (list, tuple)): if len(images) == 0: raise ValueError("Expected at least one temporal frame.") first = images[0] if isinstance(first, (list, tuple, np.ndarray)) and not hasattr(first, "shape"): return list(images) return [images] return [images] def _stack_temporal_frames(frames: list, input_data_format: ChannelDimension) -> np.ndarray: arrays = [to_numpy_array(frame) for frame in frames] if len(arrays) == 1: array = arrays[0] if array.ndim == 4: if input_data_format == ChannelDimension.FIRST: return array return np.moveaxis(array, -1, 0) if array.ndim == 3: if input_data_format == ChannelDimension.FIRST: return array[:, np.newaxis, ...] return np.moveaxis(array, -1, 0)[:, np.newaxis, ...] raise ValueError(f"Unsupported frame shape {array.shape}.") if input_data_format == ChannelDimension.LAST: stacked = np.stack(arrays, axis=0) return np.moveaxis(stacked, -1, 0) stacked = np.stack(arrays, axis=1) return stacked class PrithviImageProcessor(BaseImageProcessor): """ Image processor for Prithvi-EO-2.0 spatiotemporal HLS encoders. Accepts a single array shaped `(C, T, H, W)`, `(T, H, W, C)`, or a list of `T` frames. Applies HLS reflectance normalization with nodata masking. """ model_input_names = ["pixel_values", "temporal_coords", "location_coords"] def __init__( self, num_channels: int = 6, num_frames: int = 4, do_resize: bool = False, size: Optional[dict[str, int]] = None, 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, nodata_value: float = NO_DATA, nodata_replacement: float = NO_DATA_FLOAT, do_convert_rgb: bool = False, **kwargs, ): super().__init__(**kwargs) self.num_channels = num_channels self.num_frames = num_frames self.do_resize = do_resize self.size = size if size is not None else {"height": 224, "width": 224} 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.nodata_value = nodata_value self.nodata_replacement = nodata_replacement self.do_convert_rgb = do_convert_rgb @classmethod def from_config(cls, config): return cls( num_channels=config.num_channels, num_frames=config.num_frames, image_mean=config.image_mean, image_std=config.image_std, size={"height": config.image_size, "width": config.image_size}, ) def _normalize_hls(self, image: np.ndarray, input_data_format: ChannelDimension) -> np.ndarray: mean = np.asarray(self.image_mean, dtype=np.float32) std = np.asarray(self.image_std, dtype=np.float32) if input_data_format == ChannelDimension.FIRST: channels = image.shape[0] mean = mean[:channels].reshape(channels, *([1] * (image.ndim - 1))) std = std[:channels].reshape(channels, *([1] * (image.ndim - 1))) nodata_mask = image == self.nodata_value image = np.where(nodata_mask, self.nodata_replacement, image) return (image - mean) / std channels = image.shape[-1] mean = mean[:channels] std = std[:channels] nodata_mask = image == self.nodata_value image = np.where(nodata_mask, self.nodata_replacement, image) return (image - mean) / std @filter_out_non_signature_kwargs() def preprocess( self, images: ImageInput, temporal_coords: Optional[Union[list, np.ndarray]] = None, location_coords: Optional[Union[list, np.ndarray]] = None, do_resize: Optional[bool] = None, size: Optional[dict[str, int]] = 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, ): do_resize = do_resize if do_resize is not None else self.do_resize size = size if size is not None else self.size 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 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.") 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, ) temporal_batches = _as_temporal_frames(images) if not valid_images(temporal_batches): raise ValueError("Invalid image type. Must be PIL, numpy, or torch tensor.") processed_images = [] for sample in temporal_batches: if input_data_format is None: try: input_data_format = infer_channel_dimension_format(to_numpy_array(sample[0] if isinstance(sample, list) else sample)) except ValueError: input_data_format = ChannelDimension.LAST image = _stack_temporal_frames(sample if isinstance(sample, list) else [sample], input_data_format) if do_convert_rgb: image = self._convert_image_to_rgb(image) if do_rescale: image = image * rescale_factor if do_normalize: image = self._normalize_hls(image, ChannelDimension.FIRST) image = to_channel_dimension_format(image, data_format, input_channel_dim=ChannelDimension.FIRST) processed_images.append(image) data = {"pixel_values": processed_images} if temporal_coords is not None: data["temporal_coords"] = [np.asarray(temporal_coords, dtype=np.float32)] if location_coords is not None: data["location_coords"] = [np.asarray(location_coords, dtype=np.float32)] return BatchFeature(data=data, tensor_type=return_tensors) __all__ = ["PrithviImageProcessor"]