fix formatting and test files

Author: Donglai Wei

connectomics/config/auto_config.py

@@ -377,26 +377,24 @@ def auto_plan_config(
 
     # Collect manual overrides (values explicitly set in config)
     manual_overrides = {}
+    training_cfg = getattr(config.system, "training", None) if hasattr(config, "system") else None
     if hasattr(config, "data"):
-        if hasattr(config.data, "batch_size") and config.data.batch_size is not None:
-            manual_overrides["batch_size"] = config.data.batch_size
-        if hasattr(config.data, "num_workers") and config.data.num_workers is not None:
-            manual_overrides["num_workers"] = config.data.num_workers
+        if training_cfg and getattr(training_cfg, "batch_size", None) is not None:
+            manual_overrides["batch_size"] = training_cfg.batch_size
+        if training_cfg and getattr(training_cfg, "num_workers", None) is not None:
+            manual_overrides["num_workers"] = training_cfg.num_workers
         if hasattr(config.data, "patch_size") and config.data.patch_size is not None:
             manual_overrides["patch_size"] = config.data.patch_size
 
-    if hasattr(config, "training"):
-        if hasattr(config.training, "precision") and config.training.precision is not None:
-            manual_overrides["precision"] = config.training.precision
-        if (
-            hasattr(config.training, "accumulate_grad_batches")
-            and config.training.accumulate_grad_batches is not None
-        ):
-            manual_overrides["accumulate_grad_batches"] = config.training.accumulate_grad_batches
+    if hasattr(config, "optimization"):
+        if getattr(config.optimization, "precision", None) is not None:
+            manual_overrides["precision"] = config.optimization.precision
+        if getattr(config.optimization, "accumulate_grad_batches", None) is not None:
+            manual_overrides["accumulate_grad_batches"] = config.optimization.accumulate_grad_batches
 
-    if hasattr(config, "optimizer"):
-        if hasattr(config.optimizer, "lr") and config.optimizer.lr is not None:
-            manual_overrides["lr"] = config.optimizer.lr
+        opt_cfg = getattr(config.optimization, "optimizer", None)
+        if opt_cfg and getattr(opt_cfg, "lr", None) is not None:
+            manual_overrides["lr"] = opt_cfg.lr
 
     # Create planner
     planner = AutoConfigPlanner(
@@ -408,9 +406,8 @@ def auto_plan_config(
 
     # Plan
     use_mixed_precision = not (
-        hasattr(config, "training")
-        and hasattr(config.training, "precision")
-        and config.training.precision == "32"
+        hasattr(config, "optimization")
+        and getattr(config.optimization, "precision", None) == "32"
     )
 
     result = planner.plan(
@@ -423,20 +420,20 @@ def auto_plan_config(
     # Update config with planned values (if not manually overridden)
     OmegaConf.set_struct(config, False)  # Allow adding new fields
 
-    if "batch_size" not in manual_overrides:
-        config.data.batch_size = result.batch_size
-    if "num_workers" not in manual_overrides:
-        config.data.num_workers = result.num_workers
+    if "batch_size" not in manual_overrides and training_cfg is not None:
+        training_cfg.batch_size = result.batch_size
+    if "num_workers" not in manual_overrides and training_cfg is not None:
+        training_cfg.num_workers = result.num_workers
     if "patch_size" not in manual_overrides:
         config.data.patch_size = result.patch_size
 
     if "precision" not in manual_overrides:
-        config.training.precision = result.precision
+        config.optimization.precision = result.precision
     if "accumulate_grad_batches" not in manual_overrides:
-        config.training.accumulate_grad_batches = result.accumulate_grad_batches
+        config.optimization.accumulate_grad_batches = result.accumulate_grad_batches
 
-    if "lr" not in manual_overrides:
-        config.optimizer.lr = result.lr
+    if "lr" not in manual_overrides and hasattr(config, "optimization"):
+        config.optimization.optimizer.lr = result.lr
 
     OmegaConf.set_struct(config, True)  # Re-enable struct mode
 
@@ -460,7 +457,7 @@ def auto_plan_config(
     cfg = auto_plan_config(cfg, print_results=True)
 
     print("\nFinal Config Values:")
-    print(f"  batch_size: {cfg.data.batch_size}")
+    print(f"  batch_size: {cfg.system.training.batch_size}")
     print(f"  patch_size: {cfg.data.patch_size}")
     print(f"  precision: {cfg.optimization.precision}")
     print(f"  lr: {cfg.optimization.optimizer.lr}")

connectomics/config/hydra_config.py

@@ -450,9 +450,7 @@ class DataConfig:
     label_transform: LabelTransformConfig = field(default_factory=LabelTransformConfig)
 
     # Augmentation configuration (nested under data in YAML)
-    augmentation: Optional["AugmentationConfig"] = (
-        None  # Set to None for simple enabled flag, or full config for detailed control
-    )
+    augmentation: "AugmentationConfig" = field(default_factory=lambda: AugmentationConfig())
 
 
 @dataclass

connectomics/config/hydra_utils.py

@@ -86,7 +86,7 @@ def update_from_cli(cfg: Config, overrides: List[str]) -> Config:
     """
     Update config from command-line overrides.
 
-    Supports dot notation: ['data.batch_size=4', 'model.architecture=unet3d']
+    Supports dot notation: ['system.training.batch_size=4', 'model.architecture=unet3d']
 
     Args:
         cfg: Base Config object

connectomics/data/augment/monai_transforms.py

@@ -38,7 +38,11 @@ def __init__(
         self.rotate_ratio = rotate_ratio
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        if self.prob <= 0:
+            self._do_transform = False
+            return data
         d = dict(data)
+        self.randomize(None)
         if not self._do_transform:
             return d
 
@@ -53,6 +57,10 @@ def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
                     d[key] = self._apply_misalignment_translation(d[key])
         return d
 
+    def randomize(self, _: Any = None) -> None:
+        """Randomly decide whether to apply the transform."""
+        self._do_transform = self.R.rand() < self.prob
+
     def _apply_misalignment_translation(
         self, img: Union[np.ndarray, torch.Tensor]
     ) -> Union[np.ndarray, torch.Tensor]:
@@ -184,7 +192,12 @@ def __init__(
         self.num_sections = num_sections
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        if self.prob <= 0:
+            self._do_transform = False
+            return data
+
         d = dict(data)
+        self.randomize(None)
         if not self._do_transform:
             return d
 
@@ -193,29 +206,40 @@ def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
                 d[key] = self._apply_missing_section(d[key])
         return d
 
+    def randomize(self, _: Any = None) -> None:
+        """Randomly decide whether to apply the transform."""
+        self._do_transform = self.R.rand() < self.prob
+
     def _apply_missing_section(
         self, img: Union[np.ndarray, torch.Tensor]
     ) -> Union[np.ndarray, torch.Tensor]:
         """Remove random sections from volume."""
-        if img.ndim < 3 or img.shape[0] <= 3:
+        if img.ndim < 3:
             return img  # Skip 2D or very small volumes
 
         # Handle both numpy and torch tensors
         is_tensor = isinstance(img, torch.Tensor)
 
+        depth_axis = 0
+        if img.ndim >= 4 and img.shape[0] <= 3:
+            depth_axis = 1  # Channel-first; depth is axis 1
+
+        depth = img.shape[depth_axis]
+        if depth <= 3:
+            return img
+
         # Select sections to remove (avoid first and last)
-        num_to_remove = min(self.num_sections, img.shape[0] - 2)
+        num_to_remove = min(self.num_sections, depth - 2)
         indices_to_remove = self.R.choice(
-            np.arange(1, img.shape[0] - 1), size=num_to_remove, replace=False
+            np.arange(1, depth - 1), size=num_to_remove, replace=False
         )
 
         if is_tensor:
-            # Keep sections that are NOT in indices_to_remove
-            keep_mask = torch.ones(img.shape[0], dtype=torch.bool, device=img.device)
+            keep_mask = torch.ones(depth, dtype=torch.bool, device=img.device)
             keep_mask[indices_to_remove] = False
-            return img[keep_mask]
+            return torch.index_select(img, dim=depth_axis, index=keep_mask.nonzero(as_tuple=False).squeeze(-1))
         else:
-            return np.delete(img, indices_to_remove, axis=0)
+            return np.delete(img, indices_to_remove, axis=depth_axis)
 
 
 class RandMissingPartsd(RandomizableTransform, MapTransform):

connectomics/data/io/io.py

@@ -171,7 +171,7 @@ def read_image_as_volume(filename: str, drop_channel: bool = False) -> np.ndarra
     Raises:
         ValueError: If file format is not supported
     """
-    image_suffix = filename[filename.rfind(".") + 1:].lower()
+    image_suffix = filename[filename.rfind(".") + 1 :].lower()
     if image_suffix not in SUPPORTED_IMAGE_FORMATS:
         raise ValueError(
             f"Unsupported format: {image_suffix}. Supported formats: {SUPPORTED_IMAGE_FORMATS}"
@@ -281,7 +281,7 @@ def read_volume(
     if filename.endswith(".nii.gz"):
         image_suffix = "nii.gz"
     else:
-        image_suffix = filename[filename.rfind(".") + 1:].lower()
+        image_suffix = filename[filename.rfind(".") + 1 :].lower()
 
     if image_suffix in ["h5", "hdf5"]:
         data = read_hdf5(filename, dataset)
@@ -420,7 +420,7 @@ def get_vol_shape(filename: str, dataset: Optional[str] = None) -> tuple:
     if filename.endswith(".nii.gz"):
         image_suffix = "nii.gz"
     else:
-        image_suffix = filename[filename.rfind(".") + 1:].lower()
+        image_suffix = filename[filename.rfind(".") + 1 :].lower()
 
     if image_suffix in ["h5", "hdf5"]:
         # HDF5: Read shape from metadata (no data loading)

connectomics/data/io/tiles.py

@@ -160,22 +160,22 @@ def reconstruct_volume_from_tiles(
                     if is_image:  # Image data
                         result[
                             z - z0,
-                            y_actual_start - y0: y_actual_end - y0,
-                            x_actual_start - x0: x_actual_end - x0,
+                            y_actual_start - y0 : y_actual_end - y0,
+                            x_actual_start - x0 : x_actual_end - x0,
                         ] = patch[
-                            y_actual_start - y_patch_start: y_actual_end - y_patch_start,
-                            x_actual_start - x_patch_start: x_actual_end - x_patch_start,
+                            y_actual_start - y_patch_start : y_actual_end - y_patch_start,
+                            x_actual_start - x_patch_start : x_actual_end - x_patch_start,
                             0,
                         ]
                     else:  # Label data
                         result[
                             z - z0,
-                            y_actual_start - y0: y_actual_end - y0,
-                            x_actual_start - x0: x_actual_end - x0,
+                            y_actual_start - y0 : y_actual_end - y0,
+                            x_actual_start - x0 : x_actual_end - x0,
                         ] = rgb_to_seg(
                             patch[
-                                y_actual_start - y_patch_start: y_actual_end - y_patch_start,
-                                x_actual_start - x_patch_start: x_actual_end - x_patch_start,
+                                y_actual_start - y_patch_start : y_actual_end - y_patch_start,
+                                x_actual_start - x_patch_start : x_actual_end - x_patch_start,
                             ]
                         )
 

connectomics/data/process/crop.py

@@ -16,9 +16,9 @@ def crop_volume(data, sz, st=(0, 0, 0)):
     st = np.array(st).astype(np.int32)
 
     if data.ndim == 3:
-        return data[st[0]: st[0] + sz[0], st[1]: st[1] + sz[1], st[2]: st[2] + sz[2]]
+        return data[st[0] : st[0] + sz[0], st[1] : st[1] + sz[1], st[2] : st[2] + sz[2]]
     else:  # crop spatial dimensions
-        return data[:, st[0]: st[0] + sz[0], st[1]: st[1] + sz[1], st[2]: st[2] + sz[2]]
+        return data[:, st[0] : st[0] + sz[0], st[1] : st[1] + sz[1], st[2] : st[2] + sz[2]]
 
 
 def get_valid_pos_torch(mask, vol_sz, valid_ratio):
@@ -64,9 +64,9 @@ def get_valid_pos(mask, vol_sz, valid_ratio):
     if len(vol_sz) == 3:
         mask_sum = (
             mask_sum[
-                pad_sz_pre[0]: pad_sz_post[0],
-                pad_sz_pre[1]: pad_sz_post[1],
-                pad_sz_pre[2]: pad_sz_post[2],
+                pad_sz_pre[0] : pad_sz_post[0],
+                pad_sz_pre[1] : pad_sz_post[1],
+                pad_sz_pre[2] : pad_sz_post[2],
             ]
             >= valid_thres
         )
@@ -86,7 +86,7 @@ def get_valid_pos(mask, vol_sz, valid_ratio):
             )
     else:
         mask_sum = (
-            mask_sum[pad_sz_pre[0]: pad_sz_post[0], pad_sz_pre[1]: pad_sz_post[1]] >= valid_thres
+            mask_sum[pad_sz_pre[0] : pad_sz_post[0], pad_sz_pre[1] : pad_sz_post[1]] >= valid_thres
         )
         if mask_sum.max() > 0:
             yy, xx = np.meshgrid(np.arange(mask_sum.shape[0]), np.arange(mask_sum.shape[1]))