Can select and visualizing results from cropping, stretching or tiling images

app.py

@@ -18,12 +18,43 @@ from CLIP_Explainability.vit_cam import (
     vit_perword_relevance,
 )  # , interpret_vit_overlapped
 
-MAX_IMG_WIDTH = 450  # For small dialog
+MAX_IMG_WIDTH = 500
 MAX_IMG_HEIGHT = 800
 
 st.set_page_config(layout="wide")
 
 
+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")
+    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")
+    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")
+
+    # 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)
+    else:
+        ml_image_features = torch.from_numpy(ml_image_features).to(device)
+        ja_image_features = torch.from_numpy(ja_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
+    )
+
+    string_search()
+
+
 def init():
     st.session_state.current_page = 1
 
@@ -34,74 +65,51 @@ def init():
     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
-    )
+    with st.spinner("Loading models and data, please wait..."):
+        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)
+        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"
+        ja_model_name = "hakuhodo-tech/japanese-clip-vit-h-14-bert-wider"
+        ja_model_path = "./models/ViT-H-14-laion2B-s32B-b79K.bin"
 
-    st.session_state.ja_image_model, st.session_state.ja_image_preprocess = load(
-        ja_model_path, device=device, jit=False
-    )
+        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
-    )
+        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
+        )
+
+        # 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 (multiple languages)"
+        st.session_state.active_model = "M-CLIP (multiple languages)"
 
+    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
 
-    # Load the image IDs
-    st.session_state.images_info = pd.read_csv("./metadata.csv")
-    st.session_state.images_info.set_index("filename", inplace=True)
-
-    st.session_state.image_ids = list(
-        open("./images_list.txt", "r", encoding="utf-8").read().strip().split("\n")
-    )
-
-    # Load the image feature vectors
-    # ml_image_features = np.load("./multilingual_features.npy")
-    # ja_image_features = np.load("./hakuhodo_features.npy")
-    ml_image_features = np.load("./resized_ml_features.npy")
-    ja_image_features = np.load("./resized_ja_features.npy")
-    # ml_image_features = np.load("./tiled_ml_features.npy")
-    # ja_image_features = np.load("./tiled_ja_features.npy")
-
-    # Convert features to Tensors: Float32 on CPU and Float16 on GPU
-    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)
-    else:
-        ml_image_features = torch.from_numpy(ml_image_features).to(device)
-        ja_image_features = torch.from_numpy(ja_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
-    )
+    with st.spinner("Loading models and data, please wait..."):
+        load_image_features()
 
 
-if (
-    "ml_image_features" not in st.session_state
-    or "ja_image_features" not in st.session_state
-):
-    with st.spinner("Loading models and data, please wait..."):
-        init()
+if "images_info" not in st.session_state:
+    init()
 
 
 # The `encode_search_query` function takes a text description and encodes it into a feature vector using the CLIP model.
@@ -191,6 +199,7 @@ def visualize_gradcam(viz_image_id):
     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 (multiple languages)":
@@ -198,62 +207,141 @@ def visualize_gradcam(viz_image_id):
 
     orig_img_dims = image.size
 
-    altered_image = image.resize((img_dim, img_dim), Image.LANCZOS)
+    ##### 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 (multiple languages)":
-        p_image = (
-            st.session_state.ml_image_preprocess(altered_image)
-            .unsqueeze(0)
-            .to(st.session_state.device)
-        )
-
         # Sometimes used for token importance viz
         tokenized_text = st.session_state.ml_tokenizer.tokenize(
             st.session_state.search_field_value
         )
-        image_model = st.session_state.ml_image_model
-        # tokenize = st.session_state.ml_tokenizer.tokenize
 
         text_features = st.session_state.ml_model.forward(
             st.session_state.search_field_value, st.session_state.ml_tokenizer
         )
 
-        vis_t = interpret_vit(
-            p_image.type(st.session_state.ml_image_model.dtype),
-            text_features,
-            st.session_state.ml_image_model.visual,
-            st.session_state.device,
-            img_dim=img_dim,
-        )
+        image_model = st.session_state.ml_image_model
+        # tokenize = st.session_state.ml_tokenizer.tokenize
+        image_model.eval()
 
-    else:
-        p_image = (
-            st.session_state.ja_image_preprocess(altered_image)
-            .unsqueeze(0)
-            .to(st.session_state.device)
-        )
+        for altered_image in image_tiles:
+            image_model.zero_grad()
+
+            p_image = (
+                st.session_state.ml_image_preprocess(altered_image)
+                .unsqueeze(0)
+                .to(st.session_state.device)
+            )
 
+            vis_t = interpret_vit(
+                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)
+
+    else:
         # Sometimes used for token importance viz
         tokenized_text = st.session_state.ja_tokenizer.tokenize(
             st.session_state.search_field_value
         )
-        image_model = st.session_state.ja_image_model
 
         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)
 
-        vis_t = interpret_vit(
-            p_image.type(st.session_state.ja_image_model.dtype),
-            text_features,
-            st.session_state.ja_image_model.visual,
-            st.session_state.device,
-            img_dim=img_dim,
-        )
+        image_model = st.session_state.ja_image_model
+        image_model.eval()
+
+        for altered_image in image_tiles:
+            image_model.zero_grad()
+
+            p_image = (
+                st.session_state.ja_image_preprocess(altered_image)
+                .unsqueeze(0)
+                .to(st.session_state.device)
+            )
+
+            vis_t = interpret_vit(
+                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)
 
     transform = ToPILImage()
-    vis_img = transform(vis_t)
+
+    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]
@@ -262,14 +350,27 @@ def visualize_gradcam(viz_image_id):
         scale_factor = MAX_IMG_HEIGHT / orig_img_dims[1]
         scaled_dims = [int(orig_img_dims[0] * scale_factor), MAX_IMG_HEIGHT]
 
-    st.session_state.activations_image = vis_img.resize(scaled_dims)
+    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"""<div style="display: flex; flex-direction: column; align-items: center">
+        f"""<div style="display: flex; flex-direction: column; align-items: center;">
                 <img src="{dataurl}" />
             </div>"""
     )
@@ -326,7 +427,11 @@ def visualize_gradcam(viz_image_id):
         st.table(st.session_state.text_table_df)
 
 
-@st.dialog(" ", width="small")
+def format_vision_mode(mode_stub):
+    return f"Vision mode: {mode_stub.capitalize()}"
+
+
+@st.dialog(" ", width="large")
 def image_modal(vis_image_id):
     visualize_gradcam(vis_image_id)
 
@@ -363,7 +468,7 @@ st.markdown(
     unsafe_allow_html=True,
 )
 
-search_row = st.columns([45, 10, 13, 7, 25], vertical_alignment="center")
+search_row = st.columns([45, 5, 1, 15, 1, 8, 25], vertical_alignment="center")
 with search_row[0]:
     search_field = st.text_input(
         label="search",
@@ -379,8 +484,20 @@ with search_row[1]:
 with search_row[2]:
     st.empty()
 with search_row[3]:
-    st.markdown("**CLIP Model:**")
+    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.radio(
         "CLIP Model",
         options=["M-CLIP (multiple languages)", "J-CLIP (日本語)"],