Parallel inference

.gitignore

@@ -1,4 +1,5 @@
 slurm-*.out
+*.zarr
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

BANIS.py

@@ -1,13 +1,13 @@
 import argparse
 import gc
 import os
+import shutil
+
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
 from collections import defaultdict
 from datetime import datetime
 from typing import Any, Dict
-import random
 
-import numpy as np
 import pytorch_lightning as pl
 import torch
 import torchvision
@@ -23,7 +23,7 @@
 from tqdm import tqdm
 
 from data import load_data
-from inference import scale_sigmoid, predict_aff, compute_connected_component_segmentation
+from inference import AffinityPredictor, Thresholding
 from metrics import compute_metrics
 
 
@@ -35,7 +35,7 @@ class BANIS(LightningModule):
     def __init__(self, **kwargs: Any):
         super().__init__()
         self.save_hyperparameters()
-        print(f"hparams: \n{self.hparams}")
+        # print(f"hparams: \n{self.hparams}")
 
         self.model = create_mednext_v1(
             num_input_channels=self.hparams.num_input_channels,
@@ -163,8 +163,17 @@ def full_cube_inference(self, mode: str, global_step=None):
 
         img_data = zarr.open(os.path.join(seed_path, "data.zarr"), mode="r")["img"]
 
-        aff_pred = predict_aff(img_data, model=self, zarr_path=f"{self.hparams.save_dir}/pred_aff_{mode}.zarr", do_overlap=True, prediction_channels=3, divide=255,
-                                     small_size=self.hparams.small_size, compute_backend="local")
+        affinity_predictor = AffinityPredictor(
+            chunk_cube_size=3000,  # can be adjusted
+            compute_backend="local",
+            model=self,
+            small_size=self.hparams.small_size,
+            do_overlap=True,
+            prediction_channels=3,
+            divide=255,
+        )
+        affinity_predictor.img_to_aff(img_data, zarr_path=f"{self.hparams.save_dir}/pred_aff_{mode}.zarr")
+        aff_pred = zarr.open(f"{self.hparams.save_dir}/pred_aff_{mode}.zarr", mode="r")
 
         self._evaluate_thresholds(aff_pred, os.path.join(seed_path, "skeleton.pkl"), mode, global_step)
 
@@ -179,9 +188,9 @@ def _evaluate_thresholds(self, aff_pred: zarr.Array, skel_path: str, mode: str,
             torch.cuda.empty_cache()
             print(f"threshold {thr}")
 
-            pred_seg = compute_connected_component_segmentation(
-                aff_pred[:3] > thr  # hard affinities
-            )
+            postprocessor = Thresholding(3000, "local", thr)
+            postprocessor.aff_to_seg(aff_pred, f"{self.hparams.save_dir}/pred_seg_{mode}_tmp.zarr")
+            pred_seg = zarr.open(f"{self.hparams.save_dir}/pred_seg_{mode}_tmp.zarr", mode="r")
 
             metrics = compute_metrics(pred_seg, skel_path)
 
@@ -201,9 +210,11 @@ def _evaluate_thresholds(self, aff_pred: zarr.Array, skel_path: str, mode: str,
                     self.best_thr_so_far[mode] = thr
                     with open(f"{self.hparams.save_dir}/best_thr_{mode}.txt", "w") as f:
                         f.write(str(self.best_thr_so_far[mode]))
-                    seg_pred = zarr.array(pred_seg, dtype=np.uint32,
-                                          store=f"{self.hparams.save_dir}/pred_seg_{mode}.zarr",
-                                          chunks=(512, 512, 512), overwrite=True)
+                    if os.path.exists(f"{self.hparams.save_dir}/pred_seg_{mode}.zarr"):
+                        shutil.rmtree(f"{self.hparams.save_dir}/pred_seg_{mode}.zarr")
+                    os.replace(f"{self.hparams.save_dir}/pred_seg_{mode}_tmp.zarr", f"{self.hparams.save_dir}/pred_seg_{mode}.zarr")
+            else:
+                shutil.rmtree(f"{self.hparams.save_dir}/pred_seg_{mode}_tmp.zarr")
             best_voi = min(best_voi, metrics["voi_sum"])
 
         self.safe_add_scalar(f"{mode}_best_nerl", best_nerl, global_step)

README.md

@@ -38,6 +38,18 @@ python slurm_job_scheduler.py
 
 Adding an `auto_resubmit` argument to `config.yaml` allows Slurm to automatically resubmit jobs that reach the Slurm time limit (see `aff_train.sh`).
 
+## Prediction
+
+To predict segmentation from an image:
+
+```bash
+python inference --img_path /path/to/image.zarr --model_path /path/to/model.ckpt --chunk_cube_size 3000
+```
+
+The `chunk_cube_size` parameter sets the maximum cube size that can be loaded in memory.
+If you have enough memory available, set it to a bigger value, if you are tight with memory, set a lower value (in exchange for increased computation time).
+See [inference.py](inference.py) for other parameters.
+
 ## Evaluation
 
 To evaluate a predicted segmentation (`.zarr` or `.npy`):

aff_train.sh

@@ -1,14 +1,14 @@
 #!/bin/bash -l
 
-#SBATCH --nodes=2
-#SBATCH --gres=gpu:4
-#SBATCH --ntasks-per-node=4
+#SBATCH --nodes=1
+#SBATCH --gres=gpu:1
+#SBATCH --ntasks-per-node=1
 #SBATCH --time=7-00
-#SBATCH --cpus-per-task=16
-#SBATCH --mem=1000G
+#SBATCH --cpus-per-task=32
+#SBATCH --mem=500G
 #SBATCH --signal=B:USR1@300
 #SBATCH --open-mode=append
-#SBATCH --partition=p.large
+#SBATCH --partition=p.share
 
 mamba activate nisb
 

config.yaml

@@ -5,20 +5,20 @@ params:
     - 1e-2
   seed:
     - 0
-    #- 1
-    #- 2
-    #- 3
-    #- 4
+    - 1
+    - 2
+    - 3
+    - 4
   long_range:
     - 10
   batch_size:
-    - 1
+    - 8
   scheduler:
     - true
   model_id:
-    - "L"
+    - "S"
   kernel_size:
-    - 5
+    - 3
   synthetic:
     - 1.0
   drop_slice_prob:
@@ -32,33 +32,34 @@ params:
   affine:
     - 0.5
   n_steps:
-    - 1_000_000
+    - 50000
   small_size:
-    - 256
+    - 128
   data_setting:
-    #- "base"
-    #- "liconn"
-    #- "multichannel"
-    #- "neg_guidance"
-    #- "no_touch_thick"
-    #- "pos_guidance"
-    #- "slice_perturbed"
-    #- "touching_thin"
+    - "base"
+    - "liconn"
+    - "multichannel"
+    - "neg_guidance"
+    - "no_touch_thick"
+    - "pos_guidance"
+    - "slice_perturbed"
+    - "touching_thin"
     - "train_100"
   base_data_path:
     - "/cajal/nvmescratch/projects/NISB/"
   save_path:
-    #- "/cajal/scratch/projects/misc/riegerfr/aff_nis/"
-    - "/cajal/scratch/projects/misc/zuzur/xl_banis"
+    - "/cajal/scratch/projects/misc/riegerfr/aff_nis/"
   exp_name:
-    - "xl_test"
+    - "exp"
   real_data_path: #https://colab.research.google.com/github/funkelab/lsd/blob/master/lsd/tutorial/notebooks/lsd_data_download.ipynb
     - "/cajal/scratch/projects/misc/mdraw/data/funke/zebrafinch/training/"
   auto_resubmit:
-    - True
+    - False
   distributed:
-    - True
-  compile:
     - False
+  compile:
+    - True
   validate_extern:
-    - True
+    - True
+  augment:
+    - True

environment.yaml

@@ -8,6 +8,7 @@ dependencies:
   - bzip2=1.0.8
   - ca-certificates=2024.8.30
   - cython==3.0.11
+  - dask=2025.7.0
   - ld_impl_linux-64=2.43
   - libexpat=2.6.3
   - libffi=3.4.2
@@ -40,9 +41,11 @@ dependencies:
       - batchgenerators==0.25
       - certifi==2024.8.30
       - charset-normalizer==3.4.0
+      - cloud_volume==12.4.1
       - connected-components-3d==3.19.0
       - contourpy==1.3.0
       - cycler==0.12.1
+      - dask_jobqueue==0.9.0
       - dicom2nifti==2.5.0
       - fasteners==0.19
       - filelock==3.16.1
@@ -68,7 +71,9 @@ dependencies:
       - monai==1.3.2
       - mpmath==1.3.0
       - multidict==6.1.0
+      - mwatershed==0.5.3
       - networkx==3.3
+      - neuroglancer==2.40.1
       - nibabel==5.3.0
       - numba==0.60.0
       - numcodecs==0.13.1