add missing param to auto-run inference

Author: Sllambias

asparagus/modules/lightning_modules/segmentation_module.py

@@ -219,12 +219,10 @@ def on_test_epoch_start(self):
     def test_step(self, batch, batch_idx):
         x = batch["image"]
 
-        logits = self.model.predict(
-            mode=self.inference_mode,
+        logits = self.model.sliding_window_predict(
             data=x,
             patch_size=fit_patch_size_to_image_size(self.inference_patch_size, list(x.shape[2:])),
             overlap=0.5,
-            sliding_window_prediction=True,
         )
 
         src_logits = reverse_preprocessing(logits, batch["properties"])
@@ -249,12 +247,10 @@ def on_test_epoch_end(self):
 
     def predict_step(self, batch, batch_idx):
         x = batch["image"]
-        logits = self.model.predict(
-            mode=self.inference_mode,
+        logits = self.model.sliding_window_predict(
             data=x,
             patch_size=self.inference_patch_size,
             overlap=0.5,
-            sliding_window_prediction=True,
         )
         logits = reverse_preprocessing(
             array=logits,

asparagus/pipeline/run/test_seg.py

@@ -51,7 +51,6 @@ def main(cfg: DictConfig) -> None:
         ckpt_cfg.lightning._lightning_module,
         model=model,
         weights=path_store.ckpt_path,
-        inference_mode=ckpt_cfg.model.dimensions,
         inference_patch_size=ckpt_cfg.training.patch_size,
         test_output_path=output_path,
     )

asparagus/pipeline/run/train_seg.py

@@ -115,6 +115,7 @@ def main(cfg: DictConfig) -> None:
         optimizer=cfg.model.train_optim,
         learning_rate=cfg.model.train_lr,
         deep_supervision=cfg.model.deep_supervision,
+        inference_patch_size=cfg.training.patch_size,
         test_output_path=os.path.join(
             path_store.run_dir,
             "predictions",

configs/projects/development/runs_train/DEBUG_SEG.yaml

@@ -0,0 +1,30 @@
+# @package _global_
+defaults:
+  - /core/base@
+  - /default_train_seg@
+  - /hardware/cpu@hardware
+  - /model/unet_tiny@model
+  - _self_
+
+task: 
+root: base
+stem: debug
+
+model:
+  dimensions: 3D
+
+training:
+  patch_size: [32, 32, 32]
+  batch_size: 2
+  train_batches_per_epoch_per_device: 10
+  val_batches_per_epoch_per_device: 5
+  epochs: 5
+
+logger:
+  progress_bar: True
+  profile: False
+  wandb_log_model: False
+  wandb_logging: False
+  mlflow_logging: False
+  log_every_n_steps: 1
+  log_images_every_n_epoch: 999999

pyproject.toml

@@ -16,7 +16,7 @@ keywords = ['deep learning', 'medical image analysis','foundation models']
 
 
 dependencies = [
-    "gardening_tools>=0.1.1",
+    "gardening_tools>=0.2.0",
     "lightning==2.4.0",
     "nibabel>=5.3.2",
     "numpy>=1.23.1",