Create app.py

app.py

@@ -0,0 +1,60 @@
+import json
+import torch
+from PIL import Image
+from ruamel import yaml
+from albef.model import albef_model_for_vqa
+from albef.data.transforms import ALBEFTextTransform, testing_image_transform
+import gradio as gr
+
+data_dir = "./vqa_data/"
+
+config = yaml.load(open("./configs/vqa.yaml", "r"), Loader=yaml.Loader)
+model = albef_model_for_vqa(config)
+
+checkpoint_url = "https://download.pytorch.org/models/multimodal/albef/finetuned_vqa_checkpoint.pt"
+checkpoint = torch.hub.load_state_dict_from_url(checkpoint_url, map_location='cpu')
+model.load_state_dict(checkpoint)
+
+image_transform = testing_image_transform()
+question_transform = ALBEFTextTransform(add_end_token=False)
+answer_transform = ALBEFTextTransform(do_pre_process=False)
+
+answer_list = json.load(open(data_dir + "answer_list.json", "r"))
+
+def infer(image, question):
+  images = [image]
+  image_input = [image_transform(image) for image in images]
+  image_input = torch.stack(image_input, dim=0)
+  question_input = question_transform([question])
+  question_atts = (question_input != 0).type(torch.long)
+  answer_input = answer_transform(answer_list)
+  answer_atts = (answer_input != 0).type(torch.long)
+
+  answer_ids, _ = model(
+      image_input, 
+      question_input,
+      question_atts,
+      answer_input,
+      answer_atts,
+      k=1,
+      is_train=False,
+  )
+
+  predicted_answer_id = answer_ids[0]
+  predicted_answer = answer_list[predicted_answer_id]
+
+  return predicted_answer
+
+demo = gr.Interface(
+    fn=infer,
+    inputs=[gr.Image(label='image', type='pil', image_mode='RGB'), gr.Text(label='question')],
+    outputs=gr.Text(label='answer'),
+    # examples=[
+    #     ['vqav2.png', 'What sport is this?'],
+    #     ['vizwiz.jpeg', 'What piece of meat have I taken out of the freezer?'],
+    #     ['aqua.png', 'what does bol lean nonchalantly on'],
+    #     ['robotvqa.png', 'How many silver spoons are there?'],
+    # ]
+)
+
+demo.launch()