File size: 3,562 Bytes
846a3aa |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 |
"""
# Copyright (c) 2022, salesforce.com, inc.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
"""
import math
import numpy as np
import streamlit as st
from lavis.models.blip_models.blip_image_text_matching import compute_gradcam
from lavis.processors import load_processor
from PIL import Image
from app import device, load_demo_image
from app.utils import getAttMap, init_bert_tokenizer, load_blip_itm_model
def app():
model_type = st.sidebar.selectbox("Model:", ["BLIP_base", "BLIP_large"])
values = list(range(1, 12))
default_layer_num = values.index(7)
layer_num = (
st.sidebar.selectbox("Layer number", values, index=default_layer_num) - 1
)
st.markdown(
"<h1 style='text-align: center;'>Text Localization</h1>", unsafe_allow_html=True
)
vis_processor = load_processor("blip_image_eval").build(image_size=384)
text_processor = load_processor("blip_caption")
tokenizer = init_bert_tokenizer()
instructions = "Try the provided image and text or use your own ones."
file = st.file_uploader(instructions)
query = st.text_input(
"Try a different input.", "A girl playing with her dog on the beach."
)
submit_button = st.button("Submit")
col1, col2 = st.columns(2)
if file:
raw_img = Image.open(file).convert("RGB")
else:
raw_img = load_demo_image()
col1.header("Image")
w, h = raw_img.size
scaling_factor = 720 / w
resized_image = raw_img.resize((int(w * scaling_factor), int(h * scaling_factor)))
col1.image(resized_image, use_column_width=True)
col2.header("GradCam")
if submit_button:
if model_type.startswith("BLIP"):
blip_type = model_type.split("_")[1]
model = load_blip_itm_model(device, model_type=blip_type)
img = vis_processor(raw_img).unsqueeze(0).to(device)
qry = text_processor(query)
qry_tok = tokenizer(qry, return_tensors="pt").to(device)
norm_img = np.float32(resized_image) / 255
gradcam, _ = compute_gradcam(model, img, qry, qry_tok, block_num=layer_num)
avg_gradcam = getAttMap(norm_img, gradcam[0][1], blur=True)
col2.image(avg_gradcam, use_column_width=True, clamp=True)
num_cols = 4.0
num_tokens = len(qry_tok.input_ids[0]) - 2
num_rows = int(math.ceil(num_tokens / num_cols))
gradcam_iter = iter(gradcam[0][2:-1])
token_id_iter = iter(qry_tok.input_ids[0][1:-1])
for _ in range(num_rows):
with st.container():
for col in st.columns(int(num_cols)):
token_id = next(token_id_iter, None)
if not token_id:
break
gradcam_img = next(gradcam_iter)
word = tokenizer.decode([token_id])
gradcam_todraw = getAttMap(norm_img, gradcam_img, blur=True)
new_title = (
'<p style="text-align: center; font-size: 25px;">{}</p>'.format(
word
)
)
col.markdown(new_title, unsafe_allow_html=True)
# st.image(image, channels="BGR")
col.image(gradcam_todraw, use_column_width=True, clamp=True)
|