File size: 2,788 Bytes
dcaea4a | 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 | 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()
|