prithvi-eo-v2-tiny-tl/image_processing_prithvi.py

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