examples

app.py

@@ -49,6 +49,7 @@ def draw_heatmap(image, mask):
 
 # Define callable method for the demo
 def get_mask(image, model_name: str):
+    torch.manual_seed(seed=0)
     if image is None:
         return None, None, None
 
@@ -86,13 +87,15 @@ gr.Interface(
     get_mask,
     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"]),
+        gr.inputs.Dropdown(label="Model Name", choices=["DiffMask-ImageNet", "DiffMask-CIFAR-10"]),
     ],
     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),
     ],
+    examples=[["dogcat.jpeg", "DiffMask-ImageNet"], ["elephant-zebra.jpg", "DiffMask-ImageNet"],
+    ["finch.jpeg", "DiffMask-ImageNet"]],
     title="Vision DiffMask Demo",
     live=True,
 ).launch()