Create app.py

app.py

@@ -0,0 +1,126 @@
+import gradio as gr
+import torch
+import torch.nn.functional as F
+import requests
+import numpy as np
+import re
+import io
+
+from PIL import Image
+from transformers import ViltProcessor, ViltForMaskedLM
+from torchvision import transforms
+
+processor = ViltProcessor.from_pretrained("dandelin/vilt-b32-mlm")
+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
+        self.max = longer
+
+    def __call__(self, x):
+        w, h = x.size
+        scale = self.min / min(w, h)
+        if h < w:
+            newh, neww = self.min, scale * w
+        else:
+            newh, neww = scale * h, self.min
+
+        if max(newh, neww) > self.max:
+            scale = self.max / max(newh, neww)
+            newh = newh * scale
+            neww = neww * scale
+
+        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)
+
+def pixelbert_transform(size=800):
+    longer = int((1333 / 800) * size)
+    return transforms.Compose(
+        [
+            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])]),
+        ]
+    )
+
+
+def infer(url, mp_text):
+    try:
+        res = requests.get(url)
+        image = Image.open(io.BytesIO(res.content)).convert("RGB")
+        img = pixelbert_transform(size=384)(image)
+        img = img.unsqueeze(0).to(device)
+    except:
+        return False
+
+    tl = len(re.findall("\[MASK\]", mp_text))
+    inferred_token = [mp_text]
+    encoding = processor(image, mp_text, return_tensors="pt")
+
+    with torch.no_grad():
+        for i in range(tl):
+            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)
+            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_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)
+
+    return [np.array(image), output]
+
+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),
+    ]
+outputs_interface = [
+        gr.outputs.Image(label="Image"),
+        gr.outputs.Textbox(label="description"),
+    ]
+
+interface = gr.Interface(
+        fn=infer,
+        inputs=inputs_interface,
+        outputs=outputs_interface,
+        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].",
+            ],
+        ],
+    )
+
+interface.launch()
+