galileo-tiny-patch8/processing_galileo.py

# Copyright 2025 The Galileo Authors and The HuggingFace Inc. team.
"""Self-contained Galileo processor for trust_remote_code loading."""

from __future__ import annotations

import math
from typing import Any, NamedTuple, Optional, Union

import numpy as np
import torch

from transformers.feature_extraction_utils import BatchFeature
from transformers.processing_utils import ProcessorMixin
from transformers.utils import TensorType

from .modeling_galileo import GalileoConfig


class MaskedOutput(NamedTuple):
    space_time_x: torch.Tensor
    space_x: torch.Tensor
    time_x: torch.Tensor
    static_x: torch.Tensor
    space_time_mask: torch.Tensor
    space_mask: torch.Tensor
    time_mask: torch.Tensor
    static_mask: torch.Tensor
    months: torch.Tensor


class PretrainingNormalizer:
    def __init__(self, normalizing_dicts: dict):
        self.stats: dict[int, dict[str, np.ndarray]] = {}
        for key, val in normalizing_dicts.items():
            if isinstance(key, str) and key.isdigit():
                key = int(key)
            if not isinstance(key, int):
                continue
            self.stats[key] = {
                "mean": np.asarray(val["mean"], dtype=np.float32),
                "std": np.asarray(val["std"], dtype=np.float32),
            }

    def __call__(self, x: np.ndarray) -> np.ndarray:
        stats = self.stats[x.shape[-1]]
        return (x - stats["mean"]) / stats["std"]


def to_cartesian(lat: float, lon: float) -> np.ndarray:
    lat_rad = lat * math.pi / 180
    lon_rad = lon * math.pi / 180
    return np.array(
        [
            math.cos(lat_rad) * math.cos(lon_rad),
            math.cos(lat_rad) * math.sin(lon_rad),
            math.sin(lat_rad),
        ],
        dtype=np.float32,
    )


def construct_galileo_input(
    s1: torch.Tensor | None = None,
    s2: torch.Tensor | None = None,
    era5: torch.Tensor | None = None,
    tc: torch.Tensor | None = None,
    viirs: torch.Tensor | None = None,
    srtm: torch.Tensor | None = None,
    dw: torch.Tensor | None = None,
    wc: torch.Tensor | None = None,
    landscan: torch.Tensor | None = None,
    latlon: torch.Tensor | None = None,
    months: torch.Tensor | None = None,
    normalize: bool = False,
    band_config: GalileoConfig | None = None,
) -> MaskedOutput:
    band_config = band_config or GalileoConfig()
    bands = band_config.band_layout()
    space_time_bands = bands["space_time_bands"]
    space_time_groups = bands["space_time_groups"]
    time_bands = bands["time_bands"]
    time_groups = bands["time_groups"]
    space_bands = bands["space_bands"]
    space_groups = bands["space_groups"]
    static_bands = bands["static_bands"]
    static_groups = bands["static_groups"]

    space_time_inputs = [s1, s2]
    time_inputs = [era5, tc, viirs]
    space_inputs = [srtm, dw, wc]
    static_inputs = [landscan, latlon]
    devices = [
        x.device
        for x in space_time_inputs + time_inputs + space_inputs + static_inputs
        if x is not None
    ]

    if len(devices) == 0:
        raise ValueError("At least one input must be not None")
    device = devices[0]

    timesteps_list = [x.shape[2] for x in space_time_inputs if x is not None] + [
        x.shape[1] for x in time_inputs if x is not None
    ]
    height_list = [x.shape[0] for x in space_time_inputs if x is not None] + [
        x.shape[0] for x in space_inputs if x is not None
    ]
    width_list = [x.shape[1] for x in space_time_inputs if x is not None] + [
        x.shape[1] for x in space_inputs if x is not None
    ]

    t = timesteps_list[0] if timesteps_list else 1
    h, w = (height_list[0], width_list[0]) if height_list else (1, 1)

    s_t_x = torch.zeros((h, w, t, len(space_time_bands)), dtype=torch.float, device=device)
    s_t_m = torch.ones((h, w, t, len(space_time_groups)), dtype=torch.float, device=device)
    sp_x = torch.zeros((h, w, len(space_bands)), dtype=torch.float, device=device)
    sp_m = torch.ones((h, w, len(space_groups)), dtype=torch.float, device=device)
    t_x = torch.zeros((t, len(time_bands)), dtype=torch.float, device=device)
    t_m = torch.ones((t, len(time_groups)), dtype=torch.float, device=device)
    st_x = torch.zeros((len(static_bands)), dtype=torch.float, device=device)
    st_m = torch.ones((len(static_groups)), dtype=torch.float, device=device)

    for x, bands_list, group_key in zip([s1, s2], [bands["s1_bands"], bands["s2_bands"]], ["S1", "S2"]):
        if x is not None:
            indices = [idx for idx, val in enumerate(space_time_bands) if val in bands_list]
            groups_idx = [idx for idx, key in enumerate(space_time_groups) if group_key in key]
            s_t_x[:, :, :, indices] = x
            s_t_m[:, :, :, groups_idx] = 0

    for x, bands_list, group_key in zip(
        [srtm, dw, wc], [bands["srtm_bands"], bands["dw_bands"], bands["wc_bands"]], ["SRTM", "DW", "WC"]
    ):
        if x is not None:
            indices = [idx for idx, val in enumerate(space_bands) if val in bands_list]
            groups_idx = [idx for idx, key in enumerate(space_groups) if group_key in key]
            sp_x[:, :, indices] = x
            sp_m[:, :, groups_idx] = 0

    for x, bands_list, group_key in zip(
        [era5, tc, viirs], [bands["era5_bands"], bands["tc_bands"], bands["viirs_bands"]], ["ERA5", "TC", "VIIRS"]
    ):
        if x is not None:
            indices = [idx for idx, val in enumerate(time_bands) if val in bands_list]
            groups_idx = [idx for idx, key in enumerate(time_groups) if group_key in key]
            t_x[:, indices] = x
            t_m[:, groups_idx] = 0

    for x, bands_list, group_key in zip(
        [landscan, latlon], [bands["landscan_bands"], bands["location_bands"]], ["LS", "location"]
    ):
        if x is not None:
            if group_key == "location":
                x = torch.as_tensor(to_cartesian(float(x[0]), float(x[1])), device=device)
            indices = [idx for idx, val in enumerate(static_bands) if val in bands_list]
            groups_idx = [idx for idx, key in enumerate(static_groups) if group_key in key]
            st_x[indices] = x
            st_m[groups_idx] = 0

    if months is None:
        months = torch.ones((t), dtype=torch.long, device=device) * band_config.default_month

    if normalize:
        normalizer = PretrainingNormalizer(band_config.pretraining_normalizing_dict)
        s_t_x = torch.from_numpy(normalizer(s_t_x.cpu().numpy())).to(device)
        sp_x = torch.from_numpy(normalizer(sp_x.cpu().numpy())).to(device)
        t_x = torch.from_numpy(normalizer(t_x.cpu().numpy())).to(device)
        st_x = torch.from_numpy(normalizer(st_x.cpu().numpy())).to(device)

    return MaskedOutput(
        space_time_x=s_t_x,
        space_time_mask=s_t_m,
        space_x=sp_x,
        space_mask=sp_m,
        time_x=t_x,
        time_mask=t_m,
        static_x=st_x,
        static_mask=st_m,
        months=months,
    )


class GalileoProcessor(ProcessorMixin):
    attributes = []
    model_input_names = [
        "space_time_x",
        "space_x",
        "time_x",
        "static_x",
        "space_time_mask",
        "space_mask",
        "time_mask",
        "static_mask",
        "months",
    ]

    def __init__(
        self,
        normalize: bool = True,
        default_month: int = 6,
        patch_size: int = 8,
        s1_bands: Optional[list[str]] = None,
        s2_bands: Optional[list[str]] = None,
        era5_bands: Optional[list[str]] = None,
        tc_bands: Optional[list[str]] = None,
        viirs_bands: Optional[list[str]] = None,
        srtm_bands: Optional[list[str]] = None,
        dw_bands: Optional[list[str]] = None,
        wc_bands: Optional[list[str]] = None,
        landscan_bands: Optional[list[str]] = None,
        location_bands: Optional[list[str]] = None,
        space_time_band_groups: Optional[dict[str, list[str]]] = None,
        time_band_groups: Optional[dict[str, list[str]]] = None,
        space_band_groups: Optional[dict[str, list[str]]] = None,
        pretraining_normalizing_dict: Optional[dict[str, dict[str, list[float]]]] = None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.normalize = normalize
        self.default_month = default_month
        self.patch_size = patch_size
        self.band_config = GalileoConfig(
            default_month=default_month,
            s1_bands=s1_bands,
            s2_bands=s2_bands,
            era5_bands=era5_bands,
            tc_bands=tc_bands,
            viirs_bands=viirs_bands,
            srtm_bands=srtm_bands,
            dw_bands=dw_bands,
            wc_bands=wc_bands,
            landscan_bands=landscan_bands,
            location_bands=location_bands,
            space_time_band_groups=space_time_band_groups,
            time_band_groups=time_band_groups,
            space_band_groups=space_band_groups,
            pretraining_normalizing_dict=pretraining_normalizing_dict,
        )

    def __call__(
        self,
        s1: Optional[Union[torch.Tensor, np.ndarray]] = None,
        s2: Optional[Union[torch.Tensor, np.ndarray]] = None,
        era5: Optional[Union[torch.Tensor, np.ndarray]] = None,
        tc: Optional[Union[torch.Tensor, np.ndarray]] = None,
        viirs: Optional[Union[torch.Tensor, np.ndarray]] = None,
        srtm: Optional[Union[torch.Tensor, np.ndarray]] = None,
        dw: Optional[Union[torch.Tensor, np.ndarray]] = None,
        wc: Optional[Union[torch.Tensor, np.ndarray]] = None,
        landscan: Optional[Union[torch.Tensor, np.ndarray]] = None,
        latlon: Optional[Union[torch.Tensor, np.ndarray]] = None,
        months: Optional[Union[torch.Tensor, np.ndarray, list[int]]] = None,
        normalize: Optional[bool] = None,
        patch_size: Optional[int] = None,
        return_tensors: Optional[Union[str, TensorType]] = None,
        **kwargs: Any,
    ) -> BatchFeature:
        normalize = self.normalize if normalize is None else normalize
        if patch_size is not None:
            self.patch_size = patch_size

        def _to_tensor(value):
            if value is None:
                return None
            if torch.is_tensor(value):
                return value
            return torch.as_tensor(value, dtype=torch.float32)

        if months is not None and not torch.is_tensor(months):
            months = torch.as_tensor(months, dtype=torch.long)

        masked_output = construct_galileo_input(
            s1=_to_tensor(s1),
            s2=_to_tensor(s2),
            era5=_to_tensor(era5),
            tc=_to_tensor(tc),
            viirs=_to_tensor(viirs),
            srtm=_to_tensor(srtm),
            dw=_to_tensor(dw),
            wc=_to_tensor(wc),
            landscan=_to_tensor(landscan),
            latlon=_to_tensor(latlon),
            months=months,
            normalize=normalize,
            band_config=self.band_config,
        )

        if masked_output.space_time_x.dim() == 4:
            masked_output = MaskedOutput(
                space_time_x=masked_output.space_time_x.unsqueeze(0),
                space_x=masked_output.space_x.unsqueeze(0),
                time_x=masked_output.time_x.unsqueeze(0),
                static_x=masked_output.static_x.unsqueeze(0),
                space_time_mask=masked_output.space_time_mask.unsqueeze(0),
                space_mask=masked_output.space_mask.unsqueeze(0),
                time_mask=masked_output.time_mask.unsqueeze(0),
                static_mask=masked_output.static_mask.unsqueeze(0),
                months=masked_output.months.unsqueeze(0),
            )

        data = {
            "space_time_x": masked_output.space_time_x,
            "space_x": masked_output.space_x,
            "time_x": masked_output.time_x,
            "static_x": masked_output.static_x,
            "space_time_mask": masked_output.space_time_mask,
            "space_mask": masked_output.space_mask,
            "time_mask": masked_output.time_mask,
            "static_mask": masked_output.static_mask,
            "months": masked_output.months,
            "patch_size": self.patch_size,
        }

        if return_tensors == TensorType.PYTORCH:
            for key in self.model_input_names:
                value = data[key]
                if not torch.is_tensor(value):
                    data[key] = torch.as_tensor(value)
                data[key] = data[key].long() if key == "months" else data[key].float()
            if not torch.is_tensor(data["patch_size"]):
                data["patch_size"] = torch.tensor(data["patch_size"])

        return BatchFeature(data=data, tensor_type=return_tensors)


__all__ = ["GalileoProcessor"]