add aligned+recursion

app.py

@@ -419,9 +419,22 @@ def ncut_run(
                 metric="cosine" if i == 0 else recursion_metric,
             )
             logging_str += _logging_str
-            rgb = dont_use_too_much_green(rgb)
-            rgbs.append(to_pil_images(rgb))
-            inp = eigvecs.reshape(*features.shape[:3], -1)
+            
+            
+            if "AlignedThreeModelAttnNodes" == model_name:
+                # dirty patch for the alignedcut paper
+                start = time.time()
+                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)
+                rgbs.append(pil_images)
+                logging_str += f"plot time: {time.time() - start:.2f}s\n"
+            else:
+                rgb = dont_use_too_much_green(rgb)
+                rgbs.append(to_pil_images(rgb))
+                
+            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
@@ -472,15 +485,12 @@ def ncut_run(
         
         if "AlignedThreeModelAttnNodes" == model_name:
             # dirty patch for the alignedcut paper
-            # galleries = []
-            # for i_node in range(rgb.shape[1]):
-            #     _rgb = rgb[:, i_node]
-            #     galleries.append(to_pil_images(_rgb, target_size=56))
-            # return *galleries, logging_str
+            start = time.time()
             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
         
         
@@ -516,26 +526,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=20)
@@ -1018,19 +1028,28 @@ def make_dataset_images_section(advanced=False, is_random=False):
     return dataset_dropdown, num_images_slider, random_seed_slider, load_images_button
     
 
-def random_rotate_rgb_gallery(images):
-    if images is None or len(images) == 0:
-        gr.Warning("No images selected.")
-        return []
-    # read webp images
-    images = [Image.open(image[0]).convert("RGB") for image in images]
-    images = [np.array(image).astype(np.float32) for image in images]
-    images = np.stack(images)
-    images = torch.tensor(images) / 255
-    position = np.random.choice([1, 2, 4, 5, 6])
-    images = rotate_rgb_cube(images, position)
-    images = to_pil_images(images, resize=False)
-    return images
+# def random_rotate_rgb_gallery(images):
+#     if images is None or len(images) == 0:
+#         gr.Warning("No images selected.")
+#         return []
+#     # read webp images
+#     images = [Image.open(image[0]).convert("RGB") for image in images]
+#     images = [np.array(image).astype(np.float32) for image in images]
+#     images = np.stack(images)
+#     images = torch.tensor(images) / 255
+#     position = np.random.choice([1, 2, 4, 5, 6])
+#     images = rotate_rgb_cube(images, position)
+#     images = to_pil_images(images, resize=False)
+#     return images
+
+def protect_original_image_in_plot(original_image, rotated_images):
+    plot_h, plot_w = 332, 1542
+    image_h, image_w = original_image.shape[1], original_image.shape[2]
+    if not (plot_h == image_h and plot_w == image_w):
+        return rotated_images
+    protection_w = 190
+    rotated_images[:, :, :protection_w] = original_image[:, :, :protection_w]
+    return rotated_images
 
 def sequence_rotate_rgb_gallery(images):
     if images is None or len(images) == 0:
@@ -1041,8 +1060,10 @@ def sequence_rotate_rgb_gallery(images):
     images = [np.array(image).astype(np.float32) for image in images]
     images = np.stack(images)
     images = torch.tensor(images) / 255
+    original_images = images.clone()
     rotation_matrix = torch.tensor([[0, 1, 0], [0, 0, 1], [1, 0, 0]]).float()
     images = images @ rotation_matrix
+    images = protect_original_image_in_plot(original_images, images)
     images = to_pil_images(images, resize=False)
     return images
 
@@ -1055,7 +1076,9 @@ def flip_rgb_gallery(images, axis=0):
     images = [np.array(image).astype(np.float32) for image in images]
     images = np.stack(images)
     images = torch.tensor(images) / 255
+    original_images = images.clone()
     images = 1 - images
+    images = protect_original_image_in_plot(original_images, images)
     images = to_pil_images(images, resize=False)
     return images
 
@@ -1074,7 +1097,7 @@ def make_output_images_section():
     add_output_images_buttons(output_gallery)
     return output_gallery
 
-def make_parameters_section(is_lisa=False):
+def make_parameters_section(is_lisa=False, model_ratio=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()
@@ -1095,7 +1118,7 @@ def make_parameters_section(is_lisa=False):
         negative_prompt = gr.Textbox(label="Prompt (Negative)", elem_id="prompt", placeholder="e.g. 'a photo from egocentric view'", visible=False)
         node_type_dropdown = gr.Dropdown(layer_names, label="LISA (SAM) decoder: Layer and Node", value="dec_1_block", elem_id="node_type")
     else:
-        model_radio = gr.Radio(["CLIP", "DiNO", "Diffusion", "ImageNet", "MAE", "SAM"], label="Backbone", value="DiNO", elem_id="model_radio", show_label=True)
+        model_radio = gr.Radio(["CLIP", "DiNO", "Diffusion", "ImageNet", "MAE", "SAM"], label="Backbone", value="DiNO", elem_id="model_radio", show_label=True, visible=model_ratio)
         model_dropdown = gr.Dropdown(get_filtered_model_names("DiNO"), label="", value="DiNO(dino_vitb8_448)", elem_id="model_name", show_label=False)
         model_radio.change(fn=lambda x: gr.update(choices=get_filtered_model_names(x), value=get_default_model_name(x)), inputs=model_radio, outputs=[model_dropdown])
         layer_slider = gr.Slider(1, 12, step=1, label="Backbone: Layer index", value=10, elem_id="layer")
@@ -1292,6 +1315,7 @@ with demo:
                 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):
@@ -1311,14 +1335,6 @@ with demo:
                 ] = make_parameters_section()
                 num_eig_slider.visible = False
                 affinity_focal_gamma_slider.visible = False
-                # logging text box
-        with gr.Row():
-            with gr.Column(scale=5, min_width=200):
-                gr.Markdown(' ')
-            with gr.Column(scale=5, min_width=200):
-                gr.Markdown(' ')
-            with gr.Column(scale=5, min_width=200):
-                logging_text = gr.Textbox("Logging information", label="Logging", elem_id="logging", type="text", placeholder="Logging information")
         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")
@@ -1326,13 +1342,12 @@ with demo:
         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])
-        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),
             inputs=[
-                input_gallery, model_dropdown, layer_slider, num_eig_slider, node_type_dropdown, 
+                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, knn_ncut_slider, 
@@ -1457,7 +1472,6 @@ with demo:
         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, is_random=True)
                 num_images_slider.value = 100
                 
@@ -1476,7 +1490,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, positive_prompt, negative_prompt
-                ] = make_parameters_section()
+                ] = make_parameters_section(model_ratio=False)
                 model_dropdown.value = "AlignedThreeModelAttnNodes"
                 model_dropdown.visible = False
                 layer_slider.visible = False
@@ -1505,7 +1519,7 @@ with demo:
             outputs=[output_gallery, logging_text],
         )
         
-    with gr.Tab('Model Aligned (+Recursive)'): 
+    with gr.Tab('Model Aligned (+Rrecursion)'): 
         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.')
@@ -1513,29 +1527,58 @@ with demo:
         gr.Markdown('')
         gr.Markdown("To see a good pattern, you will need to load 100~1000 images. 100 images need 10sec for RTX4090. Running out of HuggingFace GPU Quota? Try [Demo](https://ncut-pytorch.readthedocs.io/en/latest/demo/) hosted at UPenn")
         gr.Markdown('---')
+
+        # 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=False, 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)
+        #     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=False, 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)
+        #     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=False, 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)    
+        gr.Markdown('### Output (Recursion #1)')
+        l1_gallery = gr.Gallery(format='png', value=[], label="Recursion #1", show_label=False, elem_id="ncut_l1", columns=[100], rows=[1], object_fit="contain", height="auto", show_fullscreen_button=True, interactive=False, preview=True)
+        add_output_images_buttons(l1_gallery)
+        gr.Markdown('### Output (Recursion #2)')
+        l2_gallery = gr.Gallery(format='png', value=[], label="Recursion #2", show_label=False, elem_id="ncut_l2", columns=[100], rows=[1], object_fit="contain", height="auto", show_fullscreen_button=True, interactive=False, preview=True)
+        add_output_images_buttons(l2_gallery)
+        gr.Markdown('### Output (Recursion #3)')
+        l3_gallery = gr.Gallery(format='png', value=[], label="Recursion #3", show_label=False, elem_id="ncut_l3", columns=[100], rows=[1], object_fit="contain", height="auto", show_fullscreen_button=True, interactive=False, preview=True)
+        add_output_images_buttons(l3_gallery)  
+    
         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, is_random=True)
                 num_images_slider.value = 100
                 
                 
             with gr.Column(scale=5, min_width=200):
-                output_gallery = make_output_images_section()
-                gr.Markdown('### TIP1: use the `full-screen` button, and use `arrow keys` to navigate')
+                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")
                 gr.Markdown('---')
                 gr.Markdown('Model: CLIP(ViT-B-16/openai), DiNOv2reg(dinov2_vitb14_reg), MAE(vit_base)')
                 gr.Markdown('Layer type: attention output (attn), without sum of residual')
-                gr.Markdown('### TIP2: for large image set, please increase the `num_sample` for t-SNE and NCUT')
-                gr.Markdown('---')
                 [
                     model_dropdown, layer_slider, node_type_dropdown, num_eig_slider, 
                     affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider, 
                     embedding_method_dropdown, embedding_metric_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
                     perplexity_slider, n_neighbors_slider, min_dist_slider, 
                     sampling_method_dropdown, positive_prompt, negative_prompt
-                ] = make_parameters_section()
+                ] = make_parameters_section(model_ratio=False)
+                num_eig_slider.visible = False
+                affinity_focal_gamma_slider.visible = False
                 model_dropdown.value = "AlignedThreeModelAttnNodes"
                 model_dropdown.visible = False
                 layer_slider.visible = False
@@ -1545,23 +1588,30 @@ with demo:
                 # logging text box
                 logging_text = gr.Textbox("Logging information", label="Logging", elem_id="logging", type="text", placeholder="Logging information")
                 
-        clear_images_button.click(lambda x: ([], []), outputs=[input_gallery, output_gallery])
+        clear_images_button.click(lambda x: ([],), outputs=[input_gallery])
         
-        false_placeholder = gr.Checkbox(label="False", value=False, elem_id="false_placeholder", 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
         no_prompt = gr.Textbox("", label="", elem_id="empty_placeholder", type="text", placeholder="", visible=False)
         
         submit_button.click(
-            run_fn,
+            partial(run_fn, n_ret=3),
             inputs=[
-                input_gallery, model_dropdown, layer_slider, num_eig_slider, node_type_dropdown, 
+                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, knn_ncut_slider, 
                 embedding_method_dropdown, embedding_metric_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
-                perplexity_slider, n_neighbors_slider, min_dist_slider, sampling_method_dropdown
+                perplexity_slider, n_neighbors_slider, min_dist_slider, sampling_method_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=galleries + [logging_text],
-            outputs=[output_gallery, logging_text],
+            outputs=[l1_gallery, l2_gallery, l3_gallery, logging_text],
         )