Spaces:
Runtime error
Runtime error
File size: 1,688 Bytes
52f9197 be66f33 52f9197 be66f33 b1602ac be66f33 52f9197 7387897 52f9197 7387897 52f9197 7387897 be66f33 7387897 52f9197 be66f33 7387897 be66f33 7387897 52f9197 7387897 52f9197 7387897 be66f33 7387897 be66f33 7387897 be66f33 7387897 52f9197 be66f33 7387897 be66f33 52f9197 |
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 |
from functools import cache
from pickle import load
import gradio as gr
import matplotlib.pyplot as plt
import numpy as np
import torch
from msma import build_model, config_presets
@cache
def load_model(preset="edm2-img64-s-fid", device='cpu'):
return build_model(preset, device)
@cache
def load_reference_scores(model_dir):
with np.load(f"{model_dir}/refscores.npz", "rb") as f:
ref_nll = f["arr_0"]
return ref_nll
def compute_gmm_likelihood(x_score, model_dir):
with open(f"{model_dir}/gmm.pkl", "rb") as f:
clf = load(f)
nll = -clf.score(x_score)
ref_nll = load_reference_scores(model_dir)
percentile = (ref_nll < nll).mean() * 100
return nll, percentile, ref_nll
def plot_against_reference(nll, ref_nll):
fig, ax = plt.subplots()
ax.hist(ref_nll, label="Reference Scores")
ax.axvline(nll, label='Image Score', c='red', ls="--")
plt.legend()
fig.tight_layout()
return fig
def run_inference(img, preset="edm2-img64-s-fid", device="cuda"):
img = torch.from_numpy(img).permute(2,0,1).unsqueeze(0)
img = torch.nn.functional.interpolate(img, size=64, mode='bilinear')
model = load_model(preset=preset, device=device)
x = model(img.cuda())
x = x.square().sum(dim=(2, 3, 4)) ** 0.5
nll, pct, ref_nll = compute_gmm_likelihood(x.cpu(), model_dir=f"models/{preset}")
plot = plot_against_reference(nll, ref_nll)
outstr = f"Anomaly score: {nll:.3f} / {pct:.2f} percentile"
return outstr, plot
demo = gr.Interface(
fn=run_inference,
inputs=["image"],
outputs=["text", gr.Plot(label="Comparing to Imagenette")],
)
if __name__ == "__main__":
demo.launch()
|