Modify mask decoder for two-class mask output.

SuperMedIntel · Mar 6, 2024 · 04434fd · 04434fd
1 parent bb57efb
commit 04434fd
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/models/sam/modeling/mask_decoder.py b/models/sam/modeling/mask_decoder.py
@@ -99,8 +99,8 @@ def forward(
         )
 
         # Select the correct mask or masks for output
-        if multimask_output:
-            mask_slice = slice(1, None)
+        if multimask_output: # for REFUGE dataset
+            mask_slice = slice(0, 2)
         else:
             mask_slice = slice(0, 1)
         masks = masks[:, mask_slice, :, :]