format changes

app.py

@@ -11,17 +11,19 @@ from msma import ScoreFlow, config_presets
 
 
 @cache
-def load_model(modeldir, preset="edm2-img64-s-fid", device='cpu', outdir=None):
+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
 
+
 @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)
@@ -32,47 +34,53 @@ def compute_gmm_likelihood(x_score, model_dir):
 
     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="--")
+    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.5), np.quantile(h, 0.999)
+    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 = (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 fig
+    # 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"):
+
+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 = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0)
         img = img.float().to(device)
-        model = load_model(modeldir='models', preset=preset, device=device)
+        model = load_model(modeldir="models", preset=preset, device=device)
         img_likelihood = model(img).cpu().numpy()
-
-        img = torch.nn.functional.interpolate(img, size=64, mode='bilinear')
+        # 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}")
+        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)
@@ -83,11 +91,19 @@ def run_inference(input_img, preset="edm2-img64-s-fid", device="cuda"):
 
 demo = gr.Interface(
     fn=run_inference,
-    inputs=[gr.Image(type='pil', label="Input Image")],
-    outputs=["text",
-             gr.Plot(label="Anomaly Heatmap"),
-             gr.Plot(label="Comparing to Imagenette"),
-            ],
+    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__":