--- 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}