Sync updated code and configs (no weight re-upload)

README.md

@@ -1,20 +1,3 @@
----
-license: mit
-language:
-- en
-tags:
-- remote-sensing
-- earth-observation
-- vision
-- feature-extraction
-- galileo
-- sentinel-1
-- sentinel-2
-- multimodal
-library_name: transformers
-pipeline_tag: feature-extraction
----
-
 # Galileo Transformers Models
 
 Self-contained HuggingFace model checkpoints for [Galileo](https://arxiv.org/abs/2502.09356).
@@ -31,6 +14,8 @@ Each checkpoint subfolder ships remote code for model, processor, and custom pip
 
 ## Usage
 
+Galileo operates on native patch grids (default **`patch_size: 8`** in `preprocessor_config.json`). Stack shapes are `(H, W, T, C)`; no fixed 224×224 resize is applied.
+
 ```python
 from transformers import pipeline
 import numpy as np
@@ -43,15 +28,15 @@ pipe = pipeline(
     trust_remote_code=True,
 )
 
-# 10-band Sentinel-2 stack (B2,B3,B4,B5,B6,B7,B8,B8A,B11,B12)
-s2 = np.random.randn(32, 32, 1, 10).astype(np.float32)
+# 10-band Sentinel-2 stack at native spatial size
+s2 = np.random.randn(64, 64, 1, 10).astype(np.float32)
 features = pipe(s2=s2, pool=True, return_tensors=True)
 ```
 
 Sentinel-1 only:
 
 ```python
-s1 = np.random.randn(32, 32, 1, 2).astype(np.float32)
+s1 = np.random.randn(64, 64, 1, 2).astype(np.float32)
 features = pipe(s1=s1, pool=True, return_tensors=True)
 ```
 

galileo-base-patch8/config.json

@@ -43,5 +43,262 @@
         "AutoModel"
       ]
     }
+  },
+  "s1_bands": [
+    "VV",
+    "VH"
+  ],
+  "s2_bands": [
+    "B2",
+    "B3",
+    "B4",
+    "B5",
+    "B6",
+    "B7",
+    "B8",
+    "B8A",
+    "B11",
+    "B12"
+  ],
+  "era5_bands": [
+    "temperature_2m",
+    "total_precipitation_sum"
+  ],
+  "tc_bands": [
+    "def",
+    "soil",
+    "aet"
+  ],
+  "viirs_bands": [
+    "avg_rad"
+  ],
+  "srtm_bands": [
+    "elevation",
+    "slope"
+  ],
+  "dw_bands": [
+    "DW_water",
+    "DW_trees",
+    "DW_grass",
+    "DW_flooded_vegetation",
+    "DW_crops",
+    "DW_shrub_and_scrub",
+    "DW_built",
+    "DW_bare",
+    "DW_snow_and_ice"
+  ],
+  "wc_bands": [
+    "WC_temporarycrops",
+    "WC_maize",
+    "WC_wintercereals",
+    "WC_springcereals",
+    "WC_irrigation"
+  ],
+  "landscan_bands": [
+    "b1"
+  ],
+  "location_bands": [
+    "x",
+    "y",
+    "z"
+  ],
+  "space_time_band_groups": {
+    "S1": [
+      "VV",
+      "VH"
+    ],
+    "S2_RGB": [
+      "B2",
+      "B3",
+      "B4"
+    ],
+    "S2_Red_Edge": [
+      "B5",
+      "B6",
+      "B7"
+    ],
+    "S2_NIR_10m": [
+      "B8"
+    ],
+    "S2_NIR_20m": [
+      "B8A"
+    ],
+    "S2_SWIR": [
+      "B11",
+      "B12"
+    ],
+    "NDVI": [
+      "NDVI"
+    ]
+  },
+  "time_band_groups": {
+    "ERA5": [
+      "temperature_2m",
+      "total_precipitation_sum"
+    ],
+    "TC": [
+      "def",
+      "soil",
+      "aet"
+    ],
+    "VIIRS": [
+      "avg_rad"
+    ]
+  },
+  "space_band_groups": {
+    "SRTM": [
+      "elevation",
+      "slope"
+    ],
+    "DW": [
+      "DW_water",
+      "DW_trees",
+      "DW_grass",
+      "DW_flooded_vegetation",
+      "DW_crops",
+      "DW_shrub_and_scrub",
+      "DW_built",
+      "DW_bare",
+      "DW_snow_and_ice"
+    ],
+    "WC": [
+      "WC_temporarycrops",
+      "WC_maize",
+      "WC_wintercereals",
+      "WC_springcereals",
+      "WC_irrigation"
+    ]
+  },
+  "pretraining_normalizing_dict": {
+    "13": {
+      "mean": [
+        -11.728724389184965,
+        -18.85558188024017,
+        1395.3408730676722,
+        1338.4026921784578,
+        1343.09883810357,
+        1543.8607982512297,
+        2186.2022069512263,
+        2525.0932853316694,
+        2410.3377187373408,
+        2750.2854646886753,
+        2234.911100061487,
+        1474.5311266077113,
+        0.2892116502999044
+      ],
+      "std": [
+        4.887145774840316,
+        5.730270320384293,
+        917.7041440370853,
+        913.2988423581528,
+        1092.678723527555,
+        1047.2206083460424,
+        1048.0101611156767,
+        1143.6903026819996,
+        1098.979177731649,
+        1204.472755085893,
+        1145.9774063078878,
+        980.2429840007796,
+        0.2720939024500081
+      ]
+    },
+    "16": {
+      "mean": [
+        673.0152819503361,
+        5.930092668915115,
+        0.10470439140978786,
+        0.23965913270066183,
+        0.08158044385860364,
+        0.04246976254259546,
+        0.11304392863520317,
+        0.17329647890362473,
+        0.0698981691616277,
+        0.12130267132802142,
+        0.04671318615236216,
+        10.973119802517362,
+        1.0927069179958768,
+        1.6991394232855903,
+        0.03720594618055555,
+        1.3671352688259548
+      ],
+      "std": [
+        983.0697298296237,
+        8.167406789813247,
+        0.18771647977504985,
+        0.2368313455675914,
+        0.08024268534756586,
+        0.04045374496146404,
+        0.11350342472061795,
+        0.1279898111718168,
+        0.12042341550438586,
+        0.13602408145504347,
+        0.043971116096060345,
+        31.255340146970997,
+        10.395974878206689,
+        12.92380617159917,
+        1.9285254295940466,
+        11.612179775408928
+      ]
+    },
+    "6": {
+      "mean": [
+        271.5674963541667,
+        0.08554303677156568,
+        657.3181260091111,
+        692.1291795806885,
+        562.781331880633,
+        1.5647115934036673
+      ],
+      "std": [
+        79.80828940314429,
+        0.11669547098151486,
+        704.0008695557707,
+        925.0116126406431,
+        453.2434022278578,
+        7.513020170832818
+      ]
+    },
+    "18": {
+      "mean": [
+        188.20315880851746,
+        0.2804946561574936,
+        0.11371652073860168,
+        0.058778801321983334,
+        0.10474256777763366,
+        0.2396918488264084,
+        0.08152248692512512,
+        0.04248040814399719,
+        0.11303179881572724,
+        0.17326324067115784,
+        0.06998309404850006,
+        0.12122812910079957,
+        0.04671641788482666,
+        10.98456594619751,
+        1.0968475807189941,
+        1.6947754135131836,
+        0.03320046615600586,
+        1.3602827312469483
+      ],
+      "std": [
+        1154.5919128300602,
+        0.5276998078079327,
+        0.7021637331734328,
+        0.36528892213195063,
+        0.17470213191865785,
+        0.20411195416718833,
+        0.0660782470089761,
+        0.03380702424871257,
+        0.09809195568521663,
+        0.11292471052124119,
+        0.09720748930233268,
+        0.12912217763726777,
+        0.0399973913151906,
+        23.725471823867462,
+        5.715238079725388,
+        9.030481416228302,
+        0.9950220242487364,
+        7.754429123862099
+      ]
+    }
   }
-}
+}

galileo-base-patch8/modeling_galileo.py

@@ -26,95 +26,11 @@ from einops import rearrange, repeat
 from torch import Tensor, vmap
 from torch.jit import Final
 
-# constants
-CONFIG_FILENAME = "config.json"
-ENCODER_FILENAME = "encoder.pt"
-BASE_GSD = 10
-
-# band information
-S1_BANDS = ["VV", "VH"]
-S2_BANDS = [
-    "B2",
-    "B3",
-    "B4",
-    "B5",
-    "B6",
-    "B7",
-    "B8",
-    "B8A",
-    "B11",
-    "B12",
-]
-ERA5_BANDS = ["temperature_2m", "total_precipitation_sum"]
-TC_BANDS = ["def", "soil", "aet"]
-VIIRS_BANDS = ["avg_rad"]
-SRTM_BANDS = ["elevation", "slope"]
-DW_BANDS = [
-    "DW_water",
-    "DW_trees",
-    "DW_grass",
-    "DW_flooded_vegetation",
-    "DW_crops",
-    "DW_shrub_and_scrub",
-    "DW_built",
-    "DW_bare",
-    "DW_snow_and_ice",
-]
-WC_BANDS = [
-    "WC_temporarycrops",
-    "WC_maize",
-    "WC_wintercereals",
-    "WC_springcereals",
-    "WC_irrigation",
-]
-STATIC_DW_BANDS = [f"{x}_static" for x in DW_BANDS]
-STATIC_WC_BANDS = [f"{x}_static" for x in WC_BANDS]
-
-LANDSCAN_BANDS = ["b1"]
-LOCATION_BANDS = ["x", "y", "z"]
-
-SPACE_TIME_BANDS = S1_BANDS + S2_BANDS + ["NDVI"]
-TIME_BANDS = ERA5_BANDS + TC_BANDS + VIIRS_BANDS
-SPACE_BANDS = SRTM_BANDS + DW_BANDS + WC_BANDS
-STATIC_BANDS = LANDSCAN_BANDS + LOCATION_BANDS + STATIC_DW_BANDS + STATIC_WC_BANDS
-
-
-SPACE_TIME_BANDS_GROUPS_IDX: OrderedDictType[str, List[int]] = OrderedDict(
-    {
-        "S1": [SPACE_TIME_BANDS.index(b) for b in S1_BANDS],
-        "S2_RGB": [SPACE_TIME_BANDS.index(b) for b in ["B2", "B3", "B4"]],
-        "S2_Red_Edge": [SPACE_TIME_BANDS.index(b) for b in ["B5", "B6", "B7"]],
-        "S2_NIR_10m": [SPACE_TIME_BANDS.index(b) for b in ["B8"]],
-        "S2_NIR_20m": [SPACE_TIME_BANDS.index(b) for b in ["B8A"]],
-        "S2_SWIR": [SPACE_TIME_BANDS.index(b) for b in ["B11", "B12"]],
-        "NDVI": [SPACE_TIME_BANDS.index("NDVI")],
-    }
-)
-
-TIME_BAND_GROUPS_IDX: OrderedDictType[str, List[int]] = OrderedDict(
-    {
-        "ERA5": [TIME_BANDS.index(b) for b in ERA5_BANDS],
-        "TC": [TIME_BANDS.index(b) for b in TC_BANDS],
-        "VIIRS": [TIME_BANDS.index(b) for b in VIIRS_BANDS],
-    }
-)
 
-SPACE_BAND_GROUPS_IDX: OrderedDictType[str, List[int]] = OrderedDict(
-    {
-        "SRTM": [SPACE_BANDS.index(b) for b in SRTM_BANDS],
-        "DW": [SPACE_BANDS.index(b) for b in DW_BANDS],
-        "WC": [SPACE_BANDS.index(b) for b in WC_BANDS],
-    }
-)
-
-STATIC_BAND_GROUPS_IDX: OrderedDictType[str, List[int]] = OrderedDict(
-    {
-        "LS": [STATIC_BANDS.index(b) for b in LANDSCAN_BANDS],
-        "location": [STATIC_BANDS.index(b) for b in LOCATION_BANDS],
-        "DW_static": [STATIC_BANDS.index(b) for b in STATIC_DW_BANDS],
-        "WC_static": [STATIC_BANDS.index(b) for b in STATIC_WC_BANDS],
-    }
-)
+def _band_group_indices(
+    bands: Sequence[str], groups: Dict[str, List[str]]
+) -> OrderedDictType[str, List[int]]:
+    return OrderedDict((name, [bands.index(b) for b in group_bands]) for name, group_bands in groups.items())
 
 
 def get_2d_sincos_pos_embed_with_resolution(
@@ -565,13 +481,16 @@ class GalileoBase(nn.Module):
         base_patch_size: int = 4,
         use_channel_embs: bool = True,
         drop_path: float = 0.0,
+        band_layout: Optional[Dict[str, Any]] = None,
     ):
         super().__init__()
 
-        self.space_time_groups = SPACE_TIME_BANDS_GROUPS_IDX
-        self.space_groups = SPACE_BAND_GROUPS_IDX
-        self.time_groups = TIME_BAND_GROUPS_IDX
-        self.static_groups = STATIC_BAND_GROUPS_IDX
+        band_layout = band_layout or GalileoConfig().band_layout()
+        self.space_time_groups = band_layout["space_time_groups"]
+        self.space_groups = band_layout["space_groups"]
+        self.time_groups = band_layout["time_groups"]
+        self.static_groups = band_layout["static_groups"]
+        self.base_gsd = band_layout["input_resolution_m"]
         self.embedding_size = embedding_size
         self.base_patch_size = base_patch_size
 
@@ -606,16 +525,16 @@ class GalileoBase(nn.Module):
         else:
             args = {"requires_grad": False}
         self.s_t_channel_embed = nn.Parameter(
-            torch.zeros(len(SPACE_TIME_BANDS_GROUPS_IDX), int(embedding_size * 0.25)), **args
+            torch.zeros(len(self.space_time_groups), int(embedding_size * 0.25)), **args
         )
         self.sp_channel_embed = nn.Parameter(
-            torch.zeros(len(SPACE_BAND_GROUPS_IDX), int(embedding_size * 0.25)), **args
+            torch.zeros(len(self.space_groups), int(embedding_size * 0.25)), **args
         )
         self.t_channel_embed = nn.Parameter(
-            torch.zeros(len(TIME_BAND_GROUPS_IDX), int(embedding_size * 0.25)), **args
+            torch.zeros(len(self.time_groups), int(embedding_size * 0.25)), **args
         )
         self.st_channel_embed = nn.Parameter(
-            torch.zeros(len(STATIC_BAND_GROUPS_IDX), int(embedding_size * 0.25)), **args
+            torch.zeros(len(self.static_groups), int(embedding_size * 0.25)), **args
         )
 
         self.apply(self._init_weights)
@@ -720,7 +639,7 @@ class GalileoBase(nn.Module):
         if patch_size is None:
             patch_size = self.base_patch_size
         token_res = input_res * patch_size
-        gsd_ratio = token_res / BASE_GSD
+        gsd_ratio = token_res / self.base_gsd
 
         assert h == w, "get_2d_sincos_pos_embed_with_resolution currently requires that h==w"
         spatial_embed = get_2d_sincos_pos_embed_with_resolution(
@@ -755,6 +674,7 @@ class Encoder(GalileoBase):
         max_sequence_length=24,
         freeze_projections: bool = False,
         drop_path: float = 0.0,
+        band_layout: Optional[Dict[str, Any]] = None,
     ):
         super().__init__(
             embedding_size,
@@ -765,6 +685,7 @@ class Encoder(GalileoBase):
             max_patch_size,
             use_channel_embs=True,
             drop_path=drop_path,
+            band_layout=band_layout,
         )
 
         self.space_time_embed = nn.ModuleDict(
@@ -1108,11 +1029,14 @@ class Encoder(GalileoBase):
         st_m: torch.Tensor,
         months: torch.Tensor,
         patch_size: int,
-        input_resolution_m: Optional[int] = BASE_GSD,
+        input_resolution_m: Optional[int] = None,
         exit_after: Optional[int] = None,
         token_exit_cfg: Optional[Dict] = None,
         add_layernorm_on_exit: bool = True,
     ):
+        if input_resolution_m is None:
+            input_resolution_m = self.base_gsd
+
         (
             s_t_x,
             sp_x,
@@ -1161,25 +1085,31 @@ class Encoder(GalileoBase):
         )
 
     @classmethod
-    def load_from_folder(cls, folder: Path, device: torch.device):
-        if not (folder / CONFIG_FILENAME).exists():
+    def load_from_folder(
+        cls,
+        folder: Path,
+        device: torch.device,
+        config_filename: str = "config.json",
+        encoder_filename: str = "encoder.pt",
+    ):
+        if not (folder / config_filename).exists():
             all_files_in_folder = [f.name for f in folder.glob("*")]
             raise ValueError(
-                f"Expected {CONFIG_FILENAME} in {folder}, found {all_files_in_folder}"
+                f"Expected {config_filename} in {folder}, found {all_files_in_folder}"
             )
-        if not (folder / ENCODER_FILENAME).exists():
+        if not (folder / encoder_filename).exists():
             all_files_in_folder = [f.name for f in folder.glob("*")]
             raise ValueError(
-                f"Expected {ENCODER_FILENAME} in {folder}, found {all_files_in_folder}"
+                f"Expected {encoder_filename} in {folder}, found {all_files_in_folder}"
             )
 
-        with (folder / CONFIG_FILENAME).open("r") as f:
+        with (folder / config_filename).open("r") as f:
             config = json.load(f)
             model_config = config["model"]
             encoder_config = model_config["encoder"]
         encoder = cls(**encoder_config)
 
-        state_dict = torch.load(folder / ENCODER_FILENAME, map_location=device)
+        state_dict = torch.load(folder / encoder_filename, map_location=device)
         for key in list(state_dict.keys()):
             # this cleans the state dict, which occasionally had an extra
             # ".backbone" included in the key names
@@ -1189,6 +1119,141 @@ class Encoder(GalileoBase):
 logger = logging.get_logger(__name__)
 
 
+def _default_pretraining_normalizing_dict() -> Dict[str, Dict[str, List[float]]]:
+    return {
+        "13": {
+            "mean": [
+                -11.728724389184965,
+                -18.85558188024017,
+                1395.3408730676722,
+                1338.4026921784578,
+                1343.09883810357,
+                1543.8607982512297,
+                2186.2022069512263,
+                2525.0932853316694,
+                2410.3377187373408,
+                2750.2854646886753,
+                2234.911100061487,
+                1474.5311266077113,
+                0.2892116502999044,
+            ],
+            "std": [
+                4.887145774840316,
+                5.730270320384293,
+                917.7041440370853,
+                913.2988423581528,
+                1092.678723527555,
+                1047.2206083460424,
+                1048.0101611156767,
+                1143.6903026819996,
+                1098.979177731649,
+                1204.472755085893,
+                1145.9774063078878,
+                980.2429840007796,
+                0.2720939024500081,
+            ],
+        },
+        "16": {
+            "mean": [
+                673.0152819503361,
+                5.930092668915115,
+                0.10470439140978786,
+                0.23965913270066183,
+                0.08158044385860364,
+                0.04246976254259546,
+                0.11304392863520317,
+                0.17329647890362473,
+                0.0698981691616277,
+                0.12130267132802142,
+                0.04671318615236216,
+                10.973119802517362,
+                1.0927069179958768,
+                1.6991394232855903,
+                0.03720594618055555,
+                1.3671352688259548,
+            ],
+            "std": [
+                983.0697298296237,
+                8.167406789813247,
+                0.18771647977504985,
+                0.2368313455675914,
+                0.08024268534756586,
+                0.04045374496146404,
+                0.11350342472061795,
+                0.1279898111718168,
+                0.12042341550438586,
+                0.13602408145504347,
+                0.043971116096060345,
+                31.255340146970997,
+                10.395974878206689,
+                12.92380617159917,
+                1.9285254295940466,
+                11.612179775408928,
+            ],
+        },
+        "6": {
+            "mean": [
+                271.5674963541667,
+                0.08554303677156568,
+                657.3181260091111,
+                692.1291795806885,
+                562.781331880633,
+                1.5647115934036673,
+            ],
+            "std": [
+                79.80828940314429,
+                0.11669547098151486,
+                704.0008695557707,
+                925.0116126406431,
+                453.2434022278578,
+                7.513020170832818,
+            ],
+        },
+        "18": {
+            "mean": [
+                188.20315880851746,
+                0.2804946561574936,
+                0.11371652073860168,
+                0.058778801321983334,
+                0.10474256777763366,
+                0.2396918488264084,
+                0.08152248692512512,
+                0.04248040814399719,
+                0.11303179881572724,
+                0.17326324067115784,
+                0.06998309404850006,
+                0.12122812910079957,
+                0.04671641788482666,
+                10.98456594619751,
+                1.0968475807189941,
+                1.6947754135131836,
+                0.03320046615600586,
+                1.3602827312469483,
+            ],
+            "std": [
+                1154.5919128300602,
+                0.5276998078079327,
+                0.7021637331734328,
+                0.36528892213195063,
+                0.17470213191865785,
+                0.20411195416718833,
+                0.0660782470089761,
+                0.03380702424871257,
+                0.09809195568521663,
+                0.11292471052124119,
+                0.09720748930233268,
+                0.12912217763726777,
+                0.0399973913151906,
+                23.725471823867462,
+                5.715238079725388,
+                9.030481416228302,
+                0.9950220242487364,
+                7.754429123862099,
+            ],
+        },
+    }
+
+
 class GalileoConfig(PreTrainedConfig):
     model_type = "galileo"
 
@@ -1206,6 +1271,20 @@ class GalileoConfig(PreTrainedConfig):
         default_month: int = 6,
         global_pool: bool = True,
         input_resolution_m: int = 10,
+        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)
@@ -1221,6 +1300,89 @@ class GalileoConfig(PreTrainedConfig):
         self.default_month = default_month
         self.global_pool = global_pool
         self.input_resolution_m = input_resolution_m
+        self.s1_bands = s1_bands if s1_bands is not None else ["VV", "VH"]
+        self.s2_bands = s2_bands if s2_bands is not None else [
+            "B2", "B3", "B4", "B5", "B6", "B7", "B8", "B8A", "B11", "B12",
+        ]
+        self.era5_bands = era5_bands if era5_bands is not None else [
+            "temperature_2m", "total_precipitation_sum",
+        ]
+        self.tc_bands = tc_bands if tc_bands is not None else ["def", "soil", "aet"]
+        self.viirs_bands = viirs_bands if viirs_bands is not None else ["avg_rad"]
+        self.srtm_bands = srtm_bands if srtm_bands is not None else ["elevation", "slope"]
+        self.dw_bands = dw_bands if dw_bands is not None else [
+            "DW_water", "DW_trees", "DW_grass", "DW_flooded_vegetation", "DW_crops",
+            "DW_shrub_and_scrub", "DW_built", "DW_bare", "DW_snow_and_ice",
+        ]
+        self.wc_bands = wc_bands if wc_bands is not None else [
+            "WC_temporarycrops", "WC_maize", "WC_wintercereals", "WC_springcereals", "WC_irrigation",
+        ]
+        self.landscan_bands = landscan_bands if landscan_bands is not None else ["b1"]
+        self.location_bands = location_bands if location_bands is not None else ["x", "y", "z"]
+        self.space_time_band_groups = space_time_band_groups if space_time_band_groups is not None else {
+            "S1": ["VV", "VH"],
+            "S2_RGB": ["B2", "B3", "B4"],
+            "S2_Red_Edge": ["B5", "B6", "B7"],
+            "S2_NIR_10m": ["B8"],
+            "S2_NIR_20m": ["B8A"],
+            "S2_SWIR": ["B11", "B12"],
+            "NDVI": ["NDVI"],
+        }
+        self.time_band_groups = time_band_groups if time_band_groups is not None else {
+            "ERA5": ["temperature_2m", "total_precipitation_sum"],
+            "TC": ["def", "soil", "aet"],
+            "VIIRS": ["avg_rad"],
+        }
+        self.space_band_groups = space_band_groups if space_band_groups is not None else {
+            "SRTM": ["elevation", "slope"],
+            "DW": [
+                "DW_water", "DW_trees", "DW_grass", "DW_flooded_vegetation", "DW_crops",
+                "DW_shrub_and_scrub", "DW_built", "DW_bare", "DW_snow_and_ice",
+            ],
+            "WC": [
+                "WC_temporarycrops", "WC_maize", "WC_wintercereals", "WC_springcereals", "WC_irrigation",
+            ],
+        }
+        self.pretraining_normalizing_dict = (
+            pretraining_normalizing_dict
+            if pretraining_normalizing_dict is not None
+            else _default_pretraining_normalizing_dict()
+        )
+
+    def band_layout(self) -> Dict[str, Any]:
+        space_time_bands = self.s1_bands + self.s2_bands + ["NDVI"]
+        time_bands = self.era5_bands + self.tc_bands + self.viirs_bands
+        space_bands = self.srtm_bands + self.dw_bands + self.wc_bands
+        static_dw_bands = [f"{band}_static" for band in self.dw_bands]
+        static_wc_bands = [f"{band}_static" for band in self.wc_bands]
+        static_bands = self.landscan_bands + self.location_bands + static_dw_bands + static_wc_bands
+        static_band_groups = {
+            "LS": self.landscan_bands,
+            "location": self.location_bands,
+            "DW_static": static_dw_bands,
+            "WC_static": static_wc_bands,
+        }
+        return {
+            "s1_bands": self.s1_bands,
+            "s2_bands": self.s2_bands,
+            "era5_bands": self.era5_bands,
+            "tc_bands": self.tc_bands,
+            "viirs_bands": self.viirs_bands,
+            "srtm_bands": self.srtm_bands,
+            "dw_bands": self.dw_bands,
+            "wc_bands": self.wc_bands,
+            "landscan_bands": self.landscan_bands,
+            "location_bands": self.location_bands,
+            "space_time_bands": space_time_bands,
+            "time_bands": time_bands,
+            "space_bands": space_bands,
+            "static_bands": static_bands,
+            "space_time_groups": _band_group_indices(space_time_bands, self.space_time_band_groups),
+            "time_groups": _band_group_indices(time_bands, self.time_band_groups),
+            "space_groups": _band_group_indices(space_bands, self.space_band_groups),
+            "static_groups": _band_group_indices(static_bands, static_band_groups),
+            "input_resolution_m": self.input_resolution_m,
+        }
 
 
 class GalileoPreTrainedModel(PreTrainedModel):
@@ -1245,6 +1407,7 @@ class GalileoEncoderModel(GalileoPreTrainedModel):
             max_sequence_length=config.max_sequence_length,
             freeze_projections=config.freeze_projections,
             drop_path=config.drop_path,
+            band_layout=config.band_layout(),
         )
         self.add_pooling_layer = add_pooling_layer and config.global_pool
         self.post_init()

galileo-base-patch8/preprocessor_config.json

@@ -5,5 +5,262 @@
   "patch_size": 8,
   "auto_map": {
     "AutoProcessor": "processing_galileo.GalileoProcessor"
+  },
+  "s1_bands": [
+    "VV",
+    "VH"
+  ],
+  "s2_bands": [
+    "B2",
+    "B3",
+    "B4",
+    "B5",
+    "B6",
+    "B7",
+    "B8",
+    "B8A",
+    "B11",
+    "B12"
+  ],
+  "era5_bands": [
+    "temperature_2m",
+    "total_precipitation_sum"
+  ],
+  "tc_bands": [
+    "def",
+    "soil",
+    "aet"
+  ],
+  "viirs_bands": [
+    "avg_rad"
+  ],
+  "srtm_bands": [
+    "elevation",
+    "slope"
+  ],
+  "dw_bands": [
+    "DW_water",
+    "DW_trees",
+    "DW_grass",
+    "DW_flooded_vegetation",
+    "DW_crops",
+    "DW_shrub_and_scrub",
+    "DW_built",
+    "DW_bare",
+    "DW_snow_and_ice"
+  ],
+  "wc_bands": [
+    "WC_temporarycrops",
+    "WC_maize",
+    "WC_wintercereals",
+    "WC_springcereals",
+    "WC_irrigation"
+  ],
+  "landscan_bands": [
+    "b1"
+  ],
+  "location_bands": [
+    "x",
+    "y",
+    "z"
+  ],
+  "space_time_band_groups": {
+    "S1": [
+      "VV",
+      "VH"
+    ],
+    "S2_RGB": [
+      "B2",
+      "B3",
+      "B4"
+    ],
+    "S2_Red_Edge": [
+      "B5",
+      "B6",
+      "B7"
+    ],
+    "S2_NIR_10m": [
+      "B8"
+    ],
+    "S2_NIR_20m": [
+      "B8A"
+    ],
+    "S2_SWIR": [
+      "B11",
+      "B12"
+    ],
+    "NDVI": [
+      "NDVI"
+    ]
+  },
+  "time_band_groups": {
+    "ERA5": [
+      "temperature_2m",
+      "total_precipitation_sum"
+    ],
+    "TC": [
+      "def",
+      "soil",
+      "aet"
+    ],
+    "VIIRS": [
+      "avg_rad"
+    ]
+  },
+  "space_band_groups": {
+    "SRTM": [
+      "elevation",
+      "slope"
+    ],
+    "DW": [
+      "DW_water",
+      "DW_trees",
+      "DW_grass",
+      "DW_flooded_vegetation",
+      "DW_crops",
+      "DW_shrub_and_scrub",
+      "DW_built",
+      "DW_bare",
+      "DW_snow_and_ice"
+    ],
+    "WC": [
+      "WC_temporarycrops",
+      "WC_maize",
+      "WC_wintercereals",
+      "WC_springcereals",
+      "WC_irrigation"
+    ]
+  },
+  "pretraining_normalizing_dict": {
+    "13": {
+      "mean": [
+        -11.728724389184965,
+        -18.85558188024017,
+        1395.3408730676722,
+        1338.4026921784578,
+        1343.09883810357,
+        1543.8607982512297,
+        2186.2022069512263,
+        2525.0932853316694,
+        2410.3377187373408,
+        2750.2854646886753,
+        2234.911100061487,
+        1474.5311266077113,
+        0.2892116502999044
+      ],
+      "std": [
+        4.887145774840316,
+        5.730270320384293,
+        917.7041440370853,
+        913.2988423581528,
+        1092.678723527555,
+        1047.2206083460424,
+        1048.0101611156767,
+        1143.6903026819996,
+        1098.979177731649,
+        1204.472755085893,
+        1145.9774063078878,
+        980.2429840007796,
+        0.2720939024500081
+      ]
+    },
+    "16": {
+      "mean": [
+        673.0152819503361,
+        5.930092668915115,
+        0.10470439140978786,
+        0.23965913270066183,
+        0.08158044385860364,
+        0.04246976254259546,
+        0.11304392863520317,
+        0.17329647890362473,
+        0.0698981691616277,
+        0.12130267132802142,
+        0.04671318615236216,
+        10.973119802517362,
+        1.0927069179958768,
+        1.6991394232855903,
+        0.03720594618055555,
+        1.3671352688259548
+      ],
+      "std": [
+        983.0697298296237,
+        8.167406789813247,
+        0.18771647977504985,
+        0.2368313455675914,
+        0.08024268534756586,
+        0.04045374496146404,
+        0.11350342472061795,
+        0.1279898111718168,
+        0.12042341550438586,
+        0.13602408145504347,
+        0.043971116096060345,
+        31.255340146970997,
+        10.395974878206689,
+        12.92380617159917,
+        1.9285254295940466,
+        11.612179775408928
+      ]
+    },
+    "6": {
+      "mean": [
+        271.5674963541667,
+        0.08554303677156568,
+        657.3181260091111,
+        692.1291795806885,
+        562.781331880633,
+        1.5647115934036673
+      ],
+      "std": [
+        79.80828940314429,
+        0.11669547098151486,
+        704.0008695557707,
+        925.0116126406431,
+        453.2434022278578,
+        7.513020170832818
+      ]
+    },
+    "18": {
+      "mean": [
+        188.20315880851746,
+        0.2804946561574936,
+        0.11371652073860168,
+        0.058778801321983334,
+        0.10474256777763366,
+        0.2396918488264084,
+        0.08152248692512512,
+        0.04248040814399719,
+        0.11303179881572724,
+        0.17326324067115784,
+        0.06998309404850006,
+        0.12122812910079957,
+        0.04671641788482666,
+        10.98456594619751,
+        1.0968475807189941,
+        1.6947754135131836,
+        0.03320046615600586,
+        1.3602827312469483
+      ],
+      "std": [
+        1154.5919128300602,
+        0.5276998078079327,
+        0.7021637331734328,
+        0.36528892213195063,
+        0.17470213191865785,
+        0.20411195416718833,
+        0.0660782470089761,
+        0.03380702424871257,
+        0.09809195568521663,
+        0.11292471052124119,
+        0.09720748930233268,
+        0.12912217763726777,
+        0.0399973913151906,
+        23.725471823867462,
+        5.715238079725388,
+        9.030481416228302,
+        0.9950220242487364,
+        7.754429123862099
+      ]
+    }
   }
-}
+}

galileo-base-patch8/processing_galileo.py

@@ -4,7 +4,6 @@
 from __future__ import annotations
 
 import math
-from collections import OrderedDict
 from typing import Any, NamedTuple, Optional, Union
 
 import numpy as np
@@ -14,191 +13,7 @@ from transformers.feature_extraction_utils import BatchFeature
 from transformers.processing_utils import ProcessorMixin
 from transformers.utils import TensorType
 
-
-S1_BANDS = ["VV", "VH"]
-S2_BANDS = ["B2", "B3", "B4", "B5", "B6", "B7", "B8", "B8A", "B11", "B12"]
-ERA5_BANDS = ["temperature_2m", "total_precipitation_sum"]
-TC_BANDS = ["def", "soil", "aet"]
-VIIRS_BANDS = ["avg_rad"]
-SRTM_BANDS = ["elevation", "slope"]
-DW_BANDS = [
-    "DW_water", "DW_trees", "DW_grass", "DW_flooded_vegetation", "DW_crops",
-    "DW_shrub_and_scrub", "DW_built", "DW_bare", "DW_snow_and_ice",
-]
-WC_BANDS = [
-    "WC_temporarycrops", "WC_maize", "WC_wintercereals", "WC_springcereals", "WC_irrigation",
-]
-LANDSCAN_BANDS = ["b1"]
-LOCATION_BANDS = ["x", "y", "z"]
-STATIC_DW_BANDS = [f"{x}_static" for x in DW_BANDS]
-STATIC_WC_BANDS = [f"{x}_static" for x in WC_BANDS]
-
-SPACE_TIME_BANDS = S1_BANDS + S2_BANDS + ["NDVI"]
-TIME_BANDS = ERA5_BANDS + TC_BANDS + VIIRS_BANDS
-SPACE_BANDS = SRTM_BANDS + DW_BANDS + WC_BANDS
-STATIC_BANDS = LANDSCAN_BANDS + LOCATION_BANDS + STATIC_DW_BANDS + STATIC_WC_BANDS
-
-SPACE_TIME_BANDS_GROUPS_IDX = OrderedDict({
-    "S1": [SPACE_TIME_BANDS.index(b) for b in S1_BANDS],
-    "S2_RGB": [SPACE_TIME_BANDS.index(b) for b in ["B2", "B3", "B4"]],
-    "S2_Red_Edge": [SPACE_TIME_BANDS.index(b) for b in ["B5", "B6", "B7"]],
-    "S2_NIR_10m": [SPACE_TIME_BANDS.index(b) for b in ["B8"]],
-    "S2_NIR_20m": [SPACE_TIME_BANDS.index(b) for b in ["B8A"]],
-    "S2_SWIR": [SPACE_TIME_BANDS.index(b) for b in ["B11", "B12"]],
-    "NDVI": [SPACE_TIME_BANDS.index("NDVI")],
-})
-TIME_BAND_GROUPS_IDX = OrderedDict({
-    "ERA5": [TIME_BANDS.index(b) for b in ERA5_BANDS],
-    "TC": [TIME_BANDS.index(b) for b in TC_BANDS],
-    "VIIRS": [TIME_BANDS.index(b) for b in VIIRS_BANDS],
-})
-SPACE_BAND_GROUPS_IDX = OrderedDict({
-    "SRTM": [SPACE_BANDS.index(b) for b in SRTM_BANDS],
-    "DW": [SPACE_BANDS.index(b) for b in DW_BANDS],
-    "WC": [SPACE_BANDS.index(b) for b in WC_BANDS],
-})
-STATIC_BAND_GROUPS_IDX = OrderedDict({
-    "LS": [STATIC_BANDS.index(b) for b in LANDSCAN_BANDS],
-    "location": [STATIC_BANDS.index(b) for b in LOCATION_BANDS],
-    "DW_static": [STATIC_BANDS.index(b) for b in STATIC_DW_BANDS],
-    "WC_static": [STATIC_BANDS.index(b) for b in STATIC_WC_BANDS],
-})
-
-
-DEFAULT_MONTH = 5
-
-PRETRAINING_NORMALIZING_DICT = {
-    "13": {
-        "mean": [
-            -11.728724389184965,
-            -18.85558188024017,
-            1395.3408730676722,
-            1338.4026921784578,
-            1343.09883810357,
-            1543.8607982512297,
-            2186.2022069512263,
-            2525.0932853316694,
-            2410.3377187373408,
-            2750.2854646886753,
-            2234.911100061487,
-            1474.5311266077113,
-            0.2892116502999044,
-        ],
-        "std": [
-            4.887145774840316,
-            5.730270320384293,
-            917.7041440370853,
-            913.2988423581528,
-            1092.678723527555,
-            1047.2206083460424,
-            1048.0101611156767,
-            1143.6903026819996,
-            1098.979177731649,
-            1204.472755085893,
-            1145.9774063078878,
-            980.2429840007796,
-            0.2720939024500081,
-        ],
-    },
-    "16": {
-        "mean": [
-            673.0152819503361,
-            5.930092668915115,
-            0.10470439140978786,
-            0.23965913270066183,
-            0.08158044385860364,
-            0.04246976254259546,
-            0.11304392863520317,
-            0.17329647890362473,
-            0.0698981691616277,
-            0.12130267132802142,
-            0.04671318615236216,
-            10.973119802517362,
-            1.0927069179958768,
-            1.6991394232855903,
-            0.03720594618055555,
-            1.3671352688259548,
-        ],
-        "std": [
-            983.0697298296237,
-            8.167406789813247,
-            0.18771647977504985,
-            0.2368313455675914,
-            0.08024268534756586,
-            0.04045374496146404,
-            0.11350342472061795,
-            0.1279898111718168,
-            0.12042341550438586,
-            0.13602408145504347,
-            0.043971116096060345,
-            31.255340146970997,
-            10.395974878206689,
-            12.92380617159917,
-            1.9285254295940466,
-            11.612179775408928,
-        ],
-    },
-    "6": {
-        "mean": [
-            271.5674963541667,
-            0.08554303677156568,
-            657.3181260091111,
-            692.1291795806885,
-            562.781331880633,
-            1.5647115934036673,
-        ],
-        "std": [
-            79.80828940314429,
-            0.11669547098151486,
-            704.0008695557707,
-            925.0116126406431,
-            453.2434022278578,
-            7.513020170832818,
-        ],
-    },
-    "18": {
-        "mean": [
-            188.20315880851746,
-            0.2804946561574936,
-            0.11371652073860168,
-            0.058778801321983334,
-            0.10474256777763366,
-            0.2396918488264084,
-            0.08152248692512512,
-            0.04248040814399719,
-            0.11303179881572724,
-            0.17326324067115784,
-            0.06998309404850006,
-            0.12122812910079957,
-            0.04671641788482666,
-            10.98456594619751,
-            1.0968475807189941,
-            1.6947754135131836,
-            0.03320046615600586,
-            1.3602827312469483,
-        ],
-        "std": [
-            1154.5919128300602,
-            0.5276998078079327,
-            0.7021637331734328,
-            0.36528892213195063,
-            0.17470213191865785,
-            0.20411195416718833,
-            0.0660782470089761,
-            0.03380702424871257,
-            0.09809195568521663,
-            0.11292471052124119,
-            0.09720748930233268,
-            0.12912217763726777,
-            0.0399973913151906,
-            23.725471823867462,
-            5.715238079725388,
-            9.030481416228302,
-            0.9950220242487364,
-            7.754429123862099,
-        ],
-    },
-}
+from .modeling_galileo import GalileoConfig
 
 
 class MaskedOutput(NamedTuple):
@@ -257,7 +72,19 @@ def construct_galileo_input(
     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]
@@ -285,56 +112,56 @@ def construct_galileo_input(
     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_BANDS_GROUPS_IDX)), 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_BAND_GROUPS_IDX)), 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_BAND_GROUPS_IDX)), dtype=torch.float, device=device)
-    st_x = torch.zeros((len(STATIC_BANDS)), dtype=torch.float, device=device)
-    st_m = torch.ones((len(STATIC_BAND_GROUPS_IDX)), dtype=torch.float, device=device)
+    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], [S1_BANDS, S2_BANDS], ["S1", "S2"]):
+    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_BANDS_GROUPS_IDX) if group_key in key]
+            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], [SRTM_BANDS, DW_BANDS, WC_BANDS], ["SRTM", "DW", "WC"]
+        [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_BAND_GROUPS_IDX) if group_key in key]
+            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], [ERA5_BANDS, TC_BANDS, VIIRS_BANDS], ["ERA5", "TC", "VIIRS"]
+        [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_BAND_GROUPS_IDX) if group_key in key]
+            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], [LANDSCAN_BANDS, LOCATION_BANDS], ["LS", "location"]
+        [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_BAND_GROUPS_IDX) if group_key in key]
+            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) * DEFAULT_MONTH
+        months = torch.ones((t), dtype=torch.long, device=device) * band_config.default_month
 
     if normalize:
-        normalizer = PretrainingNormalizer(PRETRAINING_NORMALIZING_DICT)
+        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)
@@ -367,11 +194,48 @@ class GalileoProcessor(ProcessorMixin):
         "months",
     ]
 
-    def __init__(self, normalize: bool = True, default_month: int = 6, patch_size: int = 8, **kwargs):
+    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,
@@ -418,6 +282,7 @@ class GalileoProcessor(ProcessorMixin):
             latlon=_to_tensor(latlon),
             months=months,
             normalize=normalize,
+            band_config=self.band_config,
         )
 
         if masked_output.space_time_x.dim() == 4:

galileo-nano-patch8/config.json

@@ -43,5 +43,262 @@
         "AutoModel"
       ]
     }
+  },
+  "s1_bands": [
+    "VV",
+    "VH"
+  ],
+  "s2_bands": [
+    "B2",
+    "B3",
+    "B4",
+    "B5",
+    "B6",
+    "B7",
+    "B8",
+    "B8A",
+    "B11",
+    "B12"
+  ],
+  "era5_bands": [
+    "temperature_2m",
+    "total_precipitation_sum"
+  ],
+  "tc_bands": [
+    "def",
+    "soil",
+    "aet"
+  ],
+  "viirs_bands": [
+    "avg_rad"
+  ],
+  "srtm_bands": [
+    "elevation",
+    "slope"
+  ],
+  "dw_bands": [
+    "DW_water",
+    "DW_trees",
+    "DW_grass",
+    "DW_flooded_vegetation",
+    "DW_crops",
+    "DW_shrub_and_scrub",
+    "DW_built",
+    "DW_bare",
+    "DW_snow_and_ice"
+  ],
+  "wc_bands": [
+    "WC_temporarycrops",
+    "WC_maize",
+    "WC_wintercereals",
+    "WC_springcereals",
+    "WC_irrigation"
+  ],
+  "landscan_bands": [
+    "b1"
+  ],
+  "location_bands": [
+    "x",
+    "y",
+    "z"
+  ],
+  "space_time_band_groups": {
+    "S1": [
+      "VV",
+      "VH"
+    ],
+    "S2_RGB": [
+      "B2",
+      "B3",
+      "B4"
+    ],
+    "S2_Red_Edge": [
+      "B5",
+      "B6",
+      "B7"
+    ],
+    "S2_NIR_10m": [
+      "B8"
+    ],
+    "S2_NIR_20m": [
+      "B8A"
+    ],
+    "S2_SWIR": [
+      "B11",
+      "B12"
+    ],
+    "NDVI": [
+      "NDVI"
+    ]
+  },
+  "time_band_groups": {
+    "ERA5": [
+      "temperature_2m",
+      "total_precipitation_sum"
+    ],
+    "TC": [
+      "def",
+      "soil",
+      "aet"
+    ],
+    "VIIRS": [
+      "avg_rad"
+    ]
+  },
+  "space_band_groups": {
+    "SRTM": [
+      "elevation",
+      "slope"
+    ],
+    "DW": [
+      "DW_water",
+      "DW_trees",
+      "DW_grass",
+      "DW_flooded_vegetation",
+      "DW_crops",
+      "DW_shrub_and_scrub",
+      "DW_built",
+      "DW_bare",
+      "DW_snow_and_ice"
+    ],
+    "WC": [
+      "WC_temporarycrops",
+      "WC_maize",
+      "WC_wintercereals",
+      "WC_springcereals",
+      "WC_irrigation"
+    ]
+  },
+  "pretraining_normalizing_dict": {
+    "13": {
+      "mean": [
+        -11.728724389184965,
+        -18.85558188024017,
+        1395.3408730676722,
+        1338.4026921784578,
+        1343.09883810357,
+        1543.8607982512297,
+        2186.2022069512263,
+        2525.0932853316694,
+        2410.3377187373408,
+        2750.2854646886753,
+        2234.911100061487,
+        1474.5311266077113,
+        0.2892116502999044
+      ],
+      "std": [
+        4.887145774840316,
+        5.730270320384293,
+        917.7041440370853,
+        913.2988423581528,
+        1092.678723527555,
+        1047.2206083460424,
+        1048.0101611156767,
+        1143.6903026819996,
+        1098.979177731649,
+        1204.472755085893,
+        1145.9774063078878,
+        980.2429840007796,
+        0.2720939024500081
+      ]
+    },
+    "16": {
+      "mean": [
+        673.0152819503361,
+        5.930092668915115,
+        0.10470439140978786,
+        0.23965913270066183,
+        0.08158044385860364,
+        0.04246976254259546,
+        0.11304392863520317,
+        0.17329647890362473,
+        0.0698981691616277,
+        0.12130267132802142,
+        0.04671318615236216,
+        10.973119802517362,
+        1.0927069179958768,
+        1.6991394232855903,
+        0.03720594618055555,
+        1.3671352688259548
+      ],
+      "std": [
+        983.0697298296237,
+        8.167406789813247,
+        0.18771647977504985,
+        0.2368313455675914,
+        0.08024268534756586,
+        0.04045374496146404,
+        0.11350342472061795,
+        0.1279898111718168,
+        0.12042341550438586,
+        0.13602408145504347,
+        0.043971116096060345,
+        31.255340146970997,
+        10.395974878206689,
+        12.92380617159917,
+        1.9285254295940466,
+        11.612179775408928
+      ]
+    },
+    "6": {
+      "mean": [
+        271.5674963541667,
+        0.08554303677156568,
+        657.3181260091111,
+        692.1291795806885,
+        562.781331880633,
+        1.5647115934036673
+      ],
+      "std": [
+        79.80828940314429,
+        0.11669547098151486,
+        704.0008695557707,
+        925.0116126406431,
+        453.2434022278578,
+        7.513020170832818
+      ]
+    },
+    "18": {
+      "mean": [
+        188.20315880851746,
+        0.2804946561574936,
+        0.11371652073860168,
+        0.058778801321983334,
+        0.10474256777763366,
+        0.2396918488264084,
+        0.08152248692512512,
+        0.04248040814399719,
+        0.11303179881572724,
+        0.17326324067115784,
+        0.06998309404850006,
+        0.12122812910079957,
+        0.04671641788482666,
+        10.98456594619751,
+        1.0968475807189941,
+        1.6947754135131836,
+        0.03320046615600586,
+        1.3602827312469483
+      ],
+      "std": [
+        1154.5919128300602,
+        0.5276998078079327,
+        0.7021637331734328,
+        0.36528892213195063,
+        0.17470213191865785,
+        0.20411195416718833,
+        0.0660782470089761,
+        0.03380702424871257,
+        0.09809195568521663,
+        0.11292471052124119,
+        0.09720748930233268,
+        0.12912217763726777,
+        0.0399973913151906,
+        23.725471823867462,
+        5.715238079725388,
+        9.030481416228302,
+        0.9950220242487364,
+        7.754429123862099
+      ]
+    }
   }
-}
+}

galileo-nano-patch8/modeling_galileo.py

@@ -26,95 +26,11 @@ from einops import rearrange, repeat
 from torch import Tensor, vmap
 from torch.jit import Final
 
-# constants
-CONFIG_FILENAME = "config.json"
-ENCODER_FILENAME = "encoder.pt"
-BASE_GSD = 10
-
-# band information
-S1_BANDS = ["VV", "VH"]
-S2_BANDS = [
-    "B2",
-    "B3",
-    "B4",
-    "B5",
-    "B6",
-    "B7",
-    "B8",
-    "B8A",
-    "B11",
-    "B12",
-]
-ERA5_BANDS = ["temperature_2m", "total_precipitation_sum"]
-TC_BANDS = ["def", "soil", "aet"]
-VIIRS_BANDS = ["avg_rad"]
-SRTM_BANDS = ["elevation", "slope"]
-DW_BANDS = [
-    "DW_water",
-    "DW_trees",
-    "DW_grass",
-    "DW_flooded_vegetation",
-    "DW_crops",
-    "DW_shrub_and_scrub",
-    "DW_built",
-    "DW_bare",
-    "DW_snow_and_ice",
-]
-WC_BANDS = [
-    "WC_temporarycrops",
-    "WC_maize",
-    "WC_wintercereals",
-    "WC_springcereals",
-    "WC_irrigation",
-]
-STATIC_DW_BANDS = [f"{x}_static" for x in DW_BANDS]
-STATIC_WC_BANDS = [f"{x}_static" for x in WC_BANDS]
-
-LANDSCAN_BANDS = ["b1"]
-LOCATION_BANDS = ["x", "y", "z"]
-
-SPACE_TIME_BANDS = S1_BANDS + S2_BANDS + ["NDVI"]
-TIME_BANDS = ERA5_BANDS + TC_BANDS + VIIRS_BANDS
-SPACE_BANDS = SRTM_BANDS + DW_BANDS + WC_BANDS
-STATIC_BANDS = LANDSCAN_BANDS + LOCATION_BANDS + STATIC_DW_BANDS + STATIC_WC_BANDS
-
-
-SPACE_TIME_BANDS_GROUPS_IDX: OrderedDictType[str, List[int]] = OrderedDict(
-    {
-        "S1": [SPACE_TIME_BANDS.index(b) for b in S1_BANDS],
-        "S2_RGB": [SPACE_TIME_BANDS.index(b) for b in ["B2", "B3", "B4"]],
-        "S2_Red_Edge": [SPACE_TIME_BANDS.index(b) for b in ["B5", "B6", "B7"]],
-        "S2_NIR_10m": [SPACE_TIME_BANDS.index(b) for b in ["B8"]],
-        "S2_NIR_20m": [SPACE_TIME_BANDS.index(b) for b in ["B8A"]],
-        "S2_SWIR": [SPACE_TIME_BANDS.index(b) for b in ["B11", "B12"]],
-        "NDVI": [SPACE_TIME_BANDS.index("NDVI")],
-    }
-)
-
-TIME_BAND_GROUPS_IDX: OrderedDictType[str, List[int]] = OrderedDict(
-    {
-        "ERA5": [TIME_BANDS.index(b) for b in ERA5_BANDS],
-        "TC": [TIME_BANDS.index(b) for b in TC_BANDS],
-        "VIIRS": [TIME_BANDS.index(b) for b in VIIRS_BANDS],
-    }
-)
 
-SPACE_BAND_GROUPS_IDX: OrderedDictType[str, List[int]] = OrderedDict(
-    {
-        "SRTM": [SPACE_BANDS.index(b) for b in SRTM_BANDS],
-        "DW": [SPACE_BANDS.index(b) for b in DW_BANDS],
-        "WC": [SPACE_BANDS.index(b) for b in WC_BANDS],
-    }
-)
-
-STATIC_BAND_GROUPS_IDX: OrderedDictType[str, List[int]] = OrderedDict(
-    {
-        "LS": [STATIC_BANDS.index(b) for b in LANDSCAN_BANDS],
-        "location": [STATIC_BANDS.index(b) for b in LOCATION_BANDS],
-        "DW_static": [STATIC_BANDS.index(b) for b in STATIC_DW_BANDS],
-        "WC_static": [STATIC_BANDS.index(b) for b in STATIC_WC_BANDS],
-    }
-)
+def _band_group_indices(
+    bands: Sequence[str], groups: Dict[str, List[str]]
+) -> OrderedDictType[str, List[int]]:
+    return OrderedDict((name, [bands.index(b) for b in group_bands]) for name, group_bands in groups.items())
 
 
 def get_2d_sincos_pos_embed_with_resolution(
@@ -565,13 +481,16 @@ class GalileoBase(nn.Module):
         base_patch_size: int = 4,
         use_channel_embs: bool = True,
         drop_path: float = 0.0,
+        band_layout: Optional[Dict[str, Any]] = None,
     ):
         super().__init__()
 
-        self.space_time_groups = SPACE_TIME_BANDS_GROUPS_IDX
-        self.space_groups = SPACE_BAND_GROUPS_IDX
-        self.time_groups = TIME_BAND_GROUPS_IDX
-        self.static_groups = STATIC_BAND_GROUPS_IDX
+        band_layout = band_layout or GalileoConfig().band_layout()
+        self.space_time_groups = band_layout["space_time_groups"]
+        self.space_groups = band_layout["space_groups"]
+        self.time_groups = band_layout["time_groups"]
+        self.static_groups = band_layout["static_groups"]
+        self.base_gsd = band_layout["input_resolution_m"]
         self.embedding_size = embedding_size
         self.base_patch_size = base_patch_size
 
@@ -606,16 +525,16 @@ class GalileoBase(nn.Module):
         else:
             args = {"requires_grad": False}
         self.s_t_channel_embed = nn.Parameter(
-            torch.zeros(len(SPACE_TIME_BANDS_GROUPS_IDX), int(embedding_size * 0.25)), **args
+            torch.zeros(len(self.space_time_groups), int(embedding_size * 0.25)), **args
         )
         self.sp_channel_embed = nn.Parameter(
-            torch.zeros(len(SPACE_BAND_GROUPS_IDX), int(embedding_size * 0.25)), **args
+            torch.zeros(len(self.space_groups), int(embedding_size * 0.25)), **args
         )
         self.t_channel_embed = nn.Parameter(
-            torch.zeros(len(TIME_BAND_GROUPS_IDX), int(embedding_size * 0.25)), **args
+            torch.zeros(len(self.time_groups), int(embedding_size * 0.25)), **args
         )
         self.st_channel_embed = nn.Parameter(
-            torch.zeros(len(STATIC_BAND_GROUPS_IDX), int(embedding_size * 0.25)), **args
+            torch.zeros(len(self.static_groups), int(embedding_size * 0.25)), **args
         )
 
         self.apply(self._init_weights)
@@ -720,7 +639,7 @@ class GalileoBase(nn.Module):
         if patch_size is None:
             patch_size = self.base_patch_size
         token_res = input_res * patch_size
-        gsd_ratio = token_res / BASE_GSD
+        gsd_ratio = token_res / self.base_gsd
 
         assert h == w, "get_2d_sincos_pos_embed_with_resolution currently requires that h==w"
         spatial_embed = get_2d_sincos_pos_embed_with_resolution(
@@ -755,6 +674,7 @@ class Encoder(GalileoBase):
         max_sequence_length=24,
         freeze_projections: bool = False,
         drop_path: float = 0.0,
+        band_layout: Optional[Dict[str, Any]] = None,
     ):
         super().__init__(
             embedding_size,
@@ -765,6 +685,7 @@ class Encoder(GalileoBase):
             max_patch_size,
             use_channel_embs=True,
             drop_path=drop_path,
+            band_layout=band_layout,
         )
 
         self.space_time_embed = nn.ModuleDict(
@@ -1108,11 +1029,14 @@ class Encoder(GalileoBase):
         st_m: torch.Tensor,
         months: torch.Tensor,
         patch_size: int,
-        input_resolution_m: Optional[int] = BASE_GSD,
+        input_resolution_m: Optional[int] = None,
         exit_after: Optional[int] = None,
         token_exit_cfg: Optional[Dict] = None,
         add_layernorm_on_exit: bool = True,
     ):
+        if input_resolution_m is None:
+            input_resolution_m = self.base_gsd
+
         (
             s_t_x,
             sp_x,
@@ -1161,25 +1085,31 @@ class Encoder(GalileoBase):
         )
 
     @classmethod
-    def load_from_folder(cls, folder: Path, device: torch.device):
-        if not (folder / CONFIG_FILENAME).exists():
+    def load_from_folder(
+        cls,
+        folder: Path,
+        device: torch.device,
+        config_filename: str = "config.json",
+        encoder_filename: str = "encoder.pt",
+    ):
+        if not (folder / config_filename).exists():
             all_files_in_folder = [f.name for f in folder.glob("*")]
             raise ValueError(
-                f"Expected {CONFIG_FILENAME} in {folder}, found {all_files_in_folder}"
+                f"Expected {config_filename} in {folder}, found {all_files_in_folder}"
             )
-        if not (folder / ENCODER_FILENAME).exists():
+        if not (folder / encoder_filename).exists():
             all_files_in_folder = [f.name for f in folder.glob("*")]
             raise ValueError(
-                f"Expected {ENCODER_FILENAME} in {folder}, found {all_files_in_folder}"
+                f"Expected {encoder_filename} in {folder}, found {all_files_in_folder}"
             )
 
-        with (folder / CONFIG_FILENAME).open("r") as f:
+        with (folder / config_filename).open("r") as f:
             config = json.load(f)
             model_config = config["model"]
             encoder_config = model_config["encoder"]
         encoder = cls(**encoder_config)
 
-        state_dict = torch.load(folder / ENCODER_FILENAME, map_location=device)
+        state_dict = torch.load(folder / encoder_filename, map_location=device)
         for key in list(state_dict.keys()):
             # this cleans the state dict, which occasionally had an extra
             # ".backbone" included in the key names
@@ -1189,6 +1119,141 @@ class Encoder(GalileoBase):
 logger = logging.get_logger(__name__)
 
 
+def _default_pretraining_normalizing_dict() -> Dict[str, Dict[str, List[float]]]:
+    return {
+        "13": {
+            "mean": [
+                -11.728724389184965,
+                -18.85558188024017,
+                1395.3408730676722,
+                1338.4026921784578,
+                1343.09883810357,
+                1543.8607982512297,
+                2186.2022069512263,
+                2525.0932853316694,
+                2410.3377187373408,
+                2750.2854646886753,
+                2234.911100061487,
+                1474.5311266077113,
+                0.2892116502999044,
+            ],
+            "std": [
+                4.887145774840316,
+                5.730270320384293,
+                917.7041440370853,
+                913.2988423581528,
+                1092.678723527555,
+                1047.2206083460424,
+                1048.0101611156767,
+                1143.6903026819996,
+                1098.979177731649,
+                1204.472755085893,
+                1145.9774063078878,
+                980.2429840007796,
+                0.2720939024500081,
+            ],
+        },
+        "16": {
+            "mean": [
+                673.0152819503361,
+                5.930092668915115,
+                0.10470439140978786,
+                0.23965913270066183,
+                0.08158044385860364,
+                0.04246976254259546,
+                0.11304392863520317,
+                0.17329647890362473,
+                0.0698981691616277,
+                0.12130267132802142,
+                0.04671318615236216,
+                10.973119802517362,
+                1.0927069179958768,
+                1.6991394232855903,
+                0.03720594618055555,
+                1.3671352688259548,
+            ],
+            "std": [
+                983.0697298296237,
+                8.167406789813247,
+                0.18771647977504985,
+                0.2368313455675914,
+                0.08024268534756586,
+                0.04045374496146404,
+                0.11350342472061795,
+                0.1279898111718168,
+                0.12042341550438586,
+                0.13602408145504347,
+                0.043971116096060345,
+                31.255340146970997,
+                10.395974878206689,
+                12.92380617159917,
+                1.9285254295940466,
+                11.612179775408928,
+            ],
+        },
+        "6": {
+            "mean": [
+                271.5674963541667,
+                0.08554303677156568,
+                657.3181260091111,
+                692.1291795806885,
+                562.781331880633,
+                1.5647115934036673,
+            ],
+            "std": [
+                79.80828940314429,
+                0.11669547098151486,
+                704.0008695557707,
+                925.0116126406431,
+                453.2434022278578,
+                7.513020170832818,
+            ],
+        },
+        "18": {
+            "mean": [
+                188.20315880851746,
+                0.2804946561574936,
+                0.11371652073860168,
+                0.058778801321983334,
+                0.10474256777763366,
+                0.2396918488264084,
+                0.08152248692512512,
+                0.04248040814399719,
+                0.11303179881572724,
+                0.17326324067115784,
+                0.06998309404850006,
+                0.12122812910079957,
+                0.04671641788482666,
+                10.98456594619751,
+                1.0968475807189941,
+                1.6947754135131836,
+                0.03320046615600586,
+                1.3602827312469483,
+            ],
+            "std": [
+                1154.5919128300602,
+                0.5276998078079327,
+                0.7021637331734328,
+                0.36528892213195063,
+                0.17470213191865785,
+                0.20411195416718833,
+                0.0660782470089761,
+                0.03380702424871257,
+                0.09809195568521663,
+                0.11292471052124119,
+                0.09720748930233268,
+                0.12912217763726777,
+                0.0399973913151906,
+                23.725471823867462,
+                5.715238079725388,
+                9.030481416228302,
+                0.9950220242487364,
+                7.754429123862099,
+            ],
+        },
+    }
+
+
 class GalileoConfig(PreTrainedConfig):
     model_type = "galileo"
 
@@ -1206,6 +1271,20 @@ class GalileoConfig(PreTrainedConfig):
         default_month: int = 6,
         global_pool: bool = True,
         input_resolution_m: int = 10,
+        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)
@@ -1221,6 +1300,89 @@ class GalileoConfig(PreTrainedConfig):
         self.default_month = default_month
         self.global_pool = global_pool
         self.input_resolution_m = input_resolution_m
+        self.s1_bands = s1_bands if s1_bands is not None else ["VV", "VH"]
+        self.s2_bands = s2_bands if s2_bands is not None else [
+            "B2", "B3", "B4", "B5", "B6", "B7", "B8", "B8A", "B11", "B12",
+        ]
+        self.era5_bands = era5_bands if era5_bands is not None else [
+            "temperature_2m", "total_precipitation_sum",
+        ]
+        self.tc_bands = tc_bands if tc_bands is not None else ["def", "soil", "aet"]
+        self.viirs_bands = viirs_bands if viirs_bands is not None else ["avg_rad"]
+        self.srtm_bands = srtm_bands if srtm_bands is not None else ["elevation", "slope"]
+        self.dw_bands = dw_bands if dw_bands is not None else [
+            "DW_water", "DW_trees", "DW_grass", "DW_flooded_vegetation", "DW_crops",
+            "DW_shrub_and_scrub", "DW_built", "DW_bare", "DW_snow_and_ice",
+        ]
+        self.wc_bands = wc_bands if wc_bands is not None else [
+            "WC_temporarycrops", "WC_maize", "WC_wintercereals", "WC_springcereals", "WC_irrigation",
+        ]
+        self.landscan_bands = landscan_bands if landscan_bands is not None else ["b1"]
+        self.location_bands = location_bands if location_bands is not None else ["x", "y", "z"]
+        self.space_time_band_groups = space_time_band_groups if space_time_band_groups is not None else {
+            "S1": ["VV", "VH"],
+            "S2_RGB": ["B2", "B3", "B4"],
+            "S2_Red_Edge": ["B5", "B6", "B7"],
+            "S2_NIR_10m": ["B8"],
+            "S2_NIR_20m": ["B8A"],
+            "S2_SWIR": ["B11", "B12"],
+            "NDVI": ["NDVI"],
+        }
+        self.time_band_groups = time_band_groups if time_band_groups is not None else {
+            "ERA5": ["temperature_2m", "total_precipitation_sum"],
+            "TC": ["def", "soil", "aet"],
+            "VIIRS": ["avg_rad"],
+        }
+        self.space_band_groups = space_band_groups if space_band_groups is not None else {
+            "SRTM": ["elevation", "slope"],
+            "DW": [
+                "DW_water", "DW_trees", "DW_grass", "DW_flooded_vegetation", "DW_crops",
+                "DW_shrub_and_scrub", "DW_built", "DW_bare", "DW_snow_and_ice",
+            ],
+            "WC": [
+                "WC_temporarycrops", "WC_maize", "WC_wintercereals", "WC_springcereals", "WC_irrigation",
+            ],
+        }
+        self.pretraining_normalizing_dict = (
+            pretraining_normalizing_dict
+            if pretraining_normalizing_dict is not None
+            else _default_pretraining_normalizing_dict()
+        )
+
+    def band_layout(self) -> Dict[str, Any]:
+        space_time_bands = self.s1_bands + self.s2_bands + ["NDVI"]
+        time_bands = self.era5_bands + self.tc_bands + self.viirs_bands
+        space_bands = self.srtm_bands + self.dw_bands + self.wc_bands
+        static_dw_bands = [f"{band}_static" for band in self.dw_bands]
+        static_wc_bands = [f"{band}_static" for band in self.wc_bands]
+        static_bands = self.landscan_bands + self.location_bands + static_dw_bands + static_wc_bands
+        static_band_groups = {
+            "LS": self.landscan_bands,
+            "location": self.location_bands,
+            "DW_static": static_dw_bands,
+            "WC_static": static_wc_bands,
+        }
+        return {
+            "s1_bands": self.s1_bands,
+            "s2_bands": self.s2_bands,
+            "era5_bands": self.era5_bands,
+            "tc_bands": self.tc_bands,
+            "viirs_bands": self.viirs_bands,
+            "srtm_bands": self.srtm_bands,
+            "dw_bands": self.dw_bands,
+            "wc_bands": self.wc_bands,
+            "landscan_bands": self.landscan_bands,
+            "location_bands": self.location_bands,
+            "space_time_bands": space_time_bands,
+            "time_bands": time_bands,
+            "space_bands": space_bands,
+            "static_bands": static_bands,
+            "space_time_groups": _band_group_indices(space_time_bands, self.space_time_band_groups),
+            "time_groups": _band_group_indices(time_bands, self.time_band_groups),
+            "space_groups": _band_group_indices(space_bands, self.space_band_groups),
+            "static_groups": _band_group_indices(static_bands, static_band_groups),
+            "input_resolution_m": self.input_resolution_m,
+        }
 
 
 class GalileoPreTrainedModel(PreTrainedModel):
@@ -1245,6 +1407,7 @@ class GalileoEncoderModel(GalileoPreTrainedModel):
             max_sequence_length=config.max_sequence_length,
             freeze_projections=config.freeze_projections,
             drop_path=config.drop_path,
+            band_layout=config.band_layout(),
         )
         self.add_pooling_layer = add_pooling_layer and config.global_pool
         self.post_init()

galileo-nano-patch8/preprocessor_config.json

@@ -5,5 +5,262 @@
   "patch_size": 8,
   "auto_map": {
     "AutoProcessor": "processing_galileo.GalileoProcessor"
+  },
+  "s1_bands": [
+    "VV",
+    "VH"
+  ],
+  "s2_bands": [
+    "B2",
+    "B3",
+    "B4",
+    "B5",
+    "B6",
+    "B7",
+    "B8",
+    "B8A",
+    "B11",
+    "B12"
+  ],
+  "era5_bands": [
+    "temperature_2m",
+    "total_precipitation_sum"
+  ],
+  "tc_bands": [
+    "def",
+    "soil",
+    "aet"
+  ],
+  "viirs_bands": [
+    "avg_rad"
+  ],
+  "srtm_bands": [
+    "elevation",
+    "slope"
+  ],
+  "dw_bands": [
+    "DW_water",
+    "DW_trees",
+    "DW_grass",
+    "DW_flooded_vegetation",
+    "DW_crops",
+    "DW_shrub_and_scrub",
+    "DW_built",
+    "DW_bare",
+    "DW_snow_and_ice"
+  ],
+  "wc_bands": [
+    "WC_temporarycrops",
+    "WC_maize",
+    "WC_wintercereals",
+    "WC_springcereals",
+    "WC_irrigation"
+  ],
+  "landscan_bands": [
+    "b1"
+  ],
+  "location_bands": [
+    "x",
+    "y",
+    "z"
+  ],
+  "space_time_band_groups": {
+    "S1": [
+      "VV",
+      "VH"
+    ],
+    "S2_RGB": [
+      "B2",
+      "B3",
+      "B4"
+    ],
+    "S2_Red_Edge": [
+      "B5",
+      "B6",
+      "B7"
+    ],
+    "S2_NIR_10m": [
+      "B8"
+    ],
+    "S2_NIR_20m": [
+      "B8A"
+    ],
+    "S2_SWIR": [
+      "B11",
+      "B12"
+    ],
+    "NDVI": [
+      "NDVI"
+    ]
+  },
+  "time_band_groups": {
+    "ERA5": [
+      "temperature_2m",
+      "total_precipitation_sum"
+    ],
+    "TC": [
+      "def",
+      "soil",
+      "aet"
+    ],
+    "VIIRS": [
+      "avg_rad"
+    ]
+  },
+  "space_band_groups": {
+    "SRTM": [
+      "elevation",
+      "slope"
+    ],
+    "DW": [
+      "DW_water",
+      "DW_trees",
+      "DW_grass",
+      "DW_flooded_vegetation",
+      "DW_crops",
+      "DW_shrub_and_scrub",
+      "DW_built",
+      "DW_bare",
+      "DW_snow_and_ice"
+    ],
+    "WC": [
+      "WC_temporarycrops",
+      "WC_maize",
+      "WC_wintercereals",
+      "WC_springcereals",
+      "WC_irrigation"
+    ]
+  },
+  "pretraining_normalizing_dict": {
+    "13": {
+      "mean": [
+        -11.728724389184965,
+        -18.85558188024017,
+        1395.3408730676722,
+        1338.4026921784578,
+        1343.09883810357,
+        1543.8607982512297,
+        2186.2022069512263,
+        2525.0932853316694,
+        2410.3377187373408,
+        2750.2854646886753,
+        2234.911100061487,
+        1474.5311266077113,
+        0.2892116502999044
+      ],
+      "std": [
+        4.887145774840316,
+        5.730270320384293,
+        917.7041440370853,
+        913.2988423581528,
+        1092.678723527555,
+        1047.2206083460424,
+        1048.0101611156767,
+        1143.6903026819996,
+        1098.979177731649,
+        1204.472755085893,
+        1145.9774063078878,
+        980.2429840007796,
+        0.2720939024500081
+      ]
+    },
+    "16": {
+      "mean": [
+        673.0152819503361,
+        5.930092668915115,
+        0.10470439140978786,
+        0.23965913270066183,
+        0.08158044385860364,
+        0.04246976254259546,
+        0.11304392863520317,
+        0.17329647890362473,
+        0.0698981691616277,
+        0.12130267132802142,
+        0.04671318615236216,
+        10.973119802517362,
+        1.0927069179958768,
+        1.6991394232855903,
+        0.03720594618055555,
+        1.3671352688259548
+      ],
+      "std": [
+        983.0697298296237,
+        8.167406789813247,
+        0.18771647977504985,
+        0.2368313455675914,
+        0.08024268534756586,
+        0.04045374496146404,
+        0.11350342472061795,
+        0.1279898111718168,
+        0.12042341550438586,
+        0.13602408145504347,
+        0.043971116096060345,
+        31.255340146970997,
+        10.395974878206689,
+        12.92380617159917,
+        1.9285254295940466,
+        11.612179775408928
+      ]
+    },
+    "6": {
+      "mean": [
+        271.5674963541667,
+        0.08554303677156568,
+        657.3181260091111,
+        692.1291795806885,
+        562.781331880633,
+        1.5647115934036673
+      ],
+      "std": [
+        79.80828940314429,
+        0.11669547098151486,
+        704.0008695557707,
+        925.0116126406431,
+        453.2434022278578,
+        7.513020170832818
+      ]
+    },
+    "18": {
+      "mean": [
+        188.20315880851746,
+        0.2804946561574936,
+        0.11371652073860168,
+        0.058778801321983334,
+        0.10474256777763366,
+        0.2396918488264084,
+        0.08152248692512512,
+        0.04248040814399719,
+        0.11303179881572724,
+        0.17326324067115784,
+        0.06998309404850006,
+        0.12122812910079957,
+        0.04671641788482666,
+        10.98456594619751,
+        1.0968475807189941,
+        1.6947754135131836,
+        0.03320046615600586,
+        1.3602827312469483
+      ],
+      "std": [
+        1154.5919128300602,
+        0.5276998078079327,
+        0.7021637331734328,
+        0.36528892213195063,
+        0.17470213191865785,
+        0.20411195416718833,
+        0.0660782470089761,
+        0.03380702424871257,
+        0.09809195568521663,
+        0.11292471052124119,
+        0.09720748930233268,
+        0.12912217763726777,
+        0.0399973913151906,
+        23.725471823867462,
+        5.715238079725388,
+        9.030481416228302,
+        0.9950220242487364,
+        7.754429123862099
+      ]
+    }
   }
-}
+}

galileo-nano-patch8/processing_galileo.py

@@ -4,7 +4,6 @@
 from __future__ import annotations
 
 import math
-from collections import OrderedDict
 from typing import Any, NamedTuple, Optional, Union
 
 import numpy as np
@@ -14,191 +13,7 @@ from transformers.feature_extraction_utils import BatchFeature
 from transformers.processing_utils import ProcessorMixin
 from transformers.utils import TensorType
 
-
-S1_BANDS = ["VV", "VH"]
-S2_BANDS = ["B2", "B3", "B4", "B5", "B6", "B7", "B8", "B8A", "B11", "B12"]
-ERA5_BANDS = ["temperature_2m", "total_precipitation_sum"]
-TC_BANDS = ["def", "soil", "aet"]
-VIIRS_BANDS = ["avg_rad"]
-SRTM_BANDS = ["elevation", "slope"]
-DW_BANDS = [
-    "DW_water", "DW_trees", "DW_grass", "DW_flooded_vegetation", "DW_crops",
-    "DW_shrub_and_scrub", "DW_built", "DW_bare", "DW_snow_and_ice",
-]
-WC_BANDS = [
-    "WC_temporarycrops", "WC_maize", "WC_wintercereals", "WC_springcereals", "WC_irrigation",
-]
-LANDSCAN_BANDS = ["b1"]
-LOCATION_BANDS = ["x", "y", "z"]
-STATIC_DW_BANDS = [f"{x}_static" for x in DW_BANDS]
-STATIC_WC_BANDS = [f"{x}_static" for x in WC_BANDS]
-
-SPACE_TIME_BANDS = S1_BANDS + S2_BANDS + ["NDVI"]
-TIME_BANDS = ERA5_BANDS + TC_BANDS + VIIRS_BANDS
-SPACE_BANDS = SRTM_BANDS + DW_BANDS + WC_BANDS
-STATIC_BANDS = LANDSCAN_BANDS + LOCATION_BANDS + STATIC_DW_BANDS + STATIC_WC_BANDS
-
-SPACE_TIME_BANDS_GROUPS_IDX = OrderedDict({
-    "S1": [SPACE_TIME_BANDS.index(b) for b in S1_BANDS],
-    "S2_RGB": [SPACE_TIME_BANDS.index(b) for b in ["B2", "B3", "B4"]],
-    "S2_Red_Edge": [SPACE_TIME_BANDS.index(b) for b in ["B5", "B6", "B7"]],
-    "S2_NIR_10m": [SPACE_TIME_BANDS.index(b) for b in ["B8"]],
-    "S2_NIR_20m": [SPACE_TIME_BANDS.index(b) for b in ["B8A"]],
-    "S2_SWIR": [SPACE_TIME_BANDS.index(b) for b in ["B11", "B12"]],
-    "NDVI": [SPACE_TIME_BANDS.index("NDVI")],
-})
-TIME_BAND_GROUPS_IDX = OrderedDict({
-    "ERA5": [TIME_BANDS.index(b) for b in ERA5_BANDS],
-    "TC": [TIME_BANDS.index(b) for b in TC_BANDS],
-    "VIIRS": [TIME_BANDS.index(b) for b in VIIRS_BANDS],
-})
-SPACE_BAND_GROUPS_IDX = OrderedDict({
-    "SRTM": [SPACE_BANDS.index(b) for b in SRTM_BANDS],
-    "DW": [SPACE_BANDS.index(b) for b in DW_BANDS],
-    "WC": [SPACE_BANDS.index(b) for b in WC_BANDS],
-})
-STATIC_BAND_GROUPS_IDX = OrderedDict({
-    "LS": [STATIC_BANDS.index(b) for b in LANDSCAN_BANDS],
-    "location": [STATIC_BANDS.index(b) for b in LOCATION_BANDS],
-    "DW_static": [STATIC_BANDS.index(b) for b in STATIC_DW_BANDS],
-    "WC_static": [STATIC_BANDS.index(b) for b in STATIC_WC_BANDS],
-})
-
-
-DEFAULT_MONTH = 5
-
-PRETRAINING_NORMALIZING_DICT = {
-    "13": {
-        "mean": [
-            -11.728724389184965,
-            -18.85558188024017,
-            1395.3408730676722,
-            1338.4026921784578,
-            1343.09883810357,
-            1543.8607982512297,
-            2186.2022069512263,
-            2525.0932853316694,
-            2410.3377187373408,
-            2750.2854646886753,
-            2234.911100061487,
-            1474.5311266077113,
-            0.2892116502999044,
-        ],
-        "std": [
-            4.887145774840316,
-            5.730270320384293,
-            917.7041440370853,
-            913.2988423581528,
-            1092.678723527555,
-            1047.2206083460424,
-            1048.0101611156767,
-            1143.6903026819996,
-            1098.979177731649,
-            1204.472755085893,
-            1145.9774063078878,
-            980.2429840007796,
-            0.2720939024500081,
-        ],
-    },
-    "16": {
-        "mean": [
-            673.0152819503361,
-            5.930092668915115,
-            0.10470439140978786,
-            0.23965913270066183,
-            0.08158044385860364,
-            0.04246976254259546,
-            0.11304392863520317,
-            0.17329647890362473,
-            0.0698981691616277,
-            0.12130267132802142,
-            0.04671318615236216,
-            10.973119802517362,
-            1.0927069179958768,
-            1.6991394232855903,
-            0.03720594618055555,
-            1.3671352688259548,
-        ],
-        "std": [
-            983.0697298296237,
-            8.167406789813247,
-            0.18771647977504985,
-            0.2368313455675914,
-            0.08024268534756586,
-            0.04045374496146404,
-            0.11350342472061795,
-            0.1279898111718168,
-            0.12042341550438586,
-            0.13602408145504347,
-            0.043971116096060345,
-            31.255340146970997,
-            10.395974878206689,
-            12.92380617159917,
-            1.9285254295940466,
-            11.612179775408928,
-        ],
-    },
-    "6": {
-        "mean": [
-            271.5674963541667,
-            0.08554303677156568,
-            657.3181260091111,
-            692.1291795806885,
-            562.781331880633,
-            1.5647115934036673,
-        ],
-        "std": [
-            79.80828940314429,
-            0.11669547098151486,
-            704.0008695557707,
-            925.0116126406431,
-            453.2434022278578,
-            7.513020170832818,
-        ],
-    },
-    "18": {
-        "mean": [
-            188.20315880851746,
-            0.2804946561574936,
-            0.11371652073860168,
-            0.058778801321983334,
-            0.10474256777763366,
-            0.2396918488264084,
-            0.08152248692512512,
-            0.04248040814399719,
-            0.11303179881572724,
-            0.17326324067115784,
-            0.06998309404850006,
-            0.12122812910079957,
-            0.04671641788482666,
-            10.98456594619751,
-            1.0968475807189941,
-            1.6947754135131836,
-            0.03320046615600586,
-            1.3602827312469483,
-        ],
-        "std": [
-            1154.5919128300602,
-            0.5276998078079327,
-            0.7021637331734328,
-            0.36528892213195063,
-            0.17470213191865785,
-            0.20411195416718833,
-            0.0660782470089761,
-            0.03380702424871257,
-            0.09809195568521663,
-            0.11292471052124119,
-            0.09720748930233268,
-            0.12912217763726777,
-            0.0399973913151906,
-            23.725471823867462,
-            5.715238079725388,
-            9.030481416228302,
-            0.9950220242487364,
-            7.754429123862099,
-        ],
-    },
-}
+from .modeling_galileo import GalileoConfig
 
 
 class MaskedOutput(NamedTuple):
@@ -257,7 +72,19 @@ def construct_galileo_input(
     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]
@@ -285,56 +112,56 @@ def construct_galileo_input(
     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_BANDS_GROUPS_IDX)), 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_BAND_GROUPS_IDX)), 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_BAND_GROUPS_IDX)), dtype=torch.float, device=device)
-    st_x = torch.zeros((len(STATIC_BANDS)), dtype=torch.float, device=device)
-    st_m = torch.ones((len(STATIC_BAND_GROUPS_IDX)), dtype=torch.float, device=device)
+    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], [S1_BANDS, S2_BANDS], ["S1", "S2"]):
+    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_BANDS_GROUPS_IDX) if group_key in key]
+            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], [SRTM_BANDS, DW_BANDS, WC_BANDS], ["SRTM", "DW", "WC"]
+        [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_BAND_GROUPS_IDX) if group_key in key]
+            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], [ERA5_BANDS, TC_BANDS, VIIRS_BANDS], ["ERA5", "TC", "VIIRS"]
+        [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_BAND_GROUPS_IDX) if group_key in key]
+            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], [LANDSCAN_BANDS, LOCATION_BANDS], ["LS", "location"]
+        [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_BAND_GROUPS_IDX) if group_key in key]
+            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) * DEFAULT_MONTH
+        months = torch.ones((t), dtype=torch.long, device=device) * band_config.default_month
 
     if normalize:
-        normalizer = PretrainingNormalizer(PRETRAINING_NORMALIZING_DICT)
+        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)
@@ -367,11 +194,48 @@ class GalileoProcessor(ProcessorMixin):
         "months",
     ]
 
-    def __init__(self, normalize: bool = True, default_month: int = 6, patch_size: int = 8, **kwargs):
+    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,
@@ -418,6 +282,7 @@ class GalileoProcessor(ProcessorMixin):
             latlon=_to_tensor(latlon),
             months=months,
             normalize=normalize,
+            band_config=self.band_config,
         )
 
         if masked_output.space_time_x.dim() == 4:

galileo-tiny-patch8/config.json

@@ -43,5 +43,262 @@
         "AutoModel"
       ]
     }
+  },
+  "s1_bands": [
+    "VV",
+    "VH"
+  ],
+  "s2_bands": [
+    "B2",
+    "B3",
+    "B4",
+    "B5",
+    "B6",
+    "B7",
+    "B8",
+    "B8A",
+    "B11",
+    "B12"
+  ],
+  "era5_bands": [
+    "temperature_2m",
+    "total_precipitation_sum"
+  ],
+  "tc_bands": [
+    "def",
+    "soil",
+    "aet"
+  ],
+  "viirs_bands": [
+    "avg_rad"
+  ],
+  "srtm_bands": [
+    "elevation",
+    "slope"
+  ],
+  "dw_bands": [
+    "DW_water",
+    "DW_trees",
+    "DW_grass",
+    "DW_flooded_vegetation",
+    "DW_crops",
+    "DW_shrub_and_scrub",
+    "DW_built",
+    "DW_bare",
+    "DW_snow_and_ice"
+  ],
+  "wc_bands": [
+    "WC_temporarycrops",
+    "WC_maize",
+    "WC_wintercereals",
+    "WC_springcereals",
+    "WC_irrigation"
+  ],
+  "landscan_bands": [
+    "b1"
+  ],
+  "location_bands": [
+    "x",
+    "y",
+    "z"
+  ],
+  "space_time_band_groups": {
+    "S1": [
+      "VV",
+      "VH"
+    ],
+    "S2_RGB": [
+      "B2",
+      "B3",
+      "B4"
+    ],
+    "S2_Red_Edge": [
+      "B5",
+      "B6",
+      "B7"
+    ],
+    "S2_NIR_10m": [
+      "B8"
+    ],
+    "S2_NIR_20m": [
+      "B8A"
+    ],
+    "S2_SWIR": [
+      "B11",
+      "B12"
+    ],
+    "NDVI": [
+      "NDVI"
+    ]
+  },
+  "time_band_groups": {
+    "ERA5": [
+      "temperature_2m",
+      "total_precipitation_sum"
+    ],
+    "TC": [
+      "def",
+      "soil",
+      "aet"
+    ],
+    "VIIRS": [
+      "avg_rad"
+    ]
+  },
+  "space_band_groups": {
+    "SRTM": [
+      "elevation",
+      "slope"
+    ],
+    "DW": [
+      "DW_water",
+      "DW_trees",
+      "DW_grass",
+      "DW_flooded_vegetation",
+      "DW_crops",
+      "DW_shrub_and_scrub",
+      "DW_built",
+      "DW_bare",
+      "DW_snow_and_ice"
+    ],
+    "WC": [
+      "WC_temporarycrops",
+      "WC_maize",
+      "WC_wintercereals",
+      "WC_springcereals",
+      "WC_irrigation"
+    ]
+  },
+  "pretraining_normalizing_dict": {
+    "13": {
+      "mean": [
+        -11.728724389184965,
+        -18.85558188024017,
+        1395.3408730676722,
+        1338.4026921784578,
+        1343.09883810357,
+        1543.8607982512297,
+        2186.2022069512263,
+        2525.0932853316694,
+        2410.3377187373408,
+        2750.2854646886753,
+        2234.911100061487,
+        1474.5311266077113,
+        0.2892116502999044
+      ],
+      "std": [
+        4.887145774840316,
+        5.730270320384293,
+        917.7041440370853,
+        913.2988423581528,
+        1092.678723527555,
+        1047.2206083460424,
+        1048.0101611156767,
+        1143.6903026819996,
+        1098.979177731649,
+        1204.472755085893,
+        1145.9774063078878,
+        980.2429840007796,
+        0.2720939024500081
+      ]
+    },
+    "16": {
+      "mean": [
+        673.0152819503361,
+        5.930092668915115,
+        0.10470439140978786,
+        0.23965913270066183,
+        0.08158044385860364,
+        0.04246976254259546,
+        0.11304392863520317,
+        0.17329647890362473,
+        0.0698981691616277,
+        0.12130267132802142,
+        0.04671318615236216,
+        10.973119802517362,
+        1.0927069179958768,
+        1.6991394232855903,
+        0.03720594618055555,
+        1.3671352688259548
+      ],
+      "std": [
+        983.0697298296237,
+        8.167406789813247,
+        0.18771647977504985,
+        0.2368313455675914,
+        0.08024268534756586,
+        0.04045374496146404,
+        0.11350342472061795,
+        0.1279898111718168,
+        0.12042341550438586,
+        0.13602408145504347,
+        0.043971116096060345,
+        31.255340146970997,
+        10.395974878206689,
+        12.92380617159917,
+        1.9285254295940466,
+        11.612179775408928
+      ]
+    },
+    "6": {
+      "mean": [
+        271.5674963541667,
+        0.08554303677156568,
+        657.3181260091111,
+        692.1291795806885,
+        562.781331880633,
+        1.5647115934036673
+      ],
+      "std": [
+        79.80828940314429,
+        0.11669547098151486,
+        704.0008695557707,
+        925.0116126406431,
+        453.2434022278578,
+        7.513020170832818
+      ]
+    },
+    "18": {
+      "mean": [
+        188.20315880851746,
+        0.2804946561574936,
+        0.11371652073860168,
+        0.058778801321983334,
+        0.10474256777763366,
+        0.2396918488264084,
+        0.08152248692512512,
+        0.04248040814399719,
+        0.11303179881572724,
+        0.17326324067115784,
+        0.06998309404850006,
+        0.12122812910079957,
+        0.04671641788482666,
+        10.98456594619751,
+        1.0968475807189941,
+        1.6947754135131836,
+        0.03320046615600586,
+        1.3602827312469483
+      ],
+      "std": [
+        1154.5919128300602,
+        0.5276998078079327,
+        0.7021637331734328,
+        0.36528892213195063,
+        0.17470213191865785,
+        0.20411195416718833,
+        0.0660782470089761,
+        0.03380702424871257,
+        0.09809195568521663,
+        0.11292471052124119,
+        0.09720748930233268,
+        0.12912217763726777,
+        0.0399973913151906,
+        23.725471823867462,
+        5.715238079725388,
+        9.030481416228302,
+        0.9950220242487364,
+        7.754429123862099
+      ]
+    }
   }
-}
+}

galileo-tiny-patch8/modeling_galileo.py

@@ -26,95 +26,11 @@ from einops import rearrange, repeat
 from torch import Tensor, vmap
 from torch.jit import Final
 
-# constants
-CONFIG_FILENAME = "config.json"
-ENCODER_FILENAME = "encoder.pt"
-BASE_GSD = 10
-
-# band information
-S1_BANDS = ["VV", "VH"]
-S2_BANDS = [
-    "B2",
-    "B3",
-    "B4",
-    "B5",
-    "B6",
-    "B7",
-    "B8",
-    "B8A",
-    "B11",
-    "B12",
-]
-ERA5_BANDS = ["temperature_2m", "total_precipitation_sum"]
-TC_BANDS = ["def", "soil", "aet"]
-VIIRS_BANDS = ["avg_rad"]
-SRTM_BANDS = ["elevation", "slope"]
-DW_BANDS = [
-    "DW_water",
-    "DW_trees",
-    "DW_grass",
-    "DW_flooded_vegetation",
-    "DW_crops",
-    "DW_shrub_and_scrub",
-    "DW_built",
-    "DW_bare",
-    "DW_snow_and_ice",
-]
-WC_BANDS = [
-    "WC_temporarycrops",
-    "WC_maize",
-    "WC_wintercereals",
-    "WC_springcereals",
-    "WC_irrigation",
-]
-STATIC_DW_BANDS = [f"{x}_static" for x in DW_BANDS]
-STATIC_WC_BANDS = [f"{x}_static" for x in WC_BANDS]
-
-LANDSCAN_BANDS = ["b1"]
-LOCATION_BANDS = ["x", "y", "z"]
-
-SPACE_TIME_BANDS = S1_BANDS + S2_BANDS + ["NDVI"]
-TIME_BANDS = ERA5_BANDS + TC_BANDS + VIIRS_BANDS
-SPACE_BANDS = SRTM_BANDS + DW_BANDS + WC_BANDS
-STATIC_BANDS = LANDSCAN_BANDS + LOCATION_BANDS + STATIC_DW_BANDS + STATIC_WC_BANDS
-
-
-SPACE_TIME_BANDS_GROUPS_IDX: OrderedDictType[str, List[int]] = OrderedDict(
-    {
-        "S1": [SPACE_TIME_BANDS.index(b) for b in S1_BANDS],
-        "S2_RGB": [SPACE_TIME_BANDS.index(b) for b in ["B2", "B3", "B4"]],
-        "S2_Red_Edge": [SPACE_TIME_BANDS.index(b) for b in ["B5", "B6", "B7"]],
-        "S2_NIR_10m": [SPACE_TIME_BANDS.index(b) for b in ["B8"]],
-        "S2_NIR_20m": [SPACE_TIME_BANDS.index(b) for b in ["B8A"]],
-        "S2_SWIR": [SPACE_TIME_BANDS.index(b) for b in ["B11", "B12"]],
-        "NDVI": [SPACE_TIME_BANDS.index("NDVI")],
-    }
-)
-
-TIME_BAND_GROUPS_IDX: OrderedDictType[str, List[int]] = OrderedDict(
-    {
-        "ERA5": [TIME_BANDS.index(b) for b in ERA5_BANDS],
-        "TC": [TIME_BANDS.index(b) for b in TC_BANDS],
-        "VIIRS": [TIME_BANDS.index(b) for b in VIIRS_BANDS],
-    }
-)
 
-SPACE_BAND_GROUPS_IDX: OrderedDictType[str, List[int]] = OrderedDict(
-    {
-        "SRTM": [SPACE_BANDS.index(b) for b in SRTM_BANDS],
-        "DW": [SPACE_BANDS.index(b) for b in DW_BANDS],
-        "WC": [SPACE_BANDS.index(b) for b in WC_BANDS],
-    }
-)
-
-STATIC_BAND_GROUPS_IDX: OrderedDictType[str, List[int]] = OrderedDict(
-    {
-        "LS": [STATIC_BANDS.index(b) for b in LANDSCAN_BANDS],
-        "location": [STATIC_BANDS.index(b) for b in LOCATION_BANDS],
-        "DW_static": [STATIC_BANDS.index(b) for b in STATIC_DW_BANDS],
-        "WC_static": [STATIC_BANDS.index(b) for b in STATIC_WC_BANDS],
-    }
-)
+def _band_group_indices(
+    bands: Sequence[str], groups: Dict[str, List[str]]
+) -> OrderedDictType[str, List[int]]:
+    return OrderedDict((name, [bands.index(b) for b in group_bands]) for name, group_bands in groups.items())
 
 
 def get_2d_sincos_pos_embed_with_resolution(
@@ -565,13 +481,16 @@ class GalileoBase(nn.Module):
         base_patch_size: int = 4,
         use_channel_embs: bool = True,
         drop_path: float = 0.0,
+        band_layout: Optional[Dict[str, Any]] = None,
     ):
         super().__init__()
 
-        self.space_time_groups = SPACE_TIME_BANDS_GROUPS_IDX
-        self.space_groups = SPACE_BAND_GROUPS_IDX
-        self.time_groups = TIME_BAND_GROUPS_IDX
-        self.static_groups = STATIC_BAND_GROUPS_IDX
+        band_layout = band_layout or GalileoConfig().band_layout()
+        self.space_time_groups = band_layout["space_time_groups"]
+        self.space_groups = band_layout["space_groups"]
+        self.time_groups = band_layout["time_groups"]
+        self.static_groups = band_layout["static_groups"]
+        self.base_gsd = band_layout["input_resolution_m"]
         self.embedding_size = embedding_size
         self.base_patch_size = base_patch_size
 
@@ -606,16 +525,16 @@ class GalileoBase(nn.Module):
         else:
             args = {"requires_grad": False}
         self.s_t_channel_embed = nn.Parameter(
-            torch.zeros(len(SPACE_TIME_BANDS_GROUPS_IDX), int(embedding_size * 0.25)), **args
+            torch.zeros(len(self.space_time_groups), int(embedding_size * 0.25)), **args
         )
         self.sp_channel_embed = nn.Parameter(
-            torch.zeros(len(SPACE_BAND_GROUPS_IDX), int(embedding_size * 0.25)), **args
+            torch.zeros(len(self.space_groups), int(embedding_size * 0.25)), **args
         )
         self.t_channel_embed = nn.Parameter(
-            torch.zeros(len(TIME_BAND_GROUPS_IDX), int(embedding_size * 0.25)), **args
+            torch.zeros(len(self.time_groups), int(embedding_size * 0.25)), **args
         )
         self.st_channel_embed = nn.Parameter(
-            torch.zeros(len(STATIC_BAND_GROUPS_IDX), int(embedding_size * 0.25)), **args
+            torch.zeros(len(self.static_groups), int(embedding_size * 0.25)), **args
         )
 
         self.apply(self._init_weights)
@@ -720,7 +639,7 @@ class GalileoBase(nn.Module):
         if patch_size is None:
             patch_size = self.base_patch_size
         token_res = input_res * patch_size
-        gsd_ratio = token_res / BASE_GSD
+        gsd_ratio = token_res / self.base_gsd
 
         assert h == w, "get_2d_sincos_pos_embed_with_resolution currently requires that h==w"
         spatial_embed = get_2d_sincos_pos_embed_with_resolution(
@@ -755,6 +674,7 @@ class Encoder(GalileoBase):
         max_sequence_length=24,
         freeze_projections: bool = False,
         drop_path: float = 0.0,
+        band_layout: Optional[Dict[str, Any]] = None,
     ):
         super().__init__(
             embedding_size,
@@ -765,6 +685,7 @@ class Encoder(GalileoBase):
             max_patch_size,
             use_channel_embs=True,
             drop_path=drop_path,
+            band_layout=band_layout,
         )
 
         self.space_time_embed = nn.ModuleDict(
@@ -1108,11 +1029,14 @@ class Encoder(GalileoBase):
         st_m: torch.Tensor,
         months: torch.Tensor,
         patch_size: int,
-        input_resolution_m: Optional[int] = BASE_GSD,
+        input_resolution_m: Optional[int] = None,
         exit_after: Optional[int] = None,
         token_exit_cfg: Optional[Dict] = None,
         add_layernorm_on_exit: bool = True,
     ):
+        if input_resolution_m is None:
+            input_resolution_m = self.base_gsd
+
         (
             s_t_x,
             sp_x,
@@ -1161,25 +1085,31 @@ class Encoder(GalileoBase):
         )
 
     @classmethod
-    def load_from_folder(cls, folder: Path, device: torch.device):
-        if not (folder / CONFIG_FILENAME).exists():
+    def load_from_folder(
+        cls,
+        folder: Path,
+        device: torch.device,
+        config_filename: str = "config.json",
+        encoder_filename: str = "encoder.pt",
+    ):
+        if not (folder / config_filename).exists():
             all_files_in_folder = [f.name for f in folder.glob("*")]
             raise ValueError(
-                f"Expected {CONFIG_FILENAME} in {folder}, found {all_files_in_folder}"
+                f"Expected {config_filename} in {folder}, found {all_files_in_folder}"
             )
-        if not (folder / ENCODER_FILENAME).exists():
+        if not (folder / encoder_filename).exists():
             all_files_in_folder = [f.name for f in folder.glob("*")]
             raise ValueError(
-                f"Expected {ENCODER_FILENAME} in {folder}, found {all_files_in_folder}"
+                f"Expected {encoder_filename} in {folder}, found {all_files_in_folder}"
             )
 
-        with (folder / CONFIG_FILENAME).open("r") as f:
+        with (folder / config_filename).open("r") as f:
             config = json.load(f)
             model_config = config["model"]
             encoder_config = model_config["encoder"]
         encoder = cls(**encoder_config)
 
-        state_dict = torch.load(folder / ENCODER_FILENAME, map_location=device)
+        state_dict = torch.load(folder / encoder_filename, map_location=device)
         for key in list(state_dict.keys()):
             # this cleans the state dict, which occasionally had an extra
             # ".backbone" included in the key names
@@ -1189,6 +1119,141 @@ class Encoder(GalileoBase):
 logger = logging.get_logger(__name__)
 
 
+def _default_pretraining_normalizing_dict() -> Dict[str, Dict[str, List[float]]]:
+    return {
+        "13": {
+            "mean": [
+                -11.728724389184965,
+                -18.85558188024017,
+                1395.3408730676722,
+                1338.4026921784578,
+                1343.09883810357,
+                1543.8607982512297,
+                2186.2022069512263,
+                2525.0932853316694,
+                2410.3377187373408,
+                2750.2854646886753,
+                2234.911100061487,
+                1474.5311266077113,
+                0.2892116502999044,
+            ],
+            "std": [
+                4.887145774840316,
+                5.730270320384293,
+                917.7041440370853,
+                913.2988423581528,
+                1092.678723527555,
+                1047.2206083460424,
+                1048.0101611156767,
+                1143.6903026819996,
+                1098.979177731649,
+                1204.472755085893,
+                1145.9774063078878,
+                980.2429840007796,
+                0.2720939024500081,
+            ],
+        },
+        "16": {
+            "mean": [
+                673.0152819503361,
+                5.930092668915115,
+                0.10470439140978786,
+                0.23965913270066183,
+                0.08158044385860364,
+                0.04246976254259546,
+                0.11304392863520317,
+                0.17329647890362473,
+                0.0698981691616277,
+                0.12130267132802142,
+                0.04671318615236216,
+                10.973119802517362,
+                1.0927069179958768,
+                1.6991394232855903,
+                0.03720594618055555,
+                1.3671352688259548,
+            ],
+            "std": [
+                983.0697298296237,
+                8.167406789813247,
+                0.18771647977504985,
+                0.2368313455675914,
+                0.08024268534756586,
+                0.04045374496146404,
+                0.11350342472061795,
+                0.1279898111718168,
+                0.12042341550438586,
+                0.13602408145504347,
+                0.043971116096060345,
+                31.255340146970997,
+                10.395974878206689,
+                12.92380617159917,
+                1.9285254295940466,
+                11.612179775408928,
+            ],
+        },
+        "6": {
+            "mean": [
+                271.5674963541667,
+                0.08554303677156568,
+                657.3181260091111,
+                692.1291795806885,
+                562.781331880633,
+                1.5647115934036673,
+            ],
+            "std": [
+                79.80828940314429,
+                0.11669547098151486,
+                704.0008695557707,
+                925.0116126406431,
+                453.2434022278578,
+                7.513020170832818,
+            ],
+        },
+        "18": {
+            "mean": [
+                188.20315880851746,
+                0.2804946561574936,
+                0.11371652073860168,
+                0.058778801321983334,
+                0.10474256777763366,
+                0.2396918488264084,
+                0.08152248692512512,
+                0.04248040814399719,
+                0.11303179881572724,
+                0.17326324067115784,
+                0.06998309404850006,
+                0.12122812910079957,
+                0.04671641788482666,
+                10.98456594619751,
+                1.0968475807189941,
+                1.6947754135131836,
+                0.03320046615600586,
+                1.3602827312469483,
+            ],
+            "std": [
+                1154.5919128300602,
+                0.5276998078079327,
+                0.7021637331734328,
+                0.36528892213195063,
+                0.17470213191865785,
+                0.20411195416718833,
+                0.0660782470089761,
+                0.03380702424871257,
+                0.09809195568521663,
+                0.11292471052124119,
+                0.09720748930233268,
+                0.12912217763726777,
+                0.0399973913151906,
+                23.725471823867462,
+                5.715238079725388,
+                9.030481416228302,
+                0.9950220242487364,
+                7.754429123862099,
+            ],
+        },
+    }
+
+
 class GalileoConfig(PreTrainedConfig):
     model_type = "galileo"
 
@@ -1206,6 +1271,20 @@ class GalileoConfig(PreTrainedConfig):
         default_month: int = 6,
         global_pool: bool = True,
         input_resolution_m: int = 10,
+        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)
@@ -1221,6 +1300,89 @@ class GalileoConfig(PreTrainedConfig):
         self.default_month = default_month
         self.global_pool = global_pool
         self.input_resolution_m = input_resolution_m
+        self.s1_bands = s1_bands if s1_bands is not None else ["VV", "VH"]
+        self.s2_bands = s2_bands if s2_bands is not None else [
+            "B2", "B3", "B4", "B5", "B6", "B7", "B8", "B8A", "B11", "B12",
+        ]
+        self.era5_bands = era5_bands if era5_bands is not None else [
+            "temperature_2m", "total_precipitation_sum",
+        ]
+        self.tc_bands = tc_bands if tc_bands is not None else ["def", "soil", "aet"]
+        self.viirs_bands = viirs_bands if viirs_bands is not None else ["avg_rad"]
+        self.srtm_bands = srtm_bands if srtm_bands is not None else ["elevation", "slope"]
+        self.dw_bands = dw_bands if dw_bands is not None else [
+            "DW_water", "DW_trees", "DW_grass", "DW_flooded_vegetation", "DW_crops",
+            "DW_shrub_and_scrub", "DW_built", "DW_bare", "DW_snow_and_ice",
+        ]
+        self.wc_bands = wc_bands if wc_bands is not None else [
+            "WC_temporarycrops", "WC_maize", "WC_wintercereals", "WC_springcereals", "WC_irrigation",
+        ]
+        self.landscan_bands = landscan_bands if landscan_bands is not None else ["b1"]
+        self.location_bands = location_bands if location_bands is not None else ["x", "y", "z"]
+        self.space_time_band_groups = space_time_band_groups if space_time_band_groups is not None else {
+            "S1": ["VV", "VH"],
+            "S2_RGB": ["B2", "B3", "B4"],
+            "S2_Red_Edge": ["B5", "B6", "B7"],
+            "S2_NIR_10m": ["B8"],
+            "S2_NIR_20m": ["B8A"],
+            "S2_SWIR": ["B11", "B12"],
+            "NDVI": ["NDVI"],
+        }
+        self.time_band_groups = time_band_groups if time_band_groups is not None else {
+            "ERA5": ["temperature_2m", "total_precipitation_sum"],
+            "TC": ["def", "soil", "aet"],
+            "VIIRS": ["avg_rad"],
+        }
+        self.space_band_groups = space_band_groups if space_band_groups is not None else {
+            "SRTM": ["elevation", "slope"],
+            "DW": [
+                "DW_water", "DW_trees", "DW_grass", "DW_flooded_vegetation", "DW_crops",
+                "DW_shrub_and_scrub", "DW_built", "DW_bare", "DW_snow_and_ice",
+            ],
+            "WC": [
+                "WC_temporarycrops", "WC_maize", "WC_wintercereals", "WC_springcereals", "WC_irrigation",
+            ],
+        }
+        self.pretraining_normalizing_dict = (
+            pretraining_normalizing_dict
+            if pretraining_normalizing_dict is not None
+            else _default_pretraining_normalizing_dict()
+        )
+
+    def band_layout(self) -> Dict[str, Any]:
+        space_time_bands = self.s1_bands + self.s2_bands + ["NDVI"]
+        time_bands = self.era5_bands + self.tc_bands + self.viirs_bands
+        space_bands = self.srtm_bands + self.dw_bands + self.wc_bands
+        static_dw_bands = [f"{band}_static" for band in self.dw_bands]
+        static_wc_bands = [f"{band}_static" for band in self.wc_bands]
+        static_bands = self.landscan_bands + self.location_bands + static_dw_bands + static_wc_bands
+        static_band_groups = {
+            "LS": self.landscan_bands,
+            "location": self.location_bands,
+            "DW_static": static_dw_bands,
+            "WC_static": static_wc_bands,
+        }
+        return {
+            "s1_bands": self.s1_bands,
+            "s2_bands": self.s2_bands,
+            "era5_bands": self.era5_bands,
+            "tc_bands": self.tc_bands,
+            "viirs_bands": self.viirs_bands,
+            "srtm_bands": self.srtm_bands,
+            "dw_bands": self.dw_bands,
+            "wc_bands": self.wc_bands,
+            "landscan_bands": self.landscan_bands,
+            "location_bands": self.location_bands,
+            "space_time_bands": space_time_bands,
+            "time_bands": time_bands,
+            "space_bands": space_bands,
+            "static_bands": static_bands,
+            "space_time_groups": _band_group_indices(space_time_bands, self.space_time_band_groups),
+            "time_groups": _band_group_indices(time_bands, self.time_band_groups),
+            "space_groups": _band_group_indices(space_bands, self.space_band_groups),
+            "static_groups": _band_group_indices(static_bands, static_band_groups),
+            "input_resolution_m": self.input_resolution_m,
+        }
 
 
 class GalileoPreTrainedModel(PreTrainedModel):
@@ -1245,6 +1407,7 @@ class GalileoEncoderModel(GalileoPreTrainedModel):
             max_sequence_length=config.max_sequence_length,
             freeze_projections=config.freeze_projections,
             drop_path=config.drop_path,
+            band_layout=config.band_layout(),
         )
         self.add_pooling_layer = add_pooling_layer and config.global_pool
         self.post_init()

galileo-tiny-patch8/preprocessor_config.json

@@ -5,5 +5,262 @@
   "patch_size": 8,
   "auto_map": {
     "AutoProcessor": "processing_galileo.GalileoProcessor"
+  },
+  "s1_bands": [
+    "VV",
+    "VH"
+  ],
+  "s2_bands": [
+    "B2",
+    "B3",
+    "B4",
+    "B5",
+    "B6",
+    "B7",
+    "B8",
+    "B8A",
+    "B11",
+    "B12"
+  ],
+  "era5_bands": [
+    "temperature_2m",
+    "total_precipitation_sum"
+  ],
+  "tc_bands": [
+    "def",
+    "soil",
+    "aet"
+  ],
+  "viirs_bands": [
+    "avg_rad"
+  ],
+  "srtm_bands": [
+    "elevation",
+    "slope"
+  ],
+  "dw_bands": [
+    "DW_water",
+    "DW_trees",
+    "DW_grass",
+    "DW_flooded_vegetation",
+    "DW_crops",
+    "DW_shrub_and_scrub",
+    "DW_built",
+    "DW_bare",
+    "DW_snow_and_ice"
+  ],
+  "wc_bands": [
+    "WC_temporarycrops",
+    "WC_maize",
+    "WC_wintercereals",
+    "WC_springcereals",
+    "WC_irrigation"
+  ],
+  "landscan_bands": [
+    "b1"
+  ],
+  "location_bands": [
+    "x",
+    "y",
+    "z"
+  ],
+  "space_time_band_groups": {
+    "S1": [
+      "VV",
+      "VH"
+    ],
+    "S2_RGB": [
+      "B2",
+      "B3",
+      "B4"
+    ],
+    "S2_Red_Edge": [
+      "B5",
+      "B6",
+      "B7"
+    ],
+    "S2_NIR_10m": [
+      "B8"
+    ],
+    "S2_NIR_20m": [
+      "B8A"
+    ],
+    "S2_SWIR": [
+      "B11",
+      "B12"
+    ],
+    "NDVI": [
+      "NDVI"
+    ]
+  },
+  "time_band_groups": {
+    "ERA5": [
+      "temperature_2m",
+      "total_precipitation_sum"
+    ],
+    "TC": [
+      "def",
+      "soil",
+      "aet"
+    ],
+    "VIIRS": [
+      "avg_rad"
+    ]
+  },
+  "space_band_groups": {
+    "SRTM": [
+      "elevation",
+      "slope"
+    ],
+    "DW": [
+      "DW_water",
+      "DW_trees",
+      "DW_grass",
+      "DW_flooded_vegetation",
+      "DW_crops",
+      "DW_shrub_and_scrub",
+      "DW_built",
+      "DW_bare",
+      "DW_snow_and_ice"
+    ],
+    "WC": [
+      "WC_temporarycrops",
+      "WC_maize",
+      "WC_wintercereals",
+      "WC_springcereals",
+      "WC_irrigation"
+    ]
+  },
+  "pretraining_normalizing_dict": {
+    "13": {
+      "mean": [
+        -11.728724389184965,
+        -18.85558188024017,
+        1395.3408730676722,
+        1338.4026921784578,
+        1343.09883810357,
+        1543.8607982512297,
+        2186.2022069512263,
+        2525.0932853316694,
+        2410.3377187373408,
+        2750.2854646886753,
+        2234.911100061487,
+        1474.5311266077113,
+        0.2892116502999044
+      ],
+      "std": [
+        4.887145774840316,
+        5.730270320384293,
+        917.7041440370853,
+        913.2988423581528,
+        1092.678723527555,
+        1047.2206083460424,
+        1048.0101611156767,
+        1143.6903026819996,
+        1098.979177731649,
+        1204.472755085893,
+        1145.9774063078878,
+        980.2429840007796,
+        0.2720939024500081
+      ]
+    },
+    "16": {
+      "mean": [
+        673.0152819503361,
+        5.930092668915115,
+        0.10470439140978786,
+        0.23965913270066183,
+        0.08158044385860364,
+        0.04246976254259546,
+        0.11304392863520317,
+        0.17329647890362473,
+        0.0698981691616277,
+        0.12130267132802142,
+        0.04671318615236216,
+        10.973119802517362,
+        1.0927069179958768,
+        1.6991394232855903,
+        0.03720594618055555,
+        1.3671352688259548
+      ],
+      "std": [
+        983.0697298296237,
+        8.167406789813247,
+        0.18771647977504985,
+        0.2368313455675914,
+        0.08024268534756586,
+        0.04045374496146404,
+        0.11350342472061795,
+        0.1279898111718168,
+        0.12042341550438586,
+        0.13602408145504347,
+        0.043971116096060345,
+        31.255340146970997,
+        10.395974878206689,
+        12.92380617159917,
+        1.9285254295940466,
+        11.612179775408928
+      ]
+    },
+    "6": {
+      "mean": [
+        271.5674963541667,
+        0.08554303677156568,
+        657.3181260091111,
+        692.1291795806885,
+        562.781331880633,
+        1.5647115934036673
+      ],
+      "std": [
+        79.80828940314429,
+        0.11669547098151486,
+        704.0008695557707,
+        925.0116126406431,
+        453.2434022278578,
+        7.513020170832818
+      ]
+    },
+    "18": {
+      "mean": [
+        188.20315880851746,
+        0.2804946561574936,
+        0.11371652073860168,
+        0.058778801321983334,
+        0.10474256777763366,
+        0.2396918488264084,
+        0.08152248692512512,
+        0.04248040814399719,
+        0.11303179881572724,
+        0.17326324067115784,
+        0.06998309404850006,
+        0.12122812910079957,
+        0.04671641788482666,
+        10.98456594619751,
+        1.0968475807189941,
+        1.6947754135131836,
+        0.03320046615600586,
+        1.3602827312469483
+      ],
+      "std": [
+        1154.5919128300602,
+        0.5276998078079327,
+        0.7021637331734328,
+        0.36528892213195063,
+        0.17470213191865785,
+        0.20411195416718833,
+        0.0660782470089761,
+        0.03380702424871257,
+        0.09809195568521663,
+        0.11292471052124119,
+        0.09720748930233268,
+        0.12912217763726777,
+        0.0399973913151906,
+        23.725471823867462,
+        5.715238079725388,
+        9.030481416228302,
+        0.9950220242487364,
+        7.754429123862099
+      ]
+    }
   }
-}
+}

galileo-tiny-patch8/processing_galileo.py

@@ -4,7 +4,6 @@
 from __future__ import annotations
 
 import math
-from collections import OrderedDict
 from typing import Any, NamedTuple, Optional, Union
 
 import numpy as np
@@ -14,191 +13,7 @@ from transformers.feature_extraction_utils import BatchFeature
 from transformers.processing_utils import ProcessorMixin
 from transformers.utils import TensorType
 
-
-S1_BANDS = ["VV", "VH"]
-S2_BANDS = ["B2", "B3", "B4", "B5", "B6", "B7", "B8", "B8A", "B11", "B12"]
-ERA5_BANDS = ["temperature_2m", "total_precipitation_sum"]
-TC_BANDS = ["def", "soil", "aet"]
-VIIRS_BANDS = ["avg_rad"]
-SRTM_BANDS = ["elevation", "slope"]
-DW_BANDS = [
-    "DW_water", "DW_trees", "DW_grass", "DW_flooded_vegetation", "DW_crops",
-    "DW_shrub_and_scrub", "DW_built", "DW_bare", "DW_snow_and_ice",
-]
-WC_BANDS = [
-    "WC_temporarycrops", "WC_maize", "WC_wintercereals", "WC_springcereals", "WC_irrigation",
-]
-LANDSCAN_BANDS = ["b1"]
-LOCATION_BANDS = ["x", "y", "z"]
-STATIC_DW_BANDS = [f"{x}_static" for x in DW_BANDS]
-STATIC_WC_BANDS = [f"{x}_static" for x in WC_BANDS]
-
-SPACE_TIME_BANDS = S1_BANDS + S2_BANDS + ["NDVI"]
-TIME_BANDS = ERA5_BANDS + TC_BANDS + VIIRS_BANDS
-SPACE_BANDS = SRTM_BANDS + DW_BANDS + WC_BANDS
-STATIC_BANDS = LANDSCAN_BANDS + LOCATION_BANDS + STATIC_DW_BANDS + STATIC_WC_BANDS
-
-SPACE_TIME_BANDS_GROUPS_IDX = OrderedDict({
-    "S1": [SPACE_TIME_BANDS.index(b) for b in S1_BANDS],
-    "S2_RGB": [SPACE_TIME_BANDS.index(b) for b in ["B2", "B3", "B4"]],
-    "S2_Red_Edge": [SPACE_TIME_BANDS.index(b) for b in ["B5", "B6", "B7"]],
-    "S2_NIR_10m": [SPACE_TIME_BANDS.index(b) for b in ["B8"]],
-    "S2_NIR_20m": [SPACE_TIME_BANDS.index(b) for b in ["B8A"]],
-    "S2_SWIR": [SPACE_TIME_BANDS.index(b) for b in ["B11", "B12"]],
-    "NDVI": [SPACE_TIME_BANDS.index("NDVI")],
-})
-TIME_BAND_GROUPS_IDX = OrderedDict({
-    "ERA5": [TIME_BANDS.index(b) for b in ERA5_BANDS],
-    "TC": [TIME_BANDS.index(b) for b in TC_BANDS],
-    "VIIRS": [TIME_BANDS.index(b) for b in VIIRS_BANDS],
-})
-SPACE_BAND_GROUPS_IDX = OrderedDict({
-    "SRTM": [SPACE_BANDS.index(b) for b in SRTM_BANDS],
-    "DW": [SPACE_BANDS.index(b) for b in DW_BANDS],
-    "WC": [SPACE_BANDS.index(b) for b in WC_BANDS],
-})
-STATIC_BAND_GROUPS_IDX = OrderedDict({
-    "LS": [STATIC_BANDS.index(b) for b in LANDSCAN_BANDS],
-    "location": [STATIC_BANDS.index(b) for b in LOCATION_BANDS],
-    "DW_static": [STATIC_BANDS.index(b) for b in STATIC_DW_BANDS],
-    "WC_static": [STATIC_BANDS.index(b) for b in STATIC_WC_BANDS],
-})
-
-
-DEFAULT_MONTH = 5
-
-PRETRAINING_NORMALIZING_DICT = {
-    "13": {
-        "mean": [
-            -11.728724389184965,
-            -18.85558188024017,
-            1395.3408730676722,
-            1338.4026921784578,
-            1343.09883810357,
-            1543.8607982512297,
-            2186.2022069512263,
-            2525.0932853316694,
-            2410.3377187373408,
-            2750.2854646886753,
-            2234.911100061487,
-            1474.5311266077113,
-            0.2892116502999044,
-        ],
-        "std": [
-            4.887145774840316,
-            5.730270320384293,
-            917.7041440370853,
-            913.2988423581528,
-            1092.678723527555,
-            1047.2206083460424,
-            1048.0101611156767,
-            1143.6903026819996,
-            1098.979177731649,
-            1204.472755085893,
-            1145.9774063078878,
-            980.2429840007796,
-            0.2720939024500081,
-        ],
-    },
-    "16": {
-        "mean": [
-            673.0152819503361,
-            5.930092668915115,
-            0.10470439140978786,
-            0.23965913270066183,
-            0.08158044385860364,
-            0.04246976254259546,
-            0.11304392863520317,
-            0.17329647890362473,
-            0.0698981691616277,
-            0.12130267132802142,
-            0.04671318615236216,
-            10.973119802517362,
-            1.0927069179958768,
-            1.6991394232855903,
-            0.03720594618055555,
-            1.3671352688259548,
-        ],
-        "std": [
-            983.0697298296237,
-            8.167406789813247,
-            0.18771647977504985,
-            0.2368313455675914,
-            0.08024268534756586,
-            0.04045374496146404,
-            0.11350342472061795,
-            0.1279898111718168,
-            0.12042341550438586,
-            0.13602408145504347,
-            0.043971116096060345,
-            31.255340146970997,
-            10.395974878206689,
-            12.92380617159917,
-            1.9285254295940466,
-            11.612179775408928,
-        ],
-    },
-    "6": {
-        "mean": [
-            271.5674963541667,
-            0.08554303677156568,
-            657.3181260091111,
-            692.1291795806885,
-            562.781331880633,
-            1.5647115934036673,
-        ],
-        "std": [
-            79.80828940314429,
-            0.11669547098151486,
-            704.0008695557707,
-            925.0116126406431,
-            453.2434022278578,
-            7.513020170832818,
-        ],
-    },
-    "18": {
-        "mean": [
-            188.20315880851746,
-            0.2804946561574936,
-            0.11371652073860168,
-            0.058778801321983334,
-            0.10474256777763366,
-            0.2396918488264084,
-            0.08152248692512512,
-            0.04248040814399719,
-            0.11303179881572724,
-            0.17326324067115784,
-            0.06998309404850006,
-            0.12122812910079957,
-            0.04671641788482666,
-            10.98456594619751,
-            1.0968475807189941,
-            1.6947754135131836,
-            0.03320046615600586,
-            1.3602827312469483,
-        ],
-        "std": [
-            1154.5919128300602,
-            0.5276998078079327,
-            0.7021637331734328,
-            0.36528892213195063,
-            0.17470213191865785,
-            0.20411195416718833,
-            0.0660782470089761,
-            0.03380702424871257,
-            0.09809195568521663,
-            0.11292471052124119,
-            0.09720748930233268,
-            0.12912217763726777,
-            0.0399973913151906,
-            23.725471823867462,
-            5.715238079725388,
-            9.030481416228302,
-            0.9950220242487364,
-            7.754429123862099,
-        ],
-    },
-}
+from .modeling_galileo import GalileoConfig
 
 
 class MaskedOutput(NamedTuple):
@@ -257,7 +72,19 @@ def construct_galileo_input(
     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]
@@ -285,56 +112,56 @@ def construct_galileo_input(
     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_BANDS_GROUPS_IDX)), 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_BAND_GROUPS_IDX)), 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_BAND_GROUPS_IDX)), dtype=torch.float, device=device)
-    st_x = torch.zeros((len(STATIC_BANDS)), dtype=torch.float, device=device)
-    st_m = torch.ones((len(STATIC_BAND_GROUPS_IDX)), dtype=torch.float, device=device)
+    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], [S1_BANDS, S2_BANDS], ["S1", "S2"]):
+    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_BANDS_GROUPS_IDX) if group_key in key]
+            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], [SRTM_BANDS, DW_BANDS, WC_BANDS], ["SRTM", "DW", "WC"]
+        [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_BAND_GROUPS_IDX) if group_key in key]
+            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], [ERA5_BANDS, TC_BANDS, VIIRS_BANDS], ["ERA5", "TC", "VIIRS"]
+        [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_BAND_GROUPS_IDX) if group_key in key]
+            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], [LANDSCAN_BANDS, LOCATION_BANDS], ["LS", "location"]
+        [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_BAND_GROUPS_IDX) if group_key in key]
+            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) * DEFAULT_MONTH
+        months = torch.ones((t), dtype=torch.long, device=device) * band_config.default_month
 
     if normalize:
-        normalizer = PretrainingNormalizer(PRETRAINING_NORMALIZING_DICT)
+        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)
@@ -367,11 +194,48 @@ class GalileoProcessor(ProcessorMixin):
         "months",
     ]
 
-    def __init__(self, normalize: bool = True, default_month: int = 6, patch_size: int = 8, **kwargs):
+    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,
@@ -418,6 +282,7 @@ class GalileoProcessor(ProcessorMixin):
             latlon=_to_tensor(latlon),
             months=months,
             normalize=normalize,
+            band_config=self.band_config,
         )
 
         if masked_output.space_time_x.dim() == 4: