Minor improvements

Signed-off-by: João Lucas de Sousa Almeida <[email protected]>

Author: Joao-L-S-Almeida

terratorch/models/generic_unet_model_factory.py

@@ -9,8 +9,6 @@
 
 from terratorch.models.model import Model, ModelFactory, ModelOutput, register_factory
 
-from mmseg.models.decode_heads import ASPPHead 
-
 import importlib
 
 @register_factory
@@ -42,8 +40,13 @@ def build_model(
         if task not in ["segmentation", "regression"]:
             msg = f"SMP models can only perform pixel wise tasks, but got task {task}"
             raise Exception(msg)
+    
+        try: 
+            mmseg = importlib.import_module("mmseg.models.decode_heads")
+        except:
+            raise Exception("The module 'mmseg' is not installed or not accessible via PYTHONPATH.")
 
-        model_class = getattr(mmseg.models.decode_heads, model)
+        model_class = getattr(mmseg, model)
 
         model = model_class(
            dilations=dilations 

tests/manufactured-finetune_unet.yaml

@@ -94,6 +94,7 @@ model:
   init_args:
     model_args:
       decoder: "unet"
+      decoder_model: "ASPPHead"
       decoder_dilations: [1, 6, 12, 18]
       backbone_drop_path_rate: 0.3
       # backbone_window_size: 8