From 24d3b6b277d1050d93021576133935cf6f43f6a9 Mon Sep 17 00:00:00 2001
From: Rahul Soni <rahul.soni@bamfhealth.com>
Date: Thu, 22 Aug 2024 00:31:26 +0400
Subject: [PATCH] minor refactoring for PR-82

---
 models/bamf_nnunet_mr_breast/config/default.yml      | 4 +---
 models/bamf_nnunet_mr_breast/utils/NNUnetRunnerV2.py | 9 +++------
 2 files changed, 4 insertions(+), 9 deletions(-)

diff --git a/models/bamf_nnunet_mr_breast/config/default.yml b/models/bamf_nnunet_mr_breast/config/default.yml
index 23867b4b..3b0a704e 100644
--- a/models/bamf_nnunet_mr_breast/config/default.yml
+++ b/models/bamf_nnunet_mr_breast/config/default.yml
@@ -31,7 +31,7 @@ execute:
 - NiftiConverter
 - module: NNUnetRunnerV2
   nnunet_dataset: Dataset009_Breast
-  roi: BREAST+FGT
+  roi: BREAST,FGT
 - module: NNUnetRunnerV2
   nnunet_dataset: Dataset011_Breast
   roi: BREAST+BREAST_CARCINOMA
@@ -49,8 +49,6 @@ modules:
 
   NNUnetRunnerV2:
     in_data: nifti:mod=mr
-    nnunet_config: 3d_fullres
-    export_prob_maps: False
 
   BreastPostProcessor:
     in_breast_data: nifti:mod=seg:nnunet_task=Dataset009_Breast
diff --git a/models/bamf_nnunet_mr_breast/utils/NNUnetRunnerV2.py b/models/bamf_nnunet_mr_breast/utils/NNUnetRunnerV2.py
index 0aab9b1a..f8dc85bb 100644
--- a/models/bamf_nnunet_mr_breast/utils/NNUnetRunnerV2.py
+++ b/models/bamf_nnunet_mr_breast/utils/NNUnetRunnerV2.py
@@ -21,13 +21,11 @@
 
 @IO.ConfigInput('in_data', 'nifti:mod=mr', the="input data to run nnunet on")
 @IO.Config('nnunet_dataset', str, None, the='nnunet dataset name')
-@IO.Config('nnunet_config', str, None, the='nnunet model name (2d, 3d_lowres, 3d_fullres, 3d_cascade_fullres)')
 @IO.Config('folds', int, None, the='number of folds to run nnunet on')
 @IO.Config('roi', str, None, the='roi or comma separated list of roi the nnunet segments')
 class NNUnetRunnerV2(Module):
 
     nnunet_dataset: str
-    nnunet_config: str
     input_data_type: DataType
     folds: int                          # TODO: support optional config attributes
     roi: str
@@ -40,7 +38,6 @@ def task(self, instance: Instance, in_data: InstanceData, out_data: InstanceData
         # get the nnunet model to run
         self.v("Running nnUNet_predict.")
         self.v(f" > dataset:     {self.nnunet_dataset}")
-        self.v(f" > config:      {self.nnunet_config}")
         self.v(f" > input data:  {in_data.abspath}")
         self.v(f" > output data: {out_data.abspath}")
 
@@ -80,7 +77,7 @@ def task(self, instance: Instance, in_data: InstanceData, out_data: InstanceData
         bash_command += ["-i", str(inp_dir)]
         bash_command += ["-o", str(out_dir)]
         bash_command += ["-d", self.nnunet_dataset]
-        bash_command += ["-c", self.nnunet_config]
+        bash_command += ["-c", "3d_fullres"]
         
         # add optional arguments
         if self.folds is not None:
@@ -91,9 +88,9 @@ def task(self, instance: Instance, in_data: InstanceData, out_data: InstanceData
 
         # output meta
         meta = {
-            "model": "nnunet",
+            "model": "nnunet-v2",
             "nnunet_dataset": self.nnunet_dataset,
-            "nnunet_config": self.nnunet_config,
+            "nnunet_config": "3d_fullres",
             "roi": self.roi
         }