Merge pull request #500 from oesteban/fix/499-extract-wm

FIX: Better generalization and renaming+relocation in the API of ``extract_wm``

Author: oesteban

niworkflows/interfaces/images.py

@@ -585,22 +585,6 @@ def reorient(in_file, newpath=None):
     return out_file
 
 
-def extract_wm(in_seg, wm_label=3, newpath=None):
-    import nibabel as nb
-    import numpy as np
-    from nipype.utils.filemanip import fname_presuffix
-
-    nii = nb.load(in_seg)
-    data = np.zeros(nii.shape, dtype=np.uint8)
-    data[np.asanyarray(nii.dataobj) == wm_label] = 1
-
-    out_file = fname_presuffix(in_seg, suffix='_wm', newpath=newpath)
-    new = nb.Nifti1Image(data, nii.affine, nii.header)
-    new.set_data_dtype(np.uint8)
-    new.to_filename(out_file)
-    return out_file
-
-
 def normalize_xform(img):
     """
     Set identical, valid qform and sform matrices in an image.

niworkflows/utils/images.py

@@ -107,3 +107,23 @@ def update_header_fields(fname, **kwargs):
     for field, value in kwargs.items():
         img.header[field] = value
     overwrite_header(img, fname)
+
+
+def dseg_label(in_seg, label, newpath=None):
+    """Extract a particular label from a discrete segmentation."""
+    from pathlib import Path
+    import nibabel as nb
+    import numpy as np
+    from nipype.utils.filemanip import fname_presuffix
+
+    newpath = Path(newpath or '.')
+
+    nii = nb.load(in_seg)
+    data = np.int16(nii.dataobj) == label
+
+    out_file = fname_presuffix(in_seg, suffix='_mask',
+                               newpath=str(newpath.absolute()))
+    new = nii.__class__(data, nii.affine, nii.header)
+    new.set_data_dtype(np.uint8)
+    new.to_filename(out_file)
+    return out_file

niworkflows/utils/tests/test_images.py

@@ -2,7 +2,7 @@
 import numpy as np
 
 import pytest
-from ..images import update_header_fields, overwrite_header
+from ..images import update_header_fields, overwrite_header, dseg_label
 
 
 def random_image():
@@ -71,3 +71,18 @@ def test_overwrite_header_reject_mmap(tmp_path):
     img = nb.load(fname, mmap=True)
     with pytest.raises(ValueError):
         overwrite_header(img, fname)
+
+
+def test_dseg_label(tmp_path):
+    fname = str(tmp_path / 'test_file.nii.gz')
+
+    data = np.dstack((
+        np.zeros((20, 20), dtype='int16'),
+        np.ones((20, 20), dtype='int16'),
+        np.ones((20, 20), dtype='int16') * 2,
+        np.ones((20, 20), dtype='int16') * 3,
+    ))
+    nb.Nifti1Image(data, np.eye(4), None).to_filename(fname)
+
+    new_im = nb.load(dseg_label(fname, label=2, newpath=tmp_path))
+    assert np.all((data == 2).astype('int16') == np.int16(new_im.dataobj))