Spaces:
Sleeping
Sleeping
| import torch | |
| import clip | |
| from PIL import Image | |
| import numpy as np | |
| import cv2 | |
| import matplotlib.pyplot as plt | |
| def interpret(image, text, model, device, index=None): | |
| logits_per_image, logits_per_text = model(image, text) | |
| probs = logits_per_image.softmax(dim=-1).detach().cpu().numpy() | |
| if index is None: | |
| index = np.argmax(logits_per_image.cpu().data.numpy(), axis=-1) | |
| one_hot = np.zeros((1, logits_per_image.size()[-1]), dtype=np.float32) | |
| one_hot[0, index] = 1 | |
| one_hot = torch.from_numpy(one_hot).requires_grad_(True) | |
| one_hot = torch.sum(one_hot.cuda() * logits_per_image) | |
| model.zero_grad() | |
| one_hot.backward(retain_graph=True) | |
| image_attn_blocks = list(dict(model.visual.transformer.resblocks.named_children()).values()) | |
| num_tokens = image_attn_blocks[0].attn_probs.shape[-1] | |
| R = torch.eye(num_tokens, num_tokens, dtype=image_attn_blocks[0].attn_probs.dtype).to(device) | |
| for blk in image_attn_blocks: | |
| grad = blk.attn_grad | |
| cam = blk.attn_probs | |
| cam = cam.reshape(-1, cam.shape[-1], cam.shape[-1]) | |
| grad = grad.reshape(-1, grad.shape[-1], grad.shape[-1]) | |
| cam = grad * cam | |
| cam = cam.clamp(min=0).mean(dim=0) | |
| R += torch.matmul(cam, R) | |
| R[0, 0] = 0 | |
| image_relevance = R[0, 1:] | |
| # 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 | |
| image_relevance = image_relevance.reshape(1, 1, 7, 7) | |
| image_relevance = torch.nn.functional.interpolate(image_relevance, size=224, mode='bilinear') | |
| image_relevance = image_relevance.reshape(224, 224).cuda().data.cpu().numpy() | |
| image_relevance = (image_relevance - image_relevance.min()) / (image_relevance.max() - image_relevance.min()) | |
| image = image[0].permute(1, 2, 0).data.cpu().numpy() | |
| image = (image - image.min()) / (image.max() - image.min()) | |
| vis = show_cam_on_image(image, image_relevance) | |
| vis = np.uint8(255 * vis) | |
| vis = cv2.cvtColor(np.array(vis), cv2.COLOR_RGB2BGR) | |
| plt.imshow(vis) | |
| plt.show() | |
| print("Label probs:", probs) | |
| def main(): | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| model, preprocess = clip.load("ViT-B/32", device=device, jit=False) | |
| image = preprocess(Image.open("catdog.png")).unsqueeze(0).to(device) | |
| text = clip.tokenize(["a dog", "a cat"]).to(device) | |
| interpret(model=model, image=image, text=text, device=device, index=0) | |
| interpret(model=model, image=image, text=text, device=device, index=1) | |
| image = preprocess(Image.open("el1.png")).unsqueeze(0).to(device) | |
| text = clip.tokenize(["an elephant", "a zebra"]).to(device) | |
| interpret(model=model, image=image, text=text, device=device, index=0) | |
| interpret(model=model, image=image, text=text, device=device, index=1) | |
| image = preprocess(Image.open("el2.png")).unsqueeze(0).to(device) | |
| text = clip.tokenize(["an elephant", "a zebra"]).to(device) | |
| interpret(model=model, image=image, text=text, device=device, index=0) | |
| interpret(model=model, image=image, text=text, device=device, index=1) | |
| image = preprocess(Image.open("el3.png")).unsqueeze(0).to(device) | |
| text = clip.tokenize(["an elephant", "a zebra"]).to(device) | |
| interpret(model=model, image=image, text=text, device=device, index=0) | |
| interpret(model=model, image=image, text=text, device=device, index=1) | |
| image = preprocess(Image.open("el4.png")).unsqueeze(0).to(device) | |
| text = clip.tokenize(["an elephant", "a zebra"]).to(device) | |
| interpret(model=model, image=image, text=text, device=device, index=0) | |
| interpret(model=model, image=image, text=text, device=device, index=1) | |
| image = preprocess(Image.open("dogbird.png")).unsqueeze(0).to(device) | |
| text = clip.tokenize(["a basset hound", "a parrot"]).to(device) | |
| interpret(model=model, image=image, text=text, device=device, index=0) | |
| interpret(model=model, image=image, text=text, device=device, index=1) | |
| if __name__ == "__main__": | |
| main() | |