Leeps's picture
Stack activation layer panels vertically
be20821
Raw
History Blame Contribute Delete
19.9 kB
import concurrent.futures
import math
import os
from functools import lru_cache
from io import BytesIO
os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")
import gradio as gr
import numpy as np
import pandas as pd
import requests
import torch
import torch.nn.functional as F
import torchvision.models as models
from matplotlib import colormaps
from PIL import Image, ImageDraw, ImageFilter
torch.set_num_threads(max(1, min(4, os.cpu_count() or 1)))
APP_TITLE = "Model Layers Viewer"
MOMA_URL = "https://media.githubusercontent.com/media/MuseumofModernArt/collection/main/Artworks.csv"
HEADERS = {
"User-Agent": (
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) "
"AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0 Safari/537.36"
)
}
MOMA_SAMPLE_SIZE = 36
MOMA_CANDIDATES = 140
NOTEBOOK_LAYERS = [0, 2, 4, 7, 14, 18]
CHANNELS_PER_LAYER = 8
OVERLAY_LAYER = 14
PREVIEW_SIZE = 360
TILE_SIZE = 120
def stage_name(layer_index):
if layer_index <= 2:
return "Early vision"
if layer_index <= 7:
return "Textures and small shapes"
if layer_index <= 14:
return "Object parts and composition"
return "Compact semantic evidence"
def stage_explanation(layer_index):
if layer_index <= 2:
return "These first filters respond to edges, contrast, color changes, and simple marks."
if layer_index <= 7:
return "The model is combining simple marks into repeated textures, corners, curves, and local motifs."
if layer_index <= 14:
return "At this depth, each activation cell responds to a wider region of the original image."
return "The final feature blocks are very small spatial maps that summarize evidence for the image as a whole."
def make_fallback_images():
samples = []
size = 384
def canvas(bg):
return Image.new("RGB", (size, size), bg)
img = canvas((232, 225, 210))
draw = ImageDraw.Draw(img)
for i, color in enumerate([(199, 44, 65), (25, 93, 138), (240, 174, 54), (42, 120, 83)]):
draw.rectangle((30 + i * 38, 40 + i * 48, 250 + i * 20, 118 + i * 55), fill=color)
samples.append({"img": img, "title": "Fallback: stacked color blocks", "artist": "generated", "url": ""})
img = canvas((22, 25, 30))
draw = ImageDraw.Draw(img)
for radius, color in zip(
range(170, 10, -24),
[(231, 76, 60), (241, 196, 15), (52, 152, 219), (46, 204, 113), (155, 89, 182)],
):
draw.ellipse((192 - radius, 192 - radius, 192 + radius, 192 + radius), outline=color, width=14)
samples.append({"img": img, "title": "Fallback: concentric rings", "artist": "generated", "url": ""})
img = canvas((246, 244, 238))
draw = ImageDraw.Draw(img)
for x in range(-120, size + 120, 34):
draw.line((x, 0, x + 190, size), fill=(35, 63, 99), width=9)
draw.line((x + 16, 0, x + 206, size), fill=(218, 83, 44), width=4)
samples.append({"img": img, "title": "Fallback: diagonal line field", "artist": "generated", "url": ""})
img = canvas((245, 245, 242))
draw = ImageDraw.Draw(img)
points = [(40, 280), (85, 110), (150, 250), (205, 70), (260, 220), (340, 115)]
draw.line(points, fill=(26, 26, 26), width=12, joint="curve")
for x, y in points:
draw.ellipse((x - 18, y - 18, x + 18, y + 18), fill=(35, 125, 161))
samples.append({"img": img.filter(ImageFilter.SMOOTH), "title": "Fallback: bold path drawing", "artist": "generated", "url": ""})
return samples
@lru_cache(maxsize=1)
def load_model():
weights = models.MobileNet_V2_Weights.DEFAULT
model = models.mobilenet_v2(weights=weights)
model.eval()
return model, weights.transforms()
@lru_cache(maxsize=1)
def layer_choices():
model, _ = load_model()
return [f"{index}: {layer.__class__.__name__}" for index, layer in enumerate(model.features)]
def rgb_image(image):
if image is None:
return None
if isinstance(image, np.ndarray):
image = Image.fromarray(image)
return image.convert("RGB")
def fit_image(image, max_side=PREVIEW_SIZE):
image = image.copy().convert("RGB")
image.thumbnail((max_side, max_side), Image.Resampling.LANCZOS)
return image
def normalize_map(values):
values = values.astype(np.float32)
low = float(values.min())
high = float(values.max())
return (values - low) / (high - low + 1e-8)
def colorize(values, cmap_name="magma"):
norm = normalize_map(values)
rgba = colormaps[cmap_name](norm)
rgb = (rgba[:, :, :3] * 255).astype(np.uint8)
return Image.fromarray(rgb)
def overlay_heatmap(image, values, alpha=0.46, cmap_name="magma"):
heat = colorize(values, cmap_name).resize(image.size, Image.Resampling.BILINEAR).convert("RGBA")
base = image.convert("RGBA")
return Image.blend(base, heat, alpha=alpha).convert("RGB")
def fetch_moma_image(row):
try:
response = requests.get(row["ImageURL"], headers=HEADERS, timeout=10)
response.raise_for_status()
image = Image.open(BytesIO(response.content)).convert("RGB")
image.thumbnail((640, 640), Image.Resampling.LANCZOS)
return {
"img": image.copy(),
"title": str(row.get("Title") or "Untitled")[:80],
"artist": str(row.get("Artist") or "Unknown artist")[:80],
"url": str(row.get("ImageURL") or ""),
}
except Exception:
return None
@lru_cache(maxsize=1)
def load_moma_items():
try:
df = pd.read_csv(
MOMA_URL,
usecols=["Title", "Artist", "ImageURL"],
low_memory=False,
)
has_image = df["ImageURL"].notna() & df["ImageURL"].str.startswith("http", na=False)
candidates = df[has_image].sample(
min(MOMA_CANDIDATES, int(has_image.sum())),
random_state=42,
)
rows = [row for _, row in candidates.iterrows()]
with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor:
results = list(executor.map(fetch_moma_image, rows))
items = [item for item in results if item is not None][:MOMA_SAMPLE_SIZE]
if len(items) >= 8:
return items, "Loaded images from the MoMA collection data used in the notebook."
except Exception as exc:
return make_fallback_images(), f"MoMA images could not be loaded here ({type(exc).__name__}). Using generated fallback images."
return make_fallback_images(), "MoMA image downloads did not return enough usable images. Using generated fallback images."
def gallery_items():
items, _ = load_moma_items()
return [(item["img"], f"{item['title']}\n{item['artist']}") for item in items]
def collect_activations(image):
model, preprocess = load_model()
x = preprocess(image).unsqueeze(0)
activations = []
with torch.inference_mode():
out = x
for layer in model.features:
out = layer(out)
activations.append(out.detach().cpu())
return activations
def feature_vector(image):
model, preprocess = load_model()
x = preprocess(image).unsqueeze(0)
with torch.inference_mode():
out = model.features(x)
pooled = F.adaptive_avg_pool2d(out, (1, 1)).flatten(1)
return pooled.squeeze(0).cpu().numpy()
@lru_cache(maxsize=1)
def moma_feature_space():
items, message = load_moma_items()
features = np.vstack([feature_vector(item["img"]) for item in items])
method = "UMAP"
try:
import umap
reducer = umap.UMAP(
n_components=2,
n_neighbors=max(2, min(10, len(items) - 1)),
min_dist=0.12,
metric="cosine",
random_state=42,
)
embedding = reducer.fit_transform(features)
except Exception:
method = "PCA fallback"
centered = features - features.mean(axis=0, keepdims=True)
_, _, vt = np.linalg.svd(centered, full_matrices=False)
embedding = centered @ vt[:2].T
return items, features, embedding, method, message
def render_embedding_image(embedding, selected_index, method):
width, height = 720, 540
margin = 58
image = Image.new("RGB", (width, height), (250, 250, 248))
draw = ImageDraw.Draw(image)
x = embedding[:, 0]
y = embedding[:, 1]
x_span = float(x.max() - x.min()) or 1.0
y_span = float(y.max() - y.min()) or 1.0
xs = margin + ((x - x.min()) / x_span) * (width - 2 * margin)
ys = height - margin - ((y - y.min()) / y_span) * (height - 2 * margin)
for t in np.linspace(0, 1, 5):
gx = margin + t * (width - 2 * margin)
gy = margin + t * (height - 2 * margin)
draw.line((gx, margin, gx, height - margin), fill=(226, 228, 232), width=1)
draw.line((margin, gy, width - margin, gy), fill=(226, 228, 232), width=1)
draw.rectangle((margin, margin, width - margin, height - margin), outline=(190, 196, 205), width=1)
draw.text((margin, 22), f"{method} of MobileNetV2 activations for MoMA images", fill=(35, 39, 47))
for index, (px, py) in enumerate(zip(xs, ys)):
if index == selected_index:
continue
draw.ellipse((px - 6, py - 6, px + 6, py + 6), fill=(59, 130, 246), outline=(255, 255, 255), width=2)
px = xs[selected_index]
py = ys[selected_index]
draw.ellipse((px - 12, py - 12, px + 12, py + 12), fill=(225, 29, 72), outline=(20, 20, 20), width=2)
draw.text((margin, height - 36), "red point = selected image", fill=(75, 85, 99))
return image
def activation_energy(activation):
fmap = activation.clamp(min=0).mean(dim=1).squeeze(0).cpu().numpy()
return normalize_map(fmap)
def strongest_channels(activation, count):
fmap = activation.squeeze(0).cpu()
strengths = fmap.mean(dim=(1, 2))
channel_ids = torch.argsort(strengths, descending=True)[: min(int(count), fmap.shape[0])].tolist()
return [(channel, fmap[channel].numpy()) for channel in channel_ids]
def channel_grid_image(activation, count=CHANNELS_PER_LAYER):
selected = strongest_channels(activation, count)
cols = 4
rows = math.ceil(len(selected) / cols)
label_h = 20
grid = Image.new("RGB", (cols * TILE_SIZE, rows * (TILE_SIZE + label_h)), (248, 248, 246))
draw = ImageDraw.Draw(grid)
for position, (channel, values) in enumerate(selected):
x = (position % cols) * TILE_SIZE
y = (position // cols) * (TILE_SIZE + label_h)
tile = colorize(values, "viridis").resize((TILE_SIZE, TILE_SIZE), Image.Resampling.NEAREST)
grid.paste(tile, (x, y))
draw.text((x + 6, y + TILE_SIZE + 3), f"channel {channel}", fill=(35, 39, 47))
return grid
def fixed_layer_gallery(activations):
choices = layer_choices()
outputs = []
for layer_index in NOTEBOOK_LAYERS:
activation = activations[layer_index]
_, channels, height, width = activation.shape
caption = (
f"Layer {layer_index}: {stage_name(layer_index)}\n"
f"torch.Size([1, {channels}, {height}, {width}])\n"
f"{choices[layer_index].split(': ', 1)[1]} | {CHANNELS_PER_LAYER} channel maps"
)
outputs.append((channel_grid_image(activation), caption))
return outputs
def layer_summary():
return (
"### Pixel activations through the model\n"
"These are the same intermediate layers checked in the notebook: 0, 2, 4, 7, 14, and 18. Each panel is one layer. "
"Inside each panel are several channel maps from that layer: bright cells are places where that map activated strongly. "
"Earlier layers keep larger grids; later layers shrink to smaller, chunkier grids."
)
def analyze_image(image):
image = rgb_image(image)
if image is None:
items, _ = load_moma_items()
image = items[0]["img"]
activations = collect_activations(image)
overlay_activation = activations[OVERLAY_LAYER]
return (
image,
fixed_layer_gallery(activations),
overlay_heatmap(fit_image(image), activation_energy(overlay_activation), alpha=0.48, cmap_name="magma"),
layer_summary(),
)
def make_umap_plot(selected_index=0):
items, _, embedding, method, message = moma_feature_space()
selected_index = 0 if selected_index is None else int(selected_index)
selected_index = max(0, min(selected_index, len(items) - 1))
plot_image = render_embedding_image(embedding, selected_index, method)
item = items[selected_index]
distances = np.linalg.norm(embedding - embedding[selected_index], axis=1)
closest_ids = [idx for idx in np.argsort(distances) if idx != selected_index][:3]
farthest_ids = [idx for idx in np.argsort(distances)[::-1] if idx != selected_index][:3]
closest_gallery = [(items[idx]["img"], f"{items[idx]['title']}\n{items[idx]['artist']}") for idx in closest_ids]
farthest_gallery = [(items[idx]["img"], f"{items[idx]['title']}\n{items[idx]['artist']}") for idx in farthest_ids]
relation = (
f"### Selected MoMA image\n"
f"**{item['title']}** by {item['artist']}\n\n"
f"Nearby points have similar MobileNetV2 feature vectors; faraway points have very different vectors. "
f"This does not mean the artworks share labels. It means the neural net produced similar or different "
f"patterns of activation after processing the pixels.\n\n"
f"{message}"
)
return plot_image, relation, closest_gallery, farthest_gallery
def initialize():
items, _ = load_moma_items()
selected_index = 0
selected_image, layer_gallery, overlay, summary = analyze_image(items[selected_index]["img"])
umap_plot, relation, closest_gallery, farthest_gallery = make_umap_plot(selected_index)
return (
gallery_items(),
selected_image,
layer_gallery,
overlay,
summary,
umap_plot,
relation,
closest_gallery,
farthest_gallery,
selected_index,
)
def select_moma(evt: gr.SelectData):
index = evt.index if isinstance(evt.index, int) else 0
items, _ = load_moma_items()
index = max(0, min(index, len(items) - 1))
selected_image, layer_gallery, overlay, summary = analyze_image(items[index]["img"])
umap_plot, relation, closest_gallery, farthest_gallery = make_umap_plot(index)
return selected_image, layer_gallery, overlay, summary, umap_plot, relation, closest_gallery, farthest_gallery, index
def analyze_upload(image):
selected_image, layer_gallery, overlay, summary = analyze_image(image)
umap_plot, relation, closest_gallery, farthest_gallery = make_umap_plot(0)
relation = (
"### Uploaded image\n"
"The UMAP tab is built from the MoMA image set. Uploaded images are shown in the layer viewer, "
"but are not inserted into the precomputed MoMA map."
)
return selected_image, layer_gallery, overlay, summary, umap_plot, relation, closest_gallery, farthest_gallery, -1
def build_app():
theme = gr.themes.Soft(
primary_hue="blue",
secondary_hue="pink",
neutral_hue="slate",
radius_size="sm",
)
css = """
.activation-gallery img { object-fit: contain !important; }
.moma-gallery img { object-fit: cover !important; }
.compact-note textarea { font-size: 0.95rem !important; }
"""
with gr.Blocks(title=APP_TITLE, theme=theme, css=css) as demo:
selected_index = gr.State(0)
gr.Markdown(
"# Model Layers Viewer\n"
"Choose a MoMA image or upload your own. The app runs MobileNetV2 on CPU and shows pixel activation maps "
"at the same intermediate layers used in the workshop notebook.\n\n"
"Inspired by [Small ML: The Art of Data Discovery]"
"(https://colab.research.google.com/drive/1_8GasNHKJpO8x-AAt93D3LfrvCLGzsbj?usp=sharing), "
"a workshop by Sam Keene at [ITP Camp 2026](https://itp.nyu.edu/camp/2026/session/205)."
)
with gr.Row(equal_height=False):
with gr.Column(scale=1, min_width=310):
moma_gallery = gr.Gallery(
label="MoMA image set",
columns=3,
rows=4,
height=520,
object_fit="cover",
elem_classes=["moma-gallery"],
)
upload = gr.Image(label="Upload your own image", type="pil", sources=["upload", "clipboard"])
with gr.Column(scale=2, min_width=520):
with gr.Tabs():
with gr.Tab("Layer activations"):
with gr.Row(equal_height=False):
selected_image = gr.Image(label="Selected image", type="pil", interactive=False)
selected_overlay = gr.Image(label="One overlay: layer 14 activation over image", type="pil", interactive=False)
layer_summary_md = gr.Markdown()
layer_progression = gr.Gallery(
label="Notebook layers: activation maps across channels",
columns=1,
height=1400,
object_fit="contain",
elem_classes=["activation-gallery"],
)
with gr.Tab("UMAP relationships"):
umap_plot = gr.Image(label="MoMA images after MobileNetV2 activation extraction", type="pil", interactive=False)
relation_md = gr.Markdown()
with gr.Row(equal_height=False):
closest_gallery = gr.Gallery(
label="3 closest images in feature space",
columns=3,
rows=1,
height=300,
object_fit="cover",
elem_classes=["moma-gallery"],
)
farthest_gallery = gr.Gallery(
label="3 farthest images in feature space",
columns=3,
rows=1,
height=300,
object_fit="cover",
elem_classes=["moma-gallery"],
)
demo.load(
initialize,
inputs=None,
outputs=[
moma_gallery,
selected_image,
layer_progression,
selected_overlay,
layer_summary_md,
umap_plot,
relation_md,
closest_gallery,
farthest_gallery,
selected_index,
],
show_progress="full",
)
moma_gallery.select(
select_moma,
inputs=None,
outputs=[
selected_image,
layer_progression,
selected_overlay,
layer_summary_md,
umap_plot,
relation_md,
closest_gallery,
farthest_gallery,
selected_index,
],
show_progress="minimal",
)
upload.change(
analyze_upload,
inputs=[upload],
outputs=[
selected_image,
layer_progression,
selected_overlay,
layer_summary_md,
umap_plot,
relation_md,
closest_gallery,
farthest_gallery,
selected_index,
],
show_progress="minimal",
)
return demo
if __name__ == "__main__":
build_app().launch()