a b/ants/segmentation/kmeans.py
1
2
3
4
__all__ = ['kmeans_segmentation']
5
6
import ants
7
8
def kmeans_segmentation(image, k, kmask=None, mrf=0.1):
9
    """
10
    K-means image segmentation that is a wrapper around `ants.atropos`
11
12
    ANTsR function: `kmeansSegmentation`
13
14
    Arguments
15
    ---------
16
    image : ANTsImage
17
        input image
18
19
    k : integer
20
        integer number of classes
21
22
    kmask : ANTsImage (optional)
23
        segment inside this mask
24
25
    mrf : scalar
26
        smoothness, higher is smoother
27
28
    Returns
29
    -------
30
    ANTsImage
31
32
    Example
33
    -------
34
    >>> import ants
35
    >>> fi = ants.image_read(ants.get_ants_data('r16'), 'float')
36
    >>> fi = ants.n3_bias_field_correction(fi, 2)
37
    >>> seg = ants.kmeans_segmentation(fi, 3)
38
    """
39
    dim = image.dimension
40
    kmimage = ants.iMath(image, 'Normalize')
41
    if kmask is None:
42
        kmask = ants.get_mask(kmimage, 0.01, 1, cleanup=2)
43
    kmask = ants.iMath(kmask, 'FillHoles').threshold_image(1,2)
44
    nhood = 'x'.join(['1']*dim)
45
    mrf = '[%s,%s]' % (str(mrf), nhood)
46
    kmimage = ants.atropos(a = kmimage, m = mrf, c = '[5,0]', i = 'kmeans[%s]'%(str(k)), x = kmask)
47
    kmimage['segmentation'] = kmimage['segmentation'].clone(image.pixeltype)
48
    return kmimage