|
"""
|
|
# 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)
|
|
|
|
col.image(gradcam_todraw, use_column_width=True, clamp=True)
|
|
|
