Update app.py

app.py

@@ -42,22 +42,25 @@ def draw_heatmap(image, mask):
 # Define callable method for the demo
 def get_mask(image):
     if image is None:
-        return None
+        return None, None
 
     image = torch.from_numpy(image).permute(2, 0, 1).float() / 255
     dm_image = feature_extractor(image).unsqueeze(0)
-    mask = diffmask.get_mask(dm_image)["mask"][0].detach()
+    dm_out = diffmask.get_mask(dm_image)
+    mask = dm_out["mask"][0].detach()
+    pred = dm_out["pred_class"][0].detach()
+    pred = diffmask.model.config.id2label[pred.item()]
 
     masked_img = draw_mask(image, mask)
     heatmap = draw_heatmap(image, mask)
-    return np.hstack((masked_img, heatmap))
+    return np.hstack((masked_img, heatmap)), pred
 
 
 # Launch demo interface
 gr.Interface(
     get_mask,
     inputs=gr.inputs.Image(label="Input", shape=(224, 224), source="upload", type="numpy"),
-    outputs=[gr.outputs.Image(label="Output")],
+    outputs=[gr.outputs.Image(label="Output"), gr.outputs.Label(label="Prediction")],
     title="Vision DiffMask Demo",
     live=True,
 ).launch()