add inspect playground

app.py

@@ -2,6 +2,7 @@
 # %%
 import copy
 from datetime import datetime
+import io
 import math
 import pickle
 from functools import partial
@@ -168,8 +169,6 @@ def compute_ncut(
     logging_str += f"NCUT time: {time.time() - start:.2f}s\n"
     
     if only_eigvecs:
-        eigvecs = eigvecs.to("cpu").reshape(features.shape[:-1] + (num_eig,))
-        eigvecs = eigvecs.detach().numpy()
         return None, logging_str, eigvecs
     
     start = time.time()
@@ -285,10 +284,15 @@ def to_pil_images(images, target_size=512, resize=True, force_size=False):
     res = int(size * multiplier)
     if force_size:
         res = target_size
-    pil_images = [
-            Image.fromarray((image * 255).cpu().numpy().astype(np.uint8))
-            for image in images
-        ]
+
+    pil_images = []
+    for image in images:
+        if isinstance(image, torch.Tensor):
+            image = image.cpu().numpy()
+        if image.dtype == np.float32 or image.dtype == np.float64:
+            image = (image * 255).astype(np.uint8)
+        pil_images.append(Image.fromarray(image))
+        
     if resize:
         pil_images = [
             image.resize((res, res), Image.Resampling.NEAREST)
@@ -865,6 +869,7 @@ def ncut_run(
     # ailgnedcut
     if not directed:
         only_eigvecs = kwargs.get("only_eigvecs", False)
+        return_eigvec_and_rgb = kwargs.get("return_eigvec_and_rgb", False)
         
         rgb, _logging_str, eigvecs = compute_ncut(
             features,
@@ -886,9 +891,20 @@ def ncut_run(
             only_eigvecs=only_eigvecs,
         )
         
+        
         if only_eigvecs:
+            eigvecs = eigvecs.to("cpu").reshape(features.shape[:-1] + (num_eig,))
+            eigvecs = eigvecs.detach().numpy()
+            logging_str += _logging_str
             return eigvecs, logging_str
         
+        if return_eigvec_and_rgb:
+            eigvecs = eigvecs.to("cpu").reshape(features.shape[:-1] + (num_eig,))
+            eigvecs = eigvecs.detach().numpy()
+            rgb = rgb.cpu().numpy()
+            logging_str += _logging_str
+            return eigvecs, rgb, logging_str
+        
     if directed:
         head_index_text = kwargs.get("head_index_text", None)
         n_heads = features.shape[-2]   # (batch, h, w, n_heads, d)
@@ -1232,6 +1248,7 @@ def run_fn(
     advanced=False,
     directed=False,
     only_eigvecs=False,
+    return_eigvec_and_rgb=False,
 ):
     # print(node_type2, head_index_text, make_symmetric)
     progress=gr.Progress()
@@ -1373,6 +1390,7 @@ def run_fn(
         "head_index_text": head_index_text,
         "make_symmetric": make_symmetric,
         "only_eigvecs": only_eigvecs,
+        "return_eigvec_and_rgb": return_eigvec_and_rgb,
     }
     # print(kwargs)
     
@@ -1599,7 +1617,7 @@ def load_dataset_images(is_advanced, dataset_name, num_images=10,
         is_advanced = "Basic"
         
     if is_advanced == "Basic":
-        gr.Info(f"Loaded images from EgoExo")
+        gr.Info(f"Loaded images from EgoExo", duration=5)
         return default_images
     try:
         progress(0.5, desc="Downloading Dataset")
@@ -1644,7 +1662,7 @@ def load_dataset_images(is_advanced, dataset_name, num_images=10,
             image_idx.extend(idx.tolist())
     if not is_filter:
         if is_random:
-            if num_images < len(dataset):
+            if num_images <= len(dataset):
                 image_idx = np.random.RandomState(seed).choice(len(dataset), num_images, replace=False).tolist()
             else:
                 gr.Warning(f"Dataset has less than {num_images} images.")
@@ -1653,7 +1671,7 @@ def load_dataset_images(is_advanced, dataset_name, num_images=10,
             image_idx = list(range(num_images))
     key = 'image' if 'image' in dataset[0] else list(dataset[0].keys())[0]
     images = [dataset[i][key] for i in image_idx]
-    gr.Info(f"Loaded {len(images)} images from {dataset_name}")
+    gr.Info(f"Loaded {len(images)} images from {dataset_name}", duration=5)
     del dataset
     
     if dataset_name in CENTER_CROP_DATASETS:
@@ -2060,7 +2078,7 @@ def make_output_images_section(markdown=True, button=True):
         add_rotate_flip_buttons(output_gallery)
     return output_gallery
 
-def make_parameters_section(is_lisa=False, model_ratio=True, parameter_dropdown=True):
+def make_parameters_section(is_lisa=False, model_ratio=True, ncut_parameter_dropdown=True, tsne_parameter_dropdown=True):
     gr.Markdown("### Parameters <a style='color: #0044CC;' href='https://ncut-pytorch.readthedocs.io/en/latest/how_to_get_better_segmentation/' target='_blank'>Help</a>")
     from ncut_pytorch.backbone import list_models, get_demo_model_names
     model_names = list_models()
@@ -2089,7 +2107,7 @@ def make_parameters_section(is_lisa=False, model_ratio=True, parameter_dropdown=
         positive_prompt.visible = False
         negative_prompt = gr.Textbox(label="Prompt (Negative)", elem_id="prompt", placeholder="e.g. 'a photo from egocentric view'")
         negative_prompt.visible = False
-        node_type_dropdown = gr.Dropdown(["attn: attention output", "mlp: mlp output", "block: sum of residual"], label="Backbone: Layer type", value="block: sum of residual", elem_id="node_type", info="which feature to take from each layer?")
+        node_type_dropdown = gr.Dropdown(["attn: attention output", "mlp: mlp output", "block: sum of residual"], label="Backbone: Layer type", value="block: sum of residual", elem_id="node_type")
     num_eig_slider = gr.Slider(1, 1000, step=1, label="NCUT: Number of eigenvectors", value=100, elem_id="num_eig", info='increase for smaller clusters')
 
     def change_layer_slider(model_name):
@@ -2125,7 +2143,7 @@ def make_parameters_section(is_lisa=False, model_ratio=True, parameter_dropdown=
                 gr.Textbox(label="Prompt (Negative)", elem_id="prompt", placeholder="e.g. 'a photo from egocentric view'", visible=False))
     model_dropdown.change(fn=change_prompt_text, inputs=model_dropdown, outputs=[positive_prompt, negative_prompt])
     
-    with gr.Accordion("Advanced Parameters: NCUT", open=False, visible=parameter_dropdown):
+    with gr.Accordion("Advanced Parameters: NCUT", open=False, visible=ncut_parameter_dropdown):
         gr.Markdown("<a href='https://ncut-pytorch.readthedocs.io/en/latest/how_to_get_better_segmentation/' target='_blank'>Docs: How to Get Better Segmentation</a>")
         affinity_focal_gamma_slider = gr.Slider(0.01, 1, step=0.01, label="NCUT: Affinity focal gamma", value=0.5, elem_id="affinity_focal_gamma", info="decrease for shaper segmentation")
         num_sample_ncut_slider = gr.Slider(100, 50000, step=100, label="NCUT: num_sample", value=10000, elem_id="num_sample_ncut", info="Nyström approximation")
@@ -2136,7 +2154,7 @@ def make_parameters_section(is_lisa=False, model_ratio=True, parameter_dropdown=
         ncut_knn_slider = gr.Slider(1, 100, step=1, label="NCUT: KNN", value=10, elem_id="knn_ncut", info="Nyström approximation")
         ncut_indirect_connection = gr.Checkbox(label="indirect_connection", value=True, elem_id="ncut_indirect_connection", info="Add indirect connection to the sub-sampled graph")
         ncut_make_orthogonal = gr.Checkbox(label="make_orthogonal", value=False, elem_id="ncut_make_orthogonal", info="Apply post-hoc eigenvectors orthogonalization")
-    with gr.Accordion("Advanced Parameters: Visualization", open=False, visible=parameter_dropdown):
+    with gr.Accordion("Advanced Parameters: Visualization", open=False, visible=tsne_parameter_dropdown):
         # embedding_method_dropdown = gr.Dropdown(["tsne_3d", "umap_3d", "umap_sphere", "tsne_2d", "umap_2d"], label="Coloring method", value="tsne_3d", elem_id="embedding_method")
         embedding_method_dropdown = gr.Radio(["tsne_3d", "umap_3d", "umap_sphere", "tsne_2d", "umap_2d"], label="Coloring method", value="tsne_3d", elem_id="embedding_method")
         # embedding_metric_dropdown = gr.Dropdown(["euclidean", "cosine"], label="t-SNE/UMAP metric", value="euclidean", elem_id="embedding_metric")
@@ -2168,8 +2186,396 @@ demo = gr.Blocks(
 )
 with demo:
 
+    with gr.Tab('Hierarchical (dev)'):
+        eigvecs = gr.State(np.array([]))
+        tsne3d_rgb = gr.State(np.array([]))
+        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=100, markdown=False)
+                submit_button.value = "🔴 RUN NCUT"
+                num_images_slider.value = 100
+                
+                false_placeholder = gr.Checkbox(label="False", value=False, elem_id="false_placeholder", visible=False)
+                no_prompt = gr.Textbox("", label="", elem_id="empty_placeholder", type="text", placeholder="", visible=False)
+                
+            with gr.Column(scale=5, min_width=200):
+                # gr.Markdown("### Step 2a: Run Backbone and NCUT")
+                # with gr.Accordion(label="Backbone Parameters", visible=True, open=False):
+                output_gallery = gr.Gallery(format='png', value=[], label="NCUT spectral-tSNE", show_label=True, elem_id="ncut", columns=[3], rows=[1], object_fit="contain", height="auto", show_share_button=True, interactive=False)
+                def add_rotate_flip_buttons_with_state(output_gallery, tsne3d_rgb):
+                    with gr.Row():
+                        rotate_button = gr.Button("🔄 Rotate", elem_id="rotate_button", variant='secondary')
+                        rotate_button.click(sequence_rotate_rgb_gallery, inputs=[output_gallery], outputs=[output_gallery])
+                        def rotate_state(arr):
+                            rotation_matrix = np.array([[0, 1, 0], [0, 0, 1], [1, 0, 0]]).astype(np.float32)
+                            return arr @ rotation_matrix
+                        rotate_button.click(rotate_state, inputs=[tsne3d_rgb], outputs=[tsne3d_rgb])
+                        flip_button = gr.Button("🔃 Flip", elem_id="flip_button", variant='secondary')
+                        flip_button.click(flip_rgb_gallery, inputs=[output_gallery], outputs=[output_gallery])
+                        def flip_state(arr):
+                            return 1 - arr
+                        flip_button.click(flip_state, inputs=[tsne3d_rgb], outputs=[tsne3d_rgb])
+                    return rotate_button, flip_button
+                add_rotate_flip_buttons_with_state(output_gallery, tsne3d_rgb)
+                
+                [
+                    model_dropdown, layer_slider, node_type_dropdown, num_eig_slider, 
+                    affinity_focal_gamma_slider, num_sample_ncut_slider, ncut_knn_slider, ncut_indirect_connection, ncut_make_orthogonal, 
+                    embedding_method_dropdown, embedding_metric_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
+                    perplexity_slider, n_neighbors_slider, min_dist_slider, 
+                    sampling_method_dropdown, ncut_metric_dropdown, positive_prompt, negative_prompt
+                ] = make_parameters_section(ncut_parameter_dropdown=True, tsne_parameter_dropdown=True)
+                # submit_button = gr.Button("🔴 RUN NCUT", elem_id="run_ncut", variant='primary')
+                logging_text = gr.Textbox("Logging information", lines=3, label="Logging", elem_id="logging", type="text", placeholder="Logging information", autofocus=False, autoscroll=False)
+                
+                def __run_fn(*args, **kwargs):
+                    eigvecs, rgb, logging_str = run_fn(*args, **kwargs)
+                    rgb_gallery = to_pil_images(rgb)
+                    return eigvecs, rgb, rgb_gallery, logging_str
+                
+                submit_button.click(
+                    partial(__run_fn, n_ret=2, return_eigvec_and_rgb=True),
+                    inputs=[
+                        input_gallery, model_dropdown, layer_slider, num_eig_slider, node_type_dropdown, 
+                        positive_prompt, negative_prompt,
+                        false_placeholder, no_prompt, no_prompt, no_prompt,
+                        affinity_focal_gamma_slider, num_sample_ncut_slider, ncut_knn_slider, ncut_indirect_connection, ncut_make_orthogonal, 
+                        embedding_method_dropdown, embedding_metric_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
+                        perplexity_slider, n_neighbors_slider, min_dist_slider, sampling_method_dropdown, ncut_metric_dropdown
+                    ],
+                    outputs=[eigvecs, tsne3d_rgb, output_gallery, logging_text],
+                )
+            
+            with gr.Column(scale=5, min_width=200):
+                gr.Markdown('---')
+                gr.Markdown('<h3 style="text-align: center;">Help</h3>')
+                gr.Markdown('---')
+                with gr.Accordion("Instructions", open=True):
+                    gr.Markdown("""                                
+                                1. Load Dataset (left).
+                                2. Choose parameters (middle).
+                                3. 🔴 RUN NCUT.
+                                4. 🔴 RUN FPS+Cluster.
+                                5. Interact and Inspect (scroll down).
+                                """)
+                with gr.Accordion("Methods: NCUT spectral-TSNE", open=False):
+                    gr.Markdown("### <a style='color: #0044CC;' href='https://ncut-pytorch.readthedocs.io/en/latest/how_ncut_works/' target='_blank'>Documentation: How NCUT works</a>")
+                    gr.Markdown("""
+                                1. Run Backbone, feature extraction for each image.
+                                2. Vectorize latent-pixels, concatenate all the images.
+                                3. Run NCUT, on one big graph of all the images.
+                                4. Run spectral-tSNE on the NCUT eigenvectors.
+                                5. Plot the 3D spectral-tSNE as RGB.
+                                """)
+                with gr.Accordion("Methods: Hierarchical Structure", open=False):
+                    gr.Markdown("""
+                                1. Farthest Point Sampling (FPS) on the eigenvectors.
+                                2. spectral-tSNE (2D) on the FPS sampled points.
+                                3. Hierarchical clustering on the FPS sampled points.
+                                """)
+                gr.Markdown('---')
+                
+                run_hierarchical_button = gr.Button("🔴 RUN FPS+Cluster", elem_id="run_hierarchical", variant='primary')
+                with gr.Accordion("Hierarchical Structure Parameters:", open=True):
+                    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")
+                tsne_plot = gr.Image(label="spectral-tSNE tree", elem_id="tsne_plot", interactive=False, format='png')
+                
+                tsne_2d_points = gr.State(np.array([]))
+                edges = 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):
+                    # 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)
+                    # 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)
+                    # highlight the selected node
+                    if hightlight_idx is not None:
+                        if highlight_connections:
+                            from fps_cluster import find_connected_component
+                            _edges = edges[k:, :]
+                            connected_nodes = find_connected_component(_edges, hightlight_idx)
+                            ax.scatter(tsne_embed[connected_nodes, 0], tsne_embed[connected_nodes, 1], s=50, c=fps_tsne3d_rgb[connected_nodes], marker='D', edgecolor='deeppink', linewidth=1)
+                        # ax.scatter(tsne_embed[hightlight_idx, 0], tsne_embed[hightlight_idx, 1], s=300, c='r', marker='x')
+                        ax.scatter(tsne_embed[hightlight_idx, 0], tsne_embed[hightlight_idx, 1], s=200, c='cyan', marker='o', edgecolor='black', linewidth=1)
+                    ax.set_xticks([])
+                    ax.set_yticks([])
+                    ax.axis('off')
+                    ax.set_xlim(tsne_embed[:, 0].min()*1.1, tsne_embed[:, 0].max()*1.1)
+                    ax.set_ylim(tsne_embed[:, 1].min()*1.1, tsne_embed[:, 1].max()*1.1)
+                    
+                    # Remove the white space around the plot
+                    fig.tight_layout(pad=0)
+
+                    # Save the plot to an in-memory buffer
+                    buf = io.BytesIO()
+                    plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0)
+                    buf.seek(0)
+
+                    # Load the image into a NumPy array
+                    image = np.array(Image.open(buf))
 
+                    # Close the buffer and plot
+                    buf.close()
+                    plt.close(fig)
+                    
+                    pil_image = Image.fromarray(image)
+                    return pil_image
                                 
+                def run_fps_tsne_hierarchical(eigvecs, num_sample_fps, perplexity_tsne, tsne3d_rgb, seed=0):
+                    if len(eigvecs) == 0:
+                        gr.Warning("Please run NCUT first.")
+                        return
+                    eigvecs = torch.tensor(eigvecs)
+                    eigvecs = eigvecs.reshape(-1, eigvecs.shape[-1])
+                    gr.Info("Running FPS, t-SNE, and Hierarchical Clustering...", 3)
+                    from ncut_pytorch.ncut_pytorch import farthest_point_sampling
+                    from sklearn.manifold import TSNE
+                    from fps_cluster import build_tree
+                    
+                    torch.manual_seed(seed)
+                    np.random.seed(seed)
+                    
+                    fps_idx = farthest_point_sampling(eigvecs, num_sample_fps)
+                    fps_eigvecs = eigvecs[fps_idx]
+                    fps_eigvecs = fps_eigvecs.numpy()
+                    
+                    tsne3d_rgb = tsne3d_rgb.reshape(-1, 3)
+                    fps_tsne3d_rgb = tsne3d_rgb[fps_idx]
+                    
+                    np.random.seed(seed)
+                    tsne_embed = TSNE(
+                        n_components=2,
+                        perplexity=perplexity_tsne,
+                        metric='cosine',
+                        random_state=seed,
+                    ).fit_transform(fps_eigvecs)
+                    
+                    edges = build_tree(tsne_embed)
+                    
+                    # Plot the t-SNE points
+                    pil_image = plot_tsne_tree(tsne_embed, edges, fps_tsne3d_rgb, 0)
+                    
+                    return tsne_embed, edges, fps_eigvecs, fps_tsne3d_rgb, fps_idx, pil_image
+
+                run_hierarchical_button.click(
+                    run_fps_tsne_hierarchical,
+                    inputs=[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],
+                )
+        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:
+                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:
+                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):
+                    if len(output_gallery) == 0:
+                        return gr.update(value=None, interactive=False)
+                    image_idx = int(image_slider)
+                    image = output_gallery[image_idx][0]
+                    return gr.update(value={'image': image}, interactive=True)
+                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):
+                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. 
+                                - Only the last clicked dot will be used
+                                - Eraser is at top-right corner
+                                - Use the right-side Radio to switch tree/image
+                                2. Choose a granularity (right-side).
+                                3. 🔴 RUN Inspection.
+                                4. Output will be shown below.
+                                """)
+                with gr.Accordion("Outputs", open=True):
+                    gr.Markdown("""
+                                1. spectral-tSNE tree: ◆ means connected components to the selected point.
+                                2. Cluster Heatmap: max cosine similarity to any 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)
+                granularity_slider = gr.Slider(1, 1000, step=1, label="Cluster Granularity", 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)
+                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])
+                tsne_plot.change(updaste_tsne_plot_change_granularity,
+                                    inputs=[granularity_slider, tsne_2d_points, edges, fps_tsne_rgb],
+                                    outputs=[tsne_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)
+            
+            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])
+            
+        def make_one_output_row(i_row=1):
+            with gr.Row() as inspect_output_row:
+                with gr.Column(scale=5, min_width=200):
+                    output_tree_image = gr.Image(label=f"spectral-tSNE tree [row#{i_row}]", elem_id="output_image", interactive=False)
+                    text_block = gr.Textbox("", label="Logging", elem_id=f"logging_{i_row}", type="text", placeholder="Logging information", autofocus=False, autoscroll=False, lines=2, show_label=False)
+                with gr.Column(scale=10, min_width=200):
+                    heatmap_gallery = gr.Gallery(format='png', value=[], label=f"Cluster Heatmap [row#{i_row}]", show_label=True, elem_id="heatmap", columns=[6], rows=[1], object_fit="contain", height="auto", show_share_button=True, interactive=False)
+            return inspect_output_row, output_tree_image, heatmap_gallery, text_block
+        
+        gr.Markdown('---')
+        MAX_ROWS = 100
+        current_output_row = gr.State(MAX_ROWS-1)
+        inspect_output_rows, output_tree_images, heatmap_galleries, text_blocks = [], [], [], []
+        for i_row in range(MAX_ROWS, 0, -1):
+            inspect_output_row, output_tree_image, heatmap_gallery, text_block = make_one_output_row(i_row)
+            inspect_output_row.visible = False
+            inspect_output_rows.append(inspect_output_row)
+            output_tree_images.append(output_tree_image)
+            heatmap_galleries.append(heatmap_gallery)
+            text_blocks.append(text_block)
+            
+        
+        def relative_xy_last_positive(prompts):
+            image = prompts['image']
+            points = np.asarray(prompts['points'])
+            if points.shape[0] == 0:
+                return [], []
+            is_point = points[:, 5] == 4.0
+            points = points[is_point]
+            is_positive = points[:, 2] == 1.0
+            if is_positive.sum() == 0:
+                raise Exception("No blue point is selected.")
+            is_negative = points[:, 2] == 0.0
+            xy = points[:, :2].tolist()
+            if isinstance(image, str):
+                image = Image.open(image)
+                image = np.array(image)
+            h, w = image.shape[:2]
+            new_xy = [(x/w, y/h) for x, y in xy]
+
+            last_positive_idx = np.where(is_positive)[0][-1]
+            x, y = new_xy[last_positive_idx]
+            return x, y
+        
+        def find_closest_fps_point_for_tsne_tree_plot(tsne_prompt, tsne2d_embed):
+            x, y = relative_xy_last_positive(tsne_prompt)
+            x_vmax = tsne2d_embed[:, 0].max() * 1.1
+            x_vmin = tsne2d_embed[:, 0].min() * 1.1
+            y_vmax = tsne2d_embed[:, 1].max() * 1.1
+            y_vmin = tsne2d_embed[:, 1].min() * 1.1
+            x = x * (x_vmax - x_vmin) + x_vmin
+            y = 1 - y
+            y = y * (y_vmax - y_vmin) + y_vmin
+            dist = np.linalg.norm(tsne2d_embed - np.array([x, y]), axis=1)
+            closest_idx = np.argmin(dist)
+            return closest_idx
+        
+        def find_closest_fps_point_for_image_prompt(image_prompt, i_image, eigvecs, fps_eigvecs):
+            x, y = relative_xy_last_positive(image_prompt)
+            _eigvec = eigvecs[i_image]
+            h, w = _eigvec.shape[:2]
+            x = int(x * w)
+            y = int(y * h)
+            eigvec = _eigvec[y, x]
+            dist = np.linalg.norm(fps_eigvecs - eigvec, axis=1)
+            closest_idx = np.argmin(dist)
+            return closest_idx
+        
+        def find_closest_fps_point(prompt_radio, 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)
+            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, output_slot, tsne2d_embed, edges, fps_eigvecs, fps_tsne_rgb, fps_indices, granularity, eigvecs, i_image, tsne3d_rgb, input_gallery, max_rows=MAX_ROWS):
+            if len(tsne2d_embed) == 0:
+                raise gr.Error("Please run FPS+Cluster first.")
+            closest_idx = find_closest_fps_point(prompt_radio, 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)
+            
+            from fps_cluster import find_connected_component
+            connected_idxs = find_connected_component(edges[granularity:], closest_idx)
+            
+            logging_text = f"Clicked: idx={closest_idx}, RGB: {closest_rgb.tolist()}\n"
+            logging_text += f"Granularity: 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)
+            
+            # draw heatmap for the connected components
+            connected_eigvecs = fps_eigvecs[connected_idxs]
+            left = torch.tensor(eigvecs).float()  # B H W 3
+            right = torch.tensor(connected_eigvecs).float()
+            left = F.normalize(left, p=2, dim=-1)
+            right = F.normalize(right, p=2, dim=-1)
+            similarity = left @ right.T
+            similarity = similarity.max(dim=-1).values  # B H W
+            hot_map = matplotlib.cm.get_cmap('hot')
+            heatmap = hot_map(similarity)[..., :3]  # B H W 3
+            heatmap_images = to_pil_images(torch.tensor(heatmap), target_size=256, force_size=True)
+            # overlay input images on the heatmap
+            input_images = [x[0] for x in input_gallery]
+            if isinstance(input_images[0], str):
+                input_images = [Image.open(x) for x in input_images]
+            for i, img in enumerate(input_images):
+                _img = img.resize((256, 256)).convert('RGB')
+                _heatmap = heatmap_images[i].resize((256, 256)).convert('RGB')
+                blend = np.array(_img) * 0.5 + np.array(_heatmap) * 0.5
+                blend = Image.fromarray(blend.astype(np.uint8))
+                heatmap_images[i] = blend
+            
+            # tree_label = f"spectral-tSNE tree [row#{max_rows-output_slot}] k={granularity} idx={closest_idx} n={len(connected_idxs)}"
+            tree_label = f"spectral-tSNE tree [row#{max_rows-output_slot}]"
+            heatmap_label = f"Cluster Heatmap [row#{max_rows-output_slot}] k={granularity} idx={closest_idx} n={len(connected_idxs)}"
+            # update the output slots
+            output_rows = [gr.update() for _ in range(max_rows)]
+            output_tsne_plots = [gr.update() for _ in range(max_rows)]
+            output_heatmaps = [gr.update() for _ in range(max_rows)]
+            output_texts = [gr.update() for _ in range(max_rows)]
+            output_rows[output_slot] = gr.update(visible=True)
+            output_tsne_plots[output_slot] = gr.update(value=output_tsne_plot, label=tree_label)
+            output_heatmaps[output_slot] = gr.update(value=heatmap_images, label=heatmap_label)
+            output_texts[output_slot] = gr.update(value=logging_text)
+            gr.Info(f"Output in [row#{max_rows-output_slot}]", 3)
+            logging_text += f"\nOutput: [row#{max_rows-output_slot}]"
+            output_slot -= 1
+            if output_slot < 0:
+                output_slot = max_rows - 1
+            
+            return *output_rows, *output_tsne_plots, *output_heatmaps, *output_texts, output_slot, logging_text
+            
+            
+        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],
+            outputs=inspect_output_rows + output_tree_images + heatmap_galleries + text_blocks + [current_output_row, inspect_logging_text],
+        )
+        
     with gr.Tab('AlignedCut'):
 
         with gr.Row():
@@ -3448,7 +3854,7 @@ with demo:
         with gr.Row():
             with gr.Column(scale=5, min_width=200):
                 gr.Markdown("### Step 1: Load Images and Run NCUT")
-                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=100)
                 # submit_button.visible = False
                 num_images_slider.value = 30
                 [
@@ -3457,7 +3863,7 @@ with demo:
                     embedding_method_dropdown, embedding_metric_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
                     perplexity_slider, n_neighbors_slider, min_dist_slider, 
                     sampling_method_dropdown, ncut_metric_dropdown, positive_prompt, negative_prompt
-                ] = make_parameters_section(parameter_dropdown=False)
+                ] = make_parameters_section(ncut_parameter_dropdown=False)
                 num_eig_slider.value = 1000
                 num_eig_slider.visible = False
                 logging_text = gr.Textbox("Logging information", label="Logging", elem_id="logging", type="text", placeholder="Logging information", autofocus=False, autoscroll=False)
@@ -3585,7 +3991,7 @@ with demo:
                 pil_images = overlaied_images
             return pil_images, (y, x)
 
-        def farthest_point_sampling(
+        def _farthest_point_sampling(
             features,
             start_feature,
             num_sample=300,
@@ -3635,7 +4041,7 @@ with demo:
             num_childs = min(4, masked_eigvecs.shape[0])
             assert num_childs > 0
             
-            child_idx = farthest_point_sampling(masked_eigvecs, _picked_eigvec, num_sample=num_childs+1)
+            child_idx = _farthest_point_sampling(masked_eigvecs, _picked_eigvec, num_sample=num_childs+1)
             child_idx = np.sort(child_idx)[:-1]
             
             # convert child_idx to flat_idx
@@ -3718,7 +4124,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=100)
                 submit_button.visible = False
                 num_images_slider.value = 30
                 
@@ -3735,7 +4141,7 @@ with demo:
                         embedding_method_dropdown, embedding_metric_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
                         perplexity_slider, n_neighbors_slider, min_dist_slider, 
                         sampling_method_dropdown, ncut_metric_dropdown, positive_prompt, negative_prompt
-                    ] = make_parameters_section(parameter_dropdown=False)
+                    ] = make_parameters_section(ncut_parameter_dropdown=False, tsne_parameter_dropdown=False)
                     num_eig_slider.value = 1024
                     num_eig_slider.visible = False
                 submit_button = gr.Button("🔴 RUN NCUT", elem_id="run_ncut", variant='primary')

fps_cluster.py

@@ -0,0 +1,78 @@
+# %%
+import numpy as np
+import torch
+
+
+def build_tree(all_dots):
+    num_sample = all_dots.shape[0]
+    center = all_dots.mean(axis=0)
+    distances_to_center = np.linalg.norm(all_dots - center, axis=1)
+    start_idx = np.argmin(distances_to_center)
+    indices = [start_idx]
+    distances = [114514,]
+    A = all_dots[:, None] - all_dots[None, :]
+    A = (A ** 2).sum(-1)
+    A = np.sqrt(A)
+    A = torch.tensor(A)
+    for i in range(num_sample - 1):
+        _A = A[indices]
+        min_dist = _A.min(dim=0).values
+        next_idx = torch.argmax(min_dist).item()
+        distance = min_dist[next_idx].item()
+        indices.append(next_idx)
+        distances.append(distance)
+    indices = np.array(indices)
+    distances = np.array(distances)
+
+    levels = np.log2(distances[1] / distances)
+    levels = np.floor(levels).astype(int) + 1
+    levels[0] = 0
+
+    n_levels = levels.max() + 1
+    pi_indices = [indices[0],]
+    for i_level in range(1, n_levels):
+        current_level_indices = levels == i_level
+        prev_level_indices = levels < i_level
+        current_level_indices = indices[current_level_indices]
+        prev_level_indices = indices[prev_level_indices]
+        _A = A[prev_level_indices][:, current_level_indices]
+        _pi = _A.min(dim=0).indices
+        pi = prev_level_indices[_pi]
+        if isinstance(pi, np.int64) or isinstance(pi, int):
+            pi = [pi,]
+        if isinstance(pi, np.ndarray):
+            pi = pi.tolist()
+        pi_indices.extend(pi)
+    pi_indices = np.array(pi_indices)
+
+    edges = np.stack([indices, pi_indices], axis=1)
+    return edges
+
+
+def find_connected_component(edges, start_node):
+    # Dictionary to store adjacency list
+    adjacency_list = {}
+    for edge in edges:
+        # Unpack edge
+        a, b = edge
+        # Add the connection for both nodes
+        if a in adjacency_list:
+            adjacency_list[a].append(b)
+        else:
+            adjacency_list[a] = [b]
+        if b in adjacency_list:
+            adjacency_list[b].append(a)
+        else:
+            adjacency_list[b] = [a]
+    
+    # Use BFS to find all nodes in the connected component
+    connected_component = set()
+    queue = [start_node]
+    
+    while queue:
+        node = queue.pop(0)
+        if node not in connected_component:
+            connected_component.add(node)
+            queue.extend(adjacency_list.get(node, []))  # Add neighbors to the queue
+
+    return np.array(list(connected_component))