add tree+image click

app.py

@@ -13,6 +13,8 @@ import uuid
 import zipfile
 import multiprocessing as mp
 
+from gradio_image_prompter import ImagePrompter
+
 from einops import rearrange
 from matplotlib import pyplot as plt
 import matplotlib
@@ -2472,26 +2474,32 @@ with demo:
                     # 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, big_pil_image
+                    return tsne_embed, edges, fps_eigvecs, fps_tsne3d_rgb, fps_idx, pil_image, gr.update(value={'image': big_pil_image, 'points': []}, interactive=True)
 
-        big_tsne_plot = gr.Image(label="spectral-tSNE tree [+ Cluster Heatmap]", elem_id="big_tsne_plot", interactive=False, format='png')
+        gr.Markdown('---')
+        gr.Markdown('<h3 style="text-align: center;">↓ interactively inspect the hierarchical structure</h3>')
+        gr.Markdown('---')
+        # big_tsne_plot = gr.Image(label="spectral-tSNE tree [+ Cluster Heatmap]", elem_id="big_tsne_plot", interactive=False, format='png')
+        tsne_image_plot = ImagePrompter(show_label=True, elem_id="tsne_image_plot", interactive=False, label="spectral-tSNE tree [+ Cluster Heatmap]")
         
         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],
-            outputs=[tsne_2d_points, edges, fps_eigvecs, fps_tsne_rgb, fps_indices, tsne_plot, big_tsne_plot],
+            outputs=[tsne_2d_points, edges, fps_eigvecs, fps_tsne_rgb, fps_indices, tsne_plot, tsne_image_plot],
         )
-        gr.Markdown('---')
-        gr.Markdown('<h3 style="text-align: center;">↓ interactively inspect the hierarchical structure</h3>')
-        gr.Markdown('---')
         with gr.Row():
-            from gradio_image_prompter import ImagePrompter
             with gr.Column(scale=5, min_width=200) as tsne_select:
+                gr.Markdown('---')
+                gr.Markdown('<h3 style="text-align: center;">Please click on the image blow ↓</h3>')
+                gr.Markdown('---')
                 tsne_prompt_image = ImagePrompter(show_label=True, elem_id="tsne_prompt_image", interactive=False, label="spectral-tSNE tree")
                 # copy plot to tsne_prompt_image on change
                 # tsne_plot.change(fn=lambda x: gr.update(value={'image': x}, interactive=True),
                 #                 inputs=[tsne_plot], outputs=[tsne_prompt_image])
             with gr.Column(scale=5, min_width=200) as image_select:
+                gr.Markdown('---')
+                gr.Markdown('<h3 style="text-align: center;">Please click on the image blow ↓</h3>')
+                gr.Markdown('---')
                 image_plot = ImagePrompter(show_label=True, elem_id="image_plot", interactive=False, label="NCUT spectral-tSNE")
                 image_slider = gr.Slider(0, 100, step=1, label="Image Index", value=0, elem_id="image_slider", interactive=True)
                 def update_image_prompt(image_slider, output_gallery):
@@ -2505,11 +2513,16 @@ 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 tsne_image_select:
+                gr.Markdown('---')
+                gr.Markdown('<h3 style="text-align: center;">Please click on the image above ↑</h3>')
+                gr.Markdown('---')
+                tsne_non_prompt_image = gr.Image(label="spectral-tSNE tree", elem_id="tsne_non_prompt_image", interactive=False, format='png')
             with gr.Column(scale=5, min_width=200):
                 gr.Markdown('<h3 style="text-align: center;">Help</h3>')
                 with gr.Accordion("Instructions", open=True):
                     gr.Markdown("""
-                                1. Click one dot on the left-side image. 
+                                1. Click one dot on the image. 
                                 - Only the last clicked dot will be used
                                 - Eraser is at top-right corner
                                 - Use the right-side Radio to switch tree/image
@@ -2524,32 +2537,34 @@ 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"], label="Where to click on?", value="Tree", elem_id="prompt_radio", show_label=True)
+                prompt_radio = gr.Radio(["Tree [+Image]", "Image"], 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):
                     # Plot the t-SNE points
                     pil_image = plot_tsne_tree(tsne_embed, edges, fps_tsne_rgb, granularity)
-                    if tsne_prompt_image is None:
-                        return gr.update(value={'image': pil_image}, interactive=True)
-                    return gr.update(value={'image': pil_image, 'points': tsne_prompt_image['points']}, interactive=True)
+                    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],
-                                            outputs=[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],
-                                    outputs=[tsne_prompt_image])
+                                    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,
+                #                         inputs=[granularity_slider, tsne_2d_points, edges, fps_tsne_rgb, tsne_prompt_image],
+                #                         outputs=[tsne_non_prompt_image])
                 run_inspection_button = gr.Button("🔴 RUN Inspection", elem_id="run_inspection", variant='primary')
                 inspect_logging_text = gr.Textbox("Logging information", lines=3, label="Logging", elem_id="inspect_logging", type="text", placeholder="Logging information", autofocus=False, autoscroll=False)
                 # output_slot_radio = gr.Radio([1, 2, 3], label="Output Row", value=1, elem_id="output_slot", show_label=True)
 
                 delete_all_output_button = gr.Button("❌ Delete All Output", elem_id="delete_all_output", variant='secondary')
-                
+            
+            tsne_image_select.visible = True
+            tsne_select.visible = False
             image_select.visible = False
-            tsne_select.visible = True
-            prompt_radio.change(fn=lambda x: gr.update(visible=x=="Tree"), inputs=prompt_radio, outputs=[tsne_select])
             prompt_radio.change(fn=lambda x: gr.update(visible=x=="Image"), inputs=prompt_radio, outputs=[image_select])
+            prompt_radio.change(fn=lambda x: gr.update(visible=x=="Tree [+Image]"), inputs=prompt_radio, outputs=[tsne_image_select])
             
         MAX_ROWS = 20
         current_output_row = gr.State(0)
@@ -2634,19 +2649,21 @@ with demo:
             closest_idx = np.argmax(sim)
             return closest_idx, (_x_ratio, _y_ratio)
         
-        def find_closest_fps_point(prompt_radio, 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, i_image, 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":
                     return find_closest_fps_point_for_image_prompt(image_prompt, i_image, 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_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, 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):
             if len(tsne2d_embed) == 0:
                 raise gr.Error("Please run FPS+Cluster first.")
-            closest_idx, (_x, _y) = find_closest_fps_point(prompt_radio, 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, i_image, tsne2d_embed, eigvecs, fps_eigvecs)
             closest_rgb = fps_tsne_rgb[closest_idx]
             closest_rgb = (closest_rgb * 255).astype(np.uint8)
             
@@ -2728,7 +2745,7 @@ with demo:
             
         run_inspection_button.click(
             run_inspection,
-            inputs=[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, 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],
             outputs=inspect_output_rows + output_tree_images + heatmap_galleries + text_blocks + [current_output_row, output_row_occupy, inspect_logging_text],
         )