Show layer progression as activation maps

app.py

@@ -293,21 +293,19 @@ def strongest_channels(activation, count):
     return [(channel, fmap[channel].numpy()) for channel in channel_ids]
 
 
-def layer_overview_gallery(image, activations):
-    preview = fit_image(image)
+def layer_overview_gallery(activations):
     choices = layer_choices()
     outputs = []
 
     for index, activation in enumerate(activations):
         _, channels, height, width = activation.shape
         energy = activation_energy(activation)
-        overlay = overlay_heatmap(preview, energy, alpha=0.42, cmap_name="magma")
+        activation_map = colorize(energy, "magma").resize((TILE_SIZE, TILE_SIZE), Image.Resampling.BILINEAR)
         caption = (
-            f"Layer {index}: {stage_name(index)}\n"
-            f"{choices[index].split(': ', 1)[1]} | {channels} channels | {height}x{width}\n"
-            "Heat shows where the layer has the most total response."
+            f"Layer {index}\n"
+            f"{choices[index].split(': ', 1)[1]} | {channels} maps | {height}x{width}"
         )
-        outputs.append((overlay, caption))
+        outputs.append((activation_map, caption))
 
     return outputs
 
@@ -347,7 +345,7 @@ def analyze_image(image, layer_label=DEFAULT_LAYER, channels=DEFAULT_CHANNELS):
 
     return (
         image,
-        layer_overview_gallery(image, activations),
+        layer_overview_gallery(activations),
         feature_map_gallery(selected, channels),
         overlay_heatmap(fit_image(image), activation_energy(selected), alpha=0.48, cmap_name="magma"),
         layer_summary(layer_index, selected, channels),
@@ -477,7 +475,7 @@ def build_app():
                     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="Selected layer activation overview", type="pil", interactive=False)
+                            selected_overlay = gr.Image(label="Single selected-layer overlay", type="pil", interactive=False)
                         layer_summary_md = gr.Markdown()
                         channel_explanations = gr.Gallery(
                             label="Selected layer feature maps",
@@ -487,7 +485,7 @@ def build_app():
                             elem_classes=["activation-gallery"],
                         )
                         layer_progression = gr.Gallery(
-                            label="Layer-by-layer activation overview",
+                            label="Activation maps as the image moves through the model",
                             columns=3,
                             height=760,
                             object_fit="contain",