working hdbet test

Signed-off-by: neuronflow <[email protected]>
BrainLesion · Dec 4, 2023 · 3b9857e · 3b9857e
1 parent 7aee89a
commit 3b9857e
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/brainles_preprocessing/brain_extraction/brain_extractor.py b/brainles_preprocessing/brain_extraction/brain_extractor.py
@@ -62,6 +62,8 @@ def extract(
         log_file_path: str = None,
         # TODO convert mode to enum
         mode: str = "accurate",
+        device: int | str = 0,
+        do_tta: bool = True,
     ) -> None:
         # GPU + accurate + TTA
         """skullstrips images with HD-BET generates a skullstripped file and mask"""
@@ -72,9 +74,9 @@ def extract(
             # TODO consider postprocessing
             # postprocess=False,
             mode=mode,
-            device=0,
+            device=device,
             postprocess=False,
-            do_tta=True,
+            do_tta=do_tta,
             keep_mask=True,
             overwrite=True,
         )

diff --git a/tests/test_brain_extractor.py b/tests/test_brain_extractor.py
@@ -12,7 +12,7 @@ def setUp(self):
         test_data_dir = turbopath(__file__).parent + "/test_data"
         input_dir = test_data_dir + "/input"
         self.output_dir = test_data_dir + "/temp_output"
-        os.makedirs(self.output_dir)
+        os.makedirs(self.output_dir, exist_ok=True)
 
         self.brain_extractor = HDBetExtractor()
         self.input_image_path = input_dir + "/tcia_example_t1c.nii.gz"
@@ -31,6 +31,9 @@ def test_extract_creates_output_files(self):
             input_image_path=self.input_image_path,
             masked_image_path=self.masked_image_path,
             brain_mask_path=self.brain_mask_path,
+            mode="fast",
+            device="cpu",
+            do_tta=False,
             # TODO generate and also test for presence of log file
         )