add advanced tab for recursive

app.py

@@ -325,9 +325,9 @@ def make_cluster_plot(eigvecs, images, h=64, w=64, progess_start=0.6, advanced=F
         mask = mask[mask_sort_idx[:3]] 
         sort_values.append(mask.mean().item())
         # fps_heatmaps[idx.item()] = heatmap.cpu()
-        fps_heatmaps[idx.item()] = heatmap[mask_sort_idx[:10]].cpu()
+        fps_heatmaps[idx.item()] = heatmap[mask_sort_idx[:6]].cpu()
         top3_image_idx[idx.item()] = mask_sort_idx[:3]
-        top10_image_idx[idx.item()] = mask_sort_idx[:10]
+        top10_image_idx[idx.item()] = mask_sort_idx[:6]
     # do the sorting
     _sort_idx = torch.tensor(sort_values).argsort(descending=True)
     fps_idx = fps_idx[_sort_idx]
@@ -427,6 +427,7 @@ def ncut_run(
     alignedcut_eig_norm_plot=False,
     **kwargs,
 ):
+    advanced = kwargs.get("advanced", False)
     progress = gr.Progress()
     progress(0.2, desc="Feature Extraction")
     
@@ -483,6 +484,7 @@ def ncut_run(
     
     if recursion:
         rgbs = []
+        all_eigvecs = []
         recursion_gammas = [recursion_l1_gamma, recursion_l2_gamma, recursion_l3_gamma]
         inp = features
         progress_start = 0.4
@@ -509,6 +511,7 @@ def ncut_run(
                 progess_start=progress_start,
             )
             logging_str += _logging_str
+            all_eigvecs.append(eigvecs.cpu().clone())
             
             
             if "AlignedThreeModelAttnNodes" == model_name:
@@ -528,8 +531,42 @@ def ncut_run(
             inp = eigvecs.reshape(*features.shape[:-1], -1)
             if recursion_metric == "cosine":
                 inp = F.normalize(inp, dim=-1)
-        return rgbs[0], rgbs[1], rgbs[2], logging_str
-    
+                
+        if not advanced:
+            return rgbs[0], rgbs[1], rgbs[2], logging_str
+        if advanced:
+            cluster_plots, norm_plots = [], []
+            for i in range(3):
+                eigvecs = all_eigvecs[i]
+                # add norm plot, cluster plot
+                start = time.time()
+                progress_start = 0.6
+                progress(progress_start, desc=f"Plotting Clusters Recursion #{i+1}")
+                h, w = features.shape[1], features.shape[2]
+                if torch.cuda.is_available():
+                    images = images.cuda()
+                _images = reverse_transform_image(images, stablediffusion="stable" in model_name.lower())
+                cluster_images, eig_magnitude = make_cluster_plot(eigvecs, _images, h=h, w=w, progess_start=progress_start, advanced=advanced)
+                logging_str += f"Recursion #{i+1} plot time: {time.time() - start:.2f}s\n"    
+                
+                norm_images = []
+                vmin, vmax = eig_magnitude.min(), eig_magnitude.max()
+                eig_magnitude = (eig_magnitude - vmin) / (vmax - vmin)
+                eig_magnitude = eig_magnitude.cpu().numpy()
+                colormap = matplotlib.colormaps['Reds']
+                for i_image in range(eig_magnitude.shape[0]):
+                    norm_image = colormap(eig_magnitude[i_image])
+                    norm_images.append(torch.tensor(norm_image[..., :3]))
+                norm_images = to_pil_images(norm_images)
+                logging_str += f"Recursion #{i+1} Eigenvector Magnitude: [{vmin:.2f}, {vmax:.2f}]\n"
+                gr.Info(f"Recursion #{i+1} Eigenvector Magnitude:</br> Min: {vmin:.2f}, Max: {vmax:.2f}", duration=10)
+                
+                cluster_plots.append(cluster_images)
+                norm_plots.append(norm_images)
+            
+            return *rgbs, *norm_plots, *cluster_plots, logging_str
+            
+            
     if old_school_ncut:  # individual images
         logging_str += "Running NCut for each image independently\n"
         rgb = []
@@ -643,7 +680,7 @@ def ncut_run(
         norm_images = to_pil_images(norm_images)
         logging_str += "Eigenvector Magnitude\n"
         logging_str += f"Min: {vmin:.2f}, Max: {vmax:.2f}\n"
-        gr.Info(f"Eigenvector Magnitude:</br> Min: {vmin:.2f}, Max: {vmax:.2f}", duration=0)
+        gr.Info(f"Eigenvector Magnitude:</br> Min: {vmin:.2f}, Max: {vmax:.2f}", duration=10)
 
     return to_pil_images(rgb), cluster_images, norm_images, logging_str
 
@@ -651,26 +688,26 @@ def ncut_run(
 
 def _ncut_run(*args, **kwargs):
     n_ret = kwargs.pop("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 = ncut_run(*args, **kwargs)
-    ret = list(ret)[:n_ret] + [ret[-1]]
-    return ret
+        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
 
 if USE_HUGGINGFACE_ZEROGPU:
     @spaces.GPU(duration=30)
@@ -1407,7 +1444,7 @@ with demo:
             scroll_to_output=True,
         )
         
-    with gr.Tab('AlignedCut (+Norm Plot)', visible=False) as tab_alignedcut_norm:
+    with gr.Tab('AlignedCut (Advanced)', visible=False) as tab_alignedcut_advanced:
 
         with gr.Row():
             with gr.Column(scale=5, min_width=200):
@@ -1582,6 +1619,78 @@ with demo:
             outputs=[l1_gallery, l2_gallery, l3_gallery, logging_text],
             api_name="API_RecursiveCut"
         )
+        
+    with gr.Tab('Recursive Cut (Advanced)', visible=False) as tab_recursivecut_advanced:
+
+        with gr.Row():
+            with gr.Column(scale=5, min_width=200):
+                gr.Markdown('### Output (Recursion #1)')
+                l1_gallery = gr.Gallery(format='png', value=[], label="Recursion #1", show_label=True, elem_id="ncut_l1", columns=[3], rows=[5], object_fit="contain", height="auto", show_fullscreen_button=True, interactive=False)
+                add_output_images_buttons(l1_gallery)
+                l1_norm_gallery = gr.Gallery(value=[], label="Eigenvector Magnitude", show_label=True, elem_id="eig_norm", columns=[3], rows=[1], object_fit="contain", height="auto", show_share_button=True, preview=False, interactive=False)
+                l1_cluster_gallery = gr.Gallery(value=[], label="Clusters", show_label=True, elem_id="clusters", columns=[2], rows=[4], object_fit="contain", height=600, show_share_button=True, preview=True, interactive=False)
+            with gr.Column(scale=5, min_width=200):
+                gr.Markdown('### Output (Recursion #2)')
+                l2_gallery = gr.Gallery(format='png', value=[], label="Recursion #2", show_label=True, elem_id="ncut_l2", columns=[3], rows=[5], object_fit="contain", height="auto", show_fullscreen_button=True, interactive=False)
+                add_output_images_buttons(l2_gallery)
+                l2_norm_gallery = gr.Gallery(value=[], label="Eigenvector Magnitude", show_label=True, elem_id="eig_norm", columns=[3], rows=[1], object_fit="contain", height="auto", show_share_button=True, preview=False, interactive=False)
+                l2_cluster_gallery = gr.Gallery(value=[], label="Clusters", show_label=True, elem_id="clusters", columns=[2], rows=[4], object_fit="contain", height=600, show_share_button=True, preview=True, interactive=False)
+            with gr.Column(scale=5, min_width=200):
+                gr.Markdown('### Output (Recursion #3)')
+                l3_gallery = gr.Gallery(format='png', value=[], label="Recursion #3", show_label=True, elem_id="ncut_l3", columns=[3], rows=[5], object_fit="contain", height="auto", show_fullscreen_button=True, interactive=False)
+                add_output_images_buttons(l3_gallery)
+                l3_norm_gallery = gr.Gallery(value=[], label="Eigenvector Magnitude", show_label=True, elem_id="eig_norm", columns=[3], rows=[1], object_fit="contain", height="auto", show_share_button=True, preview=False, interactive=False)
+                l3_cluster_gallery = gr.Gallery(value=[], label="Clusters", show_label=True, elem_id="clusters", columns=[2], rows=[4], object_fit="contain", height=600, show_share_button=True, preview=True, interactive=False)
+        with gr.Row():
+            with gr.Column(scale=5, min_width=200):
+                input_gallery, submit_button, clear_images_button = make_input_images_section()
+                dataset_dropdown, num_images_slider, random_seed_slider, load_images_button = make_dataset_images_section(advanced=True)
+                num_images_slider.value = 100
+                clear_images_button.visible = False
+                logging_text = gr.Textbox("Logging information", label="Logging", elem_id="logging", type="text", placeholder="Logging information")
+                
+            with gr.Column(scale=5, min_width=200):
+                with gr.Accordion("➡️ Recursion config", open=True):
+                    l1_num_eig_slider = gr.Slider(1, 1000, step=1, label="Recursion #1: N eigenvectors", value=100, elem_id="l1_num_eig")
+                    l2_num_eig_slider = gr.Slider(1, 1000, step=1, label="Recursion #2: N eigenvectors", value=50, elem_id="l2_num_eig")
+                    l3_num_eig_slider = gr.Slider(1, 1000, step=1, label="Recursion #3: N eigenvectors", value=50, elem_id="l3_num_eig")
+                    metric_dropdown = gr.Dropdown(["euclidean", "cosine"], label="Recursion distance metric", value="cosine", elem_id="recursion_metric")
+                    l1_affinity_focal_gamma_slider = gr.Slider(0.01, 1, step=0.01, label="Recursion #1: Affinity focal gamma", value=0.5, elem_id="recursion_l1_gamma")
+                    l2_affinity_focal_gamma_slider = gr.Slider(0.01, 1, step=0.01, label="Recursion #2: Affinity focal gamma", value=0.5, elem_id="recursion_l2_gamma")
+                    l3_affinity_focal_gamma_slider = gr.Slider(0.01, 1, step=0.01, label="Recursion #3: Affinity focal gamma", value=0.5, elem_id="recursion_l3_gamma")
+                [
+                    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()
+                num_eig_slider.visible = False
+                affinity_focal_gamma_slider.visible = False
+        true_placeholder = gr.Checkbox(label="True placeholder", value=True, elem_id="true_placeholder")
+        true_placeholder.visible = False
+        false_placeholder = gr.Checkbox(label="False placeholder", value=False, elem_id="false_placeholder")
+        false_placeholder.visible = False
+        number_placeholder = gr.Number(0, label="Number placeholder", elem_id="number_placeholder")
+        number_placeholder.visible = False
+        clear_images_button.click(lambda x: ([],), outputs=[input_gallery])
+        no_prompt = gr.Textbox("", label="", elem_id="empty_placeholder", type="text", placeholder="", visible=False)
+        
+        submit_button.click(
+            partial(run_fn, n_ret=9, advanced=True),
+            inputs=[
+                input_gallery, model_dropdown, layer_slider, l1_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,
+                false_placeholder, number_placeholder, true_placeholder,
+                l2_num_eig_slider, l3_num_eig_slider, metric_dropdown, 
+                l1_affinity_focal_gamma_slider, l2_affinity_focal_gamma_slider, l3_affinity_focal_gamma_slider
+            ],
+            outputs=[l1_gallery, l2_gallery, l3_gallery, l1_norm_gallery, l2_norm_gallery, l3_norm_gallery, l1_cluster_gallery, l2_cluster_gallery, l3_cluster_gallery, logging_text],
+        )
 
         
     with gr.Tab('Video'): 
@@ -1741,7 +1850,7 @@ with demo:
             outputs=[output_gallery, logging_text],
         )
         
-    with gr.Tab('Model Aligned (+Rrecursion)', visible=False) as tab_model_aligned_recursion: 
+    with gr.Tab('Model Aligned (Advanced)', visible=False) as tab_model_aligned_advanced: 
         gr.Markdown('This page reproduce the results from the paper [AlignedCut](https://arxiv.org/abs/2406.18344)')
         gr.Markdown('---')
         gr.Markdown('**Features are aligned across models and layers.** A linear alignment transform is trained for each model/layer, learning signal comes from 1) fMRI brain activation and 2) segmentation preserving eigen-constraints.')
@@ -1965,8 +2074,9 @@ with demo:
             return gr.update()
         
         hidden_button.change(update_smile, [n_smiles], [n_smiles, hidden_button])
-        hidden_button.change(unlock_tabs_with_info, n_smiles, tab_alignedcut_norm)
-        hidden_button.change(unlock_tabs, n_smiles, tab_model_aligned_recursion)
+        hidden_button.change(unlock_tabs_with_info, n_smiles, tab_alignedcut_advanced)
+        hidden_button.change(unlock_tabs, n_smiles, tab_model_aligned_advanced)
+        hidden_button.change(unlock_tabs, n_smiles, tab_recursivecut_advanced)
                 
     with gr.Row():
         with gr.Column():