Random walk that traverses neighbors in a cluster

reduced_palette.py

class KMeansReducedPalette:
    # ... omitted

    def random_neighborhood_walk_recolor(self, image, max_steps):
        original_shape = image.shape
        image = self._preprocess(image)
        recolor = image.copy()
        centroid_idxs = self.kmeans.predict(image)

        for ci in range(self.num_colors):
            n_pixels = np.sum(centroid_idxs == ci)

            if n_pixels == 0:
                continue

            n_neighbors = self.centroid_nearest_pixels[ci].shape[0]
            neighbor_idxs = np.random.randint(min(max_steps, n_neighbors),
                                              size=(n_pixels,))
            recolor_ci = self.centroid_nearest_pixels[ci][neighbor_idxs]
            recolor[centroid_idxs == ci] = recolor_ci

        return recolor.reshape(original_shape)