fix top 5

app.py

@@ -9,8 +9,6 @@ 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"
@@ -52,46 +50,49 @@ def draw_heatmap(image, mask):
 # Define callable method for the demo
 def get_mask(image, model_name: str):
     if image is None:
-        return None, None
-    if model_name == 'DiffMask-CiFAR-10':
+        return None, None, None
+
+    if model_name == 'DiffMask-CIFAR-10':
         diffmask_model = diffmask
     elif model_name == 'DiffMask-ImageNet':
         diffmask_model = diffmask_imagenet
+
+    # Helper function to convert class index to name
+    def idx2cname(idx):
+        return diffmask_model.model.config.id2label[idx]
+
+    # Prepare image and pass through Vision DiffMask
     image = torch.from_numpy(image).permute(2, 0, 1).float() / 255
     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()]
 
+    # Get mask and apply on image
+    mask = dm_out["mask"][0].detach()
     masked_img = draw_mask(image, mask)
     heatmap = draw_heatmap(image, mask)
-    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
 
+    # Get logits and map to predictions with class names
+    n_classes = len(diffmask_model.model.config.id2label)
+    logits_orig = dm_out["logits_orig"][0].detach().softmax(dim=-1)
+    logits_mask = dm_out["logits"][0].detach().softmax(dim=-1)
+    orig_probs = {idx2cname(i): logits_orig[i].item() for i in range(n_classes)}
+    mask_probs = {idx2cname(i): logits_mask[i].item() for i in range(n_classes)}
 
+    return np.hstack((masked_img, heatmap)), orig_probs, mask_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"),
-    gr.Label(label="Original Probabilities"), gr.Label(label="Predicted Probabilities")],
+    inputs=[
+        gr.inputs.Image(label="Input", shape=(224, 224), source="upload", type="numpy"),
+        gr.inputs.Dropdown(label="Model Name", choices=["DiffMask-CIFAR-10", "DiffMask-ImageNet"]),
+    ],
+    outputs=[
+        gr.outputs.Image(label="Output"),
+        gr.outputs.Label(label="Original Prediction", num_top_classes=5),
+        gr.outputs.Label(label="Masked Prediction", num_top_classes=5),
+    ],
     title="Vision DiffMask Demo",
     live=True,
 ).launch()
-