fix formatting and e2e testing error

Author: Donglai Wei

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

@@ -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

@@ -355,9 +355,7 @@ def _objective(self, trial: optuna.Trial) -> float:
 
                     print(
                         "\n❌ Trial {count} failed during post-processing "
-                        "(volume {vol_idx}):".format(
-                            count=self.trial_count, vol_idx=vol_idx
-                        )
+                        "(volume {vol_idx}):".format(count=self.trial_count, vol_idx=vol_idx)
                     )
                     print(f"   Parameters: {postproc_params}")
                     print(f"   Error: {e}")
@@ -381,9 +379,7 @@ def _objective(self, trial: optuna.Trial) -> float:
 
                 print(
                     "\n❌ Trial {count} failed during metric computation "
-                    "(volume {vol_idx}):".format(
-                        count=self.trial_count, vol_idx=vol_idx
-                    )
+                    "(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}")

connectomics/decoding/postprocess.py

@@ -360,7 +360,7 @@ def apply_binary_postprocessing(
 
             if len(sizes) > cc_config.top_k:
                 # Get indices of top-k largest components
-                top_k_indices = np.argsort(sizes)[-cc_config.top_k:]
+                top_k_indices = np.argsort(sizes)[-cc_config.top_k :]
                 top_k_labels = label_ids[top_k_indices]
 
                 # Create mask keeping only top-k

connectomics/inference/tta.py

@@ -117,7 +117,7 @@ def apply_preprocessing(self, tensor: torch.Tensor) -> torch.Tensor:
                 pass  # Keep all channels
             elif isinstance(tta_channel, int):
                 if tta_channel != -1:
-                    tensor = tensor[:, tta_channel: tta_channel + 1, ...]
+                    tensor = tensor[:, tta_channel : tta_channel + 1, ...]
             elif isinstance(tta_channel, (list, tuple, Sequence)):
                 # Convert to list of integers (handle both int and string numbers
                 # from OmegaConf)

connectomics/models/build.py

@@ -65,7 +65,7 @@ def load_external_weights(model, cfg):
         stripped_count = 0
         for key, value in state_dict.items():
             if key.startswith(key_prefix):
-                new_key = key[len(key_prefix):]
+                new_key = key[len(key_prefix) :]
                 new_state_dict[new_key] = value
                 stripped_count += 1
             else:

connectomics/training/deep_supervision.py

@@ -123,7 +123,7 @@ def compute_multitask_loss(
             num_label_channels = 1
 
             # Extract label channels
-            task_label = labels[:, label_ch_offset: label_ch_offset + num_label_channels, ...]
+            task_label = labels[:, label_ch_offset : label_ch_offset + num_label_channels, ...]
             label_ch_offset += num_label_channels
 
             # Apply specified losses for this task
@@ -144,9 +144,7 @@ def compute_multitask_loss(
                     print(f"Loss value: {loss.item()}")
                     output_range = f"[{task_output.min():.4f}, {task_output.max():.4f}]"
                     label_range = f"[{task_label.min():.4f}, {task_label.max():.4f}]"
-                    print(
-                        f"Output shape: {task_output.shape}, range: {output_range}"
-                    )
+                    print(f"Output shape: {task_output.shape}, range: {output_range}")
                     print(f"Label shape: {task_label.shape}, range: {label_range}")
                     print(f"Output contains NaN: {torch.isnan(task_output).any()}")
                     print(f"Label contains NaN: {torch.isnan(task_label).any()}")
@@ -242,9 +240,7 @@ def compute_loss_for_scale(
                         print(f"Loss value: {loss.item()}")
                         output_range = f"[{task_output.min():.4f}, {task_output.max():.4f}]"
                         target_range = f"[{task_target.min():.4f}, {task_target.max():.4f}]"
-                        print(
-                            f"Output shape: {task_output.shape}, range: {output_range}"
-                        )
+                        print(f"Output shape: {task_output.shape}, range: {output_range}")
                         print(f"Target shape: {task_target.shape}, range: {target_range}")
                         if self.debug_on_nan:
                             print("\nEntering debugger...")
@@ -282,9 +278,7 @@ def compute_loss_for_scale(
                     print(f"Scale: {scale_idx}, Weight: {weight}")
                     out_range = f"[{output.min():.4f}, {output.max():.4f}]"
                     tgt_range = f"[{target.min():.4f}, {target.max():.4f}]"
-                    print(
-                        f"Output shape: {output.shape}, range: {out_range}"
-                    )
+                    print(f"Output shape: {output.shape}, range: {out_range}")
                     print(f"Target shape: {target.shape}, range: {tgt_range}")
                     print(f"Output contains NaN: {torch.isnan(output).any()}")
                     print(f"Target contains NaN: {torch.isnan(target).any()}")
@@ -394,9 +388,7 @@ def compute_standard_loss(
                     print(f"Loss index: {i}, Weight: {weight}")
                     out_range = f"[{outputs.min():.4f}, {outputs.max():.4f}]"
                     label_range = f"[{labels.min():.4f}, {labels.max():.4f}]"
-                    print(
-                        f"Output shape: {outputs.shape}, range: {out_range}"
-                    )
+                    print(f"Output shape: {outputs.shape}, range: {out_range}")
                     print(f"Label shape: {labels.shape}, range: {label_range}")
                     print(f"Output contains NaN: {torch.isnan(outputs).any()}")
                     print(f"Label contains NaN: {torch.isnan(labels).any()}")

connectomics/training/lit/config.py

@@ -137,9 +137,7 @@ def create_datamodule(
                 try:
                     size_mb = DATASETS[dataset_name]["size_mb"]
                     prompt = f"   Download {dataset_name} dataset (~{size_mb} MB)? [Y/n]: "
-                    response = (
-                        input(prompt).strip().lower()
-                    )
+                    response = input(prompt).strip().lower()
                     if response in ["", "y", "yes"]:
                         if download_dataset(dataset_name, base_dir=PathLib.cwd()):
                             print("✅ Data downloaded successfully!")