test: tried merging radio_choices.change

app.py

@@ -28,7 +28,9 @@ def preload_activation(image_name):
 def get_activation_distribution(image_name: str, model_type: str):
     activation = get_data(image_name, model_type)[0]
 
-    noisy_features_indices = (sae_data_dict["mean_acts"]["imagenet"] > 0.1).nonzero()[0].tolist()
+    noisy_features_indices = (
+        (sae_data_dict["mean_acts"]["imagenet"] > 0.1).nonzero()[0].tolist()
+    )
     activation[:, noisy_features_indices] = 0
 
     return activation
@@ -52,18 +54,31 @@ def highlight_grid(evt: gr.EventData, image_name):
 
     highlighted_image = image.copy()
     draw = ImageDraw.Draw(highlighted_image)
-    box = [grid_x * cell_width, grid_y * cell_height, (grid_x + 1) * cell_width, (grid_y + 1) * cell_height]
+    box = [
+        grid_x * cell_width,
+        grid_y * cell_height,
+        (grid_x + 1) * cell_width,
+        (grid_y + 1) * cell_height,
+    ]
     draw.rectangle(box, outline="red", width=3)
 
     return highlighted_image
 
 
 def load_image(img_name):
-    return Image.open(data_dict[img_name]["image_path"]).resize((IMAGE_SIZE, IMAGE_SIZE))
+    return Image.open(data_dict[img_name]["image_path"]).resize(
+        (IMAGE_SIZE, IMAGE_SIZE)
+    )
 
 
 def plot_activations(
-    all_activation, tile_activations=None, grid_x=None, grid_y=None, top_k=5, colors=("blue", "cyan"), model_name="CLIP"
+    all_activation,
+    tile_activations=None,
+    grid_x=None,
+    grid_y=None,
+    top_k=5,
+    colors=("blue", "cyan"),
+    model_name="CLIP",
 ):
     fig = go.Figure()
 
@@ -94,10 +109,14 @@ def plot_activations(
         return fig
 
     label = f"{model_name.split('-')[-0]} Image-level"
-    fig = _add_scatter_with_annotation(fig, all_activation, model_name, colors[0], label)
+    fig = _add_scatter_with_annotation(
+        fig, all_activation, model_name, colors[0], label
+    )
     if tile_activations is not None:
         label = f"{model_name.split('-')[-0]} Tile ({grid_x}, {grid_y})"
-        fig = _add_scatter_with_annotation(fig, tile_activations, model_name, colors[1], label)
+        fig = _add_scatter_with_annotation(
+            fig, tile_activations, model_name, colors[1], label
+        )
 
     fig.update_layout(
         title="Activation Distribution",
@@ -105,7 +124,9 @@ def plot_activations(
         yaxis_title="Activation Value",
         template="plotly_white",
     )
-    fig.update_layout(legend=dict(orientation="h", yanchor="middle", y=0.5, xanchor="center", x=0.5))
+    fig.update_layout(
+        legend=dict(orientation="h", yanchor="middle", y=0.5, xanchor="center", x=0.5)
+    )
 
     return fig
 
@@ -126,12 +147,20 @@ def get_activations(evt: gr.EventData, selected_image: str, model_name: str, col
             tile_activations = activation[token_idx]
 
     fig = plot_activations(
-        all_activation, tile_activations, grid_x, grid_y, top_k=5, model_name=model_name, colors=colors
+        all_activation,
+        tile_activations,
+        grid_x,
+        grid_y,
+        top_k=5,
+        model_name=model_name,
+        colors=colors,
     )
     return fig
 
 
-def plot_activation_distribution(evt: gr.EventData, selected_image: str, model_name: str):
+def plot_activation_distribution(
+    evt: gr.EventData, selected_image: str, model_name: str
+):
     fig = make_subplots(
         rows=2,
         cols=1,
@@ -139,8 +168,12 @@ def plot_activation_distribution(evt: gr.EventData, selected_image: str, model_n
         subplot_titles=["CLIP Activation", f"{model_name} Activation"],
     )
 
-    fig_clip = get_activations(evt, selected_image, "CLIP", colors=("#00b4d8", "#90e0ef"))
-    fig_maple = get_activations(evt, selected_image, model_name, colors=("#ff5a5f", "#ffcad4"))
+    fig_clip = get_activations(
+        evt, selected_image, "CLIP", colors=("#00b4d8", "#90e0ef")
+    )
+    fig_maple = get_activations(
+        evt, selected_image, model_name, colors=("#ff5a5f", "#ffcad4")
+    )
 
     def _attach_fig(fig, sub_fig, row, col, yref):
         for trace in sub_fig.data:
@@ -178,7 +211,9 @@ def get_segmask(selected_image, slider_value, model_type):
         mask = torch.Tensor(temp[1:,].reshape(14, 14)).view(1, 1, 14, 14)
     except Exception as e:
         print(sae_act.shape, slider_value)
-    mask = torch.nn.functional.interpolate(mask, (image.height, image.width))[0][0].numpy()
+    mask = torch.nn.functional.interpolate(mask, (image.height, image.width))[0][
+        0
+    ].numpy()
     mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-10)
 
     base_opacity = 30
@@ -201,7 +236,11 @@ def get_top_images(slider_value, toggle_btn):
             os.path.join(dataset_path, "caltech101", f"{slider_value}.jpg"),
         ]
         top_images = [
-            Image.open(path) if os.path.exists(path) else Image.new("RGB", (256, 256), (255, 255, 255))
+            (
+                Image.open(path)
+                if os.path.exists(path)
+                else Image.new("RGB", (256, 256), (255, 255, 255))
+            )
             for path in top_image_paths
         ]
         return top_images
@@ -230,9 +269,19 @@ def show_activation_heatmap(selected_image, slider_value, model_type, toggle_btn
 
 
 def show_activation_heatmap_clip(selected_image, slider_value, toggle_btn):
-    rgba_overlay, top_images, act_values = show_activation_heatmap(selected_image, slider_value, "CLIP", toggle_btn)
+    rgba_overlay, top_images, act_values = show_activation_heatmap(
+        selected_image, slider_value, "CLIP", toggle_btn
+    )
     sleep(0.1)
-    return (rgba_overlay, top_images[0], top_images[1], top_images[2], act_values[0], act_values[1], act_values[2])
+    return (
+        rgba_overlay,
+        top_images[0],
+        top_images[1],
+        top_images[2],
+        act_values[0],
+        act_values[1],
+        act_values[2],
+    )
 
 
 def show_activation_heatmap_maple(selected_image, slider_value, model_name):
@@ -251,11 +300,15 @@ def get_init_radio_options(selected_image, model_name):
         top_neurons = list(np.argsort(activations)[::-1][:top_k])
         for top_neuron in top_neurons:
             neuron_dict[top_neuron] = activations[top_neuron]
-        sorted_dict = dict(sorted(neuron_dict.items(), key=lambda item: item[1], reverse=True))
+        sorted_dict = dict(
+            sorted(neuron_dict.items(), key=lambda item: item[1], reverse=True)
+        )
         return sorted_dict
 
     clip_neuron_dict = _get_top_actvation(selected_image, "CLIP", clip_neuron_dict)
-    maple_neuron_dict = _get_top_actvation(selected_image, model_name, maple_neuron_dict)
+    maple_neuron_dict = _get_top_actvation(
+        selected_image, model_name, maple_neuron_dict
+    )
 
     radio_choices = get_radio_names(clip_neuron_dict, maple_neuron_dict)
 
@@ -270,7 +323,9 @@ def get_radio_names(clip_neuron_dict, maple_neuron_dict):
     clip_only_keys = list(set(clip_keys) - (set(maple_keys)))
     maple_only_keys = list(set(maple_keys) - (set(clip_keys)))
 
-    common_keys.sort(key=lambda x: max(clip_neuron_dict[x], maple_neuron_dict[x]), reverse=True)
+    common_keys.sort(
+        key=lambda x: max(clip_neuron_dict[x], maple_neuron_dict[x]), reverse=True
+    )
     clip_only_keys.sort(reverse=True)
     maple_only_keys.sort(reverse=True)
 
@@ -301,13 +356,17 @@ def update_radio_options(evt: gr.EventData, selected_image, model_name):
                 tile_activations = all_activation[token_idx]
                 _sort_and_save_top_k(tile_activations, neuron_dict)
 
-        sorted_dict = dict(sorted(neuron_dict.items(), key=lambda item: item[1], reverse=True))
+        sorted_dict = dict(
+            sorted(neuron_dict.items(), key=lambda item: item[1], reverse=True)
+        )
         return sorted_dict
 
     clip_neuron_dict = {}
     maple_neuron_dict = {}
     clip_neuron_dict = _get_top_actvation(evt, selected_image, "CLIP", clip_neuron_dict)
-    maple_neuron_dict = _get_top_actvation(evt, selected_image, model_name, maple_neuron_dict)
+    maple_neuron_dict = _get_top_actvation(
+        evt, selected_image, model_name, maple_neuron_dict
+    )
 
     clip_keys = list(clip_neuron_dict.keys())
     maple_keys = list(maple_neuron_dict.keys())
@@ -316,7 +375,9 @@ def update_radio_options(evt: gr.EventData, selected_image, model_name):
     clip_only_keys = list(set(clip_keys) - (set(maple_keys)))
     maple_only_keys = list(set(maple_keys) - (set(clip_keys)))
 
-    common_keys.sort(key=lambda x: max(clip_neuron_dict[x], maple_neuron_dict[x]), reverse=True)
+    common_keys.sort(
+        key=lambda x: max(clip_neuron_dict[x], maple_neuron_dict[x]), reverse=True
+    )
     clip_only_keys.sort(reverse=True)
     maple_only_keys.sort(reverse=True)
 
@@ -325,7 +386,9 @@ def update_radio_options(evt: gr.EventData, selected_image, model_name):
     out.extend([f"CLIP-{i}" for i in clip_only_keys[:5]])
     out.extend([f"MaPLE-{i}" for i in maple_only_keys[:5]])
 
-    radio_choices = gr.Radio(choices=out, label="Top activating SAE latent", value=out[0])
+    radio_choices = gr.Radio(
+        choices=out, label="Top activating SAE latent", value=out[0]
+    )
     sleep(0.1)
     return radio_choices
 
@@ -347,6 +410,35 @@ def get_data(image_name, model_name):
     return out
 
 
+def update_all(selected_image, slider_value, toggle_btn, model_name):
+    (
+        seg_mask_display,
+        top_image_1,
+        top_image_2,
+        top_image_3,
+        act_value_1,
+        act_value_2,
+        act_value_3,
+    ) = show_activation_heatmap_clip(selected_image, slider_value, toggle_btn)
+    seg_mask_display_maple = show_activation_heatmap_maple(
+        selected_image, slider_value, model_name
+    )
+    markdown_display, markdown_display_2 = update_markdown(slider_value)
+
+    return (
+        seg_mask_display,
+        seg_mask_display_maple,
+        top_image_1,
+        top_image_2,
+        top_image_3,
+        act_value_1,
+        act_value_2,
+        act_value_3,
+        markdown_display,
+        markdown_display_2,
+    )
+
+
 def load_all_data(image_root, pkl_root):
     image_files = glob(f"{image_root}/*")
     data_dict = {}
@@ -387,33 +479,59 @@ with gr.Blocks(
         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)
+            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
+                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]
             )
-            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)"
+                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
             )
-            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
+                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
+                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=load_image, inputs=[image_selector], outputs=image_display)
             model_selector.change(
-                fn=plot_activation_distribution, inputs=[image_selector, model_selector], outputs=neuron_plot
+                fn=plot_activation_distribution,
+                inputs=[image_selector, model_selector],
+                outputs=neuron_plot,
             )
 
     with gr.Row():
@@ -421,82 +539,106 @@ with gr.Blocks(
             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")
+            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])
+            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)
+            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]
+                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}")
+            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)
+            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)
+            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)
+            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
+                fn=update_radio_options,
+                inputs=[image_selector, model_selector],
+                outputs=[radio_choices],
             )
 
             model_selector.change(
-                fn=update_radio_options, inputs=[image_selector, model_selector], outputs=[radio_choices], queue=True
+                fn=update_radio_options,
+                inputs=[image_selector, model_selector],
+                outputs=[radio_choices],
             )
 
             image_selector.select(
-                fn=update_radio_options, inputs=[image_selector, model_selector], outputs=[radio_choices], queue=True
+                fn=update_radio_options,
+                inputs=[image_selector, model_selector],
+                outputs=[radio_choices],
             )
 
         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,
+            fn=update_all,
+            inputs=[image_selector, radio_choices, toggle_btn, model_selector],
+            outputs=[
+                seg_mask_display,
+                seg_mask_display_maple,
+                top_image_1,
+                top_image_2,
+                top_image_3,
+                act_value_1,
+                act_value_2,
+                act_value_3,
+                markdown_display,
+                markdown_display_2,
+            ],
         )
 
-        # 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,
+            outputs=[
+                seg_mask_display,
+                top_image_1,
+                top_image_2,
+                top_image_3,
+                act_value_1,
+                act_value_2,
+                act_value_3,
+            ],
         )
 
     # Launch the app