File size: 2,307 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 |
"""
# 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 numpy as np
import streamlit as st
import torch
from lavis.models import BlipBase, load_model
from matplotlib import pyplot as plt
from PIL import Image
from scipy.ndimage import filters
from skimage import transform as skimage_transform
def resize_img(raw_img):
w, h = raw_img.size
scaling_factor = 240 / w
resized_image = raw_img.resize((int(w * scaling_factor), int(h * scaling_factor)))
return resized_image
def read_img(filepath):
raw_image = Image.open(filepath).convert("RGB")
return raw_image
@st.cache(
hash_funcs={
torch.nn.parameter.Parameter: lambda parameter: parameter.data.detach()
.cpu()
.numpy()
},
allow_output_mutation=True,
)
def load_model_cache(name, model_type, is_eval, device):
return load_model(name, model_type, is_eval, device)
@st.cache(allow_output_mutation=True)
def init_bert_tokenizer():
tokenizer = BlipBase.init_tokenizer()
return tokenizer
def getAttMap(img, attMap, blur=True, overlap=True):
attMap -= attMap.min()
if attMap.max() > 0:
attMap /= attMap.max()
attMap = skimage_transform.resize(attMap, (img.shape[:2]), order=3, mode="constant")
if blur:
attMap = filters.gaussian_filter(attMap, 0.02 * max(img.shape[:2]))
attMap -= attMap.min()
attMap /= attMap.max()
cmap = plt.get_cmap("jet")
attMapV = cmap(attMap)
attMapV = np.delete(attMapV, 3, 2)
if overlap:
attMap = (
1 * (1 - attMap**0.7).reshape(attMap.shape + (1,)) * img
+ (attMap**0.7).reshape(attMap.shape + (1,)) * attMapV
)
return attMap
@st.cache(
hash_funcs={
torch.nn.parameter.Parameter: lambda parameter: parameter.data.detach()
.cpu()
.numpy()
},
allow_output_mutation=True,
)
def load_blip_itm_model(device, model_type="base"):
model = load_model(
"blip_image_text_matching", model_type, is_eval=True, device=device
)
return model
|