imagenet

app.py

@@ -22,7 +22,8 @@ feature_extractor = UnNest(feature_extractor)
 # Load Vision DiffMask
 diffmask = ImageInterpretationNet.load_from_checkpoint('checkpoints/diffmask.ckpt')
 diffmask.set_vision_transformer(vit)
-
+diffmask_imagenet = ImageInterpretationNet.load_from_checkpoint('checkpoints/diffmask_imagenet.ckpt')
+diffmask_imagenet.set_vision_transformer(vit)
 
 # Define mask plotting functions
 def draw_mask(image, mask):
@@ -40,27 +41,31 @@ def draw_heatmap(image, mask):
 
 
 # Define callable method for the demo
-def get_mask(image):
+def get_mask(image, model_name: str):
     if image is None:
         return None, None
-
+    if model_name == 'DiffMask-CiFAR-10':
+        diffmask_model = diffmask
+    elif model_name == 'DiffMask-ImageNet':
+        diffmask_model = diffmask_imagenet
     image = torch.from_numpy(image).permute(2, 0, 1).float() / 255
     dm_image = feature_extractor(image).unsqueeze(0)
-    dm_out = diffmask.get_mask(dm_image)
+    dm_out = model.get_mask(dm_image)
     mask = dm_out["mask"][0].detach()
     pred = dm_out["pred_class"][0].detach()
-    pred = diffmask.model.config.id2label[pred.item()]
+    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
 
-
 # Launch demo interface
 gr.Interface(
     get_mask,
-    inputs=gr.inputs.Image(label="Input", shape=(224, 224), source="upload", type="numpy"),
+    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")],
     title="Vision DiffMask Demo",
     live=True,
 ).launch()
+

checkpoints/diffmask_imagenet.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c21d3e65b049984911907d004747e2162cba8a77ac10008bb9b3d8612b745a4f
+size 16607641