diff --git a/skimage/segmentation/km_segmentation.pyx b/skimage/segmentation/km_segmentation.pyx index 53237618..2f1dae5d 100644 --- a/skimage/segmentation/km_segmentation.pyx +++ b/skimage/segmentation/km_segmentation.pyx @@ -16,40 +16,56 @@ def km_segmentation(image, n_segments=100, ratio=10., max_iter=100, sigma=1.0): means_color = image[means_y, means_x, :] cdef np.ndarray[dtype=np.float_t, ndim=2] means = np.dstack([means_y, means_x, means_color]).reshape(-1, 5) + cdef np.float_t* current_mean + cdef np.float_t* mean_entry n_means = means.shape[0] # we do the scaling of ratio in the same way as in the SLIC paper # so the values have the same meaning ratio = (ratio / float(step)) ** 2 print(ratio) - cdef np.ndarray[dtype=np.float_t, ndim=3] image_yx = np.dstack([grid_y, grid_x, image / ratio]) + cdef np.ndarray[dtype=np.float_t, ndim=3] image_yx = np.dstack([grid_y, grid_x, image / ratio]).copy("C") cdef int i, k, x, y, x_min, x_max, y_min, y_max cdef float dist_mean cdef np.ndarray[dtype=np.int_t, ndim=2] nearest_mean = np.zeros((height, width), dtype=np.int) cdef np.ndarray[dtype=np.float_t, ndim=2] distance = np.ones((height, width), dtype=np.float) * np.inf + cdef np.float_t* image_p = image_yx.data + cdef np.float_t* distance_p = distance.data + cdef np.float_t* current_pixel + cdef float tmp for i in xrange(max_iter): print("iteration %d" % i) nearest_mean_old = nearest_mean.copy() + # we construct a new means every iteration, adjust pointer + current_mean = means.data # assign pixels to means for k in xrange(n_means): # compute windows: - y_min = int(max(means[k, 0] - 2 * step, 0)) - y_max = int(min(means[k, 0] + 2 * step, height)) - x_min = int(max(means[k, 1] - 2 * step, 0)) - x_max = int(min(means[k, 1] + 2 * step, height)) - for x in xrange(x_min, x_max): - for y in xrange(y_min, y_max): + y_min = int(max(current_mean[0] - 2 * step, 0)) + y_max = int(min(current_mean[0] + 2 * step, height)) + x_min = int(max(current_mean[1] - 2 * step, 0)) + x_max = int(min(current_mean[1] + 2 * step, height)) + for y in xrange(y_min, y_max): + current_pixel = &image_p[5 * (y * width + x_min)] + for x in xrange(x_min, x_max): + mean_entry = current_mean dist_mean = 0 for c in range(5): - dist_mean += (image_yx[y, x, c] - means[k, c]) ** 2 + # you would think the compiler can optimize this itself. + # mine can't (with O2) + tmp = current_pixel[0] - mean_entry[0] + current_pixel += 1 + mean_entry += 1 + dist_mean += tmp * tmp if distance[y, x] > dist_mean: nearest_mean[y, x] = k distance[y, x] = dist_mean + current_mean += 5 if (nearest_mean == nearest_mean_old).all(): break # recompute means: means_list = [np.bincount(nearest_mean.ravel(), image_yx[:, :, j].ravel()) for j in xrange(5)] in_mean = np.bincount(nearest_mean.ravel()) - means = (np.vstack(means_list) / in_mean).T + means = (np.vstack(means_list) / in_mean).T.copy("C") return nearest_mean