From 38fbbd457721f64a15849b32f3203274f95c253d Mon Sep 17 00:00:00 2001 From: GabrieleTi Date: Wed, 20 May 2026 09:13:07 +0200 Subject: [PATCH 1/6] Fix cnn and add SSL4EO encoder configurations --- src/models/components/geo_encoders/cnn_encoder.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/models/components/geo_encoders/cnn_encoder.py b/src/models/components/geo_encoders/cnn_encoder.py index 4ece11d..a61b218 100644 --- a/src/models/components/geo_encoders/cnn_encoder.py +++ b/src/models/components/geo_encoders/cnn_encoder.py @@ -23,17 +23,19 @@ def __init__( backbone="resnet", pretrained_cnn="imagenet", resnet_version=18, - freezing_strategy="all", geo_data_name="s2", input_n_bands: int | None = None, - output_dim=512, ) -> None: super().__init__() + # Backbone configurations self.backbone = backbone - self.pretrained_cnn = pretrained_cnn - self.resnet_version = resnet_version - self.freezing_strategy = freezing_strategy + if self.backbone == "resnet": + assert resnet_version in [18, 34, 50], f"Unsupported resnet version: {resnet_version}" + self.resnet_version = resnet_version + + assert pretrained_cnn in ["imagenet", "IMAGENET1K_V1", 'SSL4EO_RGB_MOCO', None], f"Unsupported pretrained_cnn: {pretrained_cnn}" + self.pretrained_cnn = pretrained_cnn self.allowed_geo_data_names = ["s2", "aef", "tessera"] assert geo_data_name in self.allowed_geo_data_names From 769a4013bfcb6d2b41b72ab64c13eea202317980 Mon Sep 17 00:00:00 2001 From: GabrieleTi Date: Wed, 20 May 2026 09:13:24 +0200 Subject: [PATCH 2/6] Fix cnn and add SSL4EO encoder configurations --- .../components/geo_encoders/cnn_encoder.py | 133 +++++++++--------- 1 file changed, 64 insertions(+), 69 deletions(-) diff --git a/src/models/components/geo_encoders/cnn_encoder.py b/src/models/components/geo_encoders/cnn_encoder.py index a61b218..918b755 100644 --- a/src/models/components/geo_encoders/cnn_encoder.py +++ b/src/models/components/geo_encoders/cnn_encoder.py @@ -3,9 +3,16 @@ import torch import torchvision.models as models from torch import nn +from torchgeo.models import resnet50, ResNet50_Weights, ResNet18_Weights, resnet18 from src.models.components.geo_encoders.base_geo_encoder import BaseGeoEncoder +from src.utils.errors import IllegalArgumentCombination +RN_DIM = { + 18 : 512, + 34: 512, + 50: 2048 +} class CNNEncoder(BaseGeoEncoder): """Convolutional neural network EO encoder. Adapted from PECL. @@ -37,15 +44,15 @@ def __init__( assert pretrained_cnn in ["imagenet", "IMAGENET1K_V1", 'SSL4EO_RGB_MOCO', None], f"Unsupported pretrained_cnn: {pretrained_cnn}" self.pretrained_cnn = pretrained_cnn + self.output_dim = RN_DIM[resnet_version] + + # Input modality configurations self.allowed_geo_data_names = ["s2", "aef", "tessera"] assert geo_data_name in self.allowed_geo_data_names self.geo_data_name = geo_data_name self.set_n_input_bands(input_n_bands) - assert ( - self.input_n_bands >= 3 and type(self.input_n_bands) is int - ), f"input_n_bands must be int >=3, got {self.input_n_bands}" - self.output_dim = output_dim + assert (self.input_n_bands >= 3 and type(self.input_n_bands) is int), f"input_n_bands must be int >=3, got {self.input_n_bands}" def set_n_input_bands(self, n_bands: int | None = None) -> None: """Sets number of input bands based on geo_data_name if n_bands is None. @@ -67,76 +74,67 @@ def set_n_input_bands(self, n_bands: int | None = None) -> None: self.input_n_bands = n_bands return None - def get_backbone(self): + @override + def _setup(self) -> List[str]: """Gets backbone model given configuration stored in self. - :return: backbone model """ + trainable_modules = [] if self.backbone == "resnet": - assert self.resnet_version in [ - 18, - 34, - 50, - ], f"Unsupported resnet version: {self.resnet_version}" - assert self.pretrained_cnn in [ - "imagenet", - "IMAGENET1K_V1", - None, - ], f"Unsupported pretrained_cnn: {self.pretrained_cnn}" - if self.pretrained_cnn == "imagenet": - self.pretrained_cnn = "IMAGENET1K_V1" - if self.resnet_version == 18: - model = models.resnet18(weights=self.pretrained_cnn) - elif self.resnet_version == 34: - model = models.resnet34(weights=self.pretrained_cnn) - elif self.resnet_version == 50: - model = models.resnet50(weights=self.pretrained_cnn) + # Weights + # SSL4EO + if self.pretrained_cnn == "SSL4EO_RGB_MOCO": + if self.resnet_version == 18: + self.geo_encoder = resnet18(weights=ResNet18_Weights.SENTINEL2_RGB_MOCO) + elif self.resnet_version == 34: + raise IllegalArgumentCombination('SSL4EO_RGB_MOCO weights are not available for RN-34') + else: + self.geo_encoder = resnet50(weights=ResNet50_Weights.SENTINEL2_RGB_MOCO) + # Imagenet else: - raise ValueError(f"Unsupported resnet version: {self.resnet_version}") + if self.pretrained_cnn == "imagenet": + self.pretrained_cnn = "IMAGENET1K_V1" + elif self.pretrained_cnn == "imagenet_v2": + self.pretrained_cnn = "IMAGENET1K_V2" + + if self.resnet_version == 18: + self.geo_encoder = models.resnet18(weights=self.pretrained_cnn) + elif self.resnet_version == 34: + self.geo_encoder = models.resnet34(weights=self.pretrained_cnn) + else: + self.geo_encoder = models.resnet50(weights=self.pretrained_cnn) # Modify the first conv layer to accept input_n_bands channels - if self.pretrained_cnn is not None and self.input_n_bands != 3: - weight = model.conv1.weight.clone() if self.input_n_bands != 3: - model.conv1 = torch.nn.Conv2d( + + # Copy pre-trained weights + if self.pretrained_cnn is not None: + weight = self.geo_encoder.conv1.weight.clone() + + # Replace 1st conv layer + self.geo_encoder.conv1 = torch.nn.Conv2d( self.input_n_bands, 64, kernel_size=7, stride=2, padding=3, bias=False ) - if self.pretrained_cnn is not None: # copy pre-trained RGB bands - for i in range(self.input_n_bands): - model.conv1.weight.data[:, i, :, :] = weight[ - :, i % 3, :, : - ] # ensure this is not frozen - model.fc = nn.Linear(model.fc.in_features, self.output_dim) - - assert self.freezing_strategy in [ - "all", - "none", - ], f"Unsupported freezing_strategy: {self.freezing_strategy}" - layers_resnet = list(model.children()) - n_layers = len(layers_resnet) - for i_c, child in enumerate(layers_resnet): - if i_c == 0: # train first layer if not 3 bands (or no freezing) - train_if = self.freezing_strategy == "none" or self.input_n_bands != 3 - for param in child.parameters(): - param.requires_grad = train_if - elif i_c == n_layers - 1: # always train last layer - for param in child.parameters(): - param.requires_grad = True - else: # train other layers if no freezing - train_if = self.freezing_strategy == "none" - for param in child.parameters(): - param.requires_grad = train_if - - return model + + # Copy pre-trained RGB bands + if self.pretrained_cnn is not None: + with torch.no_grad(): + for i in range(self.input_n_bands): + self.geo_encoder.conv1.weight[:, i, :, :] = weight[:, i % 3, :, :] + + # Ensure replaced layer is not frozen + trainable_modules.append('geo_encoder.conv1') + + # I think for features fc often is replaced with identity? + self.geo_encoder.fc = nn.Identity() + + # self.geo_encoder.fc = nn.Linear(self.geo_encoder.fc.in_features, self.output_dim) + # trainable_modules.append('geo_encoder.fc') + + return trainable_modules else: raise ValueError(f"Unsupported backbone: {self.backbone}") - @override - def _setup(self) -> List[str]: - # TODO: could you make sure new layers are returned here to be added to trainable parts? - self.geo_encoder = self.get_backbone() - return [] - @override def forward( self, @@ -147,18 +145,15 @@ def forward( :param batch: input batch :return: extracted features """ - eo_data = batch.get("eo", {}) - + eo_data = batch.get("eo", KeyError(f"Batch must contain batch['eo']")) + eo_data = eo_data.get(self.geo_data_name, KeyError(f"Batch must contain batch['eo']['{self.geo_data_name}']")) dtype = self.dtype + if eo_data.dtype != dtype: eo_data = eo_data.to(dtype) - feats = self.geo_encoder(eo_data[self.geo_data_name]) - # n_nans = torch.sum(torch.isnan(feats)).item() - # assert ( - # n_nans == 0 - # ), f"CNNEncoder output contains {n_nans}/{feats.numel()} NaNs PRIOR to normalization with data min {eo_data[self.geo_data_name].min()} and max {eo_data[self.geo_data_name].max()}." + feats = self.geo_encoder(eo_data) if self.extra_projector: feats = self.extra_projector(feats) - return feats.to(dtype) + return feats \ No newline at end of file From cbce9f1aca05a140a6552c9dc7e70e9f3c890fee Mon Sep 17 00:00:00 2001 From: GabrieleTi Date: Wed, 20 May 2026 09:15:03 +0200 Subject: [PATCH 3/6] Add satclip requirement --- pyproject.toml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 4acb3ee..c14c61c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,11 @@ dependencies = [ "xgboost>=3.2.0", ] +satclip = [ + "albumentations", + "satclip @ git+https://github.com/gabrieletijunaityte/satclip", +] + [project.optional-dependencies] create-data = [ "geemap>=0.36.6", From 8b6a48dc4c1ae75041ec5b7fbc9e63d8839689f5 Mon Sep 17 00:00:00 2001 From: GabrieleTi Date: Wed, 20 May 2026 09:16:37 +0200 Subject: [PATCH 4/6] Add satclip encoder --- src/models/components/geo_encoders/satclip.py | 66 +++++++++++++++++++ 1 file changed, 66 insertions(+) create mode 100644 src/models/components/geo_encoders/satclip.py diff --git a/src/models/components/geo_encoders/satclip.py b/src/models/components/geo_encoders/satclip.py new file mode 100644 index 0000000..b27ed58 --- /dev/null +++ b/src/models/components/geo_encoders/satclip.py @@ -0,0 +1,66 @@ +from typing import Dict, List, override + +import torch +from huggingface_hub import hf_hub_download +from satclip.load import get_satclip + +from src.models.components.geo_encoders.base_geo_encoder import BaseGeoEncoder + + +class SatClipCoordinateEncoder(BaseGeoEncoder): + def __init__( + self, + geo_data_name="coords", + hf_cache_dir: str = "../.cache", + accelerator: torch.device = torch.device("cpu"), + ) -> None: + """SatClip coordinate encoder :param geo_data_name: type of geo data used for this encoder + (supports only coordinates) :param hf_cache_dir: hugging face cache directory to store data + :param accelerator: where to load model (as it is float64, mps is not supported)""" + super().__init__() + + self.allowed_geo_data_names = ["coords"] + assert ( + geo_data_name in self.allowed_geo_data_names + ), f"geo_data_name must be one of {self.allowed_geo_data_names}, got {geo_data_name}" + self.geo_data_name = geo_data_name + + self.cache_dir = hf_cache_dir + assert accelerator != torch.device("mps"), f"accelerator {accelerator} is not supported" + self.accelerator = accelerator + + @override + def _setup(self) -> List[str]: + """Setup satclip encoder from hugging face hub and set output dimension.""" + self.geo_encoder = get_satclip( + hf_hub_download( + "microsoft/SatCLIP-ViT16-L40", "satclip-vit16-l40.ckpt", cache_dir=self.cache_dir + ), + device=self.accelerator, + ) + + self.output_dim = self.geo_encoder.nnet.last_layer.dim_out + return [] + + @override + def forward( + self, + batch: Dict[str, torch.Tensor], + ) -> torch.Tensor: + """Forward pass of satclip encoder.""" + + coords = batch.get("eo", {}).get("coords") + + # Swap coordinates + coords = coords[:, [1, 0]] + + # SatClip needs float64 + dtype = self.dtype + if coords.dtype != dtype: + coords = coords.to(dtype) + + feats = self.geo_encoder(coords) + if self.extra_projector: + feats = self.extra_projector(feats) + + return feats From b68f12c562a464b239647cdbe38de917a569ef9f Mon Sep 17 00:00:00 2001 From: GabrieleTi Date: Wed, 20 May 2026 09:24:15 +0200 Subject: [PATCH 5/6] Satclip requirement as non-optional --- pyproject.toml | 3 --- 1 file changed, 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index c14c61c..9947bf9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,9 +36,6 @@ dependencies = [ "geotessera>=0.8.0", "ipykernel>=7.2.0", "xgboost>=3.2.0", -] - -satclip = [ "albumentations", "satclip @ git+https://github.com/gabrieletijunaityte/satclip", ] From bab8ed6265d3b54497b77811250455cc0c7faff8 Mon Sep 17 00:00:00 2001 From: GabrieleTi Date: Thu, 21 May 2026 11:11:29 +0200 Subject: [PATCH 6/6] Remove freezing config from CNN encoder --- configs/model/predictive_cnn_s2.yaml | 1 - 1 file changed, 1 deletion(-) diff --git a/configs/model/predictive_cnn_s2.yaml b/configs/model/predictive_cnn_s2.yaml index 52c46d9..48c9e60 100644 --- a/configs/model/predictive_cnn_s2.yaml +++ b/configs/model/predictive_cnn_s2.yaml @@ -3,7 +3,6 @@ _target_: src.models.predictive_model.PredictiveModel geo_encoder: _target_: src.models.components.geo_encoders.cnn_encoder.CNNEncoder resnet_version: 18 - freezing_strategy: none geo_data_name: s2 prediction_head: