MR/vilt-word-patch-alignment
0
1import gradio as gr2from transformers import ViltProcessor, ViltForQuestionAnswering3import torch4 5 6import gradio as gr7import torch8import copy9import time10import requests11import io12import numpy as np13import re14 15from PIL import Image16 17from vilt.config import ex18from vilt.modules import ViLTransformerSS19 20from vilt.modules.objectives import cost_matrix_cosine, ipot21from vilt.transforms import pixelbert_transform22from vilt.datamodules.datamodule_base import get_pretrained_tokenizer23 24 25@ex.automain26def main(_config):27 _config = copy.deepcopy(_config)28 29 loss_names = {30 "itm": 0,31 "mlm": 0.5,32 "mpp": 0,33 "vqa": 0,34 "imgcls": 0,35 "nlvr2": 0,36 "irtr": 0,37 "arc": 0,38 }39 tokenizer = get_pretrained_tokenizer(_config["tokenizer"])40 41 _config.update(42 {43 "loss_names": loss_names,44 }45 )46 47 model = ViLTransformerSS(_config)48 model.setup("test")49 model.eval()50 51 device = "cpu"52 model.to(device)53 54 def infer(url, mp_text, hidx):55 try:56 res = requests.get(url)57 image = Image.open(io.BytesIO(res.content)).convert("RGB")58 img = pixelbert_transform(size=384)(image)59 img = img.unsqueeze(0).to(device)60 except:61 return False62 63 batch = {"text": [""], "image": [None]}64 tl = len(re.findall("\[MASK\]", mp_text))65 inferred_token = [mp_text]66 batch["image"][0] = img67 68 with torch.no_grad():69 for i in range(tl):70 batch["text"] = inferred_token71 encoded = tokenizer(inferred_token)72 batch["text_ids"] = torch.tensor(encoded["input_ids"]).to(device)73 batch["text_labels"] = torch.tensor(encoded["input_ids"]).to(device)74 batch["text_masks"] = torch.tensor(encoded["attention_mask"]).to(device)75 encoded = encoded["input_ids"][0][1:-1]76 infer = model(batch)77 mlm_logits = model.mlm_score(infer["text_feats"])[0, 1:-1]78 mlm_values, mlm_ids = mlm_logits.softmax(dim=-1).max(dim=-1)79 mlm_values[torch.tensor(encoded) != 103] = 080 select = mlm_values.argmax().item()81 encoded[select] = mlm_ids[select].item()82 inferred_token = [tokenizer.decode(encoded)]83 84 selected_token = ""85 encoded = tokenizer(inferred_token)86 87 if hidx > 0 and hidx < len(encoded["input_ids"][0][:-1]):88 with torch.no_grad():89 batch["text"] = inferred_token90 batch["text_ids"] = torch.tensor(encoded["input_ids"]).to(device)91 batch["text_labels"] = torch.tensor(encoded["input_ids"]).to(device)92 batch["text_masks"] = torch.tensor(encoded["attention_mask"]).to(device)93 infer = model(batch)94 txt_emb, img_emb = infer["text_feats"], infer["image_feats"]95 txt_mask, img_mask = (96 infer["text_masks"].bool(),97 infer["image_masks"].bool(),98 )99 for i, _len in enumerate(txt_mask.sum(dim=1)):100 txt_mask[i, _len - 1] = False101 txt_mask[:, 0] = False102 img_mask[:, 0] = False103 txt_pad, img_pad = ~txt_mask, ~img_mask104 105 cost = cost_matrix_cosine(txt_emb.float(), img_emb.float())106 joint_pad = txt_pad.unsqueeze(-1) | img_pad.unsqueeze(-2)107 cost.masked_fill_(joint_pad, 0)108 109 txt_len = (txt_pad.size(1) - txt_pad.sum(dim=1, keepdim=False)).to(110 dtype=cost.dtype111 )112 img_len = (img_pad.size(1) - img_pad.sum(dim=1, keepdim=False)).to(113 dtype=cost.dtype114 )115 T = ipot(116 cost.detach(),117 txt_len,118 txt_pad,119 img_len,120 img_pad,121 joint_pad,122 0.1,123 1000,124 1,125 )126 127 plan = T[0]128 plan_single = plan * len(txt_emb)129 cost_ = plan_single.t()130 131 cost_ = cost_[hidx][1:].cpu()132 133 patch_index, (H, W) = infer["patch_index"]134 heatmap = torch.zeros(H, W)135 for i, pidx in enumerate(patch_index[0]):136 h, w = pidx[0].item(), pidx[1].item()137 heatmap[h, w] = cost_[i]138 139 heatmap = (heatmap - heatmap.mean()) / heatmap.std()140 heatmap = np.clip(heatmap, 1.0, 3.0)141 heatmap = (heatmap - heatmap.min()) / (heatmap.max() - heatmap.min())142 143 _w, _h = image.size144 overlay = Image.fromarray(np.uint8(heatmap * 255), "L").resize(145 (_w, _h), resample=Image.NEAREST146 )147 image_rgba = image.copy()148 image_rgba.putalpha(overlay)149 image = image_rgba150 151 selected_token = tokenizer.convert_ids_to_tokens(152 encoded["input_ids"][0][hidx]153 )154 155 return [np.array(image), inferred_token[0], selected_token]156 157 inputs = [158 gr.inputs.Textbox(159 label="Url of an image.",160 lines=5,161 ),162 gr.inputs.Textbox(label="Caption with [MASK] tokens to be filled.", lines=5),163 gr.inputs.Slider(164 minimum=0,165 maximum=38,166 step=1,167 label="Index of token for heatmap visualization (ignored if zero)",168 ),169 ]170 outputs = [171 gr.outputs.Image(label="Image"),172 gr.outputs.Textbox(label="description"),173 gr.outputs.Textbox(label="selected token"),174 ]175 176 interface = gr.Interface(177 fn=infer,178 inputs=inputs,179 outputs=outputs,180 examples=[181 [182 "https://s3.geograph.org.uk/geophotos/06/21/24/6212487_1cca7f3f_1024x1024.jpg",183 "a display of flowers growing out and over the [MASK] [MASK] in front of [MASK] on a [MASK] [MASK].",184 0,185 ],186 [187 "https://s3.geograph.org.uk/geophotos/06/21/24/6212487_1cca7f3f_1024x1024.jpg",188 "a display of flowers growing out and over the retaining wall in front of cottages on a cloudy day.",189 4,190 ],191 [192 "https://s3.geograph.org.uk/geophotos/06/21/24/6212487_1cca7f3f_1024x1024.jpg",193 "a display of flowers growing out and over the retaining wall in front of cottages on a cloudy day.",194 11,195 ],196 [197 "https://s3.geograph.org.uk/geophotos/06/21/24/6212487_1cca7f3f_1024x1024.jpg",198 "a display of flowers growing out and over the retaining wall in front of cottages on a cloudy day.",199 15,200 ],201 [202 "https://s3.geograph.org.uk/geophotos/06/21/24/6212487_1cca7f3f_1024x1024.jpg",203 "a display of flowers growing out and over the retaining wall in front of cottages on a cloudy day.",204 18,205 ],206 [207 "https://upload.wikimedia.org/wikipedia/commons/thumb/4/40/Living_Room.jpg/800px-Living_Room.jpg",208 "a room with a [MASK], a [MASK], a [MASK], and a [MASK].",209 0,210 ],211 [212 "https://upload.wikimedia.org/wikipedia/commons/thumb/4/40/Living_Room.jpg/800px-Living_Room.jpg",213 "a room with a rug, a chair, a painting, and a plant.",214 5,215 ],216 [217 "https://upload.wikimedia.org/wikipedia/commons/thumb/4/40/Living_Room.jpg/800px-Living_Room.jpg",218 "a room with a rug, a chair, a painting, and a plant.",219 8,220 ],221 [222 "https://upload.wikimedia.org/wikipedia/commons/thumb/4/40/Living_Room.jpg/800px-Living_Room.jpg",223 "a room with a rug, a chair, a painting, and a plant.",224 11,225 ],226 [227 "https://upload.wikimedia.org/wikipedia/commons/thumb/4/40/Living_Room.jpg/800px-Living_Room.jpg",228 "a room with a rug, a chair, a painting, and a plant.",229 15,230 ],231 ],232 )233 234 interface.launch(debug=True)235 236 237ex.run()