added more tabs

app.py

@@ -266,6 +266,10 @@ def make_cluster_plot(eigvecs, images, h=64, w=64, progess_start=0.6):
     magnitude = torch.norm(eigvecs, dim=-1)
     p = 0.8
     top_p_idx = magnitude.argsort(descending=True)[:int(p * magnitude.shape[0])]
+    
+    ret_magnitude = magnitude.reshape(-1, h, w)
+    
+    
     num_samples = 300
     if num_samples > top_p_idx.shape[0]:
         num_samples = top_p_idx.shape[0]
@@ -368,7 +372,7 @@ def make_cluster_plot(eigvecs, images, h=64, w=64, progess_start=0.6):
         # plt.imshow(img)
         # plt.axis("off")
         # plt.show()
-    return fig_images
+    return fig_images, ret_magnitude
 
     
 def ncut_run(
@@ -405,6 +409,9 @@ def ncut_run(
     lisa_prompt1="",
     lisa_prompt2="",
     lisa_prompt3="",
+    plot_clusters=False,
+    alignedcut_eig_norm_plot=False,
+    **kwargs,
 ):
     progress = gr.Progress()
     progress(0.2, desc="Feature Extraction")
@@ -538,63 +545,53 @@ def ncut_run(
             )
             logging_str += _logging_str
             rgb.append(_rgb[0])
+        return rgb, logging_str
     
     
-    cluster_images = None
-    if not old_school_ncut:  # ailgnedcut, joint across all images
-        rgb, _logging_str, eigvecs = compute_ncut(
-            features,
-            num_eig=num_eig,
-            num_sample_ncut=num_sample_ncut,
-            affinity_focal_gamma=affinity_focal_gamma,
-            knn_ncut=knn_ncut,
-            knn_tsne=knn_tsne,
-            num_sample_tsne=num_sample_tsne,
-            embedding_method=embedding_method,
-            embedding_metric=embedding_metric,
-            perplexity=perplexity,
-            n_neighbors=n_neighbors,
-            min_dist=min_dist,
-            sampling_method=sampling_method,
-            indirect_connection=indirect_connection,
-            make_orthogonal=make_orthogonal,
-            metric=ncut_metric,
-        )
-        logging_str += _logging_str
-        
-        if "AlignedThreeModelAttnNodes" == model_name:
-            # dirty patch for the alignedcut paper
-            start = time.time()
-            progress(0.6, desc="Plotting")
-            pil_images = []
-            for i_image in range(rgb.shape[0]):
-                _im = plot_one_image_36_grid(images[i_image], rgb[i_image])
-                pil_images.append(_im)
-            logging_str += f"plot time: {time.time() - start:.2f}s\n"
-            return pil_images, logging_str
-        
-        
-        if is_lisa == True:
-            # dirty patch for the LISA model
-            galleries = []
-            for i_prompt in range(len(lisa_prompts)):
-                _rgb = rgb[i_prompt]
-                galleries.append(to_pil_images(_rgb))
-            return *galleries, logging_str
-        
-        rgb = dont_use_too_much_green(rgb)
-        
-        if not video_output:
-            start = time.time()
-            progress_start = 0.6
-            progress(progress_start, desc="Plotting Clusters")
-            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 = make_cluster_plot(eigvecs, _images, h=h, w=w, progess_start=progress_start)
-            logging_str += f"plot time: {time.time() - start:.2f}s\n"
-        
+    
+    # ailgnedcut
+    
+    rgb, _logging_str, eigvecs = compute_ncut(
+        features,
+        num_eig=num_eig,
+        num_sample_ncut=num_sample_ncut,
+        affinity_focal_gamma=affinity_focal_gamma,
+        knn_ncut=knn_ncut,
+        knn_tsne=knn_tsne,
+        num_sample_tsne=num_sample_tsne,
+        embedding_method=embedding_method,
+        embedding_metric=embedding_metric,
+        perplexity=perplexity,
+        n_neighbors=n_neighbors,
+        min_dist=min_dist,
+        sampling_method=sampling_method,
+        indirect_connection=indirect_connection,
+        make_orthogonal=make_orthogonal,
+        metric=ncut_metric,
+    )
+    logging_str += _logging_str
+    
+    if "AlignedThreeModelAttnNodes" == model_name:
+        # dirty patch for the alignedcut paper
+        start = time.time()
+        progress(0.6, desc="Plotting")
+        pil_images = []
+        for i_image in range(rgb.shape[0]):
+            _im = plot_one_image_36_grid(images[i_image], rgb[i_image])
+            pil_images.append(_im)
+        logging_str += f"plot time: {time.time() - start:.2f}s\n"
+        return pil_images, logging_str
+    
+    
+    if is_lisa == True:
+        # dirty patch for the LISA model
+        galleries = []
+        for i_prompt in range(len(lisa_prompts)):
+            _rgb = rgb[i_prompt]
+            galleries.append(to_pil_images(_rgb))
+        return *galleries, logging_str
+    
+    rgb = dont_use_too_much_green(rgb)
     
     if video_output:
         progress(0.8, desc="Saving Video")
@@ -602,9 +599,37 @@ def ncut_run(
         video_cache.add_video(video_path)
         pil_images_to_video(to_pil_images(rgb), video_path, fps=5)
         return video_path, logging_str
-    
-    
-    return to_pil_images(rgb), cluster_images, logging_str
+        
+    cluster_images = None
+    if plot_clusters:
+        start = time.time()
+        progress_start = 0.6
+        progress(progress_start, desc="Plotting Clusters")
+        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)
+        logging_str += f"plot time: {time.time() - start:.2f}s\n"
+    
+    norm_images = None
+    if alignedcut_eig_norm_plot:
+        norm_images = []
+        # eig_magnitude = torch.clamp(eig_magnitude, 0, 1)
+        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_image = (norm_image[..., :3] * 255).astype(np.uint8)
+            norm_images.append(Image.fromarray(norm_image))
+        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)
+
+    return to_pil_images(rgb), cluster_images, norm_images, logging_str
+
 
 
 def _ncut_run(*args, **kwargs):
@@ -828,6 +853,8 @@ def run_fn(
     recursion_l2_gamma=0.5,
     recursion_l3_gamma=0.5,
     n_ret=1,
+    plot_clusters=False,
+    alignedcut_eig_norm_plot=False,
 ):
     
     progress=gr.Progress()
@@ -958,6 +985,8 @@ def run_fn(
         "lisa_prompt3": lisa_prompt3,
         "is_lisa": is_lisa,
         "n_ret": n_ret,
+        "plot_clusters": plot_clusters,
+        "alignedcut_eig_norm_plot": alignedcut_eig_norm_plot,
     }
     # print(kwargs)
     
@@ -1303,10 +1332,19 @@ def make_parameters_section(is_lisa=False, model_ratio=True):
             perplexity_slider, n_neighbors_slider, min_dist_slider, 
             sampling_method_dropdown, ncut_metric_dropdown, positive_prompt, negative_prompt]
 
+custom_css = """
+#unlock_button {
+    all: unset !important;
+}
+.form:has(#unlock_button) {
+    all: unset !important;
+}
+"""
 demo = gr.Blocks(
     theme=gr.themes.Base(spacing_size='md', text_size='lg', primary_hue='blue', neutral_hue='slate', secondary_hue='pink'),
     # fill_width=False,
     # title="ncut-pytorch",
+    css=custom_css,
 )
 with demo:
     with gr.Tab('AlignedCut'):
@@ -1336,7 +1374,7 @@ with demo:
         no_prompt = gr.Textbox("", label="", elem_id="empty_placeholder", type="text", placeholder="", visible=False)
         
         submit_button.click(
-            partial(run_fn, n_ret=2),
+            partial(run_fn, n_ret=2, plot_clusters=True),
             inputs=[
                 input_gallery, model_dropdown, layer_slider, num_eig_slider, node_type_dropdown, 
                 positive_prompt, negative_prompt,
@@ -1350,7 +1388,46 @@ with demo:
             scroll_to_output=True,
         )
         
+    with gr.Tab('AlignedCut (+Norm Plot)', visible=False) as tab_alignedcut_norm:
+
+        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()
+                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)
+                
+            with gr.Column(scale=5, min_width=200):
+                cluster_gallery = gr.Gallery(value=[], label="Clusters", show_label=True, elem_id="clusters", columns=[5], rows=[2], object_fit="contain", height="auto", show_share_button=True, preview=True, interactive=False)
+                output_gallery = make_output_images_section()
+                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)
+                [
+                    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.value = 30
+
+        clear_images_button.click(lambda x: ([], [], [], []), outputs=[input_gallery, output_gallery, cluster_gallery, norm_gallery])
+        
+        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=3, plot_clusters=True, alignedcut_eig_norm_plot=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=[output_gallery, cluster_gallery, norm_gallery, logging_text],
+            scroll_to_output=True,
+        )
         
     with gr.Tab('NCut'): 
         gr.Markdown('#### NCut (Legacy), not aligned, no Nyström approximation')
@@ -1645,7 +1722,7 @@ with demo:
             outputs=[output_gallery, logging_text],
         )
         
-    with gr.Tab('Model Aligned (+Rrecursion)'): 
+    with gr.Tab('Model Aligned (+Rrecursion)', visible=False) as tab_model_aligned_recursion: 
         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.')
@@ -1816,23 +1893,62 @@ with demo:
     
     
     with gr.Tab('📄About'):
-        gr.Markdown("**This demo is for the Python package `ncut-pytorch`, please visit the [Documentation](https://ncut-pytorch.readthedocs.io/)**")
-        gr.Markdown("**All the models and functions used for this demo are in the Python package `ncut-pytorch`**")
-        gr.Markdown("---")
-        gr.Markdown("---")
-        gr.Markdown("**Normalized Cuts**, aka. spectral clustering, is a graphical method to analyze data grouping in the affinity eigenvector space. It has been widely used for unsupervised segmentation in the 2000s.")
-        gr.Markdown("*Normalized Cuts and Image Segmentation, Jianbo Shi and Jitendra Malik, 2000*")
-        gr.Markdown("---")        
-        gr.Markdown("**We have improved NCut, with some advanced features:**")
-        gr.Markdown("- **Nyström** Normalized Cut, is a new approximation algorithm developed for large-scale graph cuts, a large-graph of million nodes can be processed in under 10s (cpu) or 2s (gpu).")
-        gr.Markdown("- **spectral-tSNE** visualization, a new method to visualize the high-dimensional eigenvector space with 3D RGB cube. Color is aligned across images, color infers distance in representation.")
-        gr.Markdown("*paper in prep, Yang 2024*")
-        gr.Markdown("*AlignedCut: Visual Concepts Discovery on Brain-Guided Universal Feature Space, Huzheng Yang, James Gee\*, and Jianbo Shi\*, 2024*")
-        gr.Markdown("---")        
-        gr.Markdown("---")        
-        gr.Markdown('<p style="text-align: center;">We thank the HuggingFace team for hosting this demo.</p>')        
+        with gr.Column():
+            gr.Markdown("**This demo is for the Python package `ncut-pytorch`, please visit the [Documentation](https://ncut-pytorch.readthedocs.io/)**")
+            gr.Markdown("**All the models and functions used for this demo are in the Python package `ncut-pytorch`**")
+            gr.Markdown("---")
+            gr.Markdown("---")
+            gr.Markdown("**Normalized Cuts**, aka. spectral clustering, is a graphical method to analyze data grouping in the affinity eigenvector space. It has been widely used for unsupervised segmentation in the 2000s.")
+            gr.Markdown("*Normalized Cuts and Image Segmentation, Jianbo Shi and Jitendra Malik, 2000*")
+            gr.Markdown("---")        
+            gr.Markdown("**We have improved NCut, with some advanced features:**")
+            gr.Markdown("- **Nyström** Normalized Cut, is a new approximation algorithm developed for large-scale graph cuts, a large-graph of million nodes can be processed in under 10s (cpu) or 2s (gpu).")
+            gr.Markdown("- **spectral-tSNE** visualization, a new method to visualize the high-dimensional eigenvector space with 3D RGB cube. Color is aligned across images, color infers distance in representation.")
+            gr.Markdown("*paper in prep, Yang 2024*")
+            gr.Markdown("*AlignedCut: Visual Concepts Discovery on Brain-Guided Universal Feature Space, Huzheng Yang, James Gee\*, and Jianbo Shi\*, 2024*")
+            gr.Markdown("---")        
+            gr.Markdown("---")        
+            gr.Markdown('<p style="text-align: center;">We thank the HuggingFace team for hosting this demo.</p>')        
         
+        # unlock the hidden tab
+        with gr.Row():
+            with gr.Column(scale=5):
+                gr.Markdown("")
+            with gr.Column(scale=5):
+                hidden_button = gr.Checkbox(label="🤗", value=False, elem_id="unlock_button", visible=True, interactive=True)
+            with gr.Column(scale=5):
+                gr.Markdown("")
+            
+        n_smiles = gr.State(0)
+        unlock_value = 6
+
+        def update_smile(n_smiles):
+            n_smiles = n_smiles + 1
+            n_smiles = unlock_value if n_smiles > unlock_value else n_smiles
+            if n_smiles == unlock_value - 2:
+                gr.Info("click one more time to unlock", 2)
+            if n_smiles == unlock_value:
+                label = "🔓 unlocked"
+                return n_smiles, gr.update(label=label, value=True, interactive=False)
+            label = ["😊"] * n_smiles
+            label = "".join(label)
+            return n_smiles, gr.update(label=label, value=False)
         
+        def unlock_tabs_with_info(n_smiles):
+            if n_smiles == unlock_value:
+                gr.Info("🔓 unlocked tabs", 2)
+                return gr.update(visible=True)
+            return gr.update()
+
+        def unlock_tabs(n_smiles):
+            if n_smiles == unlock_value:
+                return gr.update(visible=True)
+            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)
+                
     with gr.Row():
         with gr.Column():
             gr.Markdown("##### This demo is for `ncut-pytorch`, [Documentation](https://ncut-pytorch.readthedocs.io/) ")