initial commit

app.py

@@ -0,0 +1,102 @@
+import numpy as np
+import gradio as gr
+from huggingface_hub import hf_hub_download
+import omegaconf
+from hydra import utils
+import os
+import torch
+import matplotlib.pyplot as plt
+from attn_helper import VITAttentionGradRollout, overlay_attn
+import vc_models
+#import eaif_models
+import torchvision
+
+eai_filepath = vc_models.__file__.split('src')[0]
+
+MODEL_DIR=os.path.join(eai_filepath, 'src','model_ckpts')
+if not os.path.isdir(MODEL_DIR):
+    os.mkdir(MODEL_DIR)
+
+REPO_ID = "facebook/vc1-base"
+FILENAME = "config.yaml"
+MODEL_TUPLE = None
+
+def get_model():
+    global MODEL_TUPLE
+    download_bin()
+    if MODEL_TUPLE is None:
+        model_cfg = omegaconf.OmegaConf.load(
+            hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
+        )
+        model_cfg['model']['checkpoint_path'] = None
+        model_cfg['model']['checkpoint_path'] = 'model_ckpts/vc1_vitb.pth'#os.path.join(os.getcwd(),'pytorch_model.bin')
+        print(model_cfg)
+        MODEL_TUPLE = utils.instantiate(model_cfg)
+        MODEL_TUPLE[0].eval()
+    return MODEL_TUPLE#model,embedding_dim,transform,metadata
+
+def download_bin():
+
+    bin_file = 'vc1_vitb.pth' #'pytorch_model.bin'
+    bin_path = os.path.join(MODEL_DIR,bin_file)
+    print(bin_path)
+    if not os.path.isfile(bin_path):
+        #with open(bin_file,'w') as f:
+        model_bin = hf_hub_download(repo_id=REPO_ID, filename='pytorch_model.bin',local_dir=MODEL_DIR,local_dir_use_symlinks=True)
+        os.rename(model_bin, bin_path)
+        print(type(model_bin))
+#        os.rename(model_bin, bin_file)
+#             f.write(model_bin)
+
+
+def run_attn(input_img):
+    download_bin()
+    model, embedding_dim, transform, metadata = get_model()
+    print(input_img.shape)
+    if input_img.shape[0] != 3:
+        input_img = input_img.transpose(2, 0, 1)
+    print(input_img.shape)
+    if(len(input_img.shape)== 3):
+        input_img = torch.tensor(input_img).unsqueeze(0)
+    input_img = input_img.float()
+    resize_transform = torchvision.transforms.Resize((250,250))
+    input_img = resize_transform(input_img)
+    x = transform(input_img)
+    #y = x /x.max()  * 255
+    #y = y[0].int().permute(1,2,0).numpy()
+
+    attention_rollout = VITAttentionGradRollout(model,head_fusion="mean")
+    
+    y = model(x)
+    mask = attention_rollout.get_attn_mask()
+    print(input_img.shape)
+    print(mask.shape)
+    attn_img = overlay_attn(input_img[0].permute(1,2,0),mask)
+    
+    fig = plt.figure()
+    ax = fig.subplots()
+    im = ax.matshow(y.detach().numpy().reshape(16,-1))
+    plt.colorbar(im)
+
+    return attn_img, fig
+
+# with gr.Blocks() as demo:
+#     gr.Markdown("Visual Cortex Base Model")
+#     input_img = gr.Image(shape=(250,250))
+#     output_img = gr.Image(shape=(250,250))
+#     output_plot = gr.Plot()
+#     btn = gr.Button("Encode Representation")
+#     gr.Examples(["./trifinger.jpg","./rearrange.jpg","./kitchen.jpg","./ego4d.jpg"],input_img)
+#     # demo = gr.Interface(fn=run_attn, inputs=gr.Image(shape=(250,250)), title="Visual Cortex Base Model", 
+#     btn.click(fn=run_attn, inputs=input_img, 
+#                 outputs=[output_img,output_plot])
+
+
+input_img = gr.Image(shape=(250,250))
+output_img = gr.Image(shape=(250,250))
+output_plot = gr.Plot()
+
+demo = gr.Interface(fn=run_attn, title="Visual Cortex Base Model",
+                    examples=["./trifinger.jpg","./rearrange.jpg","./kitchen.jpg","./ego4d.jpg"], 
+                    inputs=input_img,outputs=[output_img,output_plot])
+demo.launch(share=True)

attn_helper.py

@@ -0,0 +1,107 @@
+import cv2
+from PIL import Image
+import numpy as np
+import torch
+
+import PIL
+
+def overlay_attn(original_image,mask):
+    # Colormap and alpha for attention mask
+    # COLORMAP_OCEAN
+    # COLORMAP_OCEAN
+    colormap_attn, alpha_attn = cv2.COLORMAP_OCEAN, 1 #0.85
+    
+    # Resize mask to original image size
+    w, h = original_image.shape[0], original_image.shape[1]
+    mask = cv2.resize(mask / mask.max(), (h, w))[..., np.newaxis]
+    
+    # Apply colormap to mask
+    cmap = cv2.applyColorMap(np.uint8(255 * mask), colormap_attn)
+
+    print(cmap.shape)
+    # Blend mask and original image
+    grayscale_img =  cv2.cvtColor(np.uint8(original_image), cv2.COLOR_RGB2GRAY)
+    alpha_blended = cv2.addWeighted(np.uint8(original_image),1, cmap, alpha_attn, 0)
+    # alpha_blended = cmap
+
+    # Save image
+    final_im = Image.fromarray(alpha_blended)
+    # final_im = final_im.crop((0,0,250,250))
+    return final_im
+
+
+
+class VITAttentionGradRollout:
+    '''
+        Expects timm ViT transformer model
+    '''
+    def __init__(self, model, head_fusion='min', discard_ratio=0):
+        self.model = model
+        self.head_fusion = head_fusion
+        self.discard_ratio = discard_ratio
+        print(list(model.blocks.children()))
+        
+        self.attentions = {}
+        for idx, module in enumerate(list(model.blocks.children())):
+            module.attn.register_forward_hook(self.get_attention(f"attn{idx}"))
+
+
+    def get_attention(self, name):
+        def hook(module, input, output):
+            with torch.no_grad():
+                input = input[0]
+                B, N, C = input.shape
+                qkv = (
+                    module.qkv(input)
+                    .detach()
+                    .reshape(B, N, 3, module.num_heads, C // module.num_heads)
+                    .permute(2, 0, 3, 1, 4)
+                )
+                q, k, _ = (
+                    qkv[0],
+                    qkv[1],
+                    qkv[2],
+                )  # make torchscript happy (cannot use tensor as tuple)
+                attn = (q @ k.transpose(-2, -1)) * module.scale
+                attn = attn.softmax(dim=-1)
+                self.attentions[name] = attn
+        return hook
+
+    def get_attn_mask(self,k=0):
+        attn_key = "attn" + str()
+        result = torch.eye(self.attentions['attn0'].size(-1)).to(self.attentions['attn0'].device)
+
+        # result = torch.eye(self.attentions['attn2'].size(-1)).to(self.attentions['attn2'].device)
+        with torch.no_grad():
+            # for attention in self.attentions.values():
+            for k in range(11, len(self.attentions.keys())):
+                attention = self.attentions[f'attn{k}']
+                if self.head_fusion == "mean":
+                    attention_heads_fused = attention.mean(axis=1)
+                elif self.head_fusion == "max":
+                    attention_heads_fused = attention.max(axis=1)[0]
+                elif self.head_fusion == "min":
+                    attention_heads_fused = attention.min(axis=1)[0]
+                else:
+                    raise "Attention head fusion type Not supported"
+
+                # Drop the lowest attentions, but
+                # don't drop the class token
+                flat = attention_heads_fused.view(attention_heads_fused.size(0), -1)
+                _, indices = flat.topk(int(flat.size(-1)*self.discard_ratio), -1, False)
+                indices = indices[indices != 0]
+                flat[0, indices] = 0
+                I = torch.eye(attention_heads_fused.size(-1)).to(attention_heads_fused.device)
+                a = (attention_heads_fused + 1.0*I)/2
+                a = a / a.sum(dim=-1).unsqueeze(-1)
+
+                result = torch.matmul(a, result)
+
+        # Look at the total attention between the class token,
+        # and the image patches
+        mask = result[0, 0 , 1 :]
+        # In case of 224x224 image, this brings us from 196 to 14
+        width = int(mask.size(-1)**0.5)
+        mask = mask.reshape(width, width).detach().cpu().numpy()
+        mask = mask / np.max(mask)
+        return mask