Spaces:
Running
Running
| from functools import cache | |
| from pickle import load | |
| import gradio as gr | |
| import matplotlib.pyplot as plt | |
| import numpy as np | |
| import PIL.Image as Image | |
| import torch | |
| from msma import ScoreFlow, config_presets | |
| def load_model(modeldir, preset="edm2-img64-s-fid", device="cpu", outdir=None): | |
| model = ScoreFlow(preset, device=device) | |
| model.flow.load_state_dict(torch.load(f"{modeldir}/{preset}/flow.pt")) | |
| return model | |
| 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 plot_heatmap(img: Image, heatmap: np.array): | |
| fig, ax = plt.subplots() | |
| cmap = plt.get_cmap("gist_heat") | |
| h = -heatmap[0, 0].copy() | |
| qmin, qmax = np.quantile(h, 0.8), np.quantile(h, 0.999) | |
| h = np.clip(h, a_min=qmin, a_max=qmax) | |
| h = (h - h.min()) / (h.max() - h.min()) | |
| h = cmap(h, bytes=True)[:, :, :3] | |
| h = Image.fromarray(h).resize(img.size, resample=Image.Resampling.BILINEAR) | |
| im = Image.blend(img, h, alpha=0.6) | |
| # im = ax.imshow(np.array(im)) | |
| # # fig.colorbar(im) | |
| # # plt.grid(False) | |
| # # plt.axis("off") | |
| # fig.tight_layout() | |
| return im | |
| def run_inference(input_img, preset="edm2-img64-s-fid", device="cuda"): | |
| # img = center_crop_imagenet(64, img) | |
| input_img = input_img.resize(size=(64, 64), resample=Image.Resampling.LANCZOS) | |
| with torch.inference_mode(): | |
| img = np.array(input_img) | |
| img = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0) | |
| img = img.float().to(device) | |
| model = load_model(modeldir="models", preset=preset, device=device) | |
| img_likelihood = model(img).cpu().numpy() | |
| # img_likelihood = model.scorenet(img).square().sum(1).sum(1).contiguous().float().cpu().unsqueeze(1).numpy() | |
| # print(img_likelihood.shape, img_likelihood.dtype) | |
| img = torch.nn.functional.interpolate(img, size=64, mode="bilinear") | |
| x = model.scorenet(img) | |
| x = x.square().sum(dim=(2, 3, 4)) ** 0.5 | |
| nll, pct, ref_nll = compute_gmm_likelihood( | |
| x.cpu(), model_dir=f"models/{preset}" | |
| ) | |
| outstr = f"Anomaly score: {nll:.3f} / {pct:.2f} percentile" | |
| histplot = plot_against_reference(nll, ref_nll) | |
| heatmapplot = plot_heatmap(input_img, img_likelihood) | |
| return outstr, heatmapplot, histplot | |
| demo = gr.Interface( | |
| fn=run_inference, | |
| inputs=[ | |
| gr.Image(type="pil", label="Input Image"), | |
| gr.Dropdown(choices=config_presets.keys(), label="Score Model"), | |
| ], | |
| outputs=[ | |
| "text", | |
| gr.Image(label="Anomaly Heatmap", min_width=64), | |
| gr.Plot(label="Comparing to Imagenette"), | |
| ], | |
| examples=[ | |
| ['goldfish.JPEG', "edm2-img64-s-fid"] | |
| ] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |