Merge pull request #38 from BrainLesion/nicer_hdbet_test

Nicer hdbet test

Author: neuronflow

brainles_preprocessing/brain_extraction/brain_extractor.py

@@ -88,7 +88,8 @@ def extract(
             + "_mask.nii.gz"
         )
 
-        copyfile(
-            src=hdbet_mask_path,
-            dst=brain_mask_path,
-        )
+        if hdbet_mask_path != brain_mask_path:
+            copyfile(
+                src=hdbet_mask_path,
+                dst=brain_mask_path,
+            )

tests/test_brain_extractor.py

@@ -16,8 +16,12 @@ def setUp(self):
 
         self.brain_extractor = HDBetExtractor()
         self.input_image_path = input_dir + "/tcia_example_t1c.nii.gz"
+        self.input_brain_mask_path = input_dir + "/bet_tcia_example_t1c_mask.nii.gz"
         self.masked_image_path = self.output_dir + "/bet_tcia_example_t1c.nii.gz"
-        self.brain_mask_path = self.output_dir + "/bet_tcia_example_t1c_masked.nii.gz"
+        self.brain_mask_path = self.output_dir + "/bet_tcia_example_t1c_mask.nii.gz"
+        self.masked_again_image_path = (
+            self.output_dir + "/bet_tcia_example_t1c_masked2.nii.gz"
+        )
 
         print(self.input_image_path)
         print(self.masked_image_path)
@@ -26,6 +30,7 @@ def tearDown(self):
         # Clean up created files if they exist
         shutil.rmtree(self.output_dir)
 
+
     def test_extract_creates_output_files(self):
         # we try to run the fastest possible skullstripping on GPU
         self.brain_extractor.extract(
@@ -47,11 +52,12 @@ def test_extract_creates_output_files(self):
         )
 
     def test_apply_mask_creates_output_file(self):
-        # self.brain_extractor.apply_mask(
-        #     self.input_image, self.mask_image, self.output_image
-        # )
-        # self.assertTrue(
-        #     os.path.exists(self.output_image_path),
-        #     "Output image file was not created in apply_mask.",
-        # )
-        ...
+        self.brain_extractor.apply_mask(
+            input_image_path=self.input_image_path,
+            mask_image_path=self.input_brain_mask_path,
+            masked_image_path=self.masked_again_image_path,
+        )
+        self.assertTrue(
+            os.path.exists(self.masked_again_image_path),
+            "Output image file was not created in apply_mask.",
+        )