update

sorghum_pipeline/data/mask_handler.py

@@ -2,6 +2,9 @@
 
 import numpy as np
 import cv2
+import logging
+
+logger = logging.getLogger(__name__)
 
 
 class MaskHandler:
@@ -16,20 +19,32 @@ class MaskHandler:
         if mask is None:
             return image
         
-        # 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 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 (must be exactly same H x W as image)
-        binary = (mask > 0).astype(np.uint8) * 255
-        
-        # Apply mask
-        return cv2.bitwise_and(image, image, mask=binary)
+        try:
+            # 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 exactly match image dimensions
+            mask_h, mask_w = mask.shape[:2]
+            if (mask_h, mask_w) != (img_h, img_w):
+                logger.debug(f"Resizing mask from {mask.shape} to ({img_h}, {img_w})")
+                mask = cv2.resize(mask, (img_w, img_h), interpolation=cv2.INTER_NEAREST)
+            
+            # Create binary mask (must be exactly same H x W as image)
+            binary = (mask > 0).astype(np.uint8) * 255
+            
+            # Verify dimensions match
+            if binary.shape != (img_h, img_w):
+                logger.error(f"Mask shape mismatch: binary {binary.shape} != image ({img_h}, {img_w})")
+                return image
+            
+            # Apply mask
+            result = cv2.bitwise_and(image, image, mask=binary)
+            return result
+        except Exception as e:
+            logger.error(f"Mask application failed: {e}")
+            return image

sorghum_pipeline/pipeline.py

@@ -79,8 +79,12 @@ class SorghumPipeline:
         """Segment using BRIA."""
         for key, pdata in plants.items():
             composite = pdata['composite']
+            logger.info(f"Composite shape: {composite.shape}")
             soft_mask = self.segmentation_manager.segment_image_soft(composite)
-            pdata['mask'] = (soft_mask * 255.0).astype(np.uint8)
+            logger.info(f"Soft mask shape: {soft_mask.shape}")
+            mask_uint8 = (soft_mask * 255.0).astype(np.uint8)
+            logger.info(f"Mask uint8 shape: {mask_uint8.shape}")
+            pdata['mask'] = mask_uint8
         return plants
     
     def _extract_features(self, plants: Dict[str, Any]) -> Dict[str, Any]: