Add heatmap viz

app.py

@@ -1,9 +1,11 @@
 import gradio as gr
 import torch
+import torch.nn.functional as F
 import requests
 import numpy as np
 import re
 import io
+import matplotlib.pyplot as plt
 
 from PIL import Image
 from transformers import ViltProcessor, ViltForMaskedLM
@@ -15,6 +17,7 @@ model = ViltForMaskedLM.from_pretrained("dandelin/vilt-b32-mlm")
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 model.to(device)
 
+
 class MinMaxResize:
     def __init__(self, shorter=800, longer=1333):
         self.min = shorter
@@ -36,7 +39,8 @@ class MinMaxResize:
         newh, neww = int(newh + 0.5), int(neww + 0.5)
         newh, neww = newh // 32 * 32, neww // 32 * 32
 
-        return x.resize((neww, newh), resample=Image.BICUBIC)
+        return x.resize((neww, newh), resample=Image.Resampling.BICUBIC)
+
 
 def pixelbert_transform(size=800):
     longer = int((1333 / 800) * size)
@@ -44,16 +48,99 @@ def pixelbert_transform(size=800):
         [
             MinMaxResize(shorter=size, longer=longer),
             transforms.ToTensor(),
-            transforms.Compose([transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])]),
+            transforms.Compose([transforms.Normalize(
+                mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])]),
         ]
     )
 
 
-def infer(url, mp_text):
+def cost_matrix_cosine(x, y, eps=1e-5):
+    """Compute cosine distnace across every pairs of x, y (batched)
+    [B, L_x, D] [B, L_y, D] -> [B, Lx, Ly]"""
+    assert x.dim() == y.dim()
+    assert x.size(0) == y.size(0)
+    assert x.size(2) == y.size(2)
+    x_norm = F.normalize(x, p=2, dim=-1, eps=eps)
+    y_norm = F.normalize(y, p=2, dim=-1, eps=eps)
+    cosine_sim = x_norm.matmul(y_norm.transpose(1, 2))
+    cosine_dist = 1 - cosine_sim
+    return cosine_dist
+
+
+@torch.no_grad()
+def ipot(C, x_len, x_pad, y_len, y_pad, joint_pad, beta, iteration, k):
+    """ [B, M, N], [B], [B, M], [B], [B, N], [B, M, N]"""
+    b, m, n = C.size()
+    sigma = torch.ones(b, m, dtype=C.dtype,
+                       device=C.device) / x_len.unsqueeze(1)
+    T = torch.ones(b, n, m, dtype=C.dtype, device=C.device)
+    A = torch.exp(-C.transpose(1, 2) / beta)
+
+    # mask padded positions
+    sigma.masked_fill_(x_pad, 0)
+    joint_pad = joint_pad.transpose(1, 2)
+    T.masked_fill_(joint_pad, 0)
+    A.masked_fill_(joint_pad, 0)
+
+    # broadcastable lengths
+    x_len = x_len.unsqueeze(1).unsqueeze(2)
+    y_len = y_len.unsqueeze(1).unsqueeze(2)
+
+    # mask to zero out padding in delta and sigma
+    x_mask = (x_pad.to(C.dtype) * 1e4).unsqueeze(1)
+    y_mask = (y_pad.to(C.dtype) * 1e4).unsqueeze(1)
+
+    for _ in range(iteration):
+        Q = A * T  # bs * n * m
+        sigma = sigma.view(b, m, 1)
+        for _ in range(k):
+            delta = 1 / (y_len * Q.matmul(sigma).view(b, 1, n) + y_mask)
+            sigma = 1 / (x_len * delta.matmul(Q) + x_mask)
+        T = delta.view(b, n, 1) * Q * sigma
+    T.masked_fill_(joint_pad, 0)
+    return T
+
+
+def get_model_embedding_and_mask(model, input_ids, pixel_values):
+
+    input_shape = input_ids.size()
+
+    text_batch_size, seq_length = input_shape
+    device = input_ids.device
+    attention_mask = torch.ones(((text_batch_size, seq_length)), device=device)
+    image_batch_size = pixel_values.shape[0]
+    image_token_type_idx = 1
+
+    if image_batch_size != text_batch_size:
+        raise ValueError(
+            "The text inputs and image inputs need to have the same batch size")
+
+    pixel_mask = torch.ones((image_batch_size, model.vilt.config.image_size,
+                            model.vilt.config.image_size), device=device)
+
+    text_embeds = model.vilt.embeddings.text_embeddings(
+        input_ids=input_ids, token_type_ids=None, inputs_embeds=None)
+
+    image_embeds, image_masks, patch_index = model.vilt.embeddings.visual_embed(
+        pixel_values=pixel_values, pixel_mask=pixel_mask, max_image_length=model.vilt.config.max_image_length
+    )
+    text_embeds = text_embeds + model.vilt.embeddings.token_type_embeddings(
+        torch.zeros_like(attention_mask, dtype=torch.long,
+                         device=text_embeds.device)
+    )
+    image_embeds = image_embeds + model.vilt.embeddings.token_type_embeddings(
+        torch.full_like(image_masks, image_token_type_idx,
+                        dtype=torch.long, device=text_embeds.device)
+    )
+
+    return text_embeds, image_embeds, attention_mask, image_masks, patch_index
+
+
+def infer(url, mp_text, hidx):
     try:
         res = requests.get(url)
         image = Image.open(io.BytesIO(res.content)).convert("RGB")
-        img = pixelbert_transform(size=384)(image)
+        img = pixelbert_transform(size=500)(image)
         img = img.unsqueeze(0).to(device)
     except:
         return False
@@ -67,69 +154,171 @@ def infer(url, mp_text):
             encoded = processor.tokenizer(inferred_token)
             input_ids = torch.tensor(encoded.input_ids)
             encoded = encoded["input_ids"][0][1:-1]
-            outputs = model(input_ids=input_ids, pixel_values=encoding.pixel_values)
+            outputs = model(input_ids=input_ids,
+                            pixel_values=encoding.pixel_values)
             mlm_logits = outputs.logits[0]  # shape (seq_len, vocab_size)
-            
+
             # only take into account text features (minus CLS and SEP token)
-            mlm_logits = mlm_logits[1 : input_ids.shape[1] - 1, :]
+            mlm_logits = mlm_logits[1: input_ids.shape[1] - 1, :]
             mlm_values, mlm_ids = mlm_logits.softmax(dim=-1).max(dim=-1)
-            
+
             # only take into account text
             mlm_values[torch.tensor(encoded) != 103] = 0
             select = mlm_values.argmax().item()
             encoded[select] = mlm_ids[select].item()
             inferred_token = [processor.decode(encoded)]
-    
+
     encoded = processor.tokenizer(inferred_token)
     output = processor.decode(encoded.input_ids[0], skip_special_tokens=True)
+    selected_token = ''
+
+    if hidx > 0 and hidx < len(encoded["input_ids"][0][:-1]):
+        input_ids = torch.tensor(encoded.input_ids)
+        outputs = model(
+            input_ids=input_ids, pixel_values=encoding.pixel_values, output_hidden_states=True)
+
+        txt_emb, img_emb, text_masks, image_masks, patch_index = get_model_embedding_and_mask(
+            model, input_ids=input_ids, pixel_values=encoding.pixel_values)
+
+        embedding_output = torch.cat([txt_emb, img_emb], dim=1)
+        attention_mask = torch.cat([text_masks, image_masks], dim=1)
+
+        extended_attention_mask = model.vilt.get_extended_attention_mask(
+            attention_mask, input_ids.size(), device=device)
+
+        encoder_outputs = model.vilt.encoder(
+            embedding_output,
+            attention_mask=extended_attention_mask,
+            head_mask=None,
+            output_attentions=False,
+            output_hidden_states=True,
+            return_dict=True,
+        )
+
+        x = encoder_outputs.hidden_states[-1]
+        x = model.vilt.layernorm(x)
+
+        txt_emb, img_emb = (
+            x[:, :txt_emb.shape[1]],
+            x[:, txt_emb.shape[1]:],
+        )
+
+        txt_mask, img_mask = (
+            text_masks.bool(),
+            image_masks.bool(),
+        )
+
+        for i, _len in enumerate(txt_mask.sum(dim=1)):
+            txt_mask[i, _len - 1] = False
+        txt_mask[:, 0] = False
+        img_mask[:, 0] = False
+        txt_pad, img_pad = ~txt_mask, ~img_mask
+        cost = cost_matrix_cosine(txt_emb.float(), img_emb.float())
+        joint_pad = txt_pad.unsqueeze(-1) | img_pad.unsqueeze(-2)
+        cost.masked_fill_(joint_pad, 0)
+
+        txt_len = (txt_pad.size(1) - txt_pad.sum(dim=1,
+                   keepdim=False)).to(dtype=cost.dtype)
+        img_len = (img_pad.size(1) - img_pad.sum(dim=1,
+                   keepdim=False)).to(dtype=cost.dtype)
+        T = ipot(cost.detach(),
+                 txt_len,
+                 txt_pad,
+                 img_len,
+                 img_pad,
+                 joint_pad,
+                 0.1,
+                 1000,
+                 1,
+                 )
+        plan = T[0]
+        plan_single = plan * len(txt_emb)
+        cost_ = plan_single.t()
+
+        cost_ = cost_[hidx][1:].cpu()
+
+        patch_index, (H, W) = patch_index
+        heatmap = torch.zeros(H, W)
+        for i, pidx in enumerate(patch_index[0]):
+            h, w = pidx[0].item(), pidx[1].item()
+            heatmap[h, w] = cost_[i]
+
+        heatmap = (heatmap - heatmap.mean()) / heatmap.std()
+        heatmap = np.clip(heatmap, 1.0, 3.0)
+        heatmap = (heatmap - heatmap.min()) / (heatmap.max() - heatmap.min())
+
+        _w, _h = image.size
+        overlay = Image.fromarray(np.uint8(heatmap * 255), "L").resize(
+            (_w, _h), resample=Image.Resampling.NEAREST
+        )
+        image_rgba = image.copy()
+        image_rgba.putalpha(overlay)
+        image = image_rgba
+
+        selected_token = processor.tokenizer.convert_ids_to_tokens(
+            encoded["input_ids"][0][hidx]
+        )
+
+    return [np.array(image), output, selected_token]
 
-    return [np.array(image), output]
 
 title = "What's in the picture ?"
 
 description = """
 Can't find your words to describe an image ? The pre-trained 
 ViLT model will help you. Give the url of an image and a caption with [MASK] tokens to be filled or play with the given examples !
+You can even see where the model focused its attention for a given word : just choose the index of the selected word with the slider.
 """
 
 
 inputs_interface = [
-        gr.inputs.Textbox(
-            label="Url of an image.",
-            lines=5,
-        ),
-        gr.inputs.Textbox(label="Caption with [MASK] tokens to be filled.", lines=5),
-    ]
+    gr.inputs.Textbox(
+        label="Url of an image.",
+        lines=5,
+    ),
+    gr.inputs.Textbox(
+        label="Caption with [MASK] tokens to be filled.", lines=5),
+    gr.inputs.Slider(
+        minimum=0,
+        maximum=38,
+        step=1,
+        label="Index of token for heatmap visualization (ignored if zero)",
+    ),
+]
 outputs_interface = [
-        gr.outputs.Image(label="Image"),
-        gr.outputs.Textbox(label="description"),
-    ]
+    gr.outputs.Image(label="Image"),
+    gr.outputs.Textbox(label="description"),
+    gr.outputs.Textbox(label="selected token"),
+]
 
 interface = gr.Interface(
-        fn=infer,
-        inputs=inputs_interface,
-        outputs=outputs_interface,
-        title=title,
-        description=description,
-        server_name="0.0.0.0",
-        server_port=8888,
-        examples=[
-            [
+    fn=infer,
+    inputs=inputs_interface,
+    outputs=outputs_interface,
+    title=title,
+    description=description,
+    server_name="0.0.0.0",
+    server_port=8888,
+    examples=[
+        [
                 "https://s3.geograph.org.uk/geophotos/06/21/24/6212487_1cca7f3f_1024x1024.jpg",
                 "a display of flowers growing out and over the [MASK] [MASK] in front of [MASK] on a [MASK] [MASK].",
-            ],
-     
-            [
-                "https://encrypted-tbn0.gstatic.com/images?q=tbn:ANd9GcT5W71UTcSBm3r5l9NzBemglq983bYvKOHRkw&usqp=CAU",
-                "An [MASK] with the [MASK] in the [MASK].",
-            ],
-            
-            [
-                "https://www.referenseo.com/wp-content/uploads/2019/03/image-attractive-960x540.jpg",
-                "An [MASK] is flying with a [MASK] over a [MASK].",
-            ],
+                0,
         ],
-    )
 
-interface.launch()
+        [
+            "https://encrypted-tbn0.gstatic.com/images?q=tbn:ANd9GcT5W71UTcSBm3r5l9NzBemglq983bYvKOHRkw&usqp=CAU",
+            "An [MASK] with the [MASK] in the [MASK].",
+            5,
+        ],
+
+        [
+            "https://www.referenseo.com/wp-content/uploads/2019/03/image-attractive-960x540.jpg",
+            "An [MASK] is flying with a [MASK] over a [MASK].",
+            2,
+        ],
+    ],
+)
+
 
+interface.launch()