Update app.py

app.py

@@ -3,13 +3,6 @@
 # - Standard TensorFlow (Keras) based
 # - Gradio 5 compatible (no theme=)
 # - CPU-friendly (disables GPU usage)
-# - Features:
-#    * Load model (.h5) or use example models
-#    * Graph visualization (nodes = layers, edges = inbound connections)
-#    * Click node -> inspect layer attributes, shapes, params
-#    * View weights (kernels as images + histogram)
-#    * Activation maps for conv layers (image input)
-#    * Simple vs Advanced explanatory text
 # ==========================================================
 
 import io
@@ -17,7 +10,7 @@ import os
 import math
 import traceback
 import warnings
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional
 
 import gradio as gr
 import numpy as np
@@ -27,11 +20,10 @@ import networkx as nx
 
 warnings.filterwarnings("ignore")
 
-# Try importing tensorflow (standard). If import fails, we show a friendly error.
+# Try importing tensorflow
 try:
     import tensorflow as tf
     from tensorflow import keras
-    # force CPU to avoid GPU surprises in Spaces
     try:
         tf.config.set_visible_devices([], "GPU")
     except Exception:
@@ -41,533 +33,181 @@ except Exception as e:
     TF_AVAILABLE = False
     TF_IMPORT_ERROR = str(e)
 
-
 # -------------------- Helpers --------------------
 
 def safe_load_keras_model(fileobj: Optional[io.BytesIO], chosen: str):
-    """
-    If fileobj provided (uploaded .h5), load that model.
-    Else create a built-in small example model depending on 'chosen'.
-    """
     if not TF_AVAILABLE:
-        raise RuntimeError("TensorFlow not available. Add 'tensorflow' to requirements.txt")
+        raise RuntimeError("TensorFlow not available")
 
     if fileobj:
-        # load uploaded .h5 bytes
         fileobj.seek(0)
         tmp_path = "/tmp/uploaded_model.h5"
         with open(tmp_path, "wb") as f:
             f.write(fileobj.read())
         model = keras.models.load_model(tmp_path)
         return model, "uploaded .h5 model"
-    else:
-        # built-in models: "small_cnn" or "toy_resnet"
-        if chosen == "small_cnn":
-            model = keras.Sequential(
-                [
-                    keras.layers.InputLayer(input_shape=(64, 64, 3)),
-                    keras.layers.Conv2D(16, 3, activation="relu", padding="same"),
-                    keras.layers.MaxPool2D(),
-                    keras.layers.Conv2D(32, 3, activation="relu", padding="same"),
-                    keras.layers.MaxPool2D(),
-                    keras.layers.Conv2D(64, 3, activation="relu", padding="same"),
-                    keras.layers.GlobalAveragePooling2D(),
-                    keras.layers.Dense(64, activation="relu"),
-                    keras.layers.Dense(10, activation="softmax"),
-                ]
-            )
-            # build model
-            model.build(input_shape=(None, 64, 64, 3))
-            return model, "Small CNN (example)"
-        elif chosen == "toy_resnet":
-            inputs = keras.Input(shape=(64, 64, 3))
-            x = keras.layers.Conv2D(32, 3, strides=1, padding="same", activation="relu")(inputs)
-            for _ in range(2):
-                y = keras.layers.Conv2D(32, 3, padding="same", activation="relu")(x)
-                y = keras.layers.Conv2D(32, 3, padding="same")(y)
-                x = keras.layers.add([x, y])
-                x = keras.layers.ReLU()(x)
-            x = keras.layers.GlobalAveragePooling2D()(x)
-            outputs = keras.layers.Dense(5, activation="softmax")(x)
-            model = keras.Model(inputs, outputs)
-            model.build(input_shape=(None, 64, 64, 3))
-            return model, "Toy ResNet-like (example)"
-        else:
-            # fallback to small_cnn
-            return safe_load_keras_model(None, "small_cnn")
-
-
-def model_summary_str(model: keras.Model) -> str:
-    """Return model.summary() as a string."""
-    if not TF_AVAILABLE:
-        return "TensorFlow not available."
+
+    if chosen == "small_cnn":
+        model = keras.Sequential([
+            keras.layers.InputLayer(input_shape=(64, 64, 3)),
+            keras.layers.Conv2D(16, 3, activation="relu", padding="same"),
+            keras.layers.MaxPool2D(),
+            keras.layers.Conv2D(32, 3, activation="relu", padding="same"),
+            keras.layers.MaxPool2D(),
+            keras.layers.Conv2D(64, 3, activation="relu", padding="same"),
+            keras.layers.GlobalAveragePooling2D(),
+            keras.layers.Dense(64, activation="relu"),
+            keras.layers.Dense(10, activation="softmax"),
+        ])
+        model.build((None, 64, 64, 3))
+        return model, "Small CNN (example)"
+
+    if chosen == "toy_resnet":
+        inputs = keras.Input(shape=(64, 64, 3))
+        x = keras.layers.Conv2D(32, 3, padding="same", activation="relu")(inputs)
+        for _ in range(2):
+            y = keras.layers.Conv2D(32, 3, padding="same", activation="relu")(x)
+            y = keras.layers.Conv2D(32, 3, padding="same")(y)
+            x = keras.layers.add([x, y])
+            x = keras.layers.ReLU()(x)
+        x = keras.layers.GlobalAveragePooling2D()(x)
+        outputs = keras.layers.Dense(5, activation="softmax")(x)
+        model = keras.Model(inputs, outputs)
+        model.build((None, 64, 64, 3))
+        return model, "Toy ResNet-like (example)"
+
+    return safe_load_keras_model(None, "small_cnn")
+
+
+def model_summary_str(model):
     stream = io.StringIO()
     model.summary(print_fn=lambda s: stream.write(s + "\n"))
     return stream.getvalue()
 
-
 # -------------------- Graph builder --------------------
 
-def build_layer_graph(model: keras.Model):
-    """
-    Build a directed graph (networkx) of layers. Node attributes include:
-      - name, class_name, inbound_layers, outbound_layers, input_shape, output_shape, params
-    """
+def build_layer_graph(model):
     G = nx.DiGraph()
-    # Keras keeps layers in model.layers
-    layers = model.layers
-    # build simple mapping from layer.name -> layer
-    name2layer = {layer.name: layer for layer in layers}
-
-    # gather inbound/outbound info from layer._inbound_nodes
-    for layer in layers:
-        node_attr = {}
-        node_attr["name"] = layer.name
-        node_attr["class_name"] = layer.__class__.__name__
-        try:
-            node_attr["input_shape"] = layer.input_shape
-        except Exception:
-            node_attr["input_shape"] = None
-        try:
-            node_attr["output_shape"] = layer.output_shape
-        except Exception:
-            node_attr["output_shape"] = None
-        try:
-            node_attr["params"] = layer.count_params()
-        except Exception:
-            node_attr["params"] = None
-
-        # inbound layer names (may be empty for InputLayer)
+    for layer in model.layers:
         inbound = []
-        try:
-            for node in getattr(layer, "_inbound_nodes", []) or []:
-                for inbound_layer in getattr(node, "inbound_layers", []) or []:
-                    if hasattr(inbound_layer, "name"):
-                        inbound.append(inbound_layer.name)
-        except Exception:
-            inbound = []
-
-        node_attr["inbound_layers"] = inbound
-        G.add_node(layer.name, **node_attr)
-
-    # add edges based on inbound lists
-    for node in G.nodes(data=True):
-        src = node[0]
-        inbound = node[1].get("inbound_layers", [])
-        for src_in in inbound:
-            if not G.has_node(src_in):
-                # sometimes inbound is a tensor name; ignore
-                continue
-            G.add_edge(src_in, src)
+        for node in getattr(layer, "_inbound_nodes", []) or []:
+            for l in getattr(node, "inbound_layers", []) or []:
+                inbound.append(l.name)
+        G.add_node(
+            layer.name,
+            class_name=layer.__class__.__name__,
+            input_shape=getattr(layer, "input_shape", None),
+            output_shape=getattr(layer, "output_shape", None),
+            params=layer.count_params(),
+            inbound_layers=inbound,
+        )
+    for n, d in G.nodes(data=True):
+        for src in d["inbound_layers"]:
+            if src in G:
+                G.add_edge(src, n)
     return G
 
 
-def nx_to_plotly_fig(G: nx.DiGraph, highlight_node: Optional[str] = None):
-    """
-    Convert networkx graph into a Plotly network figure for interactive selection.
-    Node hover shows class_name and params. Node click returns node name via customdata.
-    """
-    pos = nx.spring_layout(G, seed=42, k=0.5)
-    node_x = []
-    node_y = []
-    texts = []
-    customdata = []
-    sizes = []
-    for n, d in G.nodes(data=True):
-        x, y = pos[n]
-        node_x.append(x)
-        node_y.append(y)
-        cname = d.get("class_name", "")
-        params = d.get("params", 0)
-        texts.append(f"{n} ({cname})\nparams: {params}")
-        customdata.append(n)
-        sizes.append(20 if n != highlight_node else 36)
-
-    edge_x = []
-    edge_y = []
+def nx_to_plotly_fig(G):
+    pos = nx.spring_layout(G, seed=42)
+    edge_x, edge_y = [], []
     for u, v in G.edges():
         x0, y0 = pos[u]
         x1, y1 = pos[v]
         edge_x += [x0, x1, None]
         edge_y += [y0, y1, None]
 
-    edge_trace = go.Scatter(
-        x=edge_x,
-        y=edge_y,
-        line=dict(width=1, color="#888"),
-        hoverinfo="none",
-        mode="lines",
-    )
-
-    node_trace = go.Scatter(
-        x=node_x,
-        y=node_y,
-        mode="markers+text",
-        text=[n for n in G.nodes()],
-        textposition="top center",
-        marker=dict(size=sizes, color="#1f78b4"),
-        hoverinfo="text",
-        hovertext=texts,
-        customdata=customdata,
-    )
+    node_x, node_y, labels = [], [], []
+    for n in G.nodes():
+        x, y = pos[n]
+        node_x.append(x)
+        node_y.append(y)
+        labels.append(n)
 
-    fig = go.Figure(data=[edge_trace, node_trace])
-    fig.update_layout(
-        showlegend=False,
-        hovermode="closest",
-        margin=dict(b=20, l=5, r=5, t=40),
-        height=600,
-        clickmode="event+select",
-    )
+    fig = go.Figure()
+    fig.add_trace(go.Scatter(x=edge_x, y=edge_y, mode="lines"))
+    fig.add_trace(go.Scatter(x=node_x, y=node_y, mode="markers+text", text=labels))
+    fig.update_layout(height=600, showlegend=False)
     fig.update_xaxes(visible=False)
     fig.update_yaxes(visible=False)
     return fig
 
-
-# -------------------- Inspect layer details --------------------
-
-def get_layer_info(model: keras.Model, layer_name: str) -> Dict[str, Any]:
-    """Return layer info: class, input/output shapes, params, config"""
-    if not TF_AVAILABLE:
-        return {"error": "TensorFlow not installed."}
-    try:
-        layer = model.get_layer(layer_name)
-    except Exception as e:
-        return {"error": f"Layer not found: {e}"}
-    info = {
-        "name": layer.name,
-        "class_name": layer.__class__.__name__,
-        "input_shape": getattr(layer, "input_shape", None),
-        "output_shape": getattr(layer, "output_shape", None),
-        "params": layer.count_params() if hasattr(layer, "count_params") else None,
-        "trainable": getattr(layer, "trainable", None),
-        "config": getattr(layer, "get_config", lambda: {})(),
-    }
-    return info
-
-
-def visualize_weights(layer):
-    """
-    For Conv2D kernels, show first few filters as small images.
-    For Dense layers show weight histogram.
-    Returns: PIL image (visual collage) and histogram data (bins/counts)
-    """
-    try:
-        weights = layer.get_weights()
-    except Exception:
-        return None, None
-
-    if len(weights) == 0:
-        return None, None
-
-    # Conv2D: kernel shape (kh, kw, in_ch, out_ch)
-    w = weights[0]
-    if w.ndim == 4:
-        kh, kw, ic, oc = w.shape
-        # visualize up to 8 filters (channels)
-        nshow = min(8, oc)
-        tile_w = kw
-        tile_h = kh
-        pad = 2
-        # normalize each filter to 0..255
-        imgs = []
-        for i in range(nshow):
-            filt = w[:, :, :, i]
-            # collapse input channels by averaging
-            img_arr = filt.mean(axis=2)
-            mn, mx = img_arr.min(), img_arr.max()
-            if mx - mn > 1e-6:
-                img_norm = (img_arr - mn) / (mx - mn)
-            else:
-                img_norm = np.zeros_like(img_arr)
-            img8 = (img_norm * 255).astype("uint8")
-            imgs.append(Image.fromarray(img8).resize((tile_w * 8, tile_h * 8)))
-        # stitch horizontally
-        total_w = sum(im.width for im in imgs) + pad * (len(imgs) - 1)
-        hmax = max(im.height for im in imgs)
-        coll = Image.new("L", (total_w, hmax), color=0)
-        x = 0
-        for im in imgs:
-            coll.paste(im, (x, 0))
-            x += im.width + pad
-        return coll.convert("RGB"), np.histogram(w.flatten(), bins=50)
-    else:
-        # Dense or other: histogram
-        hist = np.histogram(w.flatten(), bins=80)
-        return None, hist
-
-
-# -------------------- Activation extraction --------------------
-
-def build_activation_model(model: keras.Model, layer_names: List[str]):
-    """
-    Create a model that returns outputs of specified layers.
-    """
-    if not TF_AVAILABLE:
-        raise RuntimeError("TensorFlow not available")
-    outputs = [model.get_layer(name).output for name in layer_names]
-    act_model = keras.Model(inputs=model.inputs, outputs=outputs)
-    return act_model
-
-
-def compute_activations(act_model: keras.Model, pil_img: Image.Image):
-    """
-    Resize image to model input (if possible) and return activations as np arrays.
-    For conv layers, they will be (H, W, C) arrays.
-    """
-    # determine required size from model input
-    try:
-        input_shape = act_model.input_shape
-    except Exception:
-        input_shape = None
-    if input_shape and len(input_shape) == 4:
-        ih, iw = input_shape[1], input_shape[2]
-    else:
-        ih, iw = 224, 224
-    img = pil_img.convert("RGB").resize((iw, ih))
-    arr = np.array(img).astype("float32") / 255.0
-    arr = np.expand_dims(arr, axis=0)
-    with np.errstate(all="ignore"):
-        outs = act_model.predict(arr)
-    # ensure list
-    if not isinstance(outs, list):
-        outs = [outs]
-    outs_np = [o.squeeze() for o in outs]
-    return outs_np
-
-
-# -------------------- GRADIO UI callbacks --------------------
-
-def load_model_callback(model_file, example_choice):
-    if not TF_AVAILABLE:
-        return {
-            "error": True,
-            "message": "TensorFlow not installed in the environment. Add 'tensorflow' to requirements.txt and redeploy."
-        }
-    try:
-        model, tag = safe_load_keras_model(model_file, example_choice)
-        summary = model_summary_str(model)
-        G = build_layer_graph(model)
-        fig = nx_to_plotly_fig(G)
-        # basic stats
-        total_params = model.count_params()
-        return {
-            "error": False,
-            "model": model,
-            "graph_fig": fig,
-            "summary": summary,
-            "tag": tag,
-            "total_params": total_params,
-            "nx_graph": G
-        }
-    except Exception as e:
-        return {"error": True, "message": f"Failed to load model: {e}\n{traceback.format_exc()}"}
-
+# -------------------- Inspect --------------------
 
 def node_inspect_callback(state, node_name):
-    """
-    state contains 'model' and 'nx_graph'
-    """
     if not state:
         return "No model loaded.", None, None
-    model = state.get("model")
-    nx_graph = state.get("nx_graph")
-    if node_name is None:
-        return "Click a node in the graph to inspect layer details.", None, None
-    try:
-        info = get_layer_info(model, node_name)
-        # weights visualization
-        layer = model.get_layer(node_name)
-        weights_img, hist = visualize_weights(layer)
-        # create a small HTML summary
-        html = f"**Layer:** {info['name']} ({info['class_name']})  \n"
-        html += f"- input shape: `{info['input_shape']}`  \n"
-        html += f"- output shape: `{info['output_shape']}`  \n"
-        html += f"- params: `{info['params']}`  \n"
-        html += f"- trainable: `{info['trainable']}`  \n"
-        return html, weights_img, hist
-    except Exception as e:
-        return f"Error inspecting node: {e}", None, None
-
-
-def activation_callback(state, uploaded_image, selected_layers_text):
-    """
-    Compute activations for selected layers (comma separated)
-    Return list of PIL preview images (for convs show channel grid; for dense show vector plot)
-    """
-    if not state or "model" not in state:
-        return None, "No model loaded."
-    try:
-        model = state["model"]
-        # parse layers
-        layer_names = [s.strip() for s in selected_layers_text.split(",") if s.strip()]
-        # validate layers
-        valid = []
-        for name in layer_names:
-            try:
-                _ = model.get_layer(name)
-                valid.append(name)
-            except Exception:
-                pass
-        if len(valid) == 0:
-            return None, "No valid layer names found. Use exact layer names from the graph or summary."
-
-        act_model = build_activation_model(model, valid)
-        activations = compute_activations(act_model, uploaded_image)
-        # build previews: for each activation, create a montage
-        previews = []
-        for act in activations:
-            if act.ndim == 3:
-                # H,W,C -> show first up to 12 channels in a grid
-                C = act.shape[2]
-                nshow = min(12, C)
-                # normalize each channel
-                imgs = []
-                for i in range(nshow):
-                    ch = act[:, :, i]
-                    mn, mx = ch.min(), ch.max()
-                    if mx - mn > 1e-6:
-                        chn = (ch - mn) / (mx - mn)
-                    else:
-                        chn = np.zeros_like(ch)
-                    im = Image.fromarray((chn * 255).astype("uint8")).resize((128, 128))
-                    imgs.append(im.convert("RGB"))
-                # make grid 3x4
-                cols = 4
-                rows = math.ceil(len(imgs) / cols)
-                w = cols * 128
-                h = rows * 128
-                collage = Image.new("RGB", (w, h), color=(0, 0, 0))
-                x = y = 0
-                for idx, im in enumerate(imgs):
-                    collage.paste(im, (x * 128, y * 128))
-                    x += 1
-                    if x >= cols:
-                        x = 0
-                        y += 1
-                previews.append(collage)
-            else:
-                # vector -> show as small bar chart image
-                vec = np.array(act).flatten()
-                # scale to 0..255
-                if vec.size > 0:
-                    mn, mx = vec.min(), vec.max()
-                    if mx - mn > 0:
-                        v = (vec - mn) / (mx - mn)
-                    else:
-                        v = np.zeros_like(vec)
-                else:
-                    v = vec
-                # make a simple plot image as grayscale
-                arr = (v.reshape(1, -1) * 255).astype("uint8")
-                im = Image.fromarray(arr).resize((512, 128)).convert("RGB")
-                previews.append(im)
-        return previews, "OK"
-    except Exception as e:
-        return None, f"Activation error: {e}\n{traceback.format_exc()}"
-
-
-# -------------------- Build UI (Gradio 5 compatible) --------------------
+    model = state["model"]
+    layer = model.get_layer(node_name)
+    weights = layer.get_weights()
+    hist_fig = None
+    img = None
+
+    if weights:
+        w = weights[0]
+        hist = np.histogram(w.flatten(), bins=50)
+        hist_fig = go.Figure(go.Bar(x=hist[1][:-1], y=hist[0]))
+
+        if w.ndim == 4:
+            ch = w[:, :, :, 0].mean(axis=2)
+            ch = (ch - ch.min()) / (ch.ptp() + 1e-6)
+            img = Image.fromarray((ch * 255).astype("uint8")).resize((256, 256))
+
+    txt = (
+        f"**Layer:** {layer.name}\n\n"
+        f"- Type: `{layer.__class__.__name__}`\n"
+        f"- Input: `{layer.input_shape}`\n"
+        f"- Output: `{layer.output_shape}`\n"
+        f"- Params: `{layer.count_params()}`"
+    )
+    return txt, img, hist_fig
+
+# -------------------- UI --------------------
 
 with gr.Blocks() as demo:
-    gr.Markdown("# 🔎 TensorFlow Computation Graph Visualizer (Advanced)\n"
-                "Load a Keras `.h5` model or pick an example. Click nodes to inspect layers, view weights and activations.")
+    gr.Markdown("# 🔎 TensorFlow Computation Graph Visualizer")
 
     with gr.Row():
         with gr.Column(scale=1):
-            model_file = gr.File(label="Upload Keras model (.h5)", file_types=[".h5"])
-            example_choice = gr.Dropdown(["small_cnn", "toy_resnet"], value="small_cnn", label="Or pick an example model")
+            model_file = gr.File(label="Upload .h5")
+            example = gr.Dropdown(["small_cnn", "toy_resnet"], value="small_cnn")
             load_btn = gr.Button("Load model")
-            summary_box = gr.Textbox(label="Model Summary", lines=12)
-            total_params_box = gr.Textbox(label="Total Parameters", lines=1)
-            error_box = gr.Markdown()
+            summary = gr.Textbox(lines=12)
+            params = gr.Textbox()
+            error = gr.Markdown()
 
         with gr.Column(scale=2):
-            graph_plot = gr.Plot(label="Computation Graph (click a node to inspect)")
-            node_info = gr.Markdown("Click a node to inspect its details here.")
-            weights_img = gr.Image(label="Weights preview (conv filters or hist)")
-            weights_hist = gr.Plot(label="Weights histogram")
+            graph_plot = gr.Plot()
+            layer_select = gr.Dropdown(label="Select layer to inspect")
+            node_info = gr.Markdown()
+            weights_img = gr.Image()
+            weights_hist = gr.Plot()
 
-    gr.Markdown("### Activations (upload an image to see intermediate maps)")
-    with gr.Row():
-        with gr.Column(scale=1):
-            act_img = gr.Image(label="Upload image for activations", type="pil")
-            layer_names_txt = gr.Textbox(label="Layer names (comma separated) e.g. conv2d,conv2d_1", value="")
-            act_btn = gr.Button("Compute activations")
-            act_msg = gr.Markdown()
-        with gr.Column(scale=2):
-            act_preview = gr.Gallery(
-            label="Activation previews",
-            elem_id="act_gallery",
-            columns=2,
-            height="auto"
-             )
-    # state store for model object & nx graph
     state = gr.State()
 
-    # load model button behavior
-    def on_load(model_file_obj, example_choice_val):
-        if not TF_AVAILABLE:
-            return None, None, "", "", gr.update(visible=True, value=f"TensorFlow import failed: {TF_IMPORT_ERROR}")
-        try:
-            res = load_model_callback(model_file_obj, example_choice_val)
-            if res.get("error"):
-                return None, None, "", "", gr.update(visible=True, value=res.get("message"))
-            model = res["model"]
-            fig = res["graph_fig"]
-            summary = res["summary"]
-            total = res["total_params"]
-            G = res["nx_graph"]
-            st = {"model": model, "nx_graph": G}
-            return st, fig, summary, str(total), gr.update(visible=False, value="")
-        except Exception as e:
-            return None, None, "", "", gr.update(visible=True, value=f"Load error: {e}\n{traceback.format_exc()}")
-
-    load_btn.click(on_load, inputs=[model_file, example_choice], outputs=[state, graph_plot, summary_box, total_params_box, error_box])
-
-    # when user clicks a node on the plotly graph, gradio returns event with clicked point customdata -> node name
-    # Use gr.Plot's 'plotly_events' to capture clicks
-    def on_node_click(evt, st):
-        # evt is list of click events from plotly; we take first if exists
-        if not st:
-            return "No model loaded.", None, None
-        try:
-            if not evt:
-                return "Click a node to inspect it.", None, None
-            # evt is a list of dicts, get 'customdata'
-            node_name = evt[0].get("customdata") or evt[0].get("pointIndex")
-            html, wimg, hist = node_inspect_callback(st, node_name)
-            hist_fig = None
-            if hist is not None:
-                # hist is a tuple (counts, bins)
-                hist_counts, hist_bins = hist
-                hist_fig = go.Figure(data=go.Bar(x=hist_bins[:-1].tolist(), y=hist_counts.tolist()))
-                hist_fig.update_layout(title="Weight histogram", height=240)
-            return html, wimg, hist_fig
-        except Exception as e:
-            return f"Node click error: {e}", None, None
-
-    graph_plot.plotly_events(on_node_click, inputs=[gr.Plot("plotly_events"), state], outputs=[node_info, weights_img, weights_hist])
-
-    # activation compute
-    def on_compute_activations(st, uploaded_image, layer_names_txt_val):
-        previews, msg = activation_callback(st, uploaded_image, layer_names_txt_val)
-        if previews is None:
-            return None, msg
-        # convert previews to displayable list
-        return previews, "Activations computed."
-
-    act_btn.click(on_compute_activations, inputs=[state, act_img, layer_names_txt], outputs=[act_preview, act_msg])
-
-    # friendly note for non-technical users
-    with gr.Accordion("Simple explanation (for non-technical viewers)", open=False):
-        gr.Markdown("""
-**Simple explanation**
-
-- Each rectangle (node) is a layer that transforms the data.
-- Edges show how data flows from one layer to the next.
-- Click any node to see what that layer does (shapes, number of parameters).
-- Upload an image and pick a layer to see the 'activation map' — where the network 'looks' for features.
-""")
+    def on_load(file, ex):
+        model, _ = safe_load_keras_model(file, ex)
+        G = build_layer_graph(model)
+        fig = nx_to_plotly_fig(G)
+        return (
+            {"model": model, "graph": G},
+            fig,
+            model_summary_str(model),
+            str(model.count_params()),
+            "",
+            list(G.nodes())
+        )
+
+    load_btn.click(
+        on_load,
+        inputs=[model_file, example],
+        outputs=[state, graph_plot, summary, params, error, layer_select]
+    )
+
+    layer_select.change(
+        node_inspect_callback,
+        inputs=[state, layer_select],
+        outputs=[node_info, weights_img, weights_hist]
+    )
 
 demo.launch()