| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
|
|
| import gradio as gr |
| import numpy as np |
| import pandas as pd |
| import plotly.graph_objects as go |
| import torch |
| from model import DynamicRoutingCapsuleNet, MatchedMLP |
| from PIL import Image |
| from safetensors.torch import load_file |
|
|
| PROJECT_DIR = Path(__file__).resolve().parent |
| ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "capsule-pocket" |
| FRAME = pd.read_parquet(PROJECT_DIR / "data" / "test.parquet") |
| REPORT = json.loads((ARTIFACT_DIR / "evaluation.json").read_text(encoding="utf-8")) |
| CAPSULE = DynamicRoutingCapsuleNet() |
| CAPSULE.load_state_dict(load_file(ARTIFACT_DIR / "capsule.safetensors")) |
| CAPSULE.eval() |
| MLP = MatchedMLP() |
| MLP.load_state_dict(load_file(ARTIFACT_DIR / "matched_mlp.safetensors")) |
| MLP.eval() |
|
|
|
|
| @torch.inference_mode() |
| def inspect_capsules( |
| index: int, |
| vertical: int, |
| horizontal: int, |
| occlude: bool, |
| ) -> tuple[Image.Image, go.Figure, dict]: |
| row = FRAME.iloc[int(index) % len(FRAME)] |
| image = torch.from_numpy(np.asarray(row["image"], dtype=np.float32) / 16).reshape( |
| 8, 8 |
| ) |
| image = torch.roll(image, (int(vertical), int(horizontal)), (0, 1)) |
| if vertical > 0: |
| image[: int(vertical)] = 0 |
| elif vertical < 0: |
| image[int(vertical) :] = 0 |
| if horizontal > 0: |
| image[:, : int(horizontal)] = 0 |
| elif horizontal < 0: |
| image[:, int(horizontal) :] = 0 |
| if occlude: |
| image[3:5, 3:5] = 0 |
| pixels = image.reshape(1, 64) |
| _, lengths = CAPSULE(pixels) |
| mlp_logits = MLP(pixels) |
| figure = go.Figure(go.Bar(x=list(range(10)), y=lengths[0].numpy())) |
| figure.update_layout( |
| template="plotly_dark", |
| title="Digit-capsule vector lengths", |
| xaxis_title="Class", |
| yaxis_title="Length", |
| ) |
| rendered = Image.fromarray( |
| image.mul(255).to(torch.uint8).numpy(), mode="L" |
| ).resize((512, 512), Image.Resampling.NEAREST) |
| metrics = { |
| "true_label": int(row["label"]), |
| "capsule_prediction": int(lengths.argmax(1)), |
| "mlp_prediction": int(mlp_logits.argmax(1)), |
| "verified_capsule_translation_accuracy": REPORT["results"][ |
| "dynamic_routing_capsule" |
| ]["one_pixel_translation"]["accuracy"], |
| } |
| return rendered, figure, metrics |
|
|
|
|
| with gr.Blocks(title="Capsule Pocket") as demo: |
| gr.Markdown( |
| "# Capsule Pocket\n" |
| "Inspect dynamic-routing capsule lengths beside an exactly parameter-matched " |
| "MLP under translation and occlusion." |
| ) |
| with gr.Row(): |
| index = gr.Slider(0, len(FRAME) - 1, value=8, step=1, label="Test digit") |
| vertical = gr.Slider(-1, 1, value=0, step=1, label="Vertical shift") |
| horizontal = gr.Slider(-1, 1, value=0, step=1, label="Horizontal shift") |
| occlude = gr.Checkbox(False, label="Center occlusion") |
| initial = inspect_capsules(8, 0, 0, False) |
| with gr.Row(): |
| image = gr.Image(value=initial[0], label="Input") |
| chart = gr.Plot(value=initial[1], label="Capsule lengths") |
| metrics = gr.JSON(value=initial[2], label="Matched prediction") |
| button = gr.Button("Route capsules", variant="primary") |
| button.click( |
| inspect_capsules, |
| inputs=[index, vertical, horizontal, occlude], |
| outputs=[image, chart, metrics], |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| demo.launch() |
|
|
|
|
