adds SMPModelFactory tests and SMPModelFactory to model.md

Signed-off-by: Pedro Henrique Conrado <[email protected]>

Author: PedroConrado

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`

terratorch/models/prithvi_model_factory.py

@@ -99,9 +99,9 @@ def build_model(
                 msg = f"Task {task} not supported. Please choose one of {SUPPORTED_TASKS}"
                 raise NotImplementedError(msg)
 
+            backbone_kwargs, kwargs = _extract_prefix_keys(kwargs, "backbone_")
             # These params are used in case we need a SMP decoder
             # but should not be used for timm encoder
-            backbone_kwargs, kwargs = _extract_prefix_keys(kwargs, "backbone_")
             output_stride = backbone_kwargs.pop("output_stride", None)
             out_channels = backbone_kwargs.pop("out_channels", None)
 
@@ -179,6 +179,7 @@ class SMPDecoderForPrithviWrapper(nn.Module):
         forward(x: List[torch.Tensor]) -> torch.Tensor:
             Forward pass for embeddings with specified indices.
     """
+
     def __init__(self, decoder, num_channels, in_index=-1) -> None:
         """
         Args:

terratorch/models/smp_model_factory.py

@@ -96,10 +96,9 @@ def build_model(
         as auxiliary decoders or modified encoders.
 
         Attributes:
-            task (str): Specifies the task for which the model is being built. Supported tasks include
-                        "segmentation" and "regression".
-            backbone (str, nn.Module): Specifies the backbone model to be used. If a string, it should be
-                        recognized by the model factory and be able to be parsed appropriately.
+            task (str): Specifies the task for which the model is being built. Supported tasks are
+                        "segmentation".
+            backbone (str): Specifies the backbone model to be used.
             decoder (str): Specifies the decoder to be used for constructing the
                         segmentation model.
             bands (list[terratorch.datasets.HLSBands | int]): A list specifying the bands that the model
@@ -116,12 +115,15 @@ def build_model(
 
         Raises:
             ValueError: If the specified decoder is not supported by SMP.
-            Exception: If the specified task is not "segmentation" or "regression".
+            Exception: If the specified task is not "segmentation"
 
         Returns:
             nn.Module: A model instance wrapped in SMPModelWrapper configured according to the specified
                     parameters and tasks.
         """
+        if task != "segmentation":
+            msg = f"SMP models can only perform segmentatio, but got task {task}"
+            raise Exception(msg)
 
         bands = [HLSBands.try_convert_to_hls_bands_enum(b) for b in bands]
         if in_channels is None:
@@ -197,8 +199,7 @@ def register_custom_encoder(encoder, params, pretrained):
     }
 
 
-# Gets class either from string or from Module reference.
-def make_smp_encoder(encoder = None):
+def make_smp_encoder(encoder=None):
     if isinstance(encoder, str):
         base_class = _get_class_from_string(encoder)
     else:

tests/test_smp_model_factory.py

@@ -0,0 +1,80 @@
+# Copyright contributors to the Terratorch project
+
+import os
+
+import pytest
+import torch
+
+from terratorch.models import SMPModelFactory
+from terratorch.models.backbones.prithvi_vit import PRETRAINED_BANDS
+
+# from terratorch.models.backbones.prithvi_vit import default_cfgs as vit_default_cfgs
+
+NUM_CHANNELS = 6
+NUM_CLASSES = 2
+EXPECTED_SEGMENTATION_OUTPUT_SHAPE = (1, NUM_CLASSES, 224, 224)
+EXPECTED_REGRESSION_OUTPUT_SHAPE = (1, 224, 224)
+EXPECTED_CLASSIFICATION_OUTPUT_SHAPE = (1, NUM_CLASSES)
+
+
+@pytest.fixture(scope="session")
+def model_factory() -> SMPModelFactory:
+    return SMPModelFactory()
+
+
+@pytest.fixture(scope="session")
+def model_input() -> torch.Tensor:
+    return torch.ones((1, NUM_CHANNELS, 224, 224))
+
+
+@pytest.mark.parametrize("backbone", ["timm-regnetx_002"])
+@pytest.mark.parametrize("model", ["Unet", "DeepLabV3"])
+def test_create_segmentation_model(backbone, model, model_factory: SMPModelFactory, model_input):
+    model = model_factory.build_model(
+        "segmentation",
+        backbone=backbone,
+        model=model,
+        in_channels=NUM_CHANNELS,
+        bands=PRETRAINED_BANDS,
+        pretrained=False,
+        num_classes=NUM_CLASSES,
+    )
+    model.eval()
+
+    with torch.no_grad():
+        assert model(model_input).output.shape == EXPECTED_SEGMENTATION_OUTPUT_SHAPE
+
+
+@pytest.mark.parametrize("backbone", ["timm-regnetx_002"])
+@pytest.mark.parametrize("model", ["Unet", "DeepLabV3"])
+def test_create_segmentation_model_no_in_channels(backbone, model, model_factory: SMPModelFactory, model_input):
+    model = model_factory.build_model(
+        "segmentation",
+        backbone=backbone,
+        model=model,
+        bands=PRETRAINED_BANDS,
+        pretrained=False,
+        num_classes=NUM_CLASSES,
+    )
+    model.eval()
+
+    with torch.no_grad():
+        assert model(model_input).output.shape == EXPECTED_SEGMENTATION_OUTPUT_SHAPE
+
+
+@pytest.mark.parametrize("backbone", ["timm-regnetx_002"])
+@pytest.mark.parametrize("model", ["Unet", "DeepLabV3"])
+def test_create_model_with_extra_bands(backbone, model, model_factory: SMPModelFactory):
+    model = model_factory.build_model(
+        "segmentation",
+        backbone=backbone,
+        model=model,
+        in_channels=NUM_CHANNELS + 1,
+        bands=[*PRETRAINED_BANDS, 7],  # add an extra band
+        pretrained=False,
+        num_classes=NUM_CLASSES,
+    )
+    model.eval()
+    model_input = torch.ones((1, NUM_CHANNELS + 1, 224, 224))
+    with torch.no_grad():
+        assert model(model_input).output.shape == EXPECTED_SEGMENTATION_OUTPUT_SHAPE