def test_kmeans():
    x = np.concatenate((np.ones(10) * 10, np.ones(10) * 50))
    centers, labels, maps = kmeans(x,
                                   k=2,
                                   scale_max=100,
                                   scale_min=0,
                                   map_keys=['a', 'b'])
    assert (len(centers) == 2)
    assert (sorted(centers) == [10, 50])
    assert (len(labels) == len(x))
    assert (len(maps) == 2)
    assert (np.allclose(maps['a'], np.concatenate(
        (np.ones(10), np.zeros(10)))))
    assert (np.allclose(maps['b'], np.concatenate(
        (np.zeros(10), np.ones(10)))))
Exemplo n.º 2
0
def init_values(data, k, scale_range, scale_sigma):
    """
    Initialize theta and labels. The means of each label are initialized using
    kmeans. The standard deviation of each label is chosen randomly from a uniform
    distribution between 0 and scale_sigma.

    Input
    -----
    data : array (n, m)
        Image to segment.

    k : int
        Number of labels.

    scale_range : tuple
        (min, max) Range of values from which start kmeans.

    scale_sigma : float
        Sigma will be initialized with a random value between 1 and scale_sigma.

    Output
    ------
    thetas : list of lists
        Contains [mu, sigma] for each label.

    labels : int array (n, m)
        Labels for each pixel. The labels are assigned using kmeans (distance of
        each pixel's intensity to label means).
    """
    scale_min, scale_max = scale_range
    mus, labels, _ = kmeans(data.ravel(), k, scale_max=scale_max, scale_min=scale_min)
    labels = labels.reshape(data.shape)
    thetas = []
    for i in range(k):
        sigma = np.random.rand() * (scale_sigma - 1) + 1
        thetas.append([mus[i], sigma])

    return thetas, labels.astype(int)
from fmri_utils.segmentation.kmeans import kmeans

subjects = ['sub-10159']  #, 'sub-10171', 'sub-10189']

# Current directory
my_dir = dirname(__file__)
# Data directory (project-red/data)
data_dir = pjoin(my_dir, '..', '..', '..', 'data')

for s in subjects:
    img = nib.load(
        pjoin(data_dir, 'ds000030', s, 'anat', '%s_T1w_brain.nii.gz' % s))
    data = img.get_data()

    # Segment
    _, klabel, kmap = kmeans(data.ravel())

    # Reshape things
    csf = kmap['csf'].reshape(data.shape)
    gm = kmap['gray'].reshape(data.shape)
    wm = kmap['white'].reshape(data.shape)

    # Save
    nifti_csf = nib.Nifti1Image(csf, affine=img.affine)
    nifti_gm = nib.Nifti1Image(gm, affine=img.affine)
    nifti_wm = nib.Nifti1Image(wm, affine=img.affine)
    nib.nifti1.save(
        nifti_csf,
        pjoin(data_dir, 'ds000030', s, 'segmentation_results',
              'kmeans_csf.nii'))
    nib.nifti1.save(
my_dir = dirname(__file__)
# Data directory (project-red/data)
data_dir = pjoin(my_dir, '..', '..', '..', 'data')
# Report figures directory (project-red/report/figures)
report_dir = pjoin(data_dir, '..', 'report', 'figures')

# Segmentation for subjects
for s in subjects:
    img = nib.load(
        pjoin(data_dir, 'ds000030', s, 'anat', '%s_T1w_brain.nii.gz' % s))
    data = img.get_data()
    n_slices = data.shape[-1]
    slice_s = data[..., np.floor(n_slices / 2).astype(int)]  # middle slice

    # Segment
    _, klabel, kmap = kmeans(slice_s.ravel())

    # Plots
    plt.figure(figsize=(15, 15))
    plt.subplot(1, 4, 1)
    plt.imshow(klabel.reshape(slice_s.shape))
    plt.title('Segmented slice')
    plt.subplot(1, 4, 2)
    plt.imshow(kmap['csf'].reshape(slice_s.shape), cmap='gray')
    plt.title('CSF')
    plt.subplot(1, 4, 3)
    plt.imshow(kmap['gray'].reshape(slice_s.shape), cmap='gray')
    plt.title('Gray matter')
    plt.subplot(1, 4, 4)
    im = plt.imshow(kmap['white'].reshape(slice_s.shape), cmap='gray')
    plt.title('White matter')
    slice_s = data[x1:x2, y1:y2, slice_n]

    # Segment
    _, labels, maps = mrf_em(slice_s,
                             0.05,
                             k=2,
                             max_iter=10,
                             scale_range=(100, 400),
                             scale_sigma=50,
                             max_label_iter=10,
                             njobs=2,
                             map_labels=['gray', 'white'])
    # Run kmeans for comparison
    _, klabels, kmaps = kmeans(slice_s.ravel(),
                               k=2,
                               max_iter=10 ^ 4,
                               scale_max=400,
                               scale_min=200,
                               map_keys=['gray', 'white'])
    klabels = klabels.reshape(slice_s.shape)

    # Plot segmentation
    fig1 = plt.figure(figsize=(10, 10))
    plt.subplot(1, 3, 1)
    plt.imshow(labels, interpolation="None")
    plt.title('MRF EM')
    plt.subplot(1, 3, 2)
    plt.imshow(slice_s, interpolation="None")
    plt.title('Original')
    plt.subplot(1, 3, 3)
    plt.imshow(np.abs(klabels - 1), interpolation="None")
    plt.title('kmeans')