diff --git a/skimage/segmentation/__init__.py b/skimage/segmentation/__init__.py index 8fa8dfb8..2de27ad1 100644 --- a/skimage/segmentation/__init__.py +++ b/skimage/segmentation/__init__.py @@ -1,5 +1,7 @@ from .random_walker_segmentation import random_walker from .felzenszwalb import felzenszwalb_segmentation +from .km_segmentation import km_segmentation from .quickshift import quickshift -__all__ = [random_walker, quickshift, felzenszwalb_segmentation] +__all__ = [random_walker, quickshift, felzenszwalb_segmentation, + km_segmentation] diff --git a/skimage/segmentation/km_segmentation.py b/skimage/segmentation/km_segmentation.py new file mode 100644 index 00000000..fa307ab3 --- /dev/null +++ b/skimage/segmentation/km_segmentation.py @@ -0,0 +1,42 @@ +import numpy as np + + +def km_segmentation(image, n_segments=100, ratio=50, max_iter=100): + # initialize on grid: + height, width = image.shape[:2] + # approximate grid size for desired n_segments + step = np.sqrt(height * width / n_segments) + grid_y, grid_x = np.mgrid[:height, :width] + means_y = grid_y[::step, ::step] + means_x = grid_x[::step, ::step] + + means_color = image[means_y, means_x, :] + means = np.dstack([means_y, means_x, means_color]).reshape(-1, 5) + image = np.dstack([grid_y, grid_x, image * ratio]) + + nearest_mean = np.zeros((height, width), dtype=np.int) + distance = np.ones((height, width), dtype=np.float) * np.inf + for i in xrange(max_iter): + print("iteration %d" % i) + nearest_mean_old = nearest_mean.copy() + # assign pixels to means + for k, mean in enumerate(means): + # compute windows: + y_min = int(max(mean[0] - 2 * step, 0)) + y_max = int(min(mean[0] + 2 * step, height)) + x_min = int(max(mean[1] - 2 * step, 0)) + x_max = int(min(mean[1] + 2 * step, height)) + search_window = image[y_min:y_max + 1, x_min:x_max + 1] + dist_mean = np.sum((search_window - mean) ** 2, axis=2) + assign = distance[y_min:y_max + 1, x_min:x_max + 1] > dist_mean + nearest_mean[y_min:y_max + 1, x_min:x_max + 1][assign] = k + distance[y_min:y_max + 1, x_min:x_max + 1][assign] = \ + dist_mean[assign] + if (nearest_mean == nearest_mean_old).all(): + break + # recompute means: + means = [np.bincount(nearest_mean.ravel(), image[:, :, j].ravel()) + for j in xrange(5)] + in_mean = np.bincount(nearest_mean.ravel()) + means = (np.vstack(means) / in_mean).T + return nearest_mean