probs

app.py

@@ -9,6 +9,8 @@ from utils.plot import smoothen, draw_mask_on_image, draw_heatmap_on_image
 import gradio as gr
 import numpy as np
 import torch
+import seaborn as sns
+import matplotlib.pyplot as plt
 
 # Load Vision Transformer
 hf_model = "tanlq/vit-base-patch16-224-in21k-finetuned-cifar10"
@@ -59,19 +61,36 @@ def get_mask(image, model_name: str):
     dm_image = feature_extractor(image).unsqueeze(0)
     dm_out = diffmask_model.get_mask(dm_image)
     mask = dm_out["mask"][0].detach()
+    logits = dm_out["logits"][0].detach().softmax(dim=-1)
+    logits_orig = dm_out["logits_orig"][0].detach().softmax(dim=-1)
+    # fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(10, 10))
+    # sns.displot(logits_orig.cpu().numpy().flatten(), kind="kde", label="Original", ax=ax)
+    top5logits_orig = logits_orig.topk(5, dim=-1)
+    idx = top5logits_orig.indices
+    # keep the top 5 classes from the indices of the top 5 logits
+    top5logits_orig = top5logits_orig.values
+    top5logits = logits[idx]
+
     pred = dm_out["pred_class"][0].detach()
     pred = diffmask_model.model.config.id2label[pred.item()]
 
     masked_img = draw_mask(image, mask)
     heatmap = draw_heatmap(image, mask)
-    return np.hstack((masked_img, heatmap)), pred
+    orig_probs = {diffmask_model.model.config.id2label[i]: top5logits_orig[i].item() for i in range(5)}
+    pred_probs = {diffmask_model.model.config.id2label[i]: top5logits[i].item() for i in range(5)}
+
+    return np.hstack((masked_img, heatmap)), pred, orig_probs, pred_probs
+
+
+
 
 # Launch demo interface
 gr.Interface(
     get_mask,
     inputs=[gr.inputs.Image(label="Input", shape=(224, 224), source="upload", type="numpy"),
     gr.inputs.Dropdown(["DiffMask-CiFAR-10", "DiffMask-ImageNet"])],
-    outputs=[gr.outputs.Image(label="Output"), gr.outputs.Label(label="Prediction")],
+    outputs=[gr.outputs.Image(label="Output"), gr.outputs.Label(label="Prediction"),
+    gr.Label(label="Original Probabilities"), gr.Label(label="Predicted Probabilities")],
     title="Vision DiffMask Demo",
     live=True,
 ).launch()

code/models/interpretation.py

@@ -277,7 +277,8 @@ class ImageInterpretationNet(pl.LightningModule):
         mask = F.interpolate(mask, scale_factor=S)
         mask = mask.reshape(B, H, W)
 
-        return {"mask": mask, "kl_div": kl_div, "pred_class": pred_class}
+        return {"mask": mask, "kl_div": kl_div, "pred_class": pred_class,
+                "logits": logits, "logits_orig": logits_orig}
 
     def forward(self, x: Tensor) -> Tensor:
         return self.model(x).logits

requirements.txt

@@ -4,4 +4,5 @@ pytorch_lightning
 torch
 torchvision
 transformers
-
+seaborn
+matplotlib