def mean_shift(label_embedding):
neigh = NearestNeighbors(radius=0.5, metric='euclidean').fit(label_embedding.copy().T)
predicted_instances = -np.ones(label_embedding.shape[1])
unlabeled = np.where(predicted_instances < 0)[0]
dimensions = label_embedding.shape[0]
while unlabeled.size > 0:
index = np.random.choice(unlabeled)
indices = set([index])
centre = label_embedding.T[index]
for i in range(100):
neighbors = neigh.radius_neighbors(centre.reshape(1, dimensions), return_distance=False)[0]
new_centre = label_embedding.T[neighbors].mean(axis=0)
if np.allclose(centre, new_centre):
break
indices.update(neighbors)
centre = new_centre
neighbors = neigh.radius_neighbors(centre.reshape(1, dimensions), return_distance=False)[0]
indices.update(neighbors)
centre_index = neigh.nearest(centre.reshape(1, dimensions))
predicted_instances[list(indices)] = centre_index
unlabeled = np.where(predicted_instances < 0)[0]
return predicted_instances
