bug fix: add masking during training

Author: Donglai Wei

connectomics/config/hydra_config.py

@@ -480,6 +480,9 @@ class DataConfig:
     cached_sampling_foreground_threshold: float = (
         0.0  # Minimum (label > 0) fraction required for training crops; 0 disables foreground sampling
     )
+    cached_sampling_crop_to_nonzero_mask: bool = (
+        False  # Restrict random crops to intersect the nonzero bounding box of the mask volume
+    )
 
     # Reject sampling configuration (for volumetric patch sampling)
     reject_sampling: Optional[Dict[str, Any]] = None  # Dict with 'size_thres' and 'p' keys

connectomics/data/dataset/dataset_volume_cached.py

@@ -165,6 +165,7 @@ def __init__(
         pad_mode: str = "reflect",
         max_attempts: int = 10,
         foreground_threshold: float = 0.05,
+        crop_to_nonzero_mask: bool = False,
     ):
         self.image_paths = image_paths
         self.label_paths = label_paths if label_paths else [None] * len(image_paths)
@@ -200,6 +201,7 @@ def __init__(
         self._d2_last_report_step = 0
         self.max_attempts = max_attempts
         self.foreground_threshold = foreground_threshold
+        self.crop_to_nonzero_mask = crop_to_nonzero_mask
 
         # Load all volumes into memory
         print(f"  Loading {len(image_paths)} volumes into memory...")
@@ -293,6 +295,17 @@ def __init__(
         ndim = len(self.patch_size)
         self.volume_sizes = [img.shape[-ndim:] for img in self.cached_images]  # (Z, Y, X) or (Y, X)
 
+        # Precompute mask bounding boxes for crop_to_nonzero_mask sampling
+        self.mask_bboxes: List[Optional[List[Tuple[int, int]]]] = []
+        if self.crop_to_nonzero_mask:
+            for mask in self.cached_masks:
+                self.mask_bboxes.append(self._compute_mask_bbox(mask))
+            print(
+                f"  [crop_to_nonzero_mask] Mask bboxes computed for {len(self.mask_bboxes)} volumes"
+            )
+        else:
+            self.mask_bboxes = [None] * len(self.cached_images)
+
         # [D2 DIAGNOSTIC] Print foreground sampling configuration
         if self.mode == "train":
             if self.foreground_threshold > 0:
@@ -459,10 +472,39 @@ def get_sampling_fingerprint(self, num_samples: int = 5) -> str:
             # Restore RNG state
             random.setstate(state)
 
+    def _compute_mask_bbox(
+        self, mask: Optional[np.ndarray]
+    ) -> Optional[List[Tuple[int, int]]]:
+        """Compute the axis-aligned bounding box of nonzero voxels in the mask.
+
+        Args:
+            mask: Mask array with shape (C, D, H, W) or (C, H, W). May be None.
+
+        Returns:
+            List of (start, end) pairs for each spatial dimension, or None if
+            the mask is None or entirely zero.
+        """
+        if mask is None:
+            return None
+        spatial = mask[0] > 0  # drop channel dim: (D, H, W) or (H, W)
+        if not spatial.any():
+            return None
+        ndim = spatial.ndim
+        bbox = []
+        for d in range(ndim):
+            axes = tuple(i for i in range(ndim) if i != d)
+            proj = spatial.any(axis=axes) if ndim > 1 else spatial
+            indices = np.where(proj)[0]
+            bbox.append((int(indices[0]), int(indices[-1]) + 1))
+        return bbox
+
     def _get_random_crop_position(self, vol_idx: int) -> Tuple[int, ...]:
         """
         Get a random crop position for training (like v1 VolumeDataset).
 
+        When crop_to_nonzero_mask is True, restricts sampling so every crop
+        intersects the nonzero bounding box of the mask.
+
         Args:
             vol_idx: Volume index
 
@@ -471,18 +513,22 @@ def _get_random_crop_position(self, vol_idx: int) -> Tuple[int, ...]:
         """
         vol_size = self.volume_sizes[vol_idx]
         patch_size = self.patch_size
+        mask_bbox = self.mask_bboxes[vol_idx] if self.crop_to_nonzero_mask else None
 
-        # Random position ensuring crop fits within volume
-        # Support both 2D and 3D
-        # Ensure max_start ensures we can crop at least patch_size
         positions = []
         for i in range(len(patch_size)):
-            max_start = max(0, vol_size[i] - patch_size[i])
-            # Safety check: if volume is smaller than patch_size, start at 0
-            # (crop_volume will handle padding)
-            if max_start < 0:
-                max_start = 0
-            positions.append(random.randint(0, max_start))
+            vol_max = max(0, vol_size[i] - patch_size[i])
+            if mask_bbox is not None:
+                # Intersection condition: crop [pos, pos+patch) overlaps [bbox_start, bbox_end)
+                # => pos < bbox_end  and  pos + patch > bbox_start
+                # => pos >= bbox_start - patch + 1  and  pos <= bbox_end - 1
+                min_start = max(0, mask_bbox[i][0] - patch_size[i] + 1)
+                max_start = min(vol_max, mask_bbox[i][1] - 1)
+                if min_start > max_start:  # bbox too small — fall back to full range
+                    min_start, max_start = 0, vol_max
+            else:
+                min_start, max_start = 0, vol_max
+            positions.append(random.randint(min_start, max_start))
         return tuple(positions)
 
     def _get_center_crop_position(self, vol_idx: int) -> Tuple[int, ...]:

connectomics/training/deep_supervision.py

@@ -114,6 +114,7 @@ def compute_multitask_loss(
         outputs: torch.Tensor,
         labels: torch.Tensor,
         stage: str = "train",
+        mask: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, Dict[str, float]]:
         """
         Compute losses for multi-task learning where different output channels
@@ -181,15 +182,26 @@ def compute_multitask_loss(
                 loss_fn = self.loss_functions[loss_idx]
                 weight = self.loss_weights[loss_idx]
 
-                # [D3] Pass foreground-weighted mask to loss functions that support it
-                # (e.g., WeightedMSELoss, WeightedMAELoss, SmoothL1Loss)
+                # [D3] Pass foreground-weighted mask (and validity mask) to loss functions
                 if _loss_supports_weight(loss_fn):
                     fg_weight = 2.0
                     loss_weight_mask = torch.ones_like(task_label)
                     loss_weight_mask[task_label > 0] = fg_weight
+                    if mask is not None:
+                        # Zero weight for voxels outside valid mask — excluded from gradient
+                        loss_weight_mask = loss_weight_mask * (mask > 0).float()
                     loss = loss_fn(task_output, task_label, weight=loss_weight_mask)
                 else:
-                    loss = loss_fn(task_output, task_label)
+                    if mask is not None:
+                        # For losses without per-voxel weight support (e.g. DiceLoss):
+                        # suppress output to large-negative (sigmoid≈0) and zero label
+                        # in invalid regions so they contribute ≈0 to numerator/denominator
+                        mask_bool = (mask > 0)
+                        task_output_m = task_output.masked_fill(~mask_bool, -20.0)
+                        task_label_m = task_label * mask_bool.float()
+                        loss = loss_fn(task_output_m, task_label_m)
+                    else:
+                        loss = loss_fn(task_output, task_label)
 
                 # Check for NaN/Inf
                 if self.enable_nan_detection and (torch.isnan(loss) or torch.isinf(loss)):
@@ -246,7 +258,12 @@ def compute_multitask_loss(
         return total_loss, loss_dict
 
     def compute_loss_for_scale(
-        self, output: torch.Tensor, target: torch.Tensor, scale_idx: int, stage: str = "train"
+        self,
+        output: torch.Tensor,
+        target: torch.Tensor,
+        scale_idx: int,
+        stage: str = "train",
+        mask: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, Dict[str, float]]:
         """
         Compute loss for a single scale with multi-task or standard loss.
@@ -288,7 +305,20 @@ def compute_loss_for_scale(
                     loss_fn = self.loss_functions[loss_idx]
                     weight = self.loss_weights[loss_idx]
 
-                    loss = loss_fn(task_output, task_target)
+                    if mask is not None and _loss_supports_weight(loss_fn):
+                        fg_weight = 2.0
+                        loss_weight_mask = torch.ones_like(task_target)
+                        loss_weight_mask[task_target > 0] = fg_weight
+                        loss_weight_mask = loss_weight_mask * (mask > 0).float()
+                        loss = loss_fn(task_output, task_target, weight=loss_weight_mask)
+                    elif mask is not None:
+                        mask_bool = (mask > 0)
+                        loss = loss_fn(
+                            task_output.masked_fill(~mask_bool, -20.0),
+                            task_target * mask_bool.float(),
+                        )
+                    else:
+                        loss = loss_fn(task_output, task_target)
 
                     # Check for NaN/Inf (only in training mode)
                     if (
@@ -376,7 +406,11 @@ def compute_loss_for_scale(
         return scale_loss, loss_dict
 
     def compute_deep_supervision_loss(
-        self, outputs: Dict[str, torch.Tensor], labels: torch.Tensor, stage: str = "train"
+        self,
+        outputs: Dict[str, torch.Tensor],
+        labels: torch.Tensor,
+        stage: str = "train",
+        mask: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, Dict[str, float]]:
         """
         Compute multi-scale loss with deep supervision.
@@ -385,6 +419,7 @@ def compute_deep_supervision_loss(
             outputs: Dictionary with 'output' and 'ds_i' keys for deep supervision
             labels: Ground truth labels
             stage: 'train' or 'val' for logging prefix
+            mask: Optional binary validity mask (B, 1, D, H, W); loss is zeroed outside mask
 
         Returns:
             Tuple of (total_loss, loss_dict)
@@ -419,9 +454,19 @@ def compute_deep_supervision_loss(
             # Match target to output size
             target = match_target_to_output(labels, output)
 
+            # Downsample mask to match this scale using nearest interpolation
+            scale_mask = None
+            if mask is not None:
+                if mask.shape[2:] != output.shape[2:]:
+                    scale_mask = F.interpolate(
+                        mask.float(), size=output.shape[2:], mode="nearest"
+                    )
+                else:
+                    scale_mask = mask
+
             # Compute loss for this scale
             scale_loss, scale_loss_dict = self.compute_loss_for_scale(
-                output, target, scale_idx, stage
+                output, target, scale_idx, stage, mask=scale_mask
             )
 
             # Accumulate with deep supervision weight
@@ -432,7 +477,11 @@ def compute_deep_supervision_loss(
         return total_loss, loss_dict
 
     def compute_standard_loss(
-        self, outputs: torch.Tensor, labels: torch.Tensor, stage: str = "train"
+        self,
+        outputs: torch.Tensor,
+        labels: torch.Tensor,
+        stage: str = "train",
+        mask: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, Dict[str, float]]:
         """
         Compute standard single-scale loss.
@@ -441,6 +490,7 @@ def compute_standard_loss(
             outputs: Model outputs (B, C, D, H, W)
             labels: Ground truth labels (B, C, D, H, W)
             stage: 'train' or 'val' for logging prefix
+            mask: Optional binary validity mask (B, 1, D, H, W); loss is zeroed outside mask
 
         Returns:
             Tuple of (total_loss, loss_dict)
@@ -451,18 +501,29 @@ def compute_standard_loss(
         # Check if multi-task learning is configured
         if self.is_multi_task:
             # Multi-task learning: apply specific losses to specific channels
-            total_loss, loss_dict = self.compute_multitask_loss(outputs, labels, stage=stage)
+            total_loss, loss_dict = self.compute_multitask_loss(
+                outputs, labels, stage=stage, mask=mask
+            )
         else:
             # Standard single-scale loss: apply all losses to all outputs
             for i, (loss_fn, weight) in enumerate(zip(self.loss_functions, self.loss_weights)):
-                # [D3] Pass foreground-weighted mask to loss functions that support it
+                # [D3] Pass foreground-weighted mask (and validity mask) to loss functions
                 if _loss_supports_weight(loss_fn):
                     fg_weight = 2.0
                     loss_weight_mask = torch.ones_like(labels)
                     loss_weight_mask[labels > 0] = fg_weight
+                    if mask is not None:
+                        loss_weight_mask = loss_weight_mask * (mask > 0).float()
                     loss = loss_fn(outputs, labels, weight=loss_weight_mask)
                 else:
-                    loss = loss_fn(outputs, labels)
+                    if mask is not None:
+                        mask_bool = (mask > 0)
+                        loss = loss_fn(
+                            outputs.masked_fill(~mask_bool, -20.0),
+                            labels * mask_bool.float(),
+                        )
+                    else:
+                        loss = loss_fn(outputs, labels)
 
                 # Check for NaN/Inf (only in training mode)
                 if (

connectomics/training/lit/data_factory.py

@@ -626,6 +626,7 @@ def _build_preloaded_nnunet_preprocess(split: str):
             pad_mode=pad_mode,
             max_attempts=cfg.data.cached_sampling_max_attempts,
             foreground_threshold=cfg.data.cached_sampling_foreground_threshold,
+            crop_to_nonzero_mask=cfg.data.cached_sampling_crop_to_nonzero_mask,
         )
 
         preloaded_num_workers = num_workers
@@ -691,6 +692,7 @@ def _build_preloaded_nnunet_preprocess(split: str):
                     pad_mode=pad_mode,
                     max_attempts=cfg.data.cached_sampling_max_attempts,
                     foreground_threshold=cfg.data.cached_sampling_foreground_threshold,
+                    crop_to_nonzero_mask=cfg.data.cached_sampling_crop_to_nonzero_mask,
                 )
             else:
                 from ...data.dataset import create_volume_dataset

connectomics/training/lit/model.py

@@ -590,16 +590,22 @@ def _compute_test_metrics(self, decoded_predictions: np.ndarray, labels: torch.T
     def training_step(self, batch: Dict[str, torch.Tensor], batch_idx: int) -> STEP_OUTPUT:
         """Training step with deep supervision support."""
         from ...utils.debug_utils import print_tensor_stats
-        
+
         images = batch['image']
         labels = batch['label']
+        raw_mask = batch.get('mask', None)
+        # Binarize mask: (B, 1, D, H, W) float, 1 = valid, 0 = ignore
+        mask = (raw_mask > 0).float() if raw_mask is not None else None
+
+        if batch_idx == 0 and self.current_epoch == 0:
+            import pdb; pdb.set_trace()
 
         # Forward pass
         outputs = self(images)
 
         # Check if model outputs deep supervision
         is_deep_supervision = isinstance(outputs, dict) and any(k.startswith('ds_') for k in outputs.keys())
-        
+
         # DEBUG: Print model output (raw logits, before any activation)
         if batch_idx == 0:
             pred_for_debug = outputs['output'] if is_deep_supervision else outputs
@@ -617,11 +623,11 @@ def training_step(self, batch: Dict[str, torch.Tensor], batch_idx: int) -> STEP_
         # Compute loss using deep supervision handler
         if is_deep_supervision:
             total_loss, loss_dict = self.deep_supervision_handler.compute_deep_supervision_loss(
-                outputs, labels, stage="train"
+                outputs, labels, stage="train", mask=mask
             )
         else:
             total_loss, loss_dict = self.deep_supervision_handler.compute_standard_loss(
-                outputs, labels, stage="train"
+                outputs, labels, stage="train", mask=mask
             )
 
         # [D1] Training diagnostics: log prediction and target statistics to detect SDT collapse
@@ -673,6 +679,8 @@ def validation_step(self, batch: Dict[str, torch.Tensor], batch_idx: int) -> STE
         """Validation step with deep supervision support."""
         images = batch['image']
         labels = batch['label']
+        raw_mask = batch.get('mask', None)
+        mask = (raw_mask > 0).float() if raw_mask is not None else None
 
         # Forward pass
         outputs = self(images)
@@ -683,11 +691,11 @@ def validation_step(self, batch: Dict[str, torch.Tensor], batch_idx: int) -> STE
         # Compute loss using deep supervision handler
         if is_deep_supervision:
             total_loss, loss_dict = self.deep_supervision_handler.compute_deep_supervision_loss(
-                outputs, labels, stage="val"
+                outputs, labels, stage="val", mask=mask
             )
         else:
             total_loss, loss_dict = self.deep_supervision_handler.compute_standard_loss(
-                outputs, labels, stage="val"
+                outputs, labels, stage="val", mask=mask
             )
 
         # Compute evaluation metrics if enabled

tutorials/vesicle_xm.yaml

@@ -72,8 +72,12 @@ data:
   iter_num_per_epoch: 1000
   cached_sampling_foreground_threshold: 0.1
   cached_sampling_max_attempts: 10
+  cached_sampling_crop_to_nonzero_mask: true
 
   patch_size: [32, 64, 64]
+  image_transform:
+    normalize: "0-1"    
+
   data_transform:
     resize: [32, 128, 128]