add playground

app.py

@@ -993,26 +993,28 @@ 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:
+        gr.Info("NCUT Run Started", 2)
+        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]]
+        gr.Info("NCUT Run Finished", 2)
+        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)
@@ -3557,8 +3559,8 @@ with demo:
             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)
+            left = F.normalize(left, p=2, dim=-1)
+            _right = F.normalize(right, p=2, dim=-1)
             heatmap = left @ _right.unsqueeze(-1)
             heatmap = heatmap.squeeze(-1)
             heatmap = 1 - heatmap
@@ -3707,7 +3709,192 @@ with demo:
                 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('PlayGround', visible=True) as test_playground_tab2:
+        eigvecs = gr.State(torch.tensor([]))
+        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)
+                submit_button.visible = False
+                num_images_slider.value = 30
+                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)
+                
+                
+            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):
+                    [
+                        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 = 1024
+                    num_eig_slider.visible = False
+                submit_button = gr.Button("🔴 RUN NCUT", elem_id="run_ncut", variant='primary')
+                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],
+                )
+                gr.Markdown("### Step 2b: 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 Image", elem_id="load_one_image_button", variant='primary')
+                gr.Markdown("### Step 2c: 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)
+                child_distance_slider.visible = False
+                overlay_image_checkbox = gr.Checkbox(label="Overlay Image", value=True, elem_id="overlay_image_checkbox", interactive=True)
+                n_eig_slider = gr.Slider(0, 1024, step=1, label="Number of Eigenvectors", value=256, elem_id="n_eig_slider", interactive=True)
+                run_button = gr.Button("🔴 RUN", elem_id="run_groupping", variant='primary')
+                current_plot = gr.Gallery(value=None, label="Current", show_label=True, elem_id="current_plot", interactive=False, rows=[1], columns=[2])
+                with gr.Row():
+                    doublue_eigs_button = gr.Button("⬇️ +eigvecs", elem_id="doublue_eigs_button", variant='secondary')
+                    half_eigs_button = gr.Button("⬆️ -eigvecs", elem_id="half_eigs_button", variant='secondary')
     
+        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=-1)
+            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
+            # gr.Info(f"heatmap vmin: {heatmap.min()}, vmax: {heatmap.max()}, mean: {heatmap.mean()}")
+            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)
+
+        @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}", 2)
+            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)
+            
+            return current_heatmap
+        
+        def doublue_eigs_wrapper(images, eigvecs, image1_slider, prompt_image1, n_eig, distance_slider, flat_idx=None, overlay_image=True):
+            n_eig = int(n_eig*2)
+            n_eig = min(n_eig, eigvecs.shape[-1])
+            n_eig = max(n_eig, 1)
+            return gr.update(value=n_eig), run_heatmap(images, eigvecs, image1_slider, prompt_image1, n_eig, distance_slider, flat_idx, overlay_image=overlay_image)
+                
+        def half_eigs_wrapper(images, eigvecs, image1_slider, prompt_image1, n_eig, distance_slider, current_idx=None, overlay_image=True):
+            n_eig = int(n_eig/2)
+            n_eig = min(n_eig, eigvecs.shape[-1])
+            n_eig = max(n_eig, 1)
+            return gr.update(value=n_eig), run_heatmap(images, eigvecs, image1_slider, prompt_image1, n_eig, distance_slider, current_idx, overlay_image=overlay_image)
+                
+        none_placeholder = gr.State(None)
+        run_button.click(
+            run_heatmap,
+            inputs=[input_gallery, eigvecs, image1_slider, prompt_image1, n_eig_slider, child_distance_slider, none_placeholder, overlay_image_checkbox],
+            outputs=[current_plot],
+        )
+        
+        doublue_eigs_button.click(
+            doublue_eigs_wrapper,
+            inputs=[input_gallery, eigvecs, image1_slider, prompt_image1, n_eig_slider, child_distance_slider, none_placeholder, overlay_image_checkbox],
+            outputs=[n_eig_slider, current_plot],
+        )
+        
+        half_eigs_button.click(
+            half_eigs_wrapper,
+            inputs=[input_gallery, eigvecs, image1_slider, prompt_image1, n_eig_slider, child_distance_slider, current_idx, overlay_image_checkbox],
+            outputs=[n_eig_slider, current_plot],
+        )
+        
     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/)**")