Switch to side-by-side view

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