update

sorghum_pipeline/data/mask_handler.py

@@ -16,18 +16,20 @@ class MaskHandler:
         if mask is None:
             return image
         
-        # Ensure mask is uint8
+        # Get image dimensions
+        img_h, img_w = image.shape[:2]
+        
+        # Ensure mask is 2D uint8
+        if mask.ndim > 2:
+            mask = mask[:, :, 0]  # Take first channel if multi-channel
         mask = mask.astype(np.uint8)
         
-        # Resize mask to match image if needed
-        if mask.shape[:2] != image.shape[:2]:
-            mask = cv2.resize(mask, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)
+        # Resize mask to exactly match image dimensions
+        if mask.shape != (img_h, img_w):
+            mask = cv2.resize(mask, (img_w, img_h), interpolation=cv2.INTER_NEAREST)
         
-        # Create binary mask
+        # Create binary mask (must be exactly same H x W as image)
         binary = (mask > 0).astype(np.uint8) * 255
         
-        # Ensure binary mask is single channel and matches image dimensions
-        if binary.shape[:2] != image.shape[:2]:
-            binary = cv2.resize(binary, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)
-        
+        # Apply mask
         return cv2.bitwise_and(image, image, mask=binary)