Diff of /ants/ops/mask_image.py [000000] .. [5d12a0]

Switch to side-by-side view

--- a
+++ b/ants/ops/mask_image.py
@@ -0,0 +1,55 @@
+__all__ = ['mask_image']
+
+import numpy as np
+
+import ants
+from ants.decorators import image_method
+
+@image_method
+def mask_image(image, mask, level=1, binarize=False):
+    """
+    Mask an input image by a mask image.  If the mask image has multiple labels,
+    it is possible to specify which label(s) to mask at.
+
+    ANTsR function: `maskImage`
+
+    Arguments
+    ---------
+    image : ANTsImage
+        Input image.
+
+    mask : ANTsImage
+        Mask or label image.
+
+    level : scalar or tuple of scalars
+        Level(s) at which to mask image. If vector or list of values, output image is non-zero at all locations where label image matches any of the levels specified.
+
+    binarize : boolean
+        whether binarize the output image
+
+    Returns
+    -------
+    ANTsImage
+
+    Example
+    -------
+    >>> import ants
+    >>> myimage = ants.image_read(ants.get_ants_data('r16'))
+    >>> mask = ants.get_mask(myimage)
+    >>> myimage_mask = ants.mask_image(myimage, mask, 3)
+    >>> seg = ants.kmeans_segmentation(myimage, 3)
+    >>> myimage_mask = ants.mask_image(myimage, seg['segmentation'], (1,3))
+    """
+    leveluse = level
+    if type(leveluse) is np.ndarray:
+        leveluse = level.tolist()
+    if type(leveluse) is int or type(leveluse) is float:
+        leveluse = [level]
+    image_out = image.clone() * 0
+    for mylevel in leveluse:
+        temp = ants.threshold_image(mask, mylevel, mylevel)
+        if binarize:
+            image_out = image_out + temp
+        else:
+            image_out = image_out + temp * image
+    return image_out