plotting a blended version of the heatmap

app.py

@@ -4,6 +4,7 @@ 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
@@ -39,39 +40,50 @@ def plot_against_reference(nll, ref_nll):
     fig.tight_layout()
     return fig
 
-
-def plot_heatmap(heatmap):
+def plot_heatmap(img: Image, heatmap: np.array):
     fig, ax = plt.subplots()
-    im = heatmap[0,0]
-    ax.imshow(im, cmap='gist_heat')
+    cmap = plt.get_cmap("gist_heat")
+    h = heatmap[0,0].copy()
+    qmin, qmax = np.quantile(h, 0.5), 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 fig
 
-# def compute_scores    
-
+def run_inference(input_img, preset="edm2-img64-s-fid", device="cuda"):
 
-def run_inference(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.nn.functional.interpolate(img, size=64, mode='bilinear')
-        img = img.to(device)
+        img = img.float().to(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')
         x = model.scorenet(img)
         x = x.square().sum(dim=(2, 3, 4)) ** 0.5
-        img_likelihood = model(img).cpu().numpy()
         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(img_likelihood)
+    heatmapplot = plot_heatmap(input_img, img_likelihood)
 
     return outstr, heatmapplot, histplot
 
 
 demo = gr.Interface(
     fn=run_inference,
-    inputs=["image"],
+    inputs=[gr.Image(type='pil', label="Input Image")],
     outputs=["text",
              gr.Plot(label="Anomaly Heatmap"),
              gr.Plot(label="Comparing to Imagenette"),