FIX: binarize mask

Author: Christian

src/zerodose/dataset.py

@@ -7,6 +7,7 @@
 
 from zerodose import utils
 from zerodose.niftyreg_wrapper import NiftyRegistration
+from zerodose.processing import Binarize
 from zerodose.processing import PadAndCropToMNI
 from zerodose.processing import ToFloat32
 
@@ -55,6 +56,7 @@ def _get_augmentation_transform_val(self, do_registration=True) -> tio.Compose:
 
         augmentations.extend(
             [
+                Binarize(include=["mask"]),
                 tio.transforms.ZNormalization(include=["mr"], masking_method="mask"),
                 PadAndCropToMNI(include=["mr", "mask"]),
                 ToFloat32(include=["mr"]),

src/zerodose/niftyreg_wrapper.py

@@ -36,8 +36,8 @@ def _read_matrix_nifty(file_name):
 
 def _register_mri_to_mni(mri_fname, ref):
     temp_aff = tempfile.NamedTemporaryFile(delete=False)
-
-    out_mri_fname = os.path.dirname(mri_fname) + "/out_mri1.nii.gz"
+    temp_out = tempfile.NamedTemporaryFile(delete=False, suffix=".nii.gz")
+    out_mri_fname = temp_out.name
     mni_template = ref
     niftyreg.main(
         [
@@ -55,7 +55,9 @@ def _register_mri_to_mni(mri_fname, ref):
     )
     affine_mat = _read_matrix_nifty(temp_aff.name)
     temp_aff.close()
+    temp_out.close()
     os.remove(temp_aff.name)
+    os.remove(temp_out.name)
     return affine_mat
 
 

src/zerodose/processing.py

@@ -66,6 +66,22 @@ def inverse(self):
         return PadAndCropToMNI(is_inverse=True)
 
 
+class Binarize(SpatialTransform):
+    """Binarize an image based on a threshhold."""
+
+    def __init__(self, threshold=0.5, **kwargs) -> None:
+        """Initialize the transform."""
+        self.threshold = threshold
+        super().__init__(**kwargs)
+
+    def apply_transform(self, subject: Subject) -> Subject:
+        """Apply the transform to the subject."""
+        for image in self.get_images(subject):
+            data = image.data
+            image.set_data(data > self.threshold)
+        return subject
+
+
 def _pad(image: tio.Image) -> None:
     data = processing._crop_mni_to_192(image.data)
     image.set_data(data)

tests/test_slow.py

@@ -24,7 +24,7 @@ def _get_mni_dir():
     return mni_dir
 
 
-def _augment_mni_img_for_tests(image_in_path, image_out_path):
+def _augment_mni_img_for_tests(image_in_path, image_out_path, binarize=False):
     rad = np.deg2rad(10)
     cos_gamma = np.cos(rad)
     sin_gamma = np.sin(rad)
@@ -51,6 +51,9 @@ def _augment_mni_img_for_tests(image_in_path, image_out_path):
     after_rot = resample_from_to(
         img_in, ((190, 200, 170), affine_mat.dot(img_in.affine))
     )
+    # if binarize:
+    #     data  = after_rot.get_fdata()
+    #     data.get
 
     ornt = np.array([[0, 1], [1, 1], [2, 1]])
 
@@ -164,7 +167,7 @@ def sbpet_outputfile():
     return fn
 
 
-@pytest.mark.slow
+# @pytest.mark.slow
 def test_syn_mni(runner, mri_mni_file, mask_mni_file, sbpet_outputfile) -> None:
     """Test the syn command with the standard model and MNI files."""
     result = runner.invoke(
@@ -199,4 +202,4 @@ def test_syn_niftyreg(mri_aug_file, mask_aug_file, sbpet_outputfile) -> None:
         device="cuda:0",
     )
 
-    assert nib.load(mri_aug_file).get_fdata().shape == start_shape
+    assert nib.load(sbpet_outputfile).get_fdata().shape == start_shape