update click

app.py

@@ -2303,11 +2303,12 @@ with demo:
                 
                 tsne_2d_points = gr.State(np.array([]))
                 edges = gr.State(np.array([]))
+                levels = gr.State(np.array([]))
                 fps_eigvecs = gr.State(np.array([]))
                 fps_indices = gr.State(np.array([]))
                 fps_tsne_rgb = gr.State(np.array([]))
                 
-                def plot_tsne_tree(tsne_embed, edges, fps_tsne3d_rgb, k, hightlight_idx=None, highlight_connections=False):
+                def plot_tsne_tree(tsne_embed, edges, levels, fps_tsne3d_rgb, k, hightlight_idx=None, highlight_connections=False):
                     # 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)
@@ -2316,15 +2317,24 @@ with demo:
                     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]
+                    #     # _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)
+                    max_level = levels.max()
+                    line_widths = levels / max_level * 3
+                    # flip line width
+                    line_widths = 3 - line_widths
                     for i_edge in range(k, len(edges)):
                         edge = edges[i_edge]
-                        # _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)
+                        level = levels[i_edge]
+                        plt.plot(tsne_embed[edge, 0], tsne_embed[edge, 1], c='gray', lw=line_widths[level], alpha=0.5)
+                        
                     # highlight the selected node
                     if hightlight_idx is not None:
                         if highlight_connections:
@@ -2477,17 +2487,17 @@ with demo:
                     tsne_embed[:, 1] = (tsne_embed[:, 1] - tsne_embed[:, 1].min()) / (tsne_embed[:, 1].max() - tsne_embed[:, 1].min()) * 2 - 1
                     
                     if tree_method == 'eigvecs':
-                        edges = build_tree(fps_eigvecs, dist='cosine')
+                        edges, levels = build_tree(fps_eigvecs, dist='cosine')
                     if tree_method == 'tsne':
-                        edges = build_tree(tsne_embed, dist='euclidean')
+                        edges, levels = build_tree(tsne_embed, dist='euclidean')
                     
                     # Plot the t-SNE points
-                    pil_image = plot_tsne_tree(tsne_embed, edges, fps_tsne3d_rgb, 0)
+                    pil_image = plot_tsne_tree(tsne_embed, edges, levels, fps_tsne3d_rgb, 0)
                     
                     # Plot the t-SNE points with image heatmaps
                     big_pil_image = plot_tsne_with_image_heatmaps(images, eigvecs, fps_eigvecs, tsne_embed, fps_tsne3d_rgb, max_display_dots)
                     
-                    return tsne_embed, edges, fps_eigvecs, fps_tsne3d_rgb, fps_idx, pil_image, gr.update(value={'image': big_pil_image, 'points': []}, interactive=True)
+                    return tsne_embed, edges, levels, fps_eigvecs, fps_tsne3d_rgb, fps_idx, pil_image, gr.update(value={'image': big_pil_image, 'points': []}, interactive=True)
 
         gr.Markdown('---')
         gr.Markdown('<h3 style="text-align: center;">↓ interactively inspect the hierarchical structure</h3>')
@@ -2498,7 +2508,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, tree_method_radio],
-            outputs=[tsne_2d_points, edges, fps_eigvecs, fps_tsne_rgb, fps_indices, tsne_plot, tsne_image_plot],
+            outputs=[tsne_2d_points, edges, levels, fps_eigvecs, fps_tsne_rgb, fps_indices, tsne_plot, tsne_image_plot],
         )
         with gr.Row():
             with gr.Column(scale=5, min_width=200) as tsne_select:
@@ -2526,6 +2536,15 @@ with demo:
                 image_slider.change(fn=update_image_prompt, inputs=[image_slider, output_gallery], outputs=[image_plot])
                 output_gallery.change(fn=update_image_prompt, inputs=[image_slider, output_gallery], outputs=[image_plot])
                 output_gallery.change(fn=lambda x: gr.update(maximum=len(x)-1, interactive=True), inputs=[output_gallery], outputs=[image_slider])
+            with gr.Column(scale=5, min_width=200) as orig_image_select:
+                gr.Markdown('---')
+                gr.Markdown('<h3 style="text-align: center;">Please click on the image blow ↓</h3>')
+                gr.Markdown('---')
+                image_plot2 = ImagePrompter(show_label=True, elem_id="image_plot", interactive=False, label="NCUT spectral-tSNE")
+                image_slider2 = gr.Slider(0, 100, step=1, label="Image Index", value=0, elem_id="image_slider", interactive=True)
+                image_slider2.change(fn=update_image_prompt, inputs=[image_slider2, input_gallery], outputs=[image_plot2])
+                output_gallery.change(fn=update_image_prompt, inputs=[image_slider2, input_gallery], outputs=[image_plot2])
+                output_gallery.change(fn=lambda x: gr.update(maximum=len(x)-1, interactive=True), inputs=[output_gallery], outputs=[image_slider2])
             with gr.Column(scale=5, min_width=200) as tsne_image_select:
                 gr.Markdown('---')
                 gr.Markdown('<h3 style="text-align: center;">Please click on the image above ↑</h3>')
@@ -2550,18 +2569,18 @@ with demo:
                                 2. Cluster Heatmap: max of N cosine similarity to N points in the connected components.
                                 """)
             with gr.Column(scale=5, min_width=200): 
-                prompt_radio = gr.Radio(["Tree [+Image]", "Image"], label="Where to click on?", value="Tree [+Image]", elem_id="prompt_radio", show_label=True)
+                prompt_radio = gr.Radio(["Tree [+Image]", "Image (NCUT)", "Image (Orig)"], label="Where to click on?", value="Tree [+Image]", elem_id="prompt_radio", show_label=True)
                 granularity_slider = gr.Slider(1, 1000, step=1, label="Cluster Granularity (k)", value=100, elem_id="granularity")
                 num_sample_fps_slider.change(fn=lambda x: gr.update(maximum=x, interactive=True), inputs=[num_sample_fps_slider], outputs=[granularity_slider])
-                def updaste_tsne_plot_change_granularity(granularity, tsne_embed, edges, fps_tsne_rgb, tsne_prompt_image):
+                def updaste_tsne_plot_change_granularity(granularity, tsne_embed, edges, levels, fps_tsne_rgb, tsne_prompt_image):
                     # Plot the t-SNE points
-                    pil_image = plot_tsne_tree(tsne_embed, edges, fps_tsne_rgb, granularity)
+                    pil_image = plot_tsne_tree(tsne_embed, edges, levels, fps_tsne_rgb, granularity)
                     return gr.update(value=pil_image, label=f"spectral-tSNE tree [k={granularity}]")
                 granularity_slider.change(updaste_tsne_plot_change_granularity,
-                                            inputs=[granularity_slider, tsne_2d_points, edges, fps_tsne_rgb, tsne_prompt_image],
+                                            inputs=[granularity_slider, tsne_2d_points, edges, levels, fps_tsne_rgb, tsne_prompt_image],
                                             outputs=[tsne_non_prompt_image])
                 tsne_plot.change(updaste_tsne_plot_change_granularity,
-                                    inputs=[granularity_slider, tsne_2d_points, edges, fps_tsne_rgb],
+                                    inputs=[granularity_slider, tsne_2d_points, edges, levels, fps_tsne_rgb],
                                     outputs=[tsne_non_prompt_image])
                 prompt_radio.change(update_image_prompt, inputs=[image_slider, output_gallery], outputs=[image_plot])
                 # prompt_radio.change(updaste_tsne_plot_change_granularity,
@@ -2576,7 +2595,9 @@ with demo:
             tsne_image_select.visible = True
             tsne_select.visible = False
             image_select.visible = False
-            prompt_radio.change(fn=lambda x: gr.update(visible=x=="Image"), inputs=prompt_radio, outputs=[image_select])
+            orig_image_select.visible = False
+            prompt_radio.change(fn=lambda x: gr.update(visible=x=="Image (NCUT)"), inputs=prompt_radio, outputs=[image_select])
+            prompt_radio.change(fn=lambda x: gr.update(visible=x=="Image (Orig)"), inputs=prompt_radio, outputs=[orig_image_select])
             prompt_radio.change(fn=lambda x: gr.update(visible=x=="Tree [+Image]"), inputs=prompt_radio, outputs=[tsne_image_select])
             
         MAX_ROWS = 20
@@ -2595,6 +2616,7 @@ with demo:
                 # output_row_occupy[i_row-1] = False
                 return output_row_occupy, gr.update(visible=False)
             delete_button.click(partial(delete_a_row, i_row=i_row), output_row_occupy, outputs=[output_row_occupy, inspect_output_row])
+            delete_button.visible = False
             return inspect_output_row, output_tree_image, heatmap_gallery, text_block
         
         gr.Markdown('---')
@@ -2662,21 +2684,23 @@ with demo:
             closest_idx = np.argmax(sim)
             return closest_idx, (_x_ratio, _y_ratio)
         
-        def find_closest_fps_point(prompt_radio, tsne_image_prompt, tsne_prompt, image_prompt, i_image, tsne2d_embed, eigvecs, fps_eigvecs):
+        def find_closest_fps_point(prompt_radio, tsne_image_prompt, tsne_prompt, image_prompt, image_prompt2, i_image, i_image2, tsne2d_embed, eigvecs, fps_eigvecs):
             try:
                 if prompt_radio == "Tree":
                     return find_closest_fps_point_for_tsne_tree_plot(tsne_prompt, tsne2d_embed)
-                if prompt_radio == "Image":
+                if prompt_radio == "Image (NCUT)":
                     return find_closest_fps_point_for_image_prompt(image_prompt, i_image, eigvecs, fps_eigvecs)
+                if prompt_radio == "Image (Orig)":
+                    return find_closest_fps_point_for_image_prompt(image_prompt2, i_image2, eigvecs, fps_eigvecs)
                 if prompt_radio == "Tree [+Image]":
                     return find_closest_fps_point_for_tsne_tree_plot(tsne_image_prompt, tsne2d_embed)
             except:
                 raise gr.Error("""No blue point is selected. <br/>Please left-click on the image to select a blue point. <br/>After reloading the image (e.g., change granularity), please use the eraser to remove the previous point, then click on the image to select a blue point.""")
         
-        def run_inspection(tsne_image_prompt, tsne_prompt, image_prompt, prompt_radio, current_output_row, tsne2d_embed, edges, fps_eigvecs, fps_tsne_rgb, fps_indices, granularity, eigvecs, i_image, tsne3d_rgb, input_gallery, output_row_occupy, max_rows=MAX_ROWS):
+        def run_inspection(tsne_image_prompt, tsne_prompt, image_prompt, image_prompt2, prompt_radio, current_output_row, tsne2d_embed, edges, levels, fps_eigvecs, fps_tsne_rgb, fps_indices, granularity, eigvecs, i_image, i_image2, tsne3d_rgb, input_gallery, output_row_occupy, max_rows=MAX_ROWS):
             if len(tsne2d_embed) == 0:
                 raise gr.Error("Please run FPS+Cluster first.")
-            closest_idx, (_x, _y) = find_closest_fps_point(prompt_radio, tsne_image_prompt, tsne_prompt, image_prompt, i_image, tsne2d_embed, eigvecs, fps_eigvecs)
+            closest_idx, (_x, _y) = find_closest_fps_point(prompt_radio, tsne_image_prompt, tsne_prompt, image_prompt, image_prompt2, i_image, i_image2, tsne2d_embed, eigvecs, fps_eigvecs)
             closest_rgb = fps_tsne_rgb[closest_idx]
             closest_rgb = (closest_rgb * 255).astype(np.uint8)
             
@@ -2686,7 +2710,7 @@ with demo:
             logging_text = f"Clicked: idx={closest_idx}, xy=[{_x:.2f}, {_y:.2f}], RGB={closest_rgb}"
             logging_text += f"\nGranularity: k={granularity}, Connected: n={len(connected_idxs)}"
             
-            output_tsne_plot = plot_tsne_tree(tsne2d_embed, edges, fps_tsne_rgb, granularity, closest_idx, highlight_connections=True)
+            output_tsne_plot = plot_tsne_tree(tsne2d_embed, edges, levels, fps_tsne_rgb, granularity, closest_idx, highlight_connections=True)
             
             # draw heatmap for the connected components
             ## cosine distance
@@ -2758,7 +2782,7 @@ with demo:
             
         run_inspection_button.click(
             run_inspection,
-            inputs=[tsne_image_plot, tsne_prompt_image, image_plot, prompt_radio, current_output_row, tsne_2d_points, edges, fps_eigvecs, fps_tsne_rgb, fps_indices, granularity_slider, eigvecs, image_slider, tsne3d_rgb, input_gallery, output_row_occupy],
+            inputs=[tsne_image_plot, tsne_prompt_image, image_plot, image_plot2, prompt_radio, current_output_row, tsne_2d_points, edges, levels, fps_eigvecs, fps_tsne_rgb, fps_indices, granularity_slider, eigvecs, image_slider, image_slider2, tsne3d_rgb, input_gallery, output_row_occupy],
             outputs=inspect_output_rows + output_tree_images + heatmap_galleries + text_blocks + [current_output_row, output_row_occupy, inspect_logging_text],
         )
         
@@ -2767,7 +2791,7 @@ with demo:
         with gr.Row():
             with gr.Column(scale=5, min_width=200):
                 gr.Markdown("### Step 1: Load Images")
-                input_gallery, submit_button, clear_images_button, dataset_dropdown, num_images_slider, random_seed_slider, load_images_button = make_input_images_section(n_example_images=10)
+                input_gallery, submit_button, clear_images_button, dataset_dropdown, num_images_slider, random_seed_slider, load_images_button = make_input_images_section(n_example_images=30)
                 submit_button.visible = False
                 num_images_slider.value = 30
                 
@@ -2877,7 +2901,10 @@ with demo:
                 right = eigvecs.reshape(-1, eigvecs.shape[-1])[flat_idx]
                 y, x = None, None
             else:
-                right, (y, x) = xy_eigvec(prompt_points, prompt_image_idx, eigvecs)
+                try:
+                    right, (y, x) = xy_eigvec(prompt_points, prompt_image_idx, eigvecs)
+                except:
+                    raise gr.Error("""No blue point is selected. <br/>Please left-click on the image to select a blue point. <br/>After reloading the image, please use the eraser (top-right) to remove the previous point, then click on the image to select a blue point.""")
             right = right[:n_eig]
             left = F.normalize(left, p=2, dim=-1)
             _right = F.normalize(right, p=2, dim=-1)