Update DOFASeg model to use a list of encoder output channels for imp…

…roved flexibility in multi-level feature extraction. Adjusted neck and decoder initialization to accommodate the new output channel structure.
NRCan · Dec 7, 2024 · 023cf2e · 023cf2e
1 parent d3de276
commit 023cf2e
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/geo_deep_learning/models/dofa/dofa_seg.py b/geo_deep_learning/models/dofa/dofa_seg.py
@@ -535,12 +535,12 @@ def __init__(self,
             self.embedding_dim = 1024
         else:
             raise ValueError(f"Unknown encoder: {encoder}")
-
+        encoder_out_channels = [64, 128, 320, 512]
         self.neck = MultiLevelNeck(in_channels=self.in_channels,
-                                   out_channels=self.embedding_dim,
+                                   out_channels=encoder_out_channels,
                                    scales=[4, 2, 1, 0.5])
 
-        self.decoder = Decoder(in_channels=self.in_channels, 
+        self.decoder = Decoder(in_channels=encoder_out_channels, 
                                embedding_dim=self.embedding_dim, 
                                num_classes=num_classes)
     def forward(self, x):