Spaces:
Runtime error
Runtime error
""" | |
# 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 os | |
import numpy as np | |
import streamlit as st | |
import torch | |
import torch.nn.functional as F | |
from app import cache_root, device | |
from app.utils import ( | |
getAttMap, | |
init_bert_tokenizer, | |
load_blip_itm_model, | |
read_img, | |
resize_img, | |
) | |
from lavis.models import load_model | |
from lavis.processors import load_processor | |
def load_feat(): | |
from lavis.common.utils import download_url | |
dirname = os.path.join(os.path.dirname(__file__), "assets") | |
filename = "path2feat_coco_train2014.pth" | |
filepath = os.path.join(dirname, filename) | |
url = "https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/path2feat_coco_train2014.pth" | |
if not os.path.exists(filepath): | |
download_url(url=url, root=dirname, filename="path2feat_coco_train2014.pth") | |
path2feat = torch.load(filepath) | |
paths = sorted(path2feat.keys()) | |
all_img_feats = torch.stack([path2feat[k] for k in paths], dim=0).to(device) | |
return path2feat, paths, all_img_feats | |
def load_feature_extractor_model(device): | |
model_url = "https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base.pth" | |
model = load_model( | |
"blip_feature_extractor", model_type="base", is_eval=True, device=device | |
) | |
model.load_from_pretrained(model_url) | |
return model | |
def app(): | |
# === layout === | |
model_type = st.sidebar.selectbox("Model:", ["BLIP_base", "BLIP_large"]) | |
file_root = os.path.join(cache_root, "coco/images/train2014/") | |
values = [12, 24, 48] | |
default_layer_num = values.index(24) | |
num_display = st.sidebar.selectbox( | |
"Number of images:", values, index=default_layer_num | |
) | |
show_gradcam = st.sidebar.selectbox("Show GradCam:", [True, False], index=1) | |
itm_ranking = st.sidebar.selectbox("Multimodal re-ranking:", [True, False], index=0) | |
# st.title('Multimodal Search') | |
st.markdown( | |
"<h1 style='text-align: center;'>Multimodal Search</h1>", unsafe_allow_html=True | |
) | |
# === event === | |
vis_processor = load_processor("blip_image_eval").build(image_size=384) | |
text_processor = load_processor("blip_caption") | |
user_question = st.text_input( | |
"Search query", "A dog running on the grass.", help="Type something to search." | |
) | |
user_question = text_processor(user_question) | |
feature_extractor = load_feature_extractor_model(device) | |
# ======= ITC ========= | |
sample = {"text_input": user_question} | |
with torch.no_grad(): | |
text_feature = feature_extractor.extract_features( | |
sample, mode="text" | |
).text_embeds_proj[0, 0] | |
path2feat, paths, all_img_feats = load_feat() | |
all_img_feats.to(device) | |
all_img_feats = F.normalize(all_img_feats, dim=1) | |
num_cols = 4 | |
num_rows = int(num_display / num_cols) | |
similarities = text_feature @ all_img_feats.T | |
indices = torch.argsort(similarities, descending=True)[:num_display] | |
top_paths = [paths[ind.detach().cpu().item()] for ind in indices] | |
sorted_similarities = [similarities[idx] for idx in indices] | |
filenames = [os.path.join(file_root, p) for p in top_paths] | |
# ========= ITM and GradCam ========== | |
bsz = 4 # max number of images to avoid cuda oom | |
if model_type.startswith("BLIP"): | |
blip_type = model_type.split("_")[1] | |
itm_model = load_blip_itm_model(device, model_type=blip_type) | |
tokenizer = init_bert_tokenizer() | |
queries_batch = [user_question] * bsz | |
queries_tok_batch = tokenizer(queries_batch, return_tensors="pt").to(device) | |
num_batches = int(num_display / bsz) | |
avg_gradcams = [] | |
all_raw_images = [] | |
itm_scores = [] | |
for i in range(num_batches): | |
filenames_in_batch = filenames[i * bsz : (i + 1) * bsz] | |
raw_images, images = read_and_process_images(filenames_in_batch, vis_processor) | |
gradcam, itm_output = compute_gradcam_batch( | |
itm_model, images, queries_batch, queries_tok_batch | |
) | |
all_raw_images.extend([resize_img(r_img) for r_img in raw_images]) | |
norm_imgs = [np.float32(r_img) / 255 for r_img in raw_images] | |
for norm_img, grad_cam in zip(norm_imgs, gradcam): | |
avg_gradcam = getAttMap(norm_img, grad_cam[0], blur=True) | |
avg_gradcams.append(avg_gradcam) | |
with torch.no_grad(): | |
itm_score = torch.nn.functional.softmax(itm_output, dim=1) | |
itm_scores.append(itm_score) | |
# ========= ITM re-ranking ========= | |
itm_scores = torch.cat(itm_scores)[:, 1] | |
if itm_ranking: | |
itm_scores_sorted, indices = torch.sort(itm_scores, descending=True) | |
avg_gradcams_sorted = [] | |
all_raw_images_sorted = [] | |
for idx in indices: | |
avg_gradcams_sorted.append(avg_gradcams[idx]) | |
all_raw_images_sorted.append(all_raw_images[idx]) | |
avg_gradcams = avg_gradcams_sorted | |
all_raw_images = all_raw_images_sorted | |
if show_gradcam: | |
images_to_show = iter(avg_gradcams) | |
else: | |
images_to_show = iter(all_raw_images) | |
for _ in range(num_rows): | |
with st.container(): | |
for col in st.columns(num_cols): | |
col.image(next(images_to_show), use_column_width=True, clamp=True) | |
def read_and_process_images(image_paths, vis_processor): | |
raw_images = [read_img(path) for path in image_paths] | |
images = [vis_processor(r_img) for r_img in raw_images] | |
images_tensors = torch.stack(images).to(device) | |
return raw_images, images_tensors | |
def compute_gradcam_batch(model, visual_input, text_input, tokenized_text, block_num=6): | |
model.text_encoder.base_model.base_model.encoder.layer[ | |
block_num | |
].crossattention.self.save_attention = True | |
output = model({"image": visual_input, "text_input": text_input}, match_head="itm") | |
loss = output[:, 1].sum() | |
model.zero_grad() | |
loss.backward() | |
with torch.no_grad(): | |
mask = tokenized_text.attention_mask.view( | |
tokenized_text.attention_mask.size(0), 1, -1, 1, 1 | |
) # (bsz,1,token_len, 1,1) | |
token_length = mask.sum() - 2 | |
token_length = token_length.cpu() | |
# grads and cams [bsz, num_head, seq_len, image_patch] | |
grads = model.text_encoder.base_model.base_model.encoder.layer[ | |
block_num | |
].crossattention.self.get_attn_gradients() | |
cams = model.text_encoder.base_model.base_model.encoder.layer[ | |
block_num | |
].crossattention.self.get_attention_map() | |
# assume using vit large with 576 num image patch | |
cams = cams[:, :, :, 1:].reshape(visual_input.size(0), 12, -1, 24, 24) * mask | |
grads = ( | |
grads[:, :, :, 1:].clamp(0).reshape(visual_input.size(0), 12, -1, 24, 24) | |
* mask | |
) | |
gradcam = cams * grads | |
# [enc token gradcam, average gradcam across token, gradcam for individual token] | |
# gradcam = torch.cat((gradcam[0:1,:], gradcam[1:token_length+1, :].sum(dim=0, keepdim=True)/token_length, gradcam[1:, :])) | |
gradcam = gradcam.mean(1).cpu().detach() | |
gradcam = ( | |
gradcam[:, 1 : token_length + 1, :].sum(dim=1, keepdim=True) / token_length | |
) | |
return gradcam, output | |