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