init

app.py

@@ -0,0 +1,30 @@
+import gradio as gr 
+import gradio.components as grc
+import torch
+from PIL import Image
+
+from lavis.models import load_model_and_preprocess
+from lavis.processors import load_processor
+
+raw_image = Image.open("merlion.png").convert("RGB")
+
+# setup device to use
+device = torch.device("cuda") if torch.cuda.is_available() else "cpu"
+caption = "merlion in Singapore"
+
+model, vis_processors, text_processors = load_model_and_preprocess("blip2_image_text_matching", "pretrain", device=device, is_eval=True)
+
+def predict(raw_image, caption):
+    raw_image = raw_image.convert("RGB")
+    img = vis_processors["eval"](raw_image).unsqueeze(0).to(device)
+    txt = text_processors["eval"](caption)
+    itm_output = model({"image": img, "text_input": txt}, match_head="itm")
+    itm_scores = torch.nn.functional.softmax(itm_output, dim=1)
+    itm_score = itm_scores[:, 1].item()
+    itc_score = model({"image": img, "text_input": txt}, match_head='itc')
+    return '%.3f' % itm_score, '%.4f' % itc_score
+
+app = gr.Interface(fn=predict, inputs=[grc.Image(type="pil"), grc.Textbox()], outputs=[grc.Text(label="itm score"), grc.Text(label="itc score")])
+app.launch()
+
+    

requirements.txt

@@ -0,0 +1,3 @@
+salesforce-lavis
+pillow
+torch