Spaces:
Runtime error
Runtime error
File size: 3,761 Bytes
716e4bc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 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 |