Inf009/image-text-matching

1import gradio as gr 2import gradio.components as grc3import torch4 5from lavis.models import load_model_and_preprocess6from lavis.processors import load_processor7 8# setup device to use9device = torch.device("cuda") if torch.cuda.is_available() else "cpu"10model, vis_processors, text_processors = load_model_and_preprocess("blip2_image_text_matching", "pretrain", device=device, is_eval=True)11 12def predict(raw_image, caption):13    raw_image = raw_image.convert("RGB")14    img = vis_processors["eval"](raw_image).unsqueeze(0).to(device)15    txt = text_processors["eval"](caption)16    itm_output = model({"image": img, "text_input": txt}, match_head="itm")17    itm_scores = torch.nn.functional.softmax(itm_output, dim=1)18    itm_score = itm_scores[:, 1].item()19    itc_score = model({"image": img, "text_input": txt}, match_head='itc')20    return '%.3f' % itm_score, '%.4f' % itc_score21 22app = gr.Interface(fn=predict, inputs=[grc.Image(type="pil"), grc.Textbox()], outputs=[grc.Text(label="itm score"), grc.Text(label="itc score")])23app.launch()24 25