from base64 import b64encode from io import BytesIO from math import ceil import clip from multilingual_clip import legacy_multilingual_clip, pt_multilingual_clip import numpy as np import pandas as pd from PIL import Image import requests import streamlit as st import torch from torchvision.transforms import ToPILImage from transformers import AutoTokenizer, AutoModel, BertTokenizer from CLIP_Explainability.clip_ import load, tokenize from CLIP_Explainability.rn_cam import ( # interpret_rn, interpret_rn_overlapped, rn_perword_relevance, ) from CLIP_Explainability.vit_cam import ( # interpret_vit, vit_perword_relevance, interpret_vit_overlapped, ) from pytorch_grad_cam.grad_cam import GradCAM RUN_LITE = False # Load vision model for CAM viz explainability for M-CLIP only MAX_IMG_WIDTH = 500 MAX_IMG_HEIGHT = 800 st.set_page_config(layout="wide") # The `find_best_matches` function compares the text feature vector to the feature vectors of all images and finds the best matches. The function returns the IDs of the best matching images. def find_best_matches(text_features, image_features, image_ids): # Compute the similarity between the search query and each image using the Cosine similarity similarities = (image_features @ text_features.T).squeeze(1) # Sort the images by their similarity score best_image_idx = (-similarities).argsort() # Return the image IDs of the best matches return [[image_ids[i], similarities[i].item()] for i in best_image_idx] # The `encode_search_query` function takes a text description and encodes it into a feature vector using the CLIP model. def encode_search_query(search_query, model_type): with torch.no_grad(): # Encode and normalize the search query using the multilingual model if model_type == "M-CLIP (multilingual ViT)": text_encoded = st.session_state.ml_model.forward( search_query, st.session_state.ml_tokenizer ) text_encoded /= text_encoded.norm(dim=-1, keepdim=True) elif model_type == "J-CLIP (日本語 ViT)": t_text = st.session_state.ja_tokenizer( search_query, padding=True, return_tensors="pt" ) text_encoded = st.session_state.ja_model.get_text_features(**t_text) text_encoded /= text_encoded.norm(dim=-1, keepdim=True) else: # model_type == legacy text_encoded = st.session_state.rn_model(search_query) text_encoded /= text_encoded.norm(dim=-1, keepdim=True) # Retrieve the feature vector return text_encoded def clip_search(search_query): if st.session_state.search_field_value != search_query: st.session_state.search_field_value = search_query model_type = st.session_state.active_model if len(search_query) >= 1: text_features = encode_search_query(search_query, model_type) # Compute the similarity between the descrption and each photo using the Cosine similarity # similarities = list((text_features @ photo_features.T).squeeze(0)) # Sort the photos by their similarity score if model_type == "M-CLIP (multilingual ViT)": matches = find_best_matches( text_features, st.session_state.ml_image_features, st.session_state.image_ids, ) elif model_type == "J-CLIP (日本語 ViT)": matches = find_best_matches( text_features, st.session_state.ja_image_features, st.session_state.image_ids, ) else: # model_type == legacy matches = find_best_matches( text_features, st.session_state.rn_image_features, st.session_state.image_ids, ) st.session_state.search_image_ids = [match[0] for match in matches] st.session_state.search_image_scores = {match[0]: match[1] for match in matches} def string_search(): if "search_field_value" in st.session_state: clip_search(st.session_state.search_field_value) def load_image_features(): # Load the image feature vectors if st.session_state.vision_mode == "tiled": ml_image_features = np.load("./image_features/tiled_ml_features.npy") ja_image_features = np.load("./image_features/tiled_ja_features.npy") rn_image_features = np.load("./image_features/tiled_rn_features.npy") elif st.session_state.vision_mode == "stretched": ml_image_features = np.load("./image_features/resized_ml_features.npy") ja_image_features = np.load("./image_features/resized_ja_features.npy") rn_image_features = np.load("./image_features/resized_rn_features.npy") else: # st.session_state.vision_mode == "cropped": ml_image_features = np.load("./image_features/cropped_ml_features.npy") ja_image_features = np.load("./image_features/cropped_ja_features.npy") rn_image_features = np.load("./image_features/cropped_rn_features.npy") # Convert features to Tensors: Float32 on CPU and Float16 on GPU device = st.session_state.device if device == "cpu": ml_image_features = torch.from_numpy(ml_image_features).float().to(device) ja_image_features = torch.from_numpy(ja_image_features).float().to(device) rn_image_features = torch.from_numpy(rn_image_features).float().to(device) else: ml_image_features = torch.from_numpy(ml_image_features).to(device) ja_image_features = torch.from_numpy(ja_image_features).to(device) rn_image_features = torch.from_numpy(rn_image_features).to(device) st.session_state.ml_image_features = ml_image_features / ml_image_features.norm( dim=-1, keepdim=True ) st.session_state.ja_image_features = ja_image_features / ja_image_features.norm( dim=-1, keepdim=True ) st.session_state.rn_image_features = rn_image_features / rn_image_features.norm( dim=-1, keepdim=True ) string_search() def init(): st.session_state.current_page = 1 device = "cuda" if torch.cuda.is_available() else "cpu" st.session_state.device = device # Load the open CLIP models with st.spinner("Loading models and data, please wait..."): ml_model_name = "M-CLIP/XLM-Roberta-Large-Vit-B-16Plus" ml_model_path = "./models/vit_b_16_plus_240-laion400m_e32-699c4b84.pt" st.session_state.ml_image_model, st.session_state.ml_image_preprocess = load( ml_model_path, device=device, jit=False ) st.session_state.ml_model = ( pt_multilingual_clip.MultilingualCLIP.from_pretrained(ml_model_name) ) st.session_state.ml_tokenizer = AutoTokenizer.from_pretrained(ml_model_name) ja_model_name = "hakuhodo-tech/japanese-clip-vit-h-14-bert-wider" ja_model_path = "./models/ViT-H-14-laion2B-s32B-b79K.bin" if not RUN_LITE: st.session_state.ja_image_model, st.session_state.ja_image_preprocess = ( load(ja_model_path, device=device, jit=False) ) st.session_state.ja_model = AutoModel.from_pretrained( ja_model_name, trust_remote_code=True ).to(device) st.session_state.ja_tokenizer = AutoTokenizer.from_pretrained( ja_model_name, trust_remote_code=True ) if not RUN_LITE: st.session_state.rn_image_model, st.session_state.rn_image_preprocess = ( clip.load("RN50x4", device=device) ) st.session_state.rn_model = legacy_multilingual_clip.load_model( "M-BERT-Base-69" ) st.session_state.rn_tokenizer = BertTokenizer.from_pretrained( "bert-base-multilingual-cased" ) # Load the image IDs st.session_state.images_info = pd.read_csv("./metadata.csv") st.session_state.images_info.set_index("filename", inplace=True) with open("./images_list.txt", "r", encoding="utf-8") as images_list: st.session_state.image_ids = list(images_list.read().strip().split("\n")) st.session_state.active_model = "M-CLIP (multilingual ViT)" st.session_state.vision_mode = "tiled" st.session_state.search_image_ids = [] st.session_state.search_image_scores = {} st.session_state.activations_image = None st.session_state.text_table_df = None with st.spinner("Loading models and data, please wait..."): load_image_features() if "images_info" not in st.session_state: init() def visualize_gradcam(viz_image_id): if "search_field_value" not in st.session_state: return header_cols = st.columns([80, 20], vertical_alignment="bottom") with header_cols[0]: st.title("Image + query details") with header_cols[1]: if st.button("Close"): st.rerun() st.markdown( f"**Query text:** {st.session_state.search_field_value} | **Image relevance:** {round(st.session_state.search_image_scores[viz_image_id], 3)}" ) with st.spinner("Calculating..."): # info_text = st.text("Calculating activation regions...") image_url = st.session_state.images_info.loc[viz_image_id]["image_url"] image_response = requests.get(image_url) image = Image.open(BytesIO(image_response.content), formats=["JPEG", "GIF"]) image = image.convert("RGB") img_dim = 224 if st.session_state.active_model == "M-CLIP (multilingual ViT)": img_dim = 240 elif st.session_state.active_model == "Legacy (multilingual ResNet)": img_dim = 288 orig_img_dims = image.size ##### If the features are based on tiled image slices tile_behavior = None if st.session_state.vision_mode == "tiled": scaled_dims = [img_dim, img_dim] if orig_img_dims[0] > orig_img_dims[1]: scale_ratio = round(orig_img_dims[0] / orig_img_dims[1]) if scale_ratio > 1: scaled_dims = [scale_ratio * img_dim, img_dim] tile_behavior = "width" elif orig_img_dims[0] < orig_img_dims[1]: scale_ratio = round(orig_img_dims[1] / orig_img_dims[0]) if scale_ratio > 1: scaled_dims = [img_dim, scale_ratio * img_dim] tile_behavior = "height" resized_image = image.resize(scaled_dims, Image.LANCZOS) if tile_behavior == "width": image_tiles = [] for x in range(0, scale_ratio): box = (x * img_dim, 0, (x + 1) * img_dim, img_dim) image_tiles.append(resized_image.crop(box)) elif tile_behavior == "height": image_tiles = [] for y in range(0, scale_ratio): box = (0, y * img_dim, img_dim, (y + 1) * img_dim) image_tiles.append(resized_image.crop(box)) else: image_tiles = [resized_image] elif st.session_state.vision_mode == "stretched": image_tiles = [image.resize((img_dim, img_dim), Image.LANCZOS)] else: # vision_mode == "cropped" if orig_img_dims[0] > orig_img_dims[1]: scale_factor = orig_img_dims[0] / orig_img_dims[1] resized_img_dims = (round(scale_factor * img_dim), img_dim) resized_img = image.resize(resized_img_dims) elif orig_img_dims[0] < orig_img_dims[1]: scale_factor = orig_img_dims[1] / orig_img_dims[0] resized_img_dims = (img_dim, round(scale_factor * img_dim)) else: resized_img_dims = (img_dim, img_dim) resized_img = image.resize(resized_img_dims) left = round((resized_img_dims[0] - img_dim) / 2) top = round((resized_img_dims[1] - img_dim) / 2) x_right = round(resized_img_dims[0] - img_dim) - left x_bottom = round(resized_img_dims[1] - img_dim) - top right = resized_img_dims[0] - x_right bottom = resized_img_dims[1] - x_bottom # Crop the center of the image image_tiles = [resized_img.crop((left, top, right, bottom))] image_visualizations = [] if st.session_state.active_model == "M-CLIP (multilingual ViT)": # Sometimes used for token importance viz tokenized_text = st.session_state.ml_tokenizer.tokenize( st.session_state.search_field_value ) text_features = st.session_state.ml_model.forward( st.session_state.search_field_value, st.session_state.ml_tokenizer ) image_model = st.session_state.ml_image_model for altered_image in image_tiles: p_image = ( st.session_state.ml_image_preprocess(altered_image) .unsqueeze(0) .to(st.session_state.device) ) vis_t = interpret_vit_overlapped( p_image.type(st.session_state.ml_image_model.dtype), text_features, image_model.visual, st.session_state.device, img_dim=img_dim, ) image_visualizations.append(vis_t) elif st.session_state.active_model == "J-CLIP (日本語 ViT)": # Sometimes used for token importance viz tokenized_text = st.session_state.ja_tokenizer.tokenize( st.session_state.search_field_value ) t_text = st.session_state.ja_tokenizer( st.session_state.search_field_value, return_tensors="pt" ) text_features = st.session_state.ja_model.get_text_features(**t_text) image_model = st.session_state.ja_image_model for altered_image in image_tiles: p_image = ( st.session_state.ja_image_preprocess(altered_image) .unsqueeze(0) .to(st.session_state.device) ) vis_t = interpret_vit_overlapped( p_image.type(st.session_state.ja_image_model.dtype), text_features, image_model.visual, st.session_state.device, img_dim=img_dim, ) image_visualizations.append(vis_t) else: # st.session_state.active_model == Legacy # Sometimes used for token importance viz tokenized_text = st.session_state.rn_tokenizer.tokenize( st.session_state.search_field_value ) text_features = st.session_state.rn_model( st.session_state.search_field_value ) image_model = st.session_state.rn_image_model for altered_image in image_tiles: p_image = ( st.session_state.rn_image_preprocess(altered_image) .unsqueeze(0) .to(st.session_state.device) ) vis_t = interpret_rn_overlapped( p_image.type(st.session_state.rn_image_model.dtype), text_features, image_model.visual, GradCAM, st.session_state.device, img_dim=img_dim, ) image_visualizations.append(vis_t) transform = ToPILImage() vis_images = [transform(vis_t) for vis_t in image_visualizations] if st.session_state.vision_mode == "cropped": resized_img.paste(vis_images[0], (left, top)) vis_images = [resized_img] if orig_img_dims[0] > orig_img_dims[1]: scale_factor = MAX_IMG_WIDTH / orig_img_dims[0] scaled_dims = [MAX_IMG_WIDTH, int(orig_img_dims[1] * scale_factor)] else: scale_factor = MAX_IMG_HEIGHT / orig_img_dims[1] scaled_dims = [int(orig_img_dims[0] * scale_factor), MAX_IMG_HEIGHT] if tile_behavior == "width": vis_image = Image.new("RGB", (len(vis_images) * img_dim, img_dim)) for x, v_img in enumerate(vis_images): vis_image.paste(v_img, (x * img_dim, 0)) st.session_state.activations_image = vis_image.resize(scaled_dims) elif tile_behavior == "height": vis_image = Image.new("RGB", (img_dim, len(vis_images) * img_dim)) for y, v_img in enumerate(vis_images): vis_image.paste(v_img, (0, y * img_dim)) st.session_state.activations_image = vis_image.resize(scaled_dims) else: st.session_state.activations_image = vis_images[0].resize(scaled_dims) image_io = BytesIO() st.session_state.activations_image.save(image_io, "PNG") dataurl = "data:image/png;base64," + b64encode(image_io.getvalue()).decode( "ascii" ) st.html( f"""
""" ) tokenized_text = [tok.replace("▁", "") for tok in tokenized_text if tok != "▁"] tokenized_text = [ tok for tok in tokenized_text if tok not in ["s", "ed", "a", "the", "an", "ing"] ] if ( len(tokenized_text) > 1 and len(tokenized_text) < 25 and st.button( "Calculate text importance (may take some time)", ) ): search_tokens = [] token_scores = [] progress_text = f"Processing {len(tokenized_text)} text tokens" progress_bar = st.progress(0.0, text=progress_text) for t, tok in enumerate(tokenized_text): token = tok if st.session_state.active_model == "Legacy (multilingual ResNet)": word_rel = rn_perword_relevance( p_image, st.session_state.search_field_value, image_model, tokenize, GradCAM, st.session_state.device, token, data_only=True, img_dim=img_dim, ) else: word_rel = vit_perword_relevance( p_image, st.session_state.search_field_value, image_model, tokenize, st.session_state.device, token, data_only=True, img_dim=img_dim, ) avg_score = np.mean(word_rel) if avg_score == 0 or np.isnan(avg_score): continue search_tokens.append(token) token_scores.append(1 / avg_score) progress_bar.progress( (t + 1) / len(tokenized_text), text=f"Processing token {t+1} of {len(tokenized_text)}", ) progress_bar.empty() normed_scores = torch.softmax(torch.tensor(token_scores), dim=0) token_scores = [f"{round(score.item() * 100, 3)}%" for score in normed_scores] st.session_state.text_table_df = pd.DataFrame( {"token": search_tokens, "importance": token_scores} ) st.markdown("**Importance of each text token to relevance score**") st.table(st.session_state.text_table_df) def format_vision_mode(mode_stub): return mode_stub.capitalize() @st.dialog(" ", width="large") def image_modal(vis_image_id): visualize_gradcam(vis_image_id) st.title("Explore Japanese visual aesthetics with CLIP models") st.markdown( """ """, unsafe_allow_html=True, ) search_row = st.columns([45, 8, 8, 10, 1, 8, 20], vertical_alignment="center") with search_row[0]: search_field = st.text_input( label="search", label_visibility="collapsed", placeholder="Type something, or click a suggested search below.", on_change=string_search, key="search_field_value", ) with search_row[1]: st.button( "Search", on_click=string_search, use_container_width=True, type="primary" ) with search_row[2]: st.markdown("**Vision mode:**") with search_row[3]: st.selectbox( "Vision mode", options=["tiled", "stretched", "cropped"], key="vision_mode", help="How to consider images that aren't square", on_change=load_image_features, format_func=format_vision_mode, label_visibility="collapsed", ) with search_row[4]: st.empty() with search_row[5]: st.markdown("**CLIP model:**") with search_row[6]: st.selectbox( "CLIP Model:", options=[ "M-CLIP (multilingual ViT)", "J-CLIP (日本語 ViT)", "Legacy (multilingual ResNet)", ], key="active_model", on_change=string_search, label_visibility="collapsed", ) canned_searches = st.columns([12, 22, 22, 22, 22], vertical_alignment="top") with canned_searches[0]: st.markdown("**Suggested searches:**") if st.session_state.active_model == "J-CLIP (日本語 ViT)": with canned_searches[1]: st.button( "間", on_click=clip_search, args=["間"], use_container_width=True, ) with canned_searches[2]: st.button("奥", on_click=clip_search, args=["奥"], use_container_width=True) with canned_searches[3]: st.button("山", on_click=clip_search, args=["山"], use_container_width=True) with canned_searches[4]: st.button( "花に酔えり 羽織着て刀 さす女", on_click=clip_search, args=["花に酔えり 羽織着て刀 さす女"], use_container_width=True, ) else: with canned_searches[1]: st.button( "negative space", on_click=clip_search, args=["negative space"], use_container_width=True, ) with canned_searches[2]: st.button("間", on_click=clip_search, args=["間"], use_container_width=True) with canned_searches[3]: st.button("음각", on_click=clip_search, args=["음각"], use_container_width=True) with canned_searches[4]: st.button( "αρνητικός χώρος", on_click=clip_search, args=["αρνητικός χώρος"], use_container_width=True, ) controls = st.columns([35, 5, 35, 5, 20], gap="large", vertical_alignment="center") with controls[0]: im_per_pg = st.columns([30, 70], vertical_alignment="center") with im_per_pg[0]: st.markdown("**Images/page:**") with im_per_pg[1]: batch_size = st.select_slider( "Images/page:", range(10, 50, 10), label_visibility="collapsed" ) with controls[1]: st.empty() with controls[2]: im_per_row = st.columns([30, 70], vertical_alignment="center") with im_per_row[0]: st.markdown("**Images/row:**") with im_per_row[1]: row_size = st.select_slider( "Images/row:", range(1, 6), value=5, label_visibility="collapsed" ) num_batches = ceil(len(st.session_state.image_ids) / batch_size) with controls[3]: st.empty() with controls[4]: pager = st.columns([40, 60], vertical_alignment="center") with pager[0]: st.markdown(f"Page **{st.session_state.current_page}** of **{num_batches}** ") with pager[1]: st.number_input( "Page", min_value=1, max_value=num_batches, step=1, label_visibility="collapsed", key="current_page", ) if len(st.session_state.search_image_ids) == 0: batch = [] else: batch = st.session_state.search_image_ids[ (st.session_state.current_page - 1) * batch_size : st.session_state.current_page * batch_size ] grid = st.columns(row_size) col = 0 for image_id in batch: with grid[col]: link_text = st.session_state.images_info.loc[image_id]["permalink"].split("/")[ 2 ] # st.image( # st.session_state.images_info.loc[image_id]["image_url"], # caption=st.session_state.images_info.loc[image_id]["caption"], # ) st.html( f"""
{st.session_state.images_info.loc[image_id]['caption']} [{round(st.session_state.search_image_scores[image_id], 3)}]
""" ) st.caption( f"""
{link_text}
""", unsafe_allow_html=True, ) if not RUN_LITE or st.session_state.active_model == "M-CLIP (multilingual ViT)": st.button( "Explain this", on_click=image_modal, args=[image_id], use_container_width=True, key=image_id, ) col = (col + 1) % row_size