| 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 LIFSpikingClassifier, MatchedDenseClassifier |
| from PIL import Image |
| from safetensors.torch import load_file |
|
|
| PROJECT_DIR = Path(__file__).resolve().parent |
| ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "spike-pocket" |
| FRAME = pd.read_parquet(PROJECT_DIR / "data" / "test.parquet") |
| REPORT = json.loads((ARTIFACT_DIR / "evaluation.json").read_text(encoding="utf-8")) |
| SPIKING = LIFSpikingClassifier() |
| SPIKING.load_state_dict(load_file(ARTIFACT_DIR / "spiking_lif.safetensors")) |
| SPIKING.eval() |
| DENSE = MatchedDenseClassifier() |
| DENSE.load_state_dict(load_file(ARTIFACT_DIR / "matched_dense.safetensors")) |
| DENSE.eval() |
|
|
|
|
| @torch.inference_mode() |
| def inspect_spikes( |
| index: int, |
| timesteps: int, |
| noise: float, |
| ) -> tuple[Image.Image, go.Figure, dict]: |
| row = FRAME.iloc[int(index) % len(FRAME)] |
| pixels = np.asarray(row["image"], dtype=np.float32) / 16 |
| rng = np.random.default_rng(int(index) + 2111) |
| corrupted = np.clip(pixels + rng.normal(0, noise, size=64), 0, 1).astype(np.float32) |
| tensor = torch.from_numpy(corrupted)[None] |
| generator = torch.Generator().manual_seed(int(index) + 30_000) |
| spike_logits, spike_rate, raster = SPIKING( |
| tensor, |
| timesteps=int(timesteps), |
| generator=generator, |
| return_raster=True, |
| ) |
| dense_logits = DENSE(tensor) |
| image = Image.fromarray( |
| (corrupted.reshape(8, 8) * 255).astype(np.uint8), mode="L" |
| ).resize((512, 512), Image.Resampling.NEAREST) |
| figure = go.Figure( |
| go.Heatmap( |
| z=raster[0].numpy().T, |
| colorscale=[[0, "#080b14"], [1, "#49e6ff"]], |
| showscale=False, |
| ) |
| ) |
| figure.update_layout( |
| template="plotly_dark", |
| title="Poisson input-spike raster", |
| xaxis_title="Timestep", |
| yaxis_title="Input pixel", |
| ) |
| metrics = { |
| "true_label": int(row["label"]), |
| "spiking_prediction": int(spike_logits.argmax(1)), |
| "dense_prediction": int(dense_logits.argmax(1)), |
| "hidden_spike_rate": float(spike_rate), |
| "verified_clean_accuracy": REPORT["results"]["spiking_lif"]["clean"][ |
| "accuracy" |
| ], |
| } |
| return image, figure, metrics |
|
|
|
|
| with gr.Blocks(title="Spike Pocket") as demo: |
| gr.Markdown( |
| "# Spike Pocket\n" |
| "Inspect Poisson event encoding and a surrogate-gradient leaky-integrate-" |
| "and-fire classifier beside its parameter-matched dense control." |
| ) |
| with gr.Row(): |
| index = gr.Slider(0, len(FRAME) - 1, value=12, step=1, label="Test digit") |
| timesteps = gr.Slider(8, 64, value=32, step=4, label="Timesteps") |
| noise = gr.Slider(0, 0.4, value=0.0, step=0.02, label="Input noise") |
| initial = inspect_spikes(12, 32, 0.0) |
| with gr.Row(): |
| image = gr.Image(value=initial[0], label="Rate-encoded digit") |
| raster = gr.Plot(value=initial[1], label="Spike raster") |
| metrics = gr.JSON(value=initial[2], label="Neuromorphic readout") |
| button = gr.Button("Simulate spikes", variant="primary") |
| button.click( |
| inspect_spikes, |
| inputs=[index, timesteps, noise], |
| outputs=[image, raster, metrics], |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| demo.launch() |
|
|
|
|
