imagenet

app.py

@@ -12,18 +12,25 @@ import torch
 
 # Load Vision Transformer
 hf_model = "tanlq/vit-base-patch16-224-in21k-finetuned-cifar10"
+hf_model_imagenet = "google/vit-base-patch16-224"
 vit = ViTForImageClassification.from_pretrained(hf_model)
+vit_imagenet = ViTForImageClassification.from_pretrained(hf_model_imagenet)
 vit.eval()
+vit_imagenet.eval()
 
 # Load Feature Extractor
 feature_extractor = ViTFeatureExtractor.from_pretrained(hf_model, return_tensors="pt")
+feature_extractor_imagenet = ViTFeatureExtractor.from_pretrained(hf_model_imagenet, return_tensors="pt")
 feature_extractor = UnNest(feature_extractor)
+feature_extractor_imagenet = UnNest(feature_extractor_imagenet)
 
 # 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)
+diffmask_imagenet.set_vision_transformer(vit_imagenet)
+diffmask.eval()
+diffmask_imagenet.eval()
 
 # Define mask plotting functions
 def draw_mask(image, mask):

flagged/log.csv

@@ -0,0 +1,2 @@
+'Input','model_name','Output','Prediction','flag','username','timestamp'
+'Input/0.png','DiffMask-ImageNet','Output/0.png','airplane','','','2022-06-30 18:09:57.950187'