fix build error and formatting

Author: Donglai Wei

connectomics/config/hydra_config.py

@@ -1100,7 +1100,7 @@ class TestDataConfig:
     cache_suffix: str = "_prediction.h5"
     # Image transformation (applied to test images during inference)
     image_transform: ImageTransformConfig = field(default_factory=ImageTransformConfig)
-    # Label transformation (optional, typically not used for test mode to preserve raw labels for evaluation)
+    # Label transformation (optional). Typically unused in test mode to preserve raw labels
     label_transform: Optional[LabelTransformConfig] = None
 
 
@@ -1125,7 +1125,7 @@ class TuneDataConfig:
     tune_resolution: Optional[List[int]] = None
     # Image transformation (applied to tune images during inference)
     image_transform: ImageTransformConfig = field(default_factory=ImageTransformConfig)
-    # Label transformation (optional, typically not used for tune mode to preserve raw labels for evaluation)
+    # Label transformation (optional). Typically unused in tune mode to preserve raw labels
     label_transform: Optional[LabelTransformConfig] = None
 
 
@@ -1209,36 +1209,6 @@ class TuneConfig:
     parameter_space: ParameterSpaceConfig = field(default_factory=ParameterSpaceConfig)
 
 
-# Allow safe loading of checkpoints with PyTorch 2.6+ weights_only defaults
-try:
-    import torch
-
-    if hasattr(torch, "serialization") and hasattr(torch.serialization, "add_safe_globals"):
-        torch.serialization.add_safe_globals(
-            [
-                ParameterConfig,
-                DecodingParameterSpace,
-                PostprocessingParameterSpace,
-                ParameterSpaceConfig,
-                # Core config dataclasses (for Lightning checkpoints)
-                Config,
-                SystemConfig,
-                SystemTrainingConfig,
-                SystemInferenceConfig,
-                ModelConfig,
-                DataConfig,
-                OptimizationConfig,
-                MonitorConfig,
-                InferenceConfig,
-                TestConfig,
-                TuneConfig,
-            ]
-        )
-except Exception:
-    # Best-effort registration; ignore if torch not available at import time
-    pass
-
-
 @dataclass
 class Config:
     """Main configuration for PyTorch Connectomics.
@@ -1289,6 +1259,36 @@ class Config:
     tune: Optional[TuneConfig] = None
 
 
+# Allow safe loading of checkpoints with PyTorch 2.6+ weights_only defaults
+try:
+    import torch
+
+    if hasattr(torch, "serialization") and hasattr(torch.serialization, "add_safe_globals"):
+        torch.serialization.add_safe_globals(
+            [
+                ParameterConfig,
+                DecodingParameterSpace,
+                PostprocessingParameterSpace,
+                ParameterSpaceConfig,
+                # Core config dataclasses (for Lightning checkpoints)
+                Config,
+                SystemConfig,
+                SystemTrainingConfig,
+                SystemInferenceConfig,
+                ModelConfig,
+                DataConfig,
+                OptimizationConfig,
+                MonitorConfig,
+                InferenceConfig,
+                TestConfig,
+                TuneConfig,
+            ]
+        )
+except Exception:
+    # Best-effort registration; ignore if torch not available at import time
+    pass
+
+
 # Utility functions for common configuration tasks
 def configure_edge_mode(
     cfg: Config, mode: str = "seg-all", thickness: int = 1, processing_mode: str = "3d"

connectomics/data/augment/build.py

@@ -408,7 +408,8 @@ def _build_eval_transforms_impl(cfg: Config, mode: str = "val", keys: list[str]
             transforms.append(NormalizeLabelsd(keys=["label"]))
 
         # Label transformations (affinity, distance transform, etc.)
-        # For test/tune modes: NEVER apply label transforms (keep raw instance labels for evaluation)
+        # For test/tune modes: NEVER apply label transforms
+        # (keep raw instance labels for evaluation)
         # For val mode: use training label_transform config
         label_cfg = None
         if mode == "val":

connectomics/data/augment/monai_transforms.py

@@ -99,11 +99,11 @@ def _apply_misalignment_translation(
 
         output = np.zeros(out_shape, img.dtype)
         if mode == "slip":
-            output = img[:, y0 : y0 + out_shape[1], x0 : x0 + out_shape[2]]
-            output[idx] = img[idx, y1 : y1 + out_shape[1], x1 : x1 + out_shape[2]]
+            output = img[:, y0:y0 + out_shape[1], x0:x0 + out_shape[2]]
+            output[idx] = img[idx, y1:y1 + out_shape[1], x1:x1 + out_shape[2]]
         else:
-            output[:idx] = img[:idx, y0 : y0 + out_shape[1], x0 : x0 + out_shape[2]]
-            output[idx:] = img[idx:, y1 : y1 + out_shape[1], x1 : x1 + out_shape[2]]
+            output[:idx] = img[:idx, y0:y0 + out_shape[1], x0:x0 + out_shape[2]]
+            output[idx:] = img[idx:, y1:y1 + out_shape[1], x1:x1 + out_shape[2]]
 
         if is_tensor:
             output = torch.from_numpy(output).to(device)
@@ -299,7 +299,7 @@ def _apply_missing_parts(
         x_start = self.R.randint(0, img.shape[2] - hole_w + 1)
 
         # Create hole (set to 0 or mean value)
-        img[section_idx, y_start : y_start + hole_h, x_start : x_start + hole_w] = 0
+        img[section_idx, y_start: y_start + hole_h, x_start: x_start + hole_w] = 0
 
         return img
 
@@ -452,24 +452,24 @@ def _apply_cut_noise(
                 noise = self.R.uniform(-self.noise_scale, self.noise_scale, noise_shape)
                 region = img[
                     :,
-                    z_start : z_start + z_len,
-                    y_start : y_start + y_len,
-                    x_start : x_start + x_len,
+                    z_start: z_start + z_len,
+                    y_start: y_start + y_len,
+                    x_start: x_start + x_len,
                 ]
                 noisy_region = np.clip(region + noise, 0, 1)
                 img[
                     :,
-                    z_start : z_start + z_len,
-                    y_start : y_start + y_len,
-                    x_start : x_start + x_len,
+                    z_start: z_start + z_len,
+                    y_start: y_start + y_len,
+                    x_start: x_start + x_len,
                 ] = noisy_region
             else:
                 # (C, H, W) - 2D with channels
                 noise_shape = (img.shape[0], y_len, x_len)
                 noise = self.R.uniform(-self.noise_scale, self.noise_scale, noise_shape)
-                region = img[:, y_start : y_start + y_len, x_start : x_start + x_len]
+                region = img[:, y_start: y_start + y_len, x_start: x_start + x_len]
                 noisy_region = np.clip(region + noise, 0, 1)
-                img[:, y_start : y_start + y_len, x_start : x_start + x_len] = noisy_region
+                img[:, y_start: y_start + y_len, x_start: x_start + x_len] = noisy_region
         elif img.ndim == 3:
             # 3D case: (Z, Y, X) or (C, H, W)
             # Heuristic: if first dim is small (<=4), assume it's channel (2D with channels)
@@ -478,29 +478,29 @@ def _apply_cut_noise(
                 # (C, H, W) - 2D with channels
                 noise_shape = (img.shape[0], y_len, x_len)
                 noise = self.R.uniform(-self.noise_scale, self.noise_scale, noise_shape)
-                region = img[:, y_start : y_start + y_len, x_start : x_start + x_len]
+                region = img[:, y_start: y_start + y_len, x_start: x_start + x_len]
                 noisy_region = np.clip(region + noise, 0, 1)
-                img[:, y_start : y_start + y_len, x_start : x_start + x_len] = noisy_region
+                img[:, y_start: y_start + y_len, x_start: x_start + x_len] = noisy_region
             else:
                 # (Z, Y, X) - 3D
                 z_len = max(1, int(self.length_ratio * img.shape[0]))  # Ensure at least 1
                 z_start = self.R.randint(0, max(1, img.shape[0] - z_len + 1))
                 noise_shape = (z_len, y_len, x_len)
                 noise = self.R.uniform(-self.noise_scale, self.noise_scale, noise_shape)
                 region = img[
-                    z_start : z_start + z_len, y_start : y_start + y_len, x_start : x_start + x_len
+                    z_start: z_start + z_len, y_start: y_start + y_len, x_start: x_start + x_len
                 ]
                 noisy_region = np.clip(region + noise, 0, 1)
                 img[
-                    z_start : z_start + z_len, y_start : y_start + y_len, x_start : x_start + x_len
+                    z_start: z_start + z_len, y_start: y_start + y_len, x_start: x_start + x_len
                 ] = noisy_region
         else:
             # 2D case: (H, W)
             noise_shape = (y_len, x_len)
             noise = self.R.uniform(-self.noise_scale, self.noise_scale, noise_shape)
-            region = img[y_start : y_start + y_len, x_start : x_start + x_len]
+            region = img[y_start: y_start + y_len, x_start: x_start + x_len]
             noisy_region = np.clip(region + noise, 0, 1)
-            img[y_start : y_start + y_len, x_start : x_start + x_len] = noisy_region
+            img[y_start: y_start + y_len, x_start: x_start + x_len] = noisy_region
 
         if is_tensor:
             img = torch.from_numpy(img).to(device)
@@ -886,7 +886,7 @@ def _find_best_paste(
                 neuron_tensor.flip(0) if neuron_tensor.ndim == 3 else neuron_tensor.flip(1)
             )
 
-        label_paste = labels[best_idx : best_idx + 1]
+        label_paste = labels[best_idx: best_idx + 1]
 
         if best_angle != 0:
             label_paste = self._rotate_3d(label_paste, best_angle)

connectomics/data/io/io.py

@@ -37,7 +37,7 @@ def read_hdf5(
     Args:
         filename: Path to the HDF5 file
         dataset: Name of the dataset to read. If None, reads the first dataset
-        slice_obj: Optional slice for partial loading (e.g., np.s_[0:10, :, :])
+        slice_obj: Optional slice for partial loading (e.g., np.s_[0:10,:,:])
 
     Returns:
         Data from the HDF5 file as numpy array
@@ -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]))

connectomics/decoding/optuna_tuner.py

@@ -175,7 +175,10 @@ def _validate_data(self):
                     )
                 if mask.shape != gt.shape:
                     raise ValueError(
-                        f"Mask shape {mask.shape} doesn't match ground truth shape {gt.shape} in volume {i}"
+                        "Mask shape {mask_shape} doesn't match ground truth shape "
+                        "{gt_shape} in volume {vol_idx}".format(
+                            mask_shape=mask.shape, gt_shape=gt.shape, vol_idx=i
+                        )
                     )
 
     def optimize(self) -> optuna.Study:
@@ -351,7 +354,10 @@ def _objective(self, trial: optuna.Trial) -> float:
                     import traceback
 
                     print(
-                        f"\n❌ Trial {self.trial_count} failed during post-processing (volume {vol_idx}):"
+                        "\n❌ Trial {count} failed during post-processing "
+                        "(volume {vol_idx}):".format(
+                            count=self.trial_count, vol_idx=vol_idx
+                        )
                     )
                     print(f"   Parameters: {postproc_params}")
                     print(f"   Error: {e}")
@@ -374,7 +380,10 @@ def _objective(self, trial: optuna.Trial) -> float:
                 import traceback
 
                 print(
-                    f"\n❌ Trial {self.trial_count} failed during metric computation (volume {vol_idx}):"
+                    "\n❌ Trial {count} failed during metric computation "
+                    "(volume {vol_idx}):".format(
+                        count=self.trial_count, vol_idx=vol_idx
+                    )
                 )
                 print(f"   Metric: {metric_name}")
                 print(f"   Segmentation shape: {segmentation.shape}, dtype: {segmentation.dtype}")
@@ -396,7 +405,13 @@ def _objective(self, trial: optuna.Trial) -> float:
             # Show per-volume and average
             vol_str = ", ".join([f"{m:.4f}" for m in volume_metrics])
             print(
-                f"Trial {self.trial_count:3d}: {metric_name}=[{vol_str}] avg={metric_value:.6f} ({direction})"
+                "Trial {count:3d}: {metric}=[{volumes}] avg={avg:.6f} ({direction})".format(
+                    count=self.trial_count,
+                    metric=metric_name,
+                    volumes=vol_str,
+                    avg=metric_value,
+                    direction=direction,
+                )
             )
 
         return metric_value
@@ -778,9 +793,8 @@ def run_tuning(model, trainer, cfg, checkpoint_path=None):
 
     print(f"✓ Loaded {len(predictions)} prediction volume(s)")
     for i, pred in enumerate(predictions):
-        print(
-            f"  Volume {i}: shape {pred.shape}, dtype {pred.dtype}, range [{pred.min():.3f}, {pred.max():.3f}]"
-        )
+        pred_range = f"[{pred.min():.3f}, {pred.max():.3f}]"
+        print(f"  Volume {i}: shape {pred.shape}, dtype {pred.dtype}, range {pred_range}")
 
     # Step 3: Load ground truth
     print("\n[3/4] Loading ground truth labels...")
@@ -816,8 +830,10 @@ def run_tuning(model, trainer, cfg, checkpoint_path=None):
     for i, gt in enumerate(ground_truth):
         unique_labels = np.unique(gt)
         n_nonzero_labels = len(unique_labels) - (1 if 0 in unique_labels else 0)
+        gt_range = f"[{gt.min()}, {gt.max()}]"
         print(
-            f"  Volume {i}: shape {gt.shape}, dtype {gt.dtype}, range [{gt.min()}, {gt.max()}], unique labels: {n_nonzero_labels}"
+            f"  Volume {i}: shape {gt.shape}, dtype {gt.dtype}, range {gt_range}, "
+            f"unique labels: {n_nonzero_labels}"
         )
 
         # Validate ground truth for this volume