Switch to side-by-side view

--- a
+++ b/ants/segmentation/atropos.py
@@ -0,0 +1,157 @@
+"""
+Atropos segmentation
+"""
+
+
+
+__all__ = ['atropos']
+
+import os
+import glob
+import warnings
+from tempfile import mktemp
+
+
+import ants
+from ants.internal import get_lib_fn, get_pointer_string, process_arguments
+
+
+def atropos(a, x, i='Kmeans[3]', m='[0.2,1x1]', c='[5,0]',
+            priorweight=0.25, **kwargs):
+    """
+    A finite mixture modeling (FMM) segmentation approach with possibilities
+    for specifying prior constraints. These prior constraints include the
+    specification of a prior label image, prior probability images (one for
+    each class), and/or an MRF prior to enforce spatial smoothing of the
+    labels. Similar algorithms include FAST and SPM. atropos can also perform
+    multivariate segmentation if you pass a list of images in: e.g. a=(img1,img2).
+
+    ANTsR function: `atropos`
+
+    Arguments
+    ---------
+    a : ANTsImage or list/tuple of ANTsImage types
+        One or more scalar images to segment. If priors are not used,
+        the intensities of the first image are used to order the classes
+        in the segmentation output, from lowest to highest intensity. Otherwise
+        the order of the classes is dictated by the order of the prior images.
+
+    x : ANTsImage
+        mask image.
+
+    i : string
+        initialization usually KMeans[N] for N classes or a list of N prior
+        probability images. See Atropos in ANTs for full set of options.
+
+    m : string
+        mrf parameters as a string, usually "[smoothingFactor,radius]" where
+        smoothingFactor determines the amount of smoothing and radius determines
+        the MRF neighborhood, as an ANTs style neighborhood vector eg "1x1x1"
+        for a 3D image. The radius must match the dimensionality of the image,
+        eg 1x1 for 2D and The default in ANTs is smoothingFactor=0.3 and
+        radius=1. See Atropos for more options.
+
+    c : string
+        convergence parameters, "[numberOfIterations,convergenceThreshold]".
+        A threshold of 0 runs the full numberOfIterations, otherwise Atropos
+        tests convergence by comparing the mean maximum posterior probability
+        over the whole region of interest defined by the mask x.
+
+    priorweight : scalar
+        usually 0 (priors used for initialization only), 0.25 or 0.5.
+
+    kwargs : keyword arguments
+        more parameters, see Atropos help in ANTs
+
+    Returns
+    -------
+    dictionary with the following key/value pairs:
+        `segmentation`: ANTsImage
+            actually segmented image
+        `probabilityimages` : list of ANTsImage types
+            one image for each segmentation class
+
+    Example
+    -------
+    >>> import ants
+    >>> img = ants.image_read(ants.get_ants_data('r16'))
+    >>> img = ants.resample_image(img, (64,64), 1, 0)
+    >>> mask = ants.get_mask(img)
+    >>> ants.atropos( a = img, m = '[0.2,1x1]', c = '[2,0]',  i = 'kmeans[3]', x = mask )
+    >>> seg2 = ants.atropos( a = img, m = '[0.2,1x1]', c = '[2,0]', i = seg['probabilityimages'], x = mask, priorweight=0.25 )
+    """
+    probs = mktemp(prefix='antsr', suffix='prob%02d.nii.gz')
+    tdir = probs.replace(os.path.basename(probs),'')
+    probsbase = os.path.basename(probs)
+    searchpattern = probsbase.replace('%02d', '*')
+
+    ct = 0
+    if isinstance(i, (list,tuple)) and (len(i) > 1):
+        while ct < len(i):
+            probchar = str(ct+1)
+            if ct < 9:
+                probchar = '0%s' % probchar
+            tempfn = probs.replace('%02d', probchar)
+            ants.image_write(i[ct], tempfn)
+            ct += 1
+        i = 'PriorProbabilityImages[%s,%s,%s]' % (str(len(i)), probs, str(priorweight))
+
+    if isinstance(a, (list,tuple)):
+        outimg = a[0].clone('unsigned int')
+    else:
+        outimg = a.clone('unsigned int')
+
+    mydim = outimg.dimension
+    outs = '[%s,%s]' % (get_pointer_string(outimg), probs)
+    mymask = x.clone('unsigned int')
+
+    if (not isinstance(a, (list,tuple))) or (len(a) == 1):
+        myargs = {
+            'd': mydim,
+            'a': a,
+            'm-MULTINAME-0': m,
+            'o': outs,
+            'c': c,
+            'm-MULTINAME-1': m,
+            'i': i,
+            'x': mymask
+        }
+        for k, v in kwargs.items():
+            myargs[k] = v
+
+    elif isinstance(a, (list, tuple)):
+        if len(a) > 6:
+            print('more than 6 input images not really supported, using first 6')
+            a = a[:6]
+        myargs = {
+            'd': mydim,
+            'm-MULTINAME-0': m,
+            'o': outs,
+            'c': c,
+            'm-MULTINAME-1': m,
+            'i': i,
+            'x': mymask
+        }
+        for aa_idx, aa in enumerate(a):
+            myargs['a-MULTINAME-%i'%aa_idx] = aa
+
+    processed_args = process_arguments(myargs)
+    libfn = get_lib_fn('Atropos')
+    retval = libfn(processed_args)
+
+    if retval != 0:
+        raise Exception(f"Atropos exited with non-zero status {retval}. Run with verbose=1 to see error messages")
+
+    probsout = glob.glob(os.path.join(tdir,'*'+searchpattern))
+
+    if probsout is None or len(probsout) == 0:
+        raise Exception('No Atropos output probability images found. Run with verbose=1 to see error messages')
+
+    probsout.sort()
+    probimgs = [ants.image_read(probsout[0])]
+    for idx in range(1, len(probsout)):
+        probimgs.append(ants.image_read(probsout[idx]))
+
+    outimg = outimg.clone('float')
+    return {'segmentation': outimg,
+            'probabilityimages': probimgs}