--- a +++ b/ants/segmentation/kmeans.py @@ -0,0 +1,48 @@ + + + +__all__ = ['kmeans_segmentation'] + +import ants + +def kmeans_segmentation(image, k, kmask=None, mrf=0.1): + """ + K-means image segmentation that is a wrapper around `ants.atropos` + + ANTsR function: `kmeansSegmentation` + + Arguments + --------- + image : ANTsImage + input image + + k : integer + integer number of classes + + kmask : ANTsImage (optional) + segment inside this mask + + mrf : scalar + smoothness, higher is smoother + + Returns + ------- + ANTsImage + + Example + ------- + >>> import ants + >>> fi = ants.image_read(ants.get_ants_data('r16'), 'float') + >>> fi = ants.n3_bias_field_correction(fi, 2) + >>> seg = ants.kmeans_segmentation(fi, 3) + """ + dim = image.dimension + kmimage = ants.iMath(image, 'Normalize') + if kmask is None: + kmask = ants.get_mask(kmimage, 0.01, 1, cleanup=2) + kmask = ants.iMath(kmask, 'FillHoles').threshold_image(1,2) + nhood = 'x'.join(['1']*dim) + mrf = '[%s,%s]' % (str(mrf), nhood) + kmimage = ants.atropos(a = kmimage, m = mrf, c = '[5,0]', i = 'kmeans[%s]'%(str(k)), x = kmask) + kmimage['segmentation'] = kmimage['segmentation'].clone(image.pixeltype) + return kmimage