init

.gitattributes

@@ -1,3 +1,4 @@
+<<<<<<< HEAD
 *.7z filter=lfs diff=lfs merge=lfs -text
 *.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
@@ -33,3 +34,23 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+=======
+chrismas-imagnet.pkl filter=lfs diff=lfs merge=lfs -text
+dog-imagenet.pkl filter=lfs diff=lfs merge=lfs -text
+dog-mmvp.pkl filter=lfs diff=lfs merge=lfs -text
+golden_gate_bridge.pkl filter=lfs diff=lfs merge=lfs -text
+hen-imagenet-r.pkl filter=lfs diff=lfs merge=lfs -text
+hen-imagenet.pkl filter=lfs diff=lfs merge=lfs -text
+kayaking-ucf.pkl filter=lfs diff=lfs merge=lfs -text
+owl-imagenet-sketch.pkl filter=lfs diff=lfs merge=lfs -text
+owl-imagenet.pkl filter=lfs diff=lfs merge=lfs -text
+paphiopedilum-micranthum.pkl filter=lfs diff=lfs merge=lfs -text
+phalaenopsis-aphrodite.pkl filter=lfs diff=lfs merge=lfs -text
+text-1.pkl filter=lfs diff=lfs merge=lfs -text
+text-2.pkl filter=lfs diff=lfs merge=lfs -text
+text-3.pkl filter=lfs diff=lfs merge=lfs -text
+vegetation-land-eurosat.pkl filter=lfs diff=lfs merge=lfs -text
+data/sae_data/mean_act_values_caltech101.pkl.gz filter=lfs diff=lfs merge=lfs -text
+data/sae_data/mean_act_values_imagenet-sketch.pkl.gz filter=lfs diff=lfs merge=lfs -text
+data/sae_data/mean_act_values_imagenet.pkl.gz filter=lfs diff=lfs merge=lfs -text
+>>>>>>> master

README.md

@@ -1,10 +1,19 @@
 ---
+<<<<<<< HEAD
 title: Patchsae Demo
-emoji: 📚
-colorFrom: blue
-colorTo: purple
+emoji: 😻
+colorFrom: red
+colorTo: gray
 sdk: gradio
 sdk_version: 5.8.0
+=======
+title: Paper14240
+emoji: 📈
+colorFrom: blue
+colorTo: pink
+sdk: gradio
+sdk_version: 5.5.0
+>>>>>>> master
 app_file: app.py
 pinned: false
 ---

app.py

@@ -0,0 +1,503 @@
+import gzip
+import os
+import pickle
+from glob import glob
+from time import sleep
+
+import gradio as gr
+import numpy as np
+import plotly.graph_objects as go
+import torch
+from PIL import Image, ImageDraw
+from plotly.subplots import make_subplots
+
+IMAGE_SIZE = 400
+DATASET_LIST = ["imagenet", "oxford_flowers", "ucf101", "caltech101", "dtd", "eurosat"]
+GRID_NUM = 14
+pkl_root = "./data/out"
+preloaded_data = {}
+
+
+def preload_activation(image_name):
+    for model in ["CLIP"] + [f"MaPLE-{ds}" for ds in DATASET_LIST]:
+        image_file = f"{pkl_root}/{model}/{image_name}.pkl.gz"
+        with gzip.open(image_file, "rb") as f:
+            preloaded_data[model] = pickle.load(f)
+
+
+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()
+    activation[:, noisy_features_indices] = 0
+
+    return activation
+
+
+def get_grid_loc(evt, image):
+    # Get click coordinates
+    x, y = evt._data["index"][0], evt._data["index"][1]
+
+    cell_width = image.width // GRID_NUM
+    cell_height = image.height // GRID_NUM
+
+    grid_x = x // cell_width
+    grid_y = y // cell_height
+    return grid_x, grid_y, cell_width, cell_height
+
+
+def highlight_grid(evt: gr.EventData, image_name):
+    image = data_dict[image_name]["image"]
+    grid_x, grid_y, cell_width, cell_height = get_grid_loc(evt, image)
+
+    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]
+    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))
+
+
+def plot_activations(
+    all_activation, tile_activations=None, grid_x=None, grid_y=None, top_k=5, colors=("blue", "cyan"), model_name="CLIP"
+):
+    fig = go.Figure()
+
+    def _add_scatter_with_annotation(fig, activations, model_name, color, label):
+        fig.add_trace(
+            go.Scatter(
+                x=np.arange(len(activations)),
+                y=activations,
+                mode="lines",
+                name=label,
+                line=dict(color=color, dash="solid"),
+                showlegend=True,
+            )
+        )
+        top_neurons = np.argsort(activations)[::-1][:top_k]
+        for idx in top_neurons:
+            fig.add_annotation(
+                x=idx,
+                y=activations[idx],
+                text=str(idx),
+                showarrow=True,
+                arrowhead=2,
+                ax=0,
+                ay=-15,
+                arrowcolor=color,
+                opacity=0.7,
+            )
+        return fig
+
+    label = f"{model_name.split('-')[-0]} Image-level"
+    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.update_layout(
+        title="Activation Distribution",
+        xaxis_title="SAE latent index",
+        yaxis_title="Activation Value",
+        template="plotly_white",
+    )
+    fig.update_layout(legend=dict(orientation="h", yanchor="middle", y=0.5, xanchor="center", x=0.5))
+
+    return fig
+
+
+def get_activations(evt: gr.EventData, selected_image: str, model_name: str, colors):
+    activation = get_activation_distribution(selected_image, model_name)
+    all_activation = activation.mean(0)
+
+    tile_activations = None
+    grid_x = None
+    grid_y = None
+
+    if evt is not None:
+        if evt._data is not None:
+            image = data_dict[selected_image]["image"]
+            grid_x, grid_y, cell_width, cell_height = get_grid_loc(evt, image)
+            token_idx = grid_y * GRID_NUM + grid_x + 1
+            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
+    )
+    return fig
+
+
+def plot_activation_distribution(evt: gr.EventData, selected_image: str, model_name: str):
+    fig = make_subplots(
+        rows=2,
+        cols=1,
+        shared_xaxes=True,
+        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"))
+
+    def _attach_fig(fig, sub_fig, row, col, yref):
+        for trace in sub_fig.data:
+            fig.add_trace(trace, row=row, col=col)
+
+        for annotation in sub_fig.layout.annotations:
+            annotation.update(yref=yref)
+            fig.add_annotation(annotation)
+        return fig
+
+    fig = _attach_fig(fig, fig_clip, row=1, col=1, yref="y1")
+    fig = _attach_fig(fig, fig_maple, row=2, col=1, yref="y2")
+
+    fig.update_xaxes(title_text="SAE Latent Index", row=2, col=1)
+    fig.update_xaxes(title_text="SAE Latent Index", row=1, col=1)
+    fig.update_yaxes(title_text="Activation Value", row=1, col=1)
+    fig.update_yaxes(title_text="Activation Value", row=2, col=1)
+    fig.update_layout(
+        # height=500,
+        # title="Activation Distributions",
+        template="plotly_white",
+        showlegend=True,
+        legend=dict(orientation="h", yanchor="bottom", y=-0.2, xanchor="center", x=0.5),
+        margin=dict(l=20, r=20, t=40, b=20),
+    )
+
+    return fig
+
+
+def get_segmask(selected_image, slider_value, model_type):
+    image = data_dict[selected_image]["image"]
+    sae_act = get_data(selected_image, model_type)[0]
+    temp = sae_act[:, slider_value]
+    try:
+        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 = (mask - mask.min()) / (mask.max() - mask.min() + 1e-10)
+
+    base_opacity = 30
+    image_array = np.array(image)[..., :3]
+    rgba_overlay = np.zeros((mask.shape[0], mask.shape[1], 4), dtype=np.uint8)
+    rgba_overlay[..., :3] = image_array[..., :3]
+
+    darkened_image = (image_array[..., :3] * (base_opacity / 255)).astype(np.uint8)
+    rgba_overlay[mask == 0, :3] = darkened_image[mask == 0]
+    rgba_overlay[..., 3] = 255  # Fully opaque
+
+    return rgba_overlay
+
+
+def get_top_images(slider_value, toggle_btn):
+    def _get_images(dataset_path):
+        top_image_paths = [
+            os.path.join(dataset_path, "imagenet", f"{slider_value}.jpg"),
+            os.path.join(dataset_path, "imagenet-sketch", f"{slider_value}.jpg"),
+            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))
+            for path in top_image_paths
+        ]
+        return top_images
+
+    if toggle_btn:
+        top_images = _get_images("./data/top_images_masked")
+    else:
+        top_images = _get_images("./data/top_images")
+    return top_images
+
+
+def show_activation_heatmap(selected_image, slider_value, model_type, toggle_btn=False):
+    slider_value = int(slider_value.split("-")[-1])
+    rgba_overlay = get_segmask(selected_image, slider_value, model_type)
+    top_images = get_top_images(slider_value, toggle_btn)
+
+    act_values = []
+    for dataset in ["imagenet", "imagenet-sketch", "caltech101"]:
+        act_value = sae_data_dict["mean_act_values"][dataset][slider_value, :5]
+        act_value = [str(round(value, 3)) for value in act_value]
+        act_value = " | ".join(act_value)
+        out = f"#### Activation values: {act_value}"
+        act_values.append(out)
+
+    return rgba_overlay, top_images, act_values
+
+
+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)
+    sleep(0.1)
+    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):
+    slider_value = int(slider_value.split("-")[-1])
+    rgba_overlay = get_segmask(selected_image, slider_value, model_name)
+    sleep(0.1)
+    return rgba_overlay
+
+
+def get_init_radio_options(selected_image, model_name):
+    clip_neuron_dict = {}
+    maple_neuron_dict = {}
+
+    def _get_top_actvation(selected_image, model_name, neuron_dict, top_k=5):
+        activations = get_activation_distribution(selected_image, model_name).mean(0)
+        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))
+        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)
+
+    radio_choices = get_radio_names(clip_neuron_dict, maple_neuron_dict)
+
+    return radio_choices
+
+
+def get_radio_names(clip_neuron_dict, maple_neuron_dict):
+    clip_keys = list(clip_neuron_dict.keys())
+    maple_keys = list(maple_neuron_dict.keys())
+
+    common_keys = list(set(clip_keys).intersection(set(maple_keys)))
+    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)
+    clip_only_keys.sort(reverse=True)
+    maple_only_keys.sort(reverse=True)
+
+    out = []
+    out.extend([f"common-{i}" for i in common_keys[:5]])
+    out.extend([f"CLIP-{i}" for i in clip_only_keys[:5]])
+    out.extend([f"MaPLE-{i}" for i in maple_only_keys[:5]])
+
+    return out
+
+
+def update_radio_options(evt: gr.EventData, selected_image, model_name):
+    def _sort_and_save_top_k(activations, neuron_dict, top_k=5):
+        top_neurons = list(np.argsort(activations)[::-1][:top_k])
+        for top_neuron in top_neurons:
+            neuron_dict[top_neuron] = activations[top_neuron]
+
+    def _get_top_actvation(evt, selected_image, model_name, neuron_dict):
+        all_activation = get_activation_distribution(selected_image, model_name)
+        image_activation = all_activation.mean(0)
+        _sort_and_save_top_k(image_activation, neuron_dict)
+
+        if evt is not None:
+            if evt._data is not None and isinstance(evt._data["index"], list):
+                image = data_dict[selected_image]["image"]
+                grid_x, grid_y, cell_width, cell_height = get_grid_loc(evt, image)
+                token_idx = grid_y * GRID_NUM + grid_x + 1
+                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))
+        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)
+
+    clip_keys = list(clip_neuron_dict.keys())
+    maple_keys = list(maple_neuron_dict.keys())
+
+    common_keys = list(set(clip_keys).intersection(set(maple_keys)))
+    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)
+    clip_only_keys.sort(reverse=True)
+    maple_only_keys.sort(reverse=True)
+
+    out = []
+    out.extend([f"common-{i}" for i in common_keys[:5]])
+    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])
+    sleep(0.1)
+    return radio_choices
+
+
+def update_markdown(option_value):
+    latent_idx = int(option_value.split("-")[-1])
+    out_1 = f"## Segmentation mask for the selected SAE latent - {latent_idx}"
+    out_2 = f"## Top reference images for the selected SAE latent - {latent_idx}"
+    return out_1, out_2
+
+
+def get_data(image_name, model_name):
+    pkl_root = "./data/out"
+    data_dir = f"{pkl_root}/{model_name}/{image_name}.pkl.gz"
+    with gzip.open(data_dir, "rb") as f:
+        data = pickle.load(f)
+        out = data
+
+    return out
+
+
+def load_all_data(image_root, pkl_root):
+    image_files = glob(f"{image_root}/*")
+    data_dict = {}
+    for image_file in image_files:
+        image_name = os.path.basename(image_file).split(".")[0]
+        if image_file not in data_dict:
+            data_dict[image_name] = {
+                "image": Image.open(image_file).resize((IMAGE_SIZE, IMAGE_SIZE)),
+                "image_path": image_file,
+            }
+
+    sae_data_dict = {}
+    with open("./data/sae_data/mean_acts.pkl", "rb") as f:
+        data = pickle.load(f)
+        sae_data_dict["mean_acts"] = data
+
+    sae_data_dict["mean_act_values"] = {}
+    for dataset in ["imagenet", "imagenet-sketch", "caltech101"]:
+        with gzip.open(f"./data/sae_data/mean_act_values_{dataset}.pkl.gz", "rb") as f:
+            data = pickle.load(f)
+            sae_data_dict["mean_act_values"][dataset] = data
+
+    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()

requirements.txt

@@ -0,0 +1,3 @@
+torch
+matplotlib
+plotly