| from __future__ import annotations |
|
|
| from pathlib import Path |
|
|
| import gradio as gr |
| import numpy as np |
| import plotly.graph_objects as go |
| import torch |
| from model import JEncoder |
| from plotly.subplots import make_subplots |
| from safetensors.torch import load_file |
|
|
| ARTIFACT_DIR = Path(__file__).resolve().parent / "artifacts" / "pocket-jepa" |
| MODEL = JEncoder() |
| MODEL.load_state_dict(load_file(ARTIFACT_DIR / "model.safetensors")) |
| MODEL.eval() |
| ATLAS = np.load(ARTIFACT_DIR / "j_space.npz") |
| EMBEDDINGS = ATLAS["embeddings"] |
| IMAGES = ATLAS["images"] |
| LABELS = ATLAS["labels"] |
|
|
|
|
| def explore(index: int, mask_size: int) -> tuple[go.Figure, dict]: |
| index = int(index) % len(IMAGES) |
| mask_size = int(mask_size) |
| image = IMAGES[index].copy() |
| masked = image.copy() |
| start = (8 - mask_size) // 2 |
| masked[start : start + mask_size, start : start + mask_size] = 0 |
| with torch.inference_mode(): |
| embedding = MODEL( |
| torch.from_numpy(masked[None].astype(np.float32)) |
| ).numpy()[0] |
| embedding /= max(np.linalg.norm(embedding), 1e-8) |
| distances = 1 - EMBEDDINGS @ embedding |
| nearest = np.argsort(distances)[:5] |
| panels = [image, masked, *IMAGES[nearest[:3]]] |
| titles = [ |
| f"Original: {LABELS[index]}", |
| "Masked context", |
| *[f"J-neighbor: {LABELS[item]}" for item in nearest[:3]], |
| ] |
| figure = make_subplots(rows=1, cols=5, subplot_titles=titles) |
| for panel_index, (panel, title) in enumerate(zip(panels, titles, strict=True)): |
| figure.add_trace( |
| go.Heatmap( |
| z=np.flipud(panel), |
| colorscale="Viridis", |
| showscale=False, |
| name=title, |
| ), |
| row=1, |
| col=panel_index + 1, |
| ) |
| figure.update_layout( |
| title="Masked prediction in learned J-space", |
| template="plotly_dark", |
| height=310, |
| margin=dict(t=70, b=20), |
| ) |
| return figure, { |
| "query_label": int(LABELS[index]), |
| "nearest_labels": LABELS[nearest].astype(int).tolist(), |
| "cosine_distances": np.round(distances[nearest], 4).tolist(), |
| } |
|
|
|
|
| with gr.Blocks(title="Pocket JEPA") as demo: |
| gr.Markdown( |
| "# Pocket JEPA\n" |
| "Hide the center of a digit and inspect the nearest complete images in " |
| "the learned joint-embedding prediction space." |
| ) |
| with gr.Row(): |
| index = gr.Slider(0, len(IMAGES) - 1, 0, step=1, label="Image") |
| mask = gr.Slider(2, 5, 3, step=1, label="Center mask") |
| run = gr.Button("Explore J-space", variant="primary") |
| gallery = gr.Plot() |
| metrics = gr.JSON() |
| run.click(explore, [index, mask], [gallery, metrics]) |
| demo.load(explore, [index, mask], [gallery, metrics]) |
|
|
|
|
| if __name__ == "__main__": |
| demo.launch() |
|
|
