Add faithful BANIS base reproduction

Author: Donglai Wei

connectomics/config/pipeline/config_io.py

@@ -370,6 +370,15 @@ def validate_config(cfg: Config) -> None:
             "data.dataloader.patch_size must be 2D or 3D "
             f"(got length {len(cfg.data.dataloader.patch_size)})"
         )
+    target_context = getattr(cfg.data.dataloader, "target_context", None) or []
+    if target_context:
+        if len(target_context) != len(cfg.data.dataloader.patch_size):
+            raise ValueError(
+                "data.dataloader.target_context must match patch_size dimensionality "
+                f"({len(target_context)} vs {len(cfg.data.dataloader.patch_size)})"
+            )
+        if any(int(v) < 0 for v in target_context):
+            raise ValueError("data.dataloader.target_context values must be non-negative")
     if cfg.data.dataloader.batch_size <= 0:
         raise ValueError("data.dataloader.batch_size must be positive")
 
@@ -417,6 +426,9 @@ def validate_config(cfg: Config) -> None:
 
     if cfg.optimization.gradient_clip_val < 0:
         raise ValueError("optimization.gradient_clip_val must be non-negative")
+    val_check_unit = str(getattr(cfg.optimization, "val_check_interval_unit", "epoch")).lower()
+    if val_check_unit not in {"epoch", "step"}:
+        raise ValueError("optimization.val_check_interval_unit must be 'epoch' or 'step'")
     if cfg.optimization.accumulate_grad_batches <= 0:
         raise ValueError("optimization.accumulate_grad_batches must be positive")
     if hasattr(cfg.optimization, "ema") and getattr(cfg.optimization.ema, "enabled", False):

connectomics/config/schema/data.py

@@ -76,6 +76,8 @@ class LabelTransformConfig:
     output_dtype: Optional[str] = "float32"
     output_key_format: str = "{key}_{task}"
     allow_missing_keys: bool = False
+    relabel_connected_components: bool = False  # Relabel disconnected same-ID crop components.
+    relabel_connectivity: int = 6  # Connectivity for relabel_connected_components in 3D.
     segment_id: Optional[List[int]] = None
     boundary_thickness: int = 1
     resolution: Optional[List[float]] = None  # Forwarded into compatible label targets.
@@ -175,6 +177,9 @@ class DataloaderConfig:
 
     batch_size: int = 4
     patch_size: List[int] = field(default_factory=lambda: [128, 128, 128])
+    target_context: List[int] = field(
+        default_factory=lambda: [0, 0, 0]
+    )  # Extra positive-side crop context used for target generation, then cropped away.
     pin_memory: bool = True
     use_preloaded_cache_train: bool = True  # Preload training volumes into memory
     use_preloaded_cache_val: bool = True  # Preload validation volumes into memory
@@ -196,6 +201,7 @@ class DataloaderConfig:
         False  # Voxel approach: center crops on random nonzero mask voxels (stronger guarantee)
     )
     reject_sampling: Optional[Dict[str, Any]] = None  # Dict with 'size_thres' and 'p' keys
+    val_random_sampling: bool = False  # If true, validation samples random patches, not center crops.
 
 
 @dataclass

connectomics/config/schema/monitor.py

@@ -18,6 +18,7 @@ class CheckpointConfig:
     save_every_n_epochs: int = 1
     save_every_n_steps: Optional[int] = None
     step_checkpoint_filename: str = "step-{step:08d}"
+    save_on_train_epoch_end: bool = True
     use_timestamp: bool = True
 
 

connectomics/config/schema/optimization.py

@@ -94,7 +94,8 @@ class OptimizationConfig:
     precision: str = "16-mixed"  # "32", "16-mixed", "bf16-mixed"
 
     # Validation scheduling
-    val_check_interval: Union[int, float] = 1.0  # Validate every N epochs
+    val_check_interval: Union[int, float] = 1.0  # Validate every N epochs or steps.
+    val_check_interval_unit: str = "epoch"  # "epoch" or "step"
 
     # Logging
     log_every_n_steps: int = 100

connectomics/data/augmentation/build.py

@@ -64,6 +64,48 @@ def _strict_binarize_mask(mask, threshold: float = 0.0):
     return (mask > threshold).astype(mask.dtype, copy=False)
 
 
+def _target_context(cfg: Config) -> tuple[int, ...]:
+    context = getattr(cfg.data.dataloader, "target_context", None) or []
+    if not context:
+        return tuple(0 for _ in cfg.data.dataloader.patch_size)
+    return tuple(int(v) for v in context)
+
+
+def _effective_patch_size(cfg: Config) -> tuple[int, ...] | None:
+    patch_size = tuple(cfg.data.dataloader.patch_size) if cfg.data.dataloader.patch_size else None
+    if patch_size is None:
+        return None
+    context = _target_context(cfg)
+    if len(context) != len(patch_size):
+        raise ValueError(
+            "data.dataloader.target_context must have the same length as patch_size: "
+            f"{context} vs {patch_size}"
+        )
+    return tuple(int(patch_size[i]) + int(context[i]) for i in range(len(patch_size)))
+
+
+def _append_banis_pre_target_transforms(transforms: list, label_cfg) -> None:
+    if bool(getattr(label_cfg, "relabel_connected_components", False)):
+        from ..processing.transforms import RelabelConnectedComponentsd
+
+        transforms.append(
+            RelabelConnectedComponentsd(
+                keys=["label"],
+                connectivity=int(getattr(label_cfg, "relabel_connectivity", 6)),
+            )
+        )
+
+
+def _append_target_context_crop(transforms: list, cfg: Config) -> None:
+    context = _target_context(cfg)
+    if not context or not any(v > 0 for v in context):
+        return
+
+    from ..processing.transforms import LeadingSpatialCropd
+
+    transforms.append(LeadingSpatialCropd(roi_size=tuple(cfg.data.dataloader.patch_size)))
+
+
 def _build_nnunet_preprocess_transform(keys, nnunet_pre_cfg, source_spacing):
     """Build NNUNetPreprocessd transform from config."""
     source_spacing = getattr(nnunet_pre_cfg, "source_spacing", None) or source_spacing
@@ -176,9 +218,7 @@ def build_train_transforms(
 
     # Ensure target patch size is respected (unless using pre-cached dataset)
     if not skip_loading:
-        patch_size = (
-            tuple(cfg.data.dataloader.patch_size) if cfg.data.dataloader.patch_size else None
-        )
+        patch_size = _effective_patch_size(cfg)
         if patch_size and all(size > 0 for size in patch_size):
             # Pad smaller volumes so random crops always succeed
             transforms.append(
@@ -208,6 +248,10 @@ def build_train_transforms(
             )
         )
 
+    label_cfg = getattr(cfg.data, "label_transform", None)
+    if "label" in keys and label_cfg is not None:
+        _append_banis_pre_target_transforms(transforms, label_cfg)
+
     # Add augmentations if enabled
     if cfg.data.augmentation is not None:
         # Pass do_2d flag to augmentation builder
@@ -218,12 +262,10 @@ def build_train_transforms(
         transforms.extend(_build_augmentations(cfg.data.augmentation, keys, do_2d=do_2d))
 
     # Label transformations (affinity, distance transform, etc.)
-    if hasattr(cfg.data, "label_transform"):
+    if label_cfg is not None:
         from ..processing.build import create_label_transform_pipeline
         from ..processing.transforms import SegErosionInstanced
 
-        label_cfg = cfg.data.label_transform
-
         # Apply instance erosion first if specified
         if hasattr(label_cfg, "erosion") and label_cfg.erosion > 0:
             transforms.append(SegErosionInstanced(keys=["label"], tsz_h=label_cfg.erosion))
@@ -235,6 +277,8 @@ def build_train_transforms(
         else:
             transforms.append(label_transform)
 
+    _append_target_context_crop(transforms, cfg)
+
     # NOTE: Do NOT squeeze labels here!
     # - DiceLoss needs (B, 1, H, W) with to_onehot_y=True
     # - CrossEntropyLoss needs (B, H, W)
@@ -472,7 +516,13 @@ def _resolve_eval_split():
             )
         )
 
-    patch_size = tuple(data_cfg.dataloader.patch_size) if data_cfg.dataloader.patch_size else None
+    patch_size = (
+        _effective_patch_size(cfg)
+        if mode == "val"
+        else tuple(data_cfg.dataloader.patch_size)
+        if data_cfg.dataloader.patch_size
+        else None
+    )
     if patch_size and all(size > 0 for size in patch_size):
         transforms.append(
             SpatialPadd(
@@ -510,12 +560,22 @@ def _resolve_eval_split():
     # Test: Skip cropping to enable sliding window inference on full volumes
     if mode == "val":
         if patch_size and all(size > 0 for size in patch_size):
-            transforms.append(
-                CenterSpatialCropd(
-                    keys=keys,
-                    roi_size=patch_size,
+            if bool(getattr(data_cfg.dataloader, "val_random_sampling", False)):
+                transforms.append(
+                    RandSpatialCropd(
+                        keys=keys,
+                        roi_size=patch_size,
+                        random_center=True,
+                        random_size=False,
+                    )
+                )
+            else:
+                transforms.append(
+                    CenterSpatialCropd(
+                        keys=keys,
+                        roi_size=patch_size,
+                    )
                 )
-            )
     # else: mode == "test" -> no cropping for sliding window inference
 
     # Normalization - use smart normalization
@@ -530,6 +590,10 @@ def _resolve_eval_split():
             )
         )
 
+    label_cfg = getattr(data_cfg, "label_transform", None)
+    if mode == "val" and "label" in keys and label_cfg is not None:
+        _append_banis_pre_target_transforms(transforms, label_cfg)
+
     # Only process labels if 'label' is in keys
     if "label" in keys:
         # Label transformations (affinity, distance transform, etc.)
@@ -558,6 +622,9 @@ def _resolve_eval_split():
             else:
                 transforms.append(label_transform)
 
+    if mode == "val":
+        _append_target_context_crop(transforms, cfg)
+
     # NOTE: Do NOT squeeze labels here!
     # - DiceLoss needs (B, 1, H, W) with to_onehot_y=True
     # - CrossEntropyLoss needs (B, H, W)
@@ -733,16 +800,6 @@ def _build_augmentations(aug_cfg: AugmentationConfig, keys: list[str], do_2d: bo
 
     # Intensity augmentations (only for images)
     if aug_cfg.intensity.enabled:
-        if aug_cfg.intensity.gaussian_noise_prob > 0:
-            transforms.append(
-                RandGaussianNoised(
-                    keys=["image"],
-                    prob=aug_cfg.intensity.gaussian_noise_prob,
-                    std=aug_cfg.intensity.gaussian_noise_std,
-                    sample_std=True,
-                )
-            )
-
         if getattr(aug_cfg.intensity, "banis_style", False):
             transforms.append(
                 RandMulAddIntensityd(
@@ -752,7 +809,26 @@ def _build_augmentations(aug_cfg: AugmentationConfig, keys: list[str], do_2d: bo
                     add_range=aug_cfg.intensity.add_range,
                 )
             )
+            if aug_cfg.intensity.gaussian_noise_prob > 0:
+                transforms.append(
+                    RandGaussianNoised(
+                        keys=["image"],
+                        prob=aug_cfg.intensity.gaussian_noise_prob,
+                        std=aug_cfg.intensity.gaussian_noise_std,
+                        sample_std=True,
+                    )
+                )
         else:
+            if aug_cfg.intensity.gaussian_noise_prob > 0:
+                transforms.append(
+                    RandGaussianNoised(
+                        keys=["image"],
+                        prob=aug_cfg.intensity.gaussian_noise_prob,
+                        std=aug_cfg.intensity.gaussian_noise_std,
+                        sample_std=True,
+                    )
+                )
+
             if aug_cfg.intensity.shift_intensity_prob > 0:
                 transforms.append(
                     RandShiftIntensityd(

connectomics/data/processing/transforms.py

@@ -470,6 +470,95 @@ def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
         return d
 
 
+class RelabelConnectedComponentsd(MapTransform):
+    """Relabel disconnected components inside each crop while preserving unlabeled voxels."""
+
+    def __init__(
+        self,
+        keys: KeysCollection,
+        connectivity: int = 6,
+        allow_missing_keys: bool = False,
+    ) -> None:
+        super().__init__(keys, allow_missing_keys)
+        self.connectivity = int(connectivity)
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        try:
+            import cc3d
+        except ModuleNotFoundError as exc:  # pragma: no cover
+            raise ModuleNotFoundError(
+                "relabel_connected_components requires cc3d. Install connected-components-3d."
+            ) from exc
+
+        d = dict(data)
+        for key in self.key_iterator(d):
+            if key not in d:
+                continue
+
+            label = d[key]
+            label_np = label.detach().cpu().numpy() if isinstance(label, torch.Tensor) else label
+            label_np = np.asarray(label_np)
+            restore_channel_dim = label_np.ndim == 4 and label_np.shape[0] == 1
+            seg = label_np[0] if restore_channel_dim else label_np
+            if seg.ndim != 3:
+                raise ValueError(
+                    "RelabelConnectedComponentsd expects a 3D label crop "
+                    f"with optional singleton channel dim, got {tuple(label_np.shape)}"
+                )
+
+            invalid = seg == -1
+            relabeled = cc3d.connected_components(
+                seg,
+                connectivity=self.connectivity,
+                out_dtype=np.uint32,
+            ).astype(np.int32)
+            relabeled[invalid] = -1
+            d[key] = relabeled[None, ...] if restore_channel_dim else relabeled
+        return d
+
+
+class LeadingSpatialCropd:
+    """Crop arrays/tensors to ``roi_size`` from the low-index spatial corner."""
+
+    def __init__(
+        self,
+        roi_size: Sequence[int],
+        keys: Optional[Sequence[str]] = None,
+        allow_missing_keys: bool = True,
+    ) -> None:
+        self.roi_size = tuple(int(v) for v in roi_size)
+        self.keys = tuple(keys) if keys is not None else None
+        self.allow_missing_keys = allow_missing_keys
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        d = dict(data)
+        keys = self.keys if self.keys is not None else tuple(d.keys())
+        spatial_ndim = len(self.roi_size)
+
+        for key in keys:
+            if key not in d:
+                if not self.allow_missing_keys:
+                    raise KeyError(key)
+                continue
+
+            value = d[key]
+            ndim = getattr(value, "ndim", None)
+            if ndim is None or int(ndim) < spatial_ndim:
+                continue
+
+            spatial_shape = tuple(int(v) for v in value.shape[-spatial_ndim:])
+            if all(spatial_shape[i] == self.roi_size[i] for i in range(spatial_ndim)):
+                continue
+
+            slices = [slice(None)] * int(ndim)
+            first_spatial_axis = int(ndim) - spatial_ndim
+            for axis, size in enumerate(self.roi_size, start=first_spatial_axis):
+                slices[axis] = slice(0, size)
+            d[key] = value[tuple(slices)]
+
+        return d
+
+
 class EnergyQuantized(MapTransform):
     """Quantize continuous energy maps using MONAI MapTransform.
 
@@ -881,6 +970,8 @@ def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
     "SegErosiond",
     "SegDilationd",
     "SegErosionInstanced",
+    "RelabelConnectedComponentsd",
+    "LeadingSpatialCropd",
     "EnergyQuantized",
     "DecodeQuantized",
     "SegSelectiond",