update app.py

app.py

@@ -1,23 +1,112 @@
 import torch
 import timm
 import gradio as gr
-
-"""
+from huggingface_hub import hf_hub_download
+import os
 from ViT.ViT_new import vit_base_patch16_224 as vit
+import torchvision.transforms as transforms
+import requests
+from PIL import Image
+import numpy as np
+import cv2
+
+
+# create heatmap from mask on image
+def show_cam_on_image(img, mask):
+    heatmap = cv2.applyColorMap(np.uint8(255 * mask), cv2.COLORMAP_JET)
+    heatmap = np.float32(heatmap) / 255
+    cam = heatmap + np.float32(img)
+    cam = cam / np.max(cam)
+    return cam
+
+start_layer = 0
+
+# rule 5 from paper
+def avg_heads(cam, grad):
+    cam = cam.reshape(-1, cam.shape[-2], cam.shape[-1])
+    grad = grad.reshape(-1, grad.shape[-2], grad.shape[-1])
+    cam = grad * cam
+    cam = cam.clamp(min=0).mean(dim=0)
+    return cam
+
+# rule 6 from paper
+def apply_self_attention_rules(R_ss, cam_ss):
+    R_ss_addition = torch.matmul(cam_ss, R_ss)
+    return R_ss_addition
+
+def generate_relevance(model, input, index=None):
+    output = model(input, register_hook=True)
+    if index == None:
+        index = np.argmax(output.cpu().data.numpy(), axis=-1)
+
+    one_hot = np.zeros((1, output.size()[-1]), dtype=np.float32)
+    one_hot[0, index] = 1
+    one_hot_vector = one_hot
+    one_hot = torch.from_numpy(one_hot).requires_grad_(True)
+    one_hot = torch.sum(one_hot * output)
+    model.zero_grad()
+    one_hot.backward(retain_graph=True)
+
+    num_tokens = model.blocks[0].attn.get_attention_map().shape[-1]
+    R = torch.eye(num_tokens, num_tokens)
+    for i,blk in enumerate(model.blocks):
+        if i < start_layer:
+            continue
+        grad = blk.attn.get_attn_gradients()
+        cam = blk.attn.get_attention_map()
+        cam = avg_heads(cam, grad)
+        R += apply_self_attention_rules(R, cam)
+    return R[0, 1:]
+
+def generate_visualization(model, original_image, class_index=None):
+    with torch.enable_grad():
+        transformer_attribution = generate_relevance(model, original_image.unsqueeze(0), index=class_index).detach()
+    transformer_attribution = transformer_attribution.reshape(1, 1, 14, 14)
+    transformer_attribution = torch.nn.functional.interpolate(transformer_attribution, scale_factor=16, mode='bilinear')
+    transformer_attribution = transformer_attribution.reshape(224, 224).data.cpu().numpy()
+    transformer_attribution = (transformer_attribution - transformer_attribution.min()) / (transformer_attribution.max() - transformer_attribution.min())
+    
+    image_transformer_attribution = original_image.permute(1, 2, 0).data.cpu().numpy()
+    image_transformer_attribution = (image_transformer_attribution - image_transformer_attribution.min()) / (image_transformer_attribution.max() - image_transformer_attribution.min())
+    vis = show_cam_on_image(image_transformer_attribution, transformer_attribution)
+    vis =  np.uint8(255 * vis)
+    vis = cv2.cvtColor(np.array(vis), cv2.COLOR_RGB2BGR)
+    return vis
+	
+model_finetuned = None
 
-model = vit(pretrained=True).cuda()
-model.eval()
-model_finetuned = vit().cuda()
-checkpoint = torch.load('ar_base.tar')
-model_finetuned.load_state_dict(checkpoint['state_dict'])
-model_finetuned.eval()
+normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+transform_224 = transforms.Compose([
+    transforms.ToTensor(),
+    normalize,
+])
 
-iface_orig = gr.Interface(
-)
-"""
+# Download human-readable labels for ImageNet.
+response = requests.get("https://git.io/JJkYN")
+labels = response.text.split("\n")
 
 def image_classifier(inp):
-    pass # image classifier model defined here
+	image = transform_224(inp)
+	print(image.shape)
+	#return model_finetuned(image.unsqueeze(0))
+	with torch.no_grad():
+		prediction = torch.nn.functional.softmax(model_finetuned(image.unsqueeze(0))[0], dim=0)
+		confidences = {labels[i]: float(prediction[i]) for i in range(1000)}    
+		heatmap = generate_visualization(model_finetuned, image)
+	return confidences, heatmap
 
-demo = gr.Interface(image_classifier, "image", "label")
-demo.launch(share=True)
+def _load_model(model_name: str):
+	global model_finetuned
+	path = hf_hub_download('Hila/RobustViT',
+						   f'{model_name}')
+						   
+	model = vit(pretrained=True)
+	model.eval()
+	model_finetuned = vit()
+	checkpoint = torch.load(path, map_location='cpu')
+	model_finetuned.load_state_dict(checkpoint['state_dict'])
+	model_finetuned.eval()
+	
+_load_model('ar_base.tar')
+demo = gr.Interface(image_classifier, gr.inputs.Image(shape=(224,224)), [gr.outputs.Label(num_top_classes=3), gr.Image(shape=(224,224))],examples=["samples/augreg_base/tank.png", "samples/augreg_base/sundial.png", "samples/augreg_base/lizard.png", "samples/augreg_base/storck.png", "samples/augreg_base/hummingbird2.png", "samples/augreg_base/hummingbird.png"], capture_session=True)
+demo.launch(debug=True)

requirements.txt

@@ -1,2 +1,4 @@
 torch==1.7.1
-timm
+timm
+torchvision==0.8.2
+opencv-python