app.py

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
import pathlib


# 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 = None

normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
transform_224 = transforms.Compose([
    transforms.ToTensor(),
    normalize,
])

# Download human-readable labels for ImageNet.
response = requests.get("https://git.io/JJkYN")
labels = response.text.split("\n")

def image_classifier(inp):
	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)
		
		prediction_orig = torch.nn.functional.softmax(model(image.unsqueeze(0))[0], dim=0)
		confidences_orig = {labels[i]: float(prediction_orig[i]) for i in range(1000)}    
		heatmap_orig = generate_visualization(model, image)
	return confidences, heatmap, confidences_orig, heatmap_orig

def _load_model(model_name: str):
	global model_finetuned, model
	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')

def _set_example_image(example: list) -> dict:
    return gr.Image.update(value=example[0])
	
def _clear_image():
	return None

demo = gr.Blocks(css='style.css')

with demo:
	
	
	with gr.Row():
		with gr.Column():
			gr.Markdown('## [Optimizing Relevance Maps of Vision Transformers Improves Robustness](https://github.com/hila-chefer/RobustViT) - Official Demo')
			# gr.Markdown('This is an official demo for [Optimizing Relevance Maps of Vision Transformers Improves Robustness](https://github.com/hila-chefer/RobustViT).')
			gr.Markdown('Select or upload an image and then click **Submit** to see the output.')
			with gr.Row():
				input_image = gr.Image(shape=(224,224))
			with gr.Row():	
				btn = gr.Button("Submit", variant="primary")
				clear_btn = gr.Button('Clear')
		with gr.Column():
			gr.Markdown('### Examples')
			gr.Markdown('#### Corrected Prediction')
			with gr.Row():
				paths = sorted(pathlib.Path('samples/corrected').rglob('*.png'))
				corrected_pred_examples = gr.Dataset(components=[input_image], headers=['header'],
											samples=[[path.as_posix()] for path in paths])
		
			gr.Markdown('#### Improved Explainability')
			with gr.Row():
				paths = sorted(pathlib.Path('samples/better_expl').rglob('*.png'))
				better_expl = gr.Dataset(components=[input_image], headers=['header'],
											samples=[[path.as_posix()] for path in paths])
	
	
	#gr.Markdown('### Results:')
		
	with gr.Row():
		with gr.Column():
			gr.Markdown('### Ours (finetuned model)')
			out1 = gr.outputs.Label(label="Our Classification", num_top_classes=3)
			out2 = gr.Image(label="Our Relevance",shape=(224,224), elem_id="expl1")
			
		with gr.Column():
			gr.Markdown('### Original model')
			out3 = gr.outputs.Label(label="Original Classification", num_top_classes=3)
			out4 = gr.Image(label="Original Relevance",shape=(224,224),elem_id="expl2")
		
	
	corrected_pred_examples.click(fn=_set_example_image, inputs=corrected_pred_examples, outputs=input_image)
	better_expl.click(fn=_set_example_image, inputs=better_expl, outputs=input_image)
	btn.click(fn=image_classifier, inputs=input_image, outputs=[out1, out2, out3, out4])
	clear_btn.click(fn=_clear_image, inputs=[], outputs=[input_image])
	

demo.launch()