fixed sam post-processing

juglab · Nov 11, 2024 · 16c5e4e · 16c5e4e
1 parent 945d1a9
commit 16c5e4e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/featureforest/postprocess/postprocess_with_sam.py b/src/featureforest/postprocess/postprocess_with_sam.py
@@ -141,9 +141,10 @@ def get_sam_mask(
             box=input_boxes[start:end],
             multimask_output=True,
         )
+        reduced_axis = (0, 1) if masks.ndim == 4 else 0
         final_mask = np.bitwise_or(
             final_mask,
-            np.bitwise_or.reduce(masks.astype(bool), axis=(0, 1))
+            np.bitwise_or.reduce(masks.astype(bool), axis=reduced_axis)
         )
 
     return final_mask