Create visuals.py

visuals.py

@@ -0,0 +1,136 @@
+# visuals.py
+import torch
+import matplotlib.pyplot as plt
+import numpy as np
+import gradio as gr
+from PIL import Image
+
+from uei_core.models import ModelPortfolio
+from uei_core.uncertainty import UncertaintyEstimator
+from uei_core.energy import EnergyProfiler
+
+device = "cpu"
+
+models = ModelPortfolio(device=device)
+unc = UncertaintyEstimator()
+energy = EnergyProfiler()
+
+
+# ------------------------------
+# 1️⃣ Plot: Uncertainty vs Energy Curve
+# ------------------------------
+def plot_unc_energy(img):
+    x = models.preprocess(img)
+
+    logits_s, E_s = energy.measure(models.infer_small, x)
+    logits_l, E_l = energy.measure(models.infer_large, x)
+
+    U_s = float(unc.estimate(logits_s))
+    U_l = float(unc.estimate(logits_l))
+
+    # Create plot
+    fig, ax = plt.subplots(figsize=(5,4), dpi=120)
+
+    xs = [E_s, E_l]
+    ys = [U_s, U_l]
+    labels = ["Small Model", "Large Model"]
+    colors = ["#1f77b4", "#ff7f0e"]
+
+    ax.scatter(xs, ys, s=150, color=colors)
+    ax.plot(xs, ys, linestyle="--", color="#888")
+
+    for i, label in enumerate(labels):
+        ax.annotate(label, (xs[i], ys[i]), textcoords="offset points",
+                    xytext=(8,5), ha='left', fontsize=10)
+
+    ax.set_xlabel("Energy (proxy units)")
+    ax.set_ylabel("Estimated Uncertainty")
+    ax.set_title("Uncertainty vs Energy")
+    ax.grid(True, alpha=0.3)
+
+    return fig
+
+
+# ------------------------------
+# 2️⃣ Plot: Layer Activation Heatmap
+# ------------------------------
+def activation_heatmap(img):
+    x = models.preprocess(img)
+
+    # Register forward hook on the first conv
+    activations = {}
+
+    def hook(module, input, output):
+        activations["feat"] = output.detach().cpu()
+
+    h = models.small.features[0].register_forward_hook(hook)
+    models.small(x)
+    h.remove()
+
+    feat = activations["feat"][0]  # first batch
+
+    # Average channels → 2D heatmap
+    heat = feat.mean(dim=0).numpy()
+
+    fig, ax = plt.subplots(figsize=(4,4), dpi=120)
+    ax.imshow(heat, cmap="viridis")
+    ax.set_title("Early Layer Activation Heatmap")
+    ax.axis("off")
+
+    return fig
+
+
+# ------------------------------
+# 3️⃣ Plot: Model Comparison Bars
+# ------------------------------
+def model_comparison(img):
+    x = models.preprocess(img)
+
+    logits_s, E_s = energy.measure(models.infer_small, x)
+    logits_l, E_l = energy.measure(models.infer_large, x)
+
+    U_s = float(unc.estimate(logits_s))
+    U_l = float(unc.estimate(logits_l))
+
+    fig, ax = plt.subplots(figsize=(6,4))
+
+    labels = ["Small Model", "Large Model"]
+    energy_vals = [E_s, E_l]
+    unc_vals = [U_s, U_l]
+
+    x_axis = np.arange(len(labels))
+    w = 0.35
+
+    ax.bar(x_axis - w/2, energy_vals, w, label="Energy", color="#2ca02c")
+    ax.bar(x_axis + w/2, unc_vals, w, label="Uncertainty", color="#d62728")
+
+    ax.set_xticks(x_axis)
+    ax.set_xticklabels(labels)
+    ax.set_title("Model Energy & Uncertainty Comparison")
+    ax.legend()
+    ax.grid(alpha=0.2)
+
+    return fig
+
+
+# ------------------------------
+#  🔥 Gradio Interface
+# ------------------------------
+def get_visual_ui():
+    with gr.Blocks() as demo:
+        gr.Markdown("## 🔍 UEI Visualization Dashboard")
+        gr.Markdown("Explore how UEI behaves internally with colorful charts")
+
+        img = gr.Image(type="pil", label="Upload Image")
+
+        with gr.Tabs():
+            with gr.Tab("Uncertainty vs Energy"):
+                gr.Plot(label="Chart").render(fn=plot_unc_energy, inputs=img)
+
+            with gr.Tab("Layer Activations"):
+                gr.Plot(label="Activation Heatmap").render(fn=activation_heatmap, inputs=img)
+
+            with gr.Tab("Model Comparison"):
+                gr.Plot(label="Energy & Uncertainty Bars").render(fn=model_comparison, inputs=img)
+
+    return demo