ds = load_small_mose2_clean()
img = ds.image
pts = ds.pts

kp = KeyPoints(pts, img.shape)
kp.set_data(img)
ps = kp.extract_patches(35)

size = ps.shape[1]
mode = 'grid'
zps = np.array([zernike(j, size, mode) for j in np.arange(0, 66)])

#m = extract_moments_rot(ps, zps)
# Do not open rotational-invariant mode
m = extract_moments(ps, zps)

X = m[:, 1:]
umap = UMAP(X,
            n_neighbors=30,
            min_dist=0.2,
            metric='correlation',
            n_epochs=300,
            init='pca',
            keep_log=False,
            num_negative_samples=6,
            local_connectivity=1.0)
X_umap = umap.fit_transform(X)

ll = interactive_clusters(X_umap, kp.pts, img)

fp = KeyPoints(pts[ll[1].ind], img.shape)
fp.set_data(img)
ppp = fp.extract_patches(60)