revert changes

app.py

@@ -372,233 +372,61 @@ def load_all_data(image_root, pkl_root):
     return data_dict, sae_data_dict
 
 
-# data_dict, sae_data_dict = load_all_data(image_root="./data/image", pkl_root=pkl_root)
-# default_image_name = "christmas-imagenet"
-
-
-# with gr.Blocks(
-#     theme=gr.themes.Citrus(),
-#     css="""
-#     .image-row .gr-image { margin: 0 !important; padding: 0 !important; }
-#     .image-row img { width: auto; height: 50px; } /* Set a uniform height for all images */
-# """,
-# ) as demo:
-#     with gr.Row():
-#         with gr.Column():
-#             # Left View: Image selection and click handling
-#             gr.Markdown("## Select input image and patch on the image")
-#             image_selector = gr.Dropdown(choices=list(data_dict.keys()), value=default_image_name, label="Select Image")
-#             image_display = gr.Image(value=data_dict[default_image_name]["image"], type="pil", interactive=True)
-
-#             # Update image display when a new image is selected
-#             image_selector.change(
-#                 fn=lambda img_name: data_dict[img_name]["image"], inputs=image_selector, outputs=image_display
-#             )
-#             image_display.select(fn=highlight_grid, inputs=[image_selector], outputs=[image_display])
-
-#         with gr.Column():
-#             gr.Markdown("## SAE latent activations of CLIP and MaPLE")
-#             model_options = [f"MaPLE-{dataset_name}" for dataset_name in DATASET_LIST]
-#             model_selector = gr.Dropdown(
-#                 choices=model_options, value=model_options[0], label="Select adapted model (MaPLe)"
-#             )
-#             init_plot = plot_activation_distribution(None, default_image_name, model_options[0])
-#             neuron_plot = gr.Plot(label="Neuron Activation", value=init_plot, show_label=False)
-
-#             image_selector.change(
-#                 fn=plot_activation_distribution, inputs=[image_selector, model_selector], outputs=neuron_plot
-#             )
-#             image_display.select(
-#                 fn=plot_activation_distribution, inputs=[image_selector, model_selector], outputs=neuron_plot
-#             )
-#             model_selector.change(fn=load_image, inputs=[image_selector], outputs=image_display)
-#             model_selector.change(
-#                 fn=plot_activation_distribution, inputs=[image_selector, model_selector], outputs=neuron_plot
-#             )
-
-#     with gr.Row():
-#         with gr.Column():
-#             radio_names = get_init_radio_options(default_image_name, model_options[0])
-
-#             feautre_idx = radio_names[0].split("-")[-1]
-#             markdown_display = gr.Markdown(f"## Segmentation mask for the selected SAE latent - {feautre_idx}")
-#             init_seg, init_tops, init_values = show_activation_heatmap(default_image_name, radio_names[0], "CLIP")
-
-#             gr.Markdown("### Localize SAE latent activation using CLIP")
-#             seg_mask_display = gr.Image(value=init_seg, type="pil", show_label=False)
-#             init_seg_maple, _, _ = show_activation_heatmap(default_image_name, radio_names[0], model_options[0])
-#             gr.Markdown("### Localize SAE latent activation using MaPLE")
-#             seg_mask_display_maple = gr.Image(value=init_seg_maple, type="pil", show_label=False)
-
-#         with gr.Column():
-#             gr.Markdown("## Top activating SAE latent index")
-
-#             radio_choices = gr.Radio(
-#                 choices=radio_names, label="Top activating SAE latent", interactive=True, value=radio_names[0]
-#             )
-#             toggle_btn = gr.Checkbox(label="Show segmentation mask", value=False)
-
-#             markdown_display_2 = gr.Markdown(f"## Top reference images for the selected SAE latent - {feautre_idx}")
-
-#             gr.Markdown("### ImageNet")
-#             top_image_1 = gr.Image(value=init_tops[0], type="pil", label="ImageNet", show_label=False)
-#             act_value_1 = gr.Markdown(init_values[0])
-
-#             gr.Markdown("### ImageNet-Sketch")
-#             top_image_2 = gr.Image(value=init_tops[1], type="pil", label="ImageNet-Sketch", show_label=False)
-#             act_value_2 = gr.Markdown(init_values[1])
-
-#             gr.Markdown("### Caltech101")
-#             top_image_3 = gr.Image(value=init_tops[2], type="pil", label="Caltech101", show_label=False)
-#             act_value_3 = gr.Markdown(init_values[2])
-
-#             image_display.select(
-#                 fn=update_radio_options, inputs=[image_selector, model_selector], outputs=[radio_choices], queue=True
-#             )
-
-#             model_selector.change(
-#                 fn=update_radio_options, inputs=[image_selector, model_selector], outputs=[radio_choices], queue=True
-#             )
-
-#             image_selector.select(
-#                 fn=update_radio_options, inputs=[image_selector, model_selector], outputs=[radio_choices], queue=True
-#             )
-
-#         radio_choices.change(
-#             fn=update_markdown,
-#             inputs=[radio_choices],
-#             outputs=[markdown_display, markdown_display_2],
-#             queue=True,
-#         )
-
-#         radio_choices.change(
-#             fn=show_activation_heatmap_clip,
-#             inputs=[image_selector, radio_choices, toggle_btn],
-#             outputs=[seg_mask_display, top_image_1, top_image_2, top_image_3, act_value_1, act_value_2, act_value_3],
-#             queue=True,
-#         )
-
-#         radio_choices.change(
-#             fn=show_activation_heatmap_maple,
-#             inputs=[image_selector, radio_choices, model_selector],
-#             outputs=[seg_mask_display_maple],
-#             queue=True,
-#         )
-
-#         # toggle_btn.change(
-#         #     fn=get_top_images,
-#         #     inputs=[radio_choices, toggle_btn],
-#         #     outputs=[top_image_1, top_image_2, top_image_3],
-#         #     queue=True,
-#         # )
-
-#         toggle_btn.change(
-#             fn=show_activation_heatmap_clip,
-#             inputs=[image_selector, radio_choices, toggle_btn],
-#             outputs=[seg_mask_display, top_image_1, top_image_2, top_image_3, act_value_1, act_value_2, act_value_3],
-#             queue=True,
-#         )
-
-#     # Launch the app
-#     demo.launch()
-
-# Precompute all necessary data and store in caches before launching the Gradio app.
-
-# Load data once at startup
 data_dict, sae_data_dict = load_all_data(image_root="./data/image", pkl_root=pkl_root)
 default_image_name = "christmas-imagenet"
-model_options = [f"MaPLE-{dataset_name}" for dataset_name in DATASET_LIST]
-default_model = model_options[0]
 
-# Precompute activation distributions for all images/models to avoid repeated I/O.
-activation_cache = {}
-for img_name in data_dict.keys():
-    for mdl in ["CLIP"] + model_options:
-        activation_cache[(img_name, mdl)] = get_activation_distribution(img_name, mdl)
-
-# Precompute initial radio options and top-neuron related info for default states.
-radio_names = get_init_radio_options(default_image_name, default_model)
-feautre_idx = radio_names[0].split("-")[-1]
-
-# Precompute initial figures and mask overlays so they don't need to be recomputed on load.
-init_plot = plot_activation_distribution(None, default_image_name, default_model)
-init_seg, init_tops, init_values = show_activation_heatmap(default_image_name, radio_names[0], "CLIP")
-init_seg_maple, _, _ = show_activation_heatmap(default_image_name, radio_names[0], default_model)
 
 with gr.Blocks(
     theme=gr.themes.Citrus(),
     css="""
     .image-row .gr-image { margin: 0 !important; padding: 0 !important; }
     .image-row img { width: auto; height: 50px; } /* Set a uniform height for all images */
-"""
+""",
 ) as demo:
     with gr.Row():
         with gr.Column():
+            # Left View: Image selection and click handling
             gr.Markdown("## Select input image and patch on the image")
-
-            # Instead of recomputing, just directly load from data_dict
             image_selector = gr.Dropdown(choices=list(data_dict.keys()), value=default_image_name, label="Select Image")
             image_display = gr.Image(value=data_dict[default_image_name]["image"], type="pil", interactive=True)
 
-            # When image changes, just display the corresponding image
-            def update_image_display(img_name):
-                return data_dict[img_name]["image"]
-
+            # Update image display when a new image is selected
             image_selector.change(
-                fn=update_image_display, 
-                inputs=image_selector, 
-                outputs=image_display
+                fn=lambda img_name: data_dict[img_name]["image"], inputs=image_selector, outputs=image_display
             )
-
-            # Highlight selected grid cell
             image_display.select(fn=highlight_grid, inputs=[image_selector], outputs=[image_display])
 
         with gr.Column():
             gr.Markdown("## SAE latent activations of CLIP and MaPLE")
-
+            model_options = [f"MaPLE-{dataset_name}" for dataset_name in DATASET_LIST]
             model_selector = gr.Dropdown(
-                choices=model_options, 
-                value=default_model, 
-                label="Select adapted model (MaPLe)"
+                choices=model_options, value=model_options[0], label="Select adapted model (MaPLe)"
             )
-
+            init_plot = plot_activation_distribution(None, default_image_name, model_options[0])
             neuron_plot = gr.Plot(label="Neuron Activation", value=init_plot, show_label=False)
 
-            # Update plot based on image/model
-            def update_plot(img_name, model_name):
-                # Use precomputed activation distributions from activation_cache
-                # to create the figure. If figure creation is expensive, consider caching plots as well.
-                return plot_activation_distribution(None, img_name, model_name)
-
             image_selector.change(
-                fn=update_plot, 
-                inputs=[image_selector, model_selector], 
-                outputs=neuron_plot
+                fn=plot_activation_distribution, inputs=[image_selector, model_selector], outputs=neuron_plot
             )
             image_display.select(
-                fn=update_plot, 
-                inputs=[image_selector, model_selector], 
-                outputs=neuron_plot
-            )
-            model_selector.change(
-                fn=lambda img_name: data_dict[img_name]["image"], 
-                inputs=[image_selector], 
-                outputs=image_display
+                fn=plot_activation_distribution, inputs=[image_selector, model_selector], outputs=neuron_plot
             )
+            model_selector.change(fn=load_image, inputs=[image_selector], outputs=image_display)
             model_selector.change(
-                fn=update_plot, 
-                inputs=[image_selector, model_selector], 
-                outputs=neuron_plot
+                fn=plot_activation_distribution, inputs=[image_selector, model_selector], outputs=neuron_plot
             )
 
     with gr.Row():
         with gr.Column():
-            # Use previously precomputed segmentation masks and tops instead of recomputing on load
+            radio_names = get_init_radio_options(default_image_name, model_options[0])
+
+            feautre_idx = radio_names[0].split("-")[-1]
             markdown_display = gr.Markdown(f"## Segmentation mask for the selected SAE latent - {feautre_idx}")
+            init_seg, init_tops, init_values = show_activation_heatmap(default_image_name, radio_names[0], "CLIP")
 
             gr.Markdown("### Localize SAE latent activation using CLIP")
             seg_mask_display = gr.Image(value=init_seg, type="pil", show_label=False)
-
+            init_seg_maple, _, _ = show_activation_heatmap(default_image_name, radio_names[0], model_options[0])
             gr.Markdown("### Localize SAE latent activation using MaPLE")
             seg_mask_display_maple = gr.Image(value=init_seg_maple, type="pil", show_label=False)
 
@@ -606,17 +434,12 @@ with gr.Blocks(
             gr.Markdown("## Top activating SAE latent index")
 
             radio_choices = gr.Radio(
-                choices=radio_names, 
-                label="Top activating SAE latent", 
-                interactive=True, 
-                value=radio_names[0]
+                choices=radio_names, label="Top activating SAE latent", interactive=True, value=radio_names[0]
             )
-
             toggle_btn = gr.Checkbox(label="Show segmentation mask", value=False)
 
             markdown_display_2 = gr.Markdown(f"## Top reference images for the selected SAE latent - {feautre_idx}")
 
-            # Display precomputed top images and values
             gr.Markdown("### ImageNet")
             top_image_1 = gr.Image(value=init_tops[0], type="pil", label="ImageNet", show_label=False)
             act_value_1 = gr.Markdown(init_values[0])
@@ -629,31 +452,18 @@ with gr.Blocks(
             top_image_3 = gr.Image(value=init_tops[2], type="pil", label="Caltech101", show_label=False)
             act_value_3 = gr.Markdown(init_values[2])
 
-            # Update radio choices when image/model changes.
-            # If expensive, this could be cached as well.
-            def on_image_or_model_change(img_name, model_name):
-                return update_radio_options(None, img_name, model_name)
-
             image_display.select(
-                fn=on_image_or_model_change, 
-                inputs=[image_selector, model_selector], 
-                outputs=[radio_choices], 
-                queue=True
+                fn=update_radio_options, inputs=[image_selector, model_selector], outputs=[radio_choices], queue=True
             )
+
             model_selector.change(
-                fn=on_image_or_model_change, 
-                inputs=[image_selector, model_selector], 
-                outputs=[radio_choices], 
-                queue=True
+                fn=update_radio_options, inputs=[image_selector, model_selector], outputs=[radio_choices], queue=True
             )
+
             image_selector.select(
-                fn=on_image_or_model_change, 
-                inputs=[image_selector, model_selector], 
-                outputs=[radio_choices], 
-                queue=True
+                fn=update_radio_options, inputs=[image_selector, model_selector], outputs=[radio_choices], queue=True
             )
 
-        # Update markdown titles dynamically based on selected radio choice
         radio_choices.change(
             fn=update_markdown,
             inputs=[radio_choices],
@@ -661,7 +471,6 @@ with gr.Blocks(
             queue=True,
         )
 
-        # Show activation heatmap for CLIP
         radio_choices.change(
             fn=show_activation_heatmap_clip,
             inputs=[image_selector, radio_choices, toggle_btn],
@@ -669,7 +478,6 @@ with gr.Blocks(
             queue=True,
         )
 
-        # Show activation heatmap for MaPLE
         radio_choices.change(
             fn=show_activation_heatmap_maple,
             inputs=[image_selector, radio_choices, model_selector],
@@ -677,7 +485,13 @@ with gr.Blocks(
             queue=True,
         )
 
-        # Toggle segmentation mask
+        # toggle_btn.change(
+        #     fn=get_top_images,
+        #     inputs=[radio_choices, toggle_btn],
+        #     outputs=[top_image_1, top_image_2, top_image_3],
+        #     queue=True,
+        # )
+
         toggle_btn.change(
             fn=show_activation_heatmap_clip,
             inputs=[image_selector, radio_choices, toggle_btn],
@@ -687,4 +501,3 @@ with gr.Blocks(
 
     # Launch the app
     demo.launch()
-