The number of channels can be estimated in different ways according to the kind of bands definition

Signed-off-by: João Lucas de Sousa Almeida <joao.l.sa.9.3@gmail.com>

Author: Joao-L-S-Almeida

terratorch/models/backbones/prithvi_swin.py

@@ -18,6 +18,7 @@
 from terratorch.datasets.utils import HLSBands
 from terratorch.models.backbones.prithvi_select_patch_embed_weights import prithvi_select_patch_embed_weights
 from terratorch.models.backbones.swin_encoder_decoder import MMSegSwinTransformer
+from terratorch.models.backbones.utils import _estimate_in_chans
 
 PRETRAINED_BANDS = [
     HLSBands.BLUE,
@@ -174,7 +175,8 @@ def _create_swin_mmseg_transformer(
     # the current swin model is not multitemporal
     if "num_frames" in kwargs:
         kwargs = {k: v for k, v in kwargs.items() if k != "num_frames"}
-    kwargs["in_chans"] = len(model_bands)
+
+    kwargs["in_chans"] = _estimate_in_chans(model_bands=model_bands)
 
     def checkpoint_filter_wrapper_fn(state_dict, model):
         return checkpoint_filter_fn(state_dict, model, pretrained_bands, model_bands)

terratorch/models/backbones/prithvi_vit.py

@@ -14,6 +14,7 @@
 from terratorch.datasets import HLSBands
 from terratorch.models.backbones.prithvi_select_patch_embed_weights import prithvi_select_patch_embed_weights
 from terratorch.models.backbones.vit_encoder_decoder import TemporalViTEncoder
+from terratorch.models.backbones.utils import _estimate_in_chans
 
 PRETRAINED_BANDS = [
     HLSBands.BLUE,
@@ -81,7 +82,7 @@ def _create_prithvi(
     if "features_only" in kwargs:
         kwargs = {k: v for k, v in kwargs.items() if k != "features_only"}
 
-    kwargs["in_chans"] = len(model_bands)
+    kwargs["in_chans"] = _estimate_in_chans(model_bands=model_bands)
 
     def checkpoint_filter_wrapper_fn(state_dict, model):
         return checkpoint_filter_fn(state_dict, model, pretrained_bands, model_bands)

terratorch/tasks/segmentation_tasks.py

@@ -254,7 +254,6 @@ def validation_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -
         x = batch["image"]
         y = batch["mask"]
         model_output: ModelOutput = self(x)
-        print(x.shape, model_output.output.shape, y.shape)        
         loss = self.val_loss_handler.compute_loss(model_output, y, self.criterion, self.aux_loss)
         self.val_loss_handler.log_loss(self.log, loss_dict=loss, batch_size=x.shape[0])
         y_hat_hard = to_segmentation_prediction(model_output)