Adds SMPModelFactory example

docs/models.md

@@ -61,6 +61,7 @@ By passing a list of bands being used to the constructor, we automatically filte
 
 ## Model Factory
 ### :::terratorch.models.PrithviModelFactory
+### :::terratorch.models.SMPModelFactory
 
 # Adding new model types
 Adding new model types is as simple as creating a new factory that produces models. See for instance the example below for a potential `SMPModelFactory`

examples/confs/smp_model_factory.yaml

@@ -0,0 +1,93 @@
+benchmark_suffix: smp_test
+experiment_name: smp_test
+backbone:
+  # backbone: resnet18
+  # backbone_args:
+  #   pretrained: False
+  #   output_stride: 2
+  #   smp_decoder_channels: 512
+  #   smp_encoder_depth: 5
+
+  backbone: swin3d.swin3d_backbone.Swin3dBackbone
+  backbone_args:
+    pretrained: False
+    output_stride: 2
+    out_channels:
+      - 192
+      - 384
+      - 768
+      - 768
+    smp_decoder_channels: 768
+    smp_encoder_depth: 5
+
+
+tasks:
+  - name: cashew
+    type: segmentation
+    loss: ce
+    model_factory: SMPModelFactory
+    bands:
+      - RED
+      - GREEN
+      - BLUE
+    num_classes: 7
+    max_epochs: 60
+    direction: max
+    datamodule:
+      class_path: terratorch.datamodules.MBeninSmallHolderCashewsNonGeoDataModule
+      init_args:
+        batch_size: 16
+        num_workers: 4
+        train_transform:
+        - class_path: albumentations.Resize
+          init_args:
+            always_apply: True
+            height: 224
+            width: 224  
+        - class_path: ToTensorV2   
+        test_transform:
+        - class_path: albumentations.Resize
+          init_args:
+            always_apply: True
+            height: 224
+            width: 224
+        - class_path: ToTensorV2
+        val_transform:
+        - class_path: albumentations.Resize
+          init_args:
+            height: 224
+            width: 224
+        - class_path: ToTensorV2
+        data_root: "/dccstor/geofm-finetuning/geobench/segmentation_v1.0"
+        bands:
+          - "RED"
+          - "GREEN"
+          - "BLUE"
+    decoder: Unet
+    decoder_args:
+      channels: 128
+    metric: val/Multiclass Jaccard Index
+
+n_trials: 16
+save_models: False
+storage_uri: /dccstor/geofm-finetuning/pedrohc/smp_test
+optimization_space:
+  model:
+    - DeepLabV3
+  lr:
+    min: 6e-5
+    max: 1e-3
+    type: real
+    log: true
+  batch_size:
+    - 8
+    - 16
+    - 32
+  decoder_channels:
+    - 32
+    - 64
+    - 128
+  head_dropout:
+    min: 0.2
+    max: 0.8
+    type: real

pyproject.toml

@@ -39,7 +39,9 @@ dependencies = [
   "geobench>=1.0.0",
   "mlflow>=2.12.1",
   # broken due to https://github.com/Lightning-AI/pytorch-lightning/issues/19977
-  "lightning>=2, <=2.2.5"
+  "lightning>=2, <=2.2.5",
+  # see issue #64
+  "albumentations<=1.4.10"
 ]
 
 [project.optional-dependencies]

terratorch/datamodules/generic_pixel_wise_data_module.py

@@ -3,11 +3,11 @@
 """
 This module contains generic data modules for instantiation at runtime.
 """
-
+import os
 from collections.abc import Callable, Iterable
 from pathlib import Path
 from typing import Any
-
+import numpy as np
 import albumentations as A
 import kornia.augmentation as K
 import torch
@@ -17,7 +17,7 @@
 from torchgeo.transforms import AugmentationSequential
 
 from terratorch.datasets import GenericNonGeoPixelwiseRegressionDataset, GenericNonGeoSegmentationDataset, HLSBands
-
+from terratorch.io.file import load_from_file_or_attribute
 
 def wrap_in_compose_is_list(transform_list):
     # set check shapes to false because of the multitemporal case
@@ -79,8 +79,8 @@ def __init__(
         test_data_root: Path,
         img_grep: str,
         label_grep: str,
-        means: list[float],
-        stds: list[float],
+        means: list[float] | str,
+        stds: list[float] | str,
         num_classes: int,
         predict_data_root: Path | None = None,
         train_label_data_root: Path | None = None,
@@ -91,9 +91,9 @@ def __init__(
         test_split: Path | None = None,
         ignore_split_file_extensions: bool = True,
         allow_substring_split_file: bool = True,
-        dataset_bands: list[HLSBands | int] | None = None,
-        predict_dataset_bands: list[HLSBands | int] | None = None,
-        output_bands: list[HLSBands | int] | None = None,
+        dataset_bands: list[HLSBands | int | tuple[int, int] | str] | None = None,
+        predict_dataset_bands: list[HLSBands | int | tuple[int, int] | str ] | None = None,
+        output_bands: list[HLSBands | int | tuple[int, int] | str] | None = None,
         constant_scale: float = 1,
         rgb_indices: list[int] | None = None,
         train_transform: A.Compose | None | list[A.BasicTransform] = None,
@@ -198,6 +198,9 @@ def __init__(
         #     K.Normalize(means, stds),
         #     data_keys=["image"],
         # )
+        means = load_from_file_or_attribute(means)
+        stds = load_from_file_or_attribute(stds)
+
         self.aug = Normalize(means, stds)
 
         # self.aug = Normalize(means, stds)
@@ -317,8 +320,8 @@ def __init__(
         train_data_root: Path,
         val_data_root: Path,
         test_data_root: Path,
-        means: list[float],
-        stds: list[float],
+        means: list[float] | str,
+        stds: list[float] | str,
         predict_data_root: Path | None = None,
         img_grep: str | None = "*",
         label_grep: str | None = "*",
@@ -330,9 +333,9 @@ def __init__(
         test_split: Path | None = None,
         ignore_split_file_extensions: bool = True,
         allow_substring_split_file: bool = True,
-        dataset_bands: list[HLSBands | int] | None = None,
-        predict_dataset_bands: list[HLSBands | int] | None = None,
-        output_bands: list[HLSBands | int] | None = None,
+        dataset_bands: list[HLSBands | int | tuple[int, int] | str ] | None = None,
+        predict_dataset_bands: list[HLSBands | int | tuple[int, int] | str ] | None = None,
+        output_bands: list[HLSBands | int | tuple[int, int] | str ] | None = None,
         constant_scale: float = 1,
         rgb_indices: list[int] | None = None,
         train_transform: A.Compose | None | list[A.BasicTransform] = None,
@@ -430,6 +433,9 @@ def __init__(
         #     K.Normalize(means, stds),
         #     data_keys=["image"],
         # )
+        means = load_from_file_or_attribute(means)
+        stds = load_from_file_or_attribute(stds)
+
         self.aug = Normalize(means, stds)
         self.no_data_replace = no_data_replace
         self.no_label_replace = no_label_replace

terratorch/datamodules/generic_scalar_label_data_module.py

@@ -22,12 +22,12 @@
     HLSBands,
 )
 
+from terratorch.io.file import load_from_file_or_attribute
 
 def wrap_in_compose_is_list(transform_list):
     # set check shapes to false because of the multitemporal case
     return A.Compose(transform_list, is_check_shapes=False) if isinstance(transform_list, Iterable) else transform_list
 
-
 class Normalize(Callable):
     def __init__(self, means, stds):
         super().__init__()
@@ -68,8 +68,8 @@ def __init__(
         train_data_root: Path,
         val_data_root: Path,
         test_data_root: Path,
-        means: list[float],
-        stds: list[float],
+        means: list[float] | str,
+        stds: list[float] | str,
         num_classes: int,
         predict_data_root: Path | None = None,
         train_split: Path | None = None,
@@ -166,6 +166,10 @@ def __init__(
         #     K.Normalize(means, stds),
         #     data_keys=["image"],
         # )
+
+        means = load_from_file_or_attribute(means)
+        stds = load_from_file_or_attribute(stds)
+
         self.aug = Normalize(means, stds)
 
         # self.aug = Normalize(means, stds)