add test playground

app.py

@@ -2,6 +2,7 @@
 # %%
 import copy
 from datetime import datetime
+import math
 import pickle
 from functools import partial
 from io import BytesIO
@@ -137,6 +138,7 @@ def compute_ncut(
     indirect_connection=True,
     make_orthogonal=False,
     progess_start=0.4,
+    only_eigvecs=False,
 ):        
     progress = gr.Progress()
     logging_str = ""
@@ -165,6 +167,10 @@ def compute_ncut(
     # print(f"NCUT time: {time.time() - start:.2f}s")
     logging_str += f"NCUT time: {time.time() - start:.2f}s\n"
     
+    if only_eigvecs:
+        eigvecs = eigvecs.to("cpu").reshape(features.shape[:-1] + (num_eig,))
+        return None, logging_str, eigvecs
+    
     start = time.time()
     progress(progess_start+0.01, desc="spectral-tSNE")
     _, rgb = eigenvector_to_rgb(
@@ -272,10 +278,12 @@ def dont_use_too_much_green(image_rgb):
     return image_rgb
 
 
-def to_pil_images(images, target_size=512, resize=True):
+def to_pil_images(images, target_size=512, resize=True, force_size=False):
     size = images[0].shape[1]
     multiplier = target_size // size
     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
@@ -855,6 +863,8 @@ def ncut_run(
     
     # ailgnedcut
     if not directed:
+        only_eigvecs = kwargs.get("only_eigvecs", False)
+        
         rgb, _logging_str, eigvecs = compute_ncut(
             features,
             num_eig=num_eig,
@@ -872,7 +882,12 @@ def ncut_run(
             indirect_connection=indirect_connection,
             make_orthogonal=make_orthogonal,
             metric=ncut_metric,
+            only_eigvecs=only_eigvecs,
         )
+        
+        if only_eigvecs:
+            return eigvecs, logging_str
+        
     if directed:
         head_index_text = kwargs.get("head_index_text", None)
         n_heads = features.shape[-2]   # (batch, h, w, n_heads, d)
@@ -978,26 +993,26 @@ def ncut_run(
 
 def _ncut_run(*args, **kwargs):
     n_ret = kwargs.get("n_ret", 1)
-    try:
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
+    # try:
+    #     if torch.cuda.is_available():
+    #         torch.cuda.empty_cache()
             
-        ret = ncut_run(*args, **kwargs)
+    #     ret = ncut_run(*args, **kwargs)
         
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
+    #     if torch.cuda.is_available():
+    #         torch.cuda.empty_cache()
         
-        ret = list(ret)[:n_ret] + [ret[-1]]
-        return ret
-    except Exception as e:
-        gr.Error(str(e))
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-        return *(None for _ in range(n_ret)), "Error: " + str(e)
+    #     ret = list(ret)[:n_ret] + [ret[-1]]
+    #     return ret
+    # except Exception as e:
+    #     gr.Error(str(e))
+    #     if torch.cuda.is_available():
+    #         torch.cuda.empty_cache()
+    #     return *(None for _ in range(n_ret)), "Error: " + str(e)
 
-    # ret = ncut_run(*args, **kwargs)
-    # ret = list(ret)[:n_ret] + [ret[-1]]
-    # return ret
+    ret = ncut_run(*args, **kwargs)
+    ret = list(ret)[:n_ret] + [ret[-1]]
+    return ret
 
 if USE_HUGGINGFACE_ZEROGPU:
     @spaces.GPU(duration=30)
@@ -1213,6 +1228,7 @@ def run_fn(
     alignedcut_eig_norm_plot=False,
     advanced=False,
     directed=False,
+    only_eigvecs=False,
 ):
     # print(node_type2, head_index_text, make_symmetric)
     progress=gr.Progress()
@@ -1353,6 +1369,7 @@ def run_fn(
         "node_type2": node_type2,
         "head_index_text": head_index_text,
         "make_symmetric": make_symmetric,
+        "only_eigvecs": only_eigvecs,
     }
     # print(kwargs)
     
@@ -1664,7 +1681,7 @@ def load_and_append(existing_images, *args, **kwargs):
     gr.Info(f"Total images: {len(existing_images)}")
     return existing_images
 
-def make_input_images_section(rows=1, cols=3, height="auto", advanced=False, is_random=False, allow_download=False, markdown=True):
+def make_input_images_section(rows=1, cols=3, height="auto", advanced=False, is_random=False, allow_download=False, markdown=True, n_example_images=100):
     if markdown:
         gr.Markdown('### Input Images')
     input_gallery = gr.Gallery(value=None, label="Input images", show_label=True, elem_id="input_images", columns=[cols], rows=[rows], object_fit="contain", height=height, type="pil", show_share_button=False,
@@ -1702,7 +1719,7 @@ def make_input_images_section(rows=1, cols=3, height="auto", advanced=False, is_
             with gr.Row():
                 button = gr.Button("Load\n"+name, elem_id=f"example-{name}", elem_classes="small-button", variant='secondary', size="sm", scale=1, min_width=60)
                 gallery = gr.Gallery(value=images, label=name, show_label=True, columns=[3], rows=[1], interactive=False, height=80, scale=8, object_fit="cover", min_width=140, allow_preview=False)
-            button.click(fn=lambda: gr.update(value=load_dataset_images(True, dataset_name, 100, is_random=True, seed=42)), outputs=[input_gallery])
+            button.click(fn=lambda: gr.update(value=load_dataset_images(True, dataset_name, n_example_images, is_random=True, seed=42)), outputs=[input_gallery])
             return gallery, button
         example_items = [
             ("EgoExo", ['./images/egoexo1.jpg', './images/egoexo3.jpg', './images/egoexo2.jpg'], "EgoExo"),
@@ -2040,7 +2057,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):
+def make_parameters_section(is_lisa=False, model_ratio=True, 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()
@@ -2105,18 +2122,18 @@ def make_parameters_section(is_lisa=False, model_ratio=True):
                 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):
+    with gr.Accordion("Advanced Parameters: NCUT", open=False, visible=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")
         # sampling_method_dropdown = gr.Dropdown(["QuickFPS", "random"], label="NCUT: Sampling method", value="QuickFPS", elem_id="sampling_method", info="Nyström approximation")
         sampling_method_dropdown = gr.Radio(["QuickFPS", "random"], label="NCUT: Sampling method", value="QuickFPS", elem_id="sampling_method")
         # ncut_metric_dropdown = gr.Dropdown(["euclidean", "cosine"], label="NCUT: Distance metric", value="cosine", elem_id="ncut_metric")
-        ncut_metric_dropdown = gr.Radio(["euclidean", "cosine"], label="NCUT: Distance metric", value="cosine", elem_id="ncut_metric")
+        ncut_metric_dropdown = gr.Radio(["euclidean", "cosine", "rbf"], label="NCUT: Distance metric", value="cosine", elem_id="ncut_metric")
         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):
+    with gr.Accordion("Advanced Parameters: Visualization", open=False, visible=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")
@@ -2147,8 +2164,9 @@ demo = gr.Blocks(
     css=custom_css,
 )
 with demo:
-    
-    
+
+
+                                
     with gr.Tab('AlignedCut'):
 
         with gr.Row():
@@ -2989,7 +3007,7 @@ with demo:
                 # sampling_method_dropdown = gr.Dropdown(["QuickFPS", "random"], label="NCUT: Sampling method", value="QuickFPS", elem_id="sampling_method", info="Nyström approximation")
                 sampling_method_dropdown = gr.Radio(["QuickFPS", "random"], label="NCUT: Sampling method", value="QuickFPS", elem_id="sampling_method")
                 # ncut_metric_dropdown = gr.Dropdown(["euclidean", "cosine"], label="NCUT: Distance metric", value="cosine", elem_id="ncut_metric")
-                ncut_metric_dropdown = gr.Radio(["euclidean", "cosine"], label="NCUT: Distance metric", value="cosine", elem_id="ncut_metric")
+                ncut_metric_dropdown = gr.Radio(["euclidean", "cosine", "rbf"], label="NCUT: Distance metric", value="cosine", elem_id="ncut_metric")
                 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=False, elem_id="ncut_indirect_connection", info="TODO: Indirect connection is not implemented for directed NCUT", interactive=False)
                 ncut_make_orthogonal = gr.Checkbox(label="make_orthogonal", value=False, elem_id="ncut_make_orthogonal", info="Apply post-hoc eigenvectors orthogonalization")
@@ -3422,6 +3440,274 @@ with demo:
                     outputs=[mask_gallery, crop_gallery])
                 
     
+    with gr.Tab('PlayGround (test)', visible=False) as test_playground_tab:
+        eigvecs = gr.State(torch.tensor([]))
+        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)
+                # submit_button.visible = False
+                num_images_slider.value = 30
+                [
+                    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(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)
+                
+                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)
+                
+                submit_button.click(
+                    partial(run_fn, n_ret=1, only_eigvecs=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, logging_text],
+                )
+                
+            with gr.Column(scale=5, min_width=200):
+                gr.Markdown("### Step 2a: Pick an Image")
+                from gradio_image_prompter import ImagePrompter
+                with gr.Row():
+                    image1_slider = gr.Slider(0, 100, step=1, label="Image#1 Index", value=0, elem_id="image1_slider", interactive=True)
+                load_one_image_button = gr.Button("🔴 Load", elem_id="load_one_image_button", variant='primary')
+                gr.Markdown("### Step 2b: Draw a Point")
+                gr.Markdown("""
+                    <h5>
+                    🖱️ Left Click: Foreground </br>
+                    </h5>
+                """)
+                prompt_image1 = ImagePrompter(show_label=False, elem_id="prompt_image1", interactive=False)
+                def update_prompt_image(original_images, index):
+                    images = original_images
+                    if images is None:
+                        return 
+                    total_len = len(images)
+                    if total_len == 0:
+                        return 
+                    if index >= total_len:
+                        index = total_len - 1
+                                            
+                    return ImagePrompter(value={'image': images[index][0], 'points': []}, interactive=True)
+                    # return gr.Image(value=images[index][0], elem_id=f"prompt_image{randint}", interactive=True)
+                load_one_image_button.click(update_prompt_image, inputs=[input_gallery, image1_slider], outputs=[prompt_image1])
+                
+            child_idx = gr.State([])
+            current_idx = gr.State(None)
+            n_eig = gr.State(64)
+            with gr.Column(scale=5, min_width=200):
+                gr.Markdown("### Step 3: Check groupping")
+                child_distance_slider = gr.Slider(0, 0.5, step=0.001, label="Child Distance", value=0.1, elem_id="child_distance_slider", interactive=True)
+                overlay_image_checkbox = gr.Checkbox(label="Overlay Image", value=True, elem_id="overlay_image_checkbox", interactive=True)
+                run_button = gr.Button("🔴 RUN", elem_id="run_groupping", variant='primary')
+                parent_plot = gr.Gallery(value=None, label="Parent", show_label=True, elem_id="parent_plot", interactive=False, rows=[1], columns=[2])
+                parent_button = gr.Button("Use Parent", elem_id="run_parent")
+                current_plot = gr.Gallery(value=None, label="Current", show_label=True, elem_id="current_plot", interactive=False, rows=[1], columns=[2])
+            with gr.Column(scale=5, min_width=200):
+                child_plots = []
+                child_buttons = []
+                for i in range(4):
+                    child_plots.append(gr.Gallery(value=None, label=f"Child {i}", show_label=True, elem_id=f"child_plot_{i}", interactive=False, rows=[1], columns=[2]))
+                    child_buttons.append(gr.Button(f"Use Child {i}", elem_id=f"run_child_{i}"))
+    
+        def relative_xy(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
+            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]
+            # print(new_xy)
+            return new_xy, is_positive
+        
+        def xy_eigvec(prompts, image_idx, eigvecs):
+            eigvec = eigvecs[image_idx]
+            xy, is_positive = relative_xy(prompts)
+            for i, (x, y) in enumerate(xy):
+                if not is_positive[i]:
+                    continue
+                x = int(x * eigvec.shape[1])
+                y = int(y * eigvec.shape[0])
+                return eigvec[y, x], (y, x)
+        
+        from ncut_pytorch.ncut_pytorch import _transform_heatmap
+        def _run_heatmap_fn(images, eigvecs, prompt_image_idx, prompt_points, n_eig, flat_idx=None, raw_heatmap=False, overlay_image=True):
+            left = eigvecs[..., :n_eig]
+            if flat_idx is not None:
+                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)
+            right = right[:n_eig]
+            left = F.normalize(left, p=2, dim=1)
+            _right = F.normalize(right, p=2, dim=0)
+            heatmap = left @ _right.unsqueeze(-1)
+            heatmap = heatmap.squeeze(-1)
+            heatmap = 1 - heatmap
+            heatmap = _transform_heatmap(heatmap)
+            if raw_heatmap:
+                return heatmap
+            # apply hot colormap and covert to PIL image 256x256
+            heatmap = heatmap.cpu().numpy()
+            hot_map = matplotlib.cm.get_cmap('hot')
+            heatmap = hot_map(heatmap)
+            pil_images = to_pil_images(torch.tensor(heatmap), target_size=256, force_size=True)
+            if overlay_image:
+                overlaied_images = []
+                for i_image in range(len(images)):
+                    rgb_image = images[i_image].resize((256, 256))
+                    rgb_image = np.array(rgb_image)
+                    heatmap_image = np.array(pil_images[i_image])[..., :3]
+                    blend_image = 0.5 * rgb_image + 0.5 * heatmap_image
+                    blend_image = Image.fromarray(blend_image.astype(np.uint8))
+                    overlaied_images.append(blend_image)
+                pil_images = overlaied_images
+            return pil_images, (y, x)
+
+        def farthest_point_sampling(
+            features,
+            start_feature,
+            num_sample=300,
+            h=9,
+        ):
+            import fpsample
+
+            h = min(h, int(np.log2(features.shape[0])))
+            
+            inp = features.cpu().numpy()
+            inp = np.concatenate([inp, start_feature[None, :]], axis=0)
+
+            kdline_fps_samples_idx = fpsample.bucket_fps_kdline_sampling(
+                inp, num_sample, h, start_idx=inp.shape[0] - 1
+            ).astype(np.int64)
+            return kdline_fps_samples_idx
+        
+        @torch.no_grad()
+        def run_heatmap(images, eigvecs, image1_slider, prompt_image1, n_eig, distance_slider, flat_idx=None, overlay_image=True):
+            gr.Info(f"current number of eigenvectors: {n_eig}")
+            images = [image[0] for image in images]
+            if isinstance(images[0], str):
+                images = [Image.open(image[0]).convert("RGB").resize((256, 256)) for image in images]
+                        
+            current_heatmap, (y, x) = _run_heatmap_fn(images, eigvecs, image1_slider, prompt_image1, n_eig, flat_idx, overlay_image=overlay_image)
+            parent_heatmap, _ = _run_heatmap_fn(images, eigvecs, image1_slider, prompt_image1, int(n_eig/2), flat_idx, overlay_image=overlay_image)
+            
+            # find childs
+            # pca_eigvecs
+            _eigvecs = eigvecs.reshape(-1, eigvecs.shape[-1])
+            u, s, v = torch.pca_lowrank(_eigvecs, q=8)
+            _n = _eigvecs.shape[0]
+            s /= math.sqrt(_n)
+            _eigvecs = u @ torch.diag(s)
+            
+            if flat_idx is None:
+                _picked_eigvec = _eigvecs.reshape(*eigvecs.shape[:-1], 8)[image1_slider, y, x]
+            else:
+                _picked_eigvec = _eigvecs[flat_idx]
+            l2_distance = torch.norm(_eigvecs - _picked_eigvec, dim=-1)
+            average_distance = l2_distance.mean()
+            distance_threshold = distance_slider * average_distance
+            distance_mask = l2_distance < distance_threshold
+            masked_eigvecs = _eigvecs[distance_mask]
+            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 = np.sort(child_idx)[:-1]
+            
+            # convert child_idx to flat_idx
+            dummy_idx = torch.zeros(_eigvecs.shape[0], dtype=torch.bool)
+            dummy_idx2 = torch.zeros(int(distance_mask.sum().item()), dtype=torch.bool)
+            dummy_idx2[child_idx] = True
+            dummy_idx[distance_mask] = dummy_idx2
+            child_idx = torch.where(dummy_idx)[0]
+            
+            
+            # current_child heatmap, for contrast
+            current_child_heatmap = _run_heatmap_fn(images,eigvecs, image1_slider, prompt_image1, int(n_eig*2), flat_idx, raw_heatmap=True, overlay_image=overlay_image)
+            
+            # child_heatmaps, contrast mean of current clicked point
+            child_heatmaps = []
+            for idx in child_idx:
+                child_heatmap = _run_heatmap_fn(images,eigvecs, image1_slider, prompt_image1, int(n_eig*2), idx, raw_heatmap=True, overlay_image=overlay_image)
+                heatmap = child_heatmap - current_child_heatmap
+                # convert [-1, 1] to [0, 1]
+                heatmap = (heatmap + 1) / 2
+                heatmap = heatmap.cpu().numpy()
+                cm = matplotlib.cm.get_cmap('bwr')
+                heatmap = cm(heatmap)
+                # bwr with contrast
+                pil_images1 = to_pil_images(torch.tensor(heatmap), resize=256)
+                # no contrast
+                pil_images2, _ = _run_heatmap_fn(images, eigvecs, image1_slider, prompt_image1, int(n_eig*2), idx, overlay_image=overlay_image)
+                
+                # combine contrast and no contrast
+                pil_images = []
+                for i in range(len(pil_images1)):
+                    pil_images.append(pil_images2[i])
+                    pil_images.append(pil_images1[i])
+                
+                                    
+                child_heatmaps.append(pil_images)
+            
+            return parent_heatmap, current_heatmap, *child_heatmaps, child_idx.tolist()
+                
+        # def debug_fn(eigvecs):
+        #     shape = eigvecs.shape
+        #     gr.Info(f"eigvecs shape: {shape}")
+        
+        # run_button.click(
+        #     debug_fn,
+        #     inputs=[eigvecs],
+        #     outputs=[],
+        # )
+        none_placeholder = gr.State(None)
+        run_button.click(
+            run_heatmap,
+            inputs=[input_gallery, eigvecs, image1_slider, prompt_image1, n_eig, child_distance_slider, none_placeholder, overlay_image_checkbox],
+            outputs=[parent_plot, current_plot, *child_plots, child_idx],
+        )
+        
+        def run_paraent(input_gallery, eigvecs, image1_slider, prompt_image1, n_eig, distance_slider, current_idx=None, overlay_image=True):
+            n_eig = int(n_eig/2)
+            return n_eig, *run_heatmap(input_gallery, eigvecs, image1_slider, prompt_image1, n_eig, distance_slider, current_idx, overlay_image)
+        
+        parent_button.click(
+            run_paraent,
+            inputs=[input_gallery, eigvecs, image1_slider, prompt_image1, n_eig, child_distance_slider, current_idx, overlay_image_checkbox],
+            outputs=[n_eig, parent_plot, current_plot, *child_plots, child_idx],
+        )
+        
+        def run_child(input_gallery, eigvecs, image1_slider, prompt_image1, n_eig, distance_slider, child_idx=[], overlay_image=True, i_child=0):
+            n_eig = int(n_eig*2)
+            flat_idx = child_idx[i_child]
+            return n_eig, flat_idx, *run_heatmap(input_gallery, eigvecs, image1_slider, prompt_image1, n_eig, distance_slider, flat_idx, overlay_image)
+
+        for i in range(4):
+            child_buttons[i].click(
+                partial(run_child, i_child=i),
+                inputs=[input_gallery, eigvecs, image1_slider, prompt_image1, n_eig, child_distance_slider, child_idx, overlay_image_checkbox],
+                outputs=[n_eig, current_idx, parent_plot, current_plot, *child_plots, child_idx],
+            )
+    
     with gr.Tab('📄About'):
         with gr.Column():
             gr.Markdown("**This demo is for Python package `ncut-pytorch`, please visit the [Documentation](https://ncut-pytorch.readthedocs.io/)**")
@@ -3481,6 +3767,7 @@ with demo:
         hidden_button.change(unlock_tabs, n_smiles, tab_recursivecut_advanced)
         hidden_button.change(unlock_tabs, n_smiles, tab_compare_models_advanced)
         hidden_button.change(unlock_tabs, n_smiles, tab_directed_ncut)
+        hidden_button.change(unlock_tabs, n_smiles, test_playground_tab)
                 
     # with gr.Row():
         # with gr.Column():
@@ -3522,3 +3809,13 @@ demo.launch(share=True)
 # # %%
 
 # %%
+
+# %%
+
+# %%
+
+# %%
+
+# %%
+
+# %%