add tree option

app.py

@@ -2297,6 +2297,7 @@ with demo:
                     num_sample_fps_slider = gr.Slider(1, 5000, step=1, label="FPS: num_sample", value=1000, elem_id="num_sample_fps")
                     tsne_perplexity_slider = gr.Slider(1, 1000, step=1, label="t-SNE: perplexity", value=500, elem_id="perplexity_tsne")
                     fps_hc_seed_slider = gr.Slider(0, 1000, step=1, label="Seed", value=0, elem_id="fps_hc_seed")
+                    tree_method_radio = gr.Radio(["eigvecs", "tsne"], label="Tree Method (input type)", value="eigvecs", elem_id="tree_method")
                 tsne_plot = gr.Image(label="spectral-tSNE tree", elem_id="tsne_plot", interactive=False, format='png')
                 
                 tsne_2d_points = gr.State(np.array([]))
@@ -2309,10 +2310,20 @@ with demo:
                     # Plot the t-SNE points
                     fig, ax = plt.subplots(1, 1, figsize=(6, 6))
                     ax.scatter(tsne_embed[:, 0], tsne_embed[:, 1], s=20, c=fps_tsne3d_rgb)
+                    # compute the length of the edges
+                    lengthes = np.linalg.norm(tsne_embed[edges[:, 0]] - tsne_embed[edges[:, 1]], axis=1)
+                    max_length = lengthes[k:].max()
+                    diag_length = np.linalg.norm(tsne_embed.max(axis=0) - tsne_embed.min(axis=0))
                     # draw the edges
                     for i_edge in range(k, len(edges)):
                         edge = edges[i_edge]
-                        ax.plot(tsne_embed[edge, 0], tsne_embed[edge, 1], 'k-', lw=1, alpha=0.7)
+                        # _do = np.clip(lengthes[i_edge] / (diag_length*0.3), 0, 1)
+                        if lengthes[i_edge] > diag_length*0.1:
+                            _do = 1.0
+                        else:
+                            _do = 0.0
+                        alpha = 0.7 * (1 - _do) + 0.0
+                        ax.plot(tsne_embed[edge, 0], tsne_embed[edge, 1], 'k-', lw=1, alpha=alpha)
                     # highlight the selected node
                     if hightlight_idx is not None:
                         if highlight_connections:
@@ -2428,7 +2439,7 @@ with demo:
                     return pil_image
                     
                 
-                def run_fps_tsne_hierarchical(image_gallery, eigvecs, num_sample_fps, perplexity_tsne, tsne3d_rgb, seed=0, max_display_dots=300):
+                def run_fps_tsne_hierarchical(image_gallery, eigvecs, num_sample_fps, perplexity_tsne, tsne3d_rgb, seed=0, tree_method='eigvecs', max_display_dots=300):
                     if len(eigvecs) == 0:
                         gr.Warning("Please run NCUT first.")
                         return
@@ -2464,9 +2475,10 @@ with demo:
                     tsne_embed[:, 0] = (tsne_embed[:, 0] - tsne_embed[:, 0].min()) / (tsne_embed[:, 0].max() - tsne_embed[:, 0].min()) * 2 - 1
                     tsne_embed[:, 1] = (tsne_embed[:, 1] - tsne_embed[:, 1].min()) / (tsne_embed[:, 1].max() - tsne_embed[:, 1].min()) * 2 - 1
                     
-                    edges = build_tree(tsne_embed)
-                    # edges = build_tree(fps_eigvecs, dist='cosine')
-                    # edges = build_tree(fps_tsne3d_rgb)
+                    if tree_method == 'eigvecs':
+                        edges = build_tree(fps_eigvecs, dist='cosine')
+                    if tree_method == 'tsne':
+                        edges = build_tree(tsne_embed, dist='euclidean')
                     
                     # Plot the t-SNE points
                     pil_image = plot_tsne_tree(tsne_embed, edges, fps_tsne3d_rgb, 0)
@@ -2484,7 +2496,7 @@ with demo:
         
         run_hierarchical_button.click(
             run_fps_tsne_hierarchical,
-            inputs=[input_gallery, eigvecs, num_sample_fps_slider, tsne_perplexity_slider, tsne3d_rgb, fps_hc_seed_slider],
+            inputs=[input_gallery, eigvecs, num_sample_fps_slider, tsne_perplexity_slider, tsne3d_rgb, fps_hc_seed_slider, tree_method_radio],
             outputs=[tsne_2d_points, edges, fps_eigvecs, fps_tsne_rgb, fps_indices, tsne_plot, tsne_image_plot],
         )
         with gr.Row():