fix formatting

Donglai Wei · Donglai Wei · commit 18ff57797b0a · 2025-12-10T00:40:18.000-05:00
diff --git a/connectomics/data/augment/monai_transforms.py b/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,31 @@ 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 +888,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)
diff --git a/connectomics/data/io/io.py b/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)
diff --git a/connectomics/data/io/tiles.py b/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,
                             ]
                         )
 
diff --git a/connectomics/data/process/crop.py b/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):
@@ -62,14 +62,11 @@ def get_valid_pos(mask, vol_sz, valid_ratio):
     pad_sz_post = data_sz - (vol_sz - pad_sz_pre - 1)
     valid_pos = np.zeros([0, 3])
     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],
-            ]
-            >= valid_thres
-        )
+        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],
+        ] >= valid_thres
         if mask_sum.max() > 0:
             zz, yy, xx = np.meshgrid(
                 np.arange(mask_sum.shape[0]),
@@ -86,7 +83,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]))
diff --git a/connectomics/decoding/postprocess.py b/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
diff --git a/connectomics/inference/tta.py b/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)
diff --git a/connectomics/models/build.py b/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:
diff --git a/connectomics/training/deep_supervision.py b/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
diff --git a/connectomics/training/lit/model.py b/connectomics/training/lit/model.py
@@ -426,8 +426,6 @@ def _compute_test_metrics(
         ):
             # Adapted Rand requires instance segmentation labels (integer labels), not binary
             # decoded_predictions should already be instance segmentation from decode_instance_*
-            # Check original shape before processing to handle batch dimension correctly
-            original_shape = decoded_predictions.shape
             pred_instance = torch.from_numpy(decoded_predictions).long()
 
             # Labels should also be instance segmentation (integer labels)
@@ -438,8 +436,10 @@ def _compute_test_metrics(
             while labels_instance.ndim > 3 and labels_instance.shape[0] == 1:
                 labels_instance = labels_instance.squeeze(0)
 
-            # Squeeze all leading dimensions of size 1 from predictions (remove batch & channel dims)
-            # Predictions can be: (B, C, Z, H, W), (B, Z, H, W), (C, Z, H, W), or (Z, H, W)
+            # Squeeze all leading dimensions of size 1 from predictions
+            # (remove batch & channel dims)
+            # Predictions can be: (B, C, Z, H, W), (B, Z, H, W),
+            # (C, Z, H, W), or (Z, H, W)
             while pred_instance.ndim > 3 and pred_instance.shape[0] == 1:
                 pred_instance = pred_instance.squeeze(0)
 
diff --git a/connectomics/training/lit/utils.py b/connectomics/training/lit/utils.py
@@ -195,9 +195,8 @@ def setup_config(args) -> Config:
         print("🔧 Fast-dev-run mode: Overriding config for debugging")
         print(f"   - num_gpus: {cfg.system.training.num_gpus} → 1")
         print(f"   - num_cpus: {cfg.system.training.num_cpus} → 1")
-        print(
-            f"   - num_workers: {cfg.system.training.num_workers} → 0 (avoid multiprocessing in debug mode)"
-        )
+        print(f"   - num_workers: {cfg.system.training.num_workers} → 0 "
+              "(avoid multiprocessing in debug mode)")
         print(
             f"   - batch_size: Controlled by PyTorch Lightning (--fast-dev-run={args.fast_dev_run})"
         )
diff --git a/connectomics/utils/visualizer.py b/connectomics/utils/visualizer.py
@@ -303,7 +303,7 @@ def _log_multi_channel_viz(
 
         # Show each output channel
         for i in range(min(output.shape[1], 12)):  # Increased limit to 12 channels
-            channel_img = output[:, i : i + 1].repeat(1, 3, 1, 1)  # Convert to RGB
+            channel_img = output[:, i:i + 1].repeat(1, 3, 1, 1)  # Convert to RGB
             writer.add_image(
                 f"{prefix}/output_channel_{i}",
                 vutils.make_grid(
@@ -317,7 +317,7 @@ def _log_multi_channel_viz(
 
         # Show each label channel
         for i in range(min(label.shape[1], 12)):  # Increased limit to 12 channels
-            channel_img = label[:, i : i + 1].repeat(1, 3, 1, 1)  # Convert to RGB
+            channel_img = label[:, i:i + 1].repeat(1, 3, 1, 1)  # Convert to RGB
             writer.add_image(
                 f"{prefix}/label_channel_{i}",
                 vutils.make_grid(
@@ -333,7 +333,7 @@ def _log_multi_channel_viz(
         if mask is not None and mask.numel() > 0:
             for i in range(min(mask.shape[1], 12)):  # Show up to 12 mask channels
                 # Show mask in cyan for better visibility
-                mask_channel = mask[:, i : i + 1]
+                mask_channel = mask[:, i:i + 1]
                 mask_rgb = torch.cat(
                     [
                         torch.zeros_like(mask_channel),  # R=0
@@ -393,7 +393,7 @@ def _normalize(self, tensor: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
             # Multi-channel: normalize each channel independently
             normalized = []
             for c in range(tensor.shape[1]):
-                channel = tensor[:, c : c + 1]
+                channel = tensor[:, c:c + 1]
                 min_val = channel.min()
                 max_val = channel.max()
                 if max_val > min_val:
diff --git a/justfile b/justfile
@@ -98,12 +98,17 @@ tensorboard-run experiment timestamp port='6006':
 
 # Launch any just command on SLURM (e.g., just slurm weilab 8 4 "train lucchi")
 # Automatically uses srun for distributed training when num_gpu > 1
-slurm partition num_cpu num_gpu cmd:
+slurm partition num_cpu num_gpu cmd constraint='':
     #!/usr/bin/env bash
     # Configure for multi-GPU training with PyTorch Lightning DDP
     # Set ntasks=num_gpu and use srun to launch DDP processes
     # SLURM will set CUDA_VISIBLE_DEVICES for each task automatically
 
+    constraint_flag=""
+    if [ -n "{{constraint}}" ]; then
+        constraint_flag="--constraint={{constraint}}"
+    fi
+
     sbatch --job-name="pytc_{{cmd}}" \
            --partition={{partition}} \
            --output=slurm_outputs/slurm-%j.out \
@@ -114,6 +119,7 @@ slurm partition num_cpu num_gpu cmd:
            --cpus-per-task={{num_cpu}} \
            --mem=32G \
            --time=48:00:00 \
+           $constraint_flag \
            --wrap="mkdir -p \$HOME/.just && export JUST_TEMPDIR=\$HOME/.just TMPDIR=\$HOME/.just && source /projects/weilab/weidf/lib/miniconda3/bin/activate pytc && cd $PWD && srun --ntasks={{num_gpu}} --ntasks-per-node={{num_gpu}} just {{cmd}}"
 
 # Launch parameter sweep from config (e.g., just sweep tutorials/sweep_example.yaml)
diff --git a/tests/unit/test_lit_utils.py b/tests/unit/test_lit_utils.py
@@ -44,7 +44,7 @@ def test_setup_config_applies_overrides_and_fast_dev_run(tmp_path):
     updated = setup_config(args)
 
     expected_dir = f"outputs/{cfg_path.stem}_{cfg.model.architecture}/checkpoints"
-    assert updated.monitor.checkpoint.dirpath == expected_dir
+    assert Path(updated.monitor.checkpoint.dirpath).as_posix() == expected_dir
     assert updated.optimization.optimizer.lr == 0.01
     assert updated.optimization.max_epochs == 5
     assert updated.system.training.batch_size == 2
diff --git a/tutorials/cellmap_cos7.yaml b/tutorials/cellmap_cos7.yaml
@@ -24,10 +24,10 @@ description: CellMap COS7 multi-organelle segmentation with MedNeXt
 # System
 system:
   training:
-    num_gpus: 1
-    num_cpus: 1
-    num_workers: 1
-    batch_size: 2                      # Per GPU batch size
+    num_gpus: 4
+    num_cpus: 8
+    num_workers: 8
+    batch_size: 8                      # Per GPU batch size
   inference:
     num_gpus: 1
     num_cpus: 1
