Legacy ResNet CAM visualization functionality

app.py

@@ -2,7 +2,8 @@ from base64 import b64encode
 from io import BytesIO
 from math import ceil
 
-from multilingual_clip import pt_multilingual_clip
+import clip
+from multilingual_clip import legacy_multilingual_clip, pt_multilingual_clip
 import numpy as np
 import pandas as pd
 from PIL import Image
@@ -10,13 +11,21 @@ import requests
 import streamlit as st
 import torch
 from torchvision.transforms import ToPILImage
-from transformers import AutoTokenizer, AutoModel
+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,
+    # interpret_vit,
     vit_perword_relevance,
-)  # , interpret_vit_overlapped
+    interpret_vit_overlapped,
+)
+
+from pytorch_grad_cam.grad_cam import GradCAM
 
 MAX_IMG_WIDTH = 500
 MAX_IMG_HEIGHT = 800
@@ -40,17 +49,20 @@ def find_best_matches(text_features, image_features, image_ids):
 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 (multiple languages)":
+        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)
-        else:  # model_type == "J-CLIP (日本語 only)"
+        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
@@ -69,18 +81,24 @@ def clip_search(search_query):
         # similarities = list((text_features @ photo_features.T).squeeze(0))
 
         # Sort the photos by their similarity score
-        if model_type == "M-CLIP (multiple languages)":
+        if model_type == "M-CLIP (multilingual ViT)":
             matches = find_best_matches(
                 text_features,
                 st.session_state.ml_image_features,
                 st.session_state.image_ids,
             )
-        else:  # model_type == "J-CLIP (日本語 only)"
+        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}
@@ -96,21 +114,26 @@ def load_image_features():
     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
@@ -118,6 +141,9 @@ def load_image_features():
     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()
 
@@ -129,10 +155,11 @@ def init():
     st.session_state.device = device
 
     # Load the open CLIP models
-    ml_model_name = "M-CLIP/XLM-Roberta-Large-Vit-B-16Plus"
-    ml_model_path = "./models/vit_b_16_plus_240-laion400m_e32-699c4b84.pt"
 
     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
         )
@@ -156,6 +183,17 @@ def init():
             ja_model_name, trust_remote_code=True
         )
 
+        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)
@@ -163,7 +201,7 @@ def init():
         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 (multilingual ViT)"
 
     st.session_state.vision_mode = "tiled"
     st.session_state.search_image_ids = []
@@ -194,195 +232,223 @@ def visualize_gradcam(viz_image_id):
         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 (multiple languages)":
-        img_dim = 240
-
-    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)]
+    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)
 
-    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))]
+            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
 
-    image_visualizations = []
+            # Crop the center of the image
+            image_tiles = [resized_img.crop((left, top, right, bottom))]
 
-    if st.session_state.active_model == "M-CLIP (multiple languages)":
-        # 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_visualizations = []
 
-        image_model = st.session_state.ml_image_model
-        # tokenize = st.session_state.ml_tokenizer.tokenize
-        image_model.eval()
-
-        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)
+        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
             )
 
-            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,
+            text_features = st.session_state.ml_model.forward(
+                st.session_state.search_field_value, st.session_state.ml_tokenizer
             )
 
-            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
-        )
-
-        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
-        image_model.eval()
-
-        for altered_image in image_tiles:
-            image_model.zero_grad()
+            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
+            )
 
-            p_image = (
-                st.session_state.ja_image_preprocess(altered_image)
-                .unsqueeze(0)
-                .to(st.session_state.device)
+            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
             )
 
-            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,
+            text_features = st.session_state.rn_model(
+                st.session_state.search_field_value
             )
 
-            image_visualizations.append(vis_t)
+            image_model = st.session_state.rn_image_model
 
-    transform = ToPILImage()
+            for altered_image in image_tiles:
+                p_image = (
+                    st.session_state.rn_image_preprocess(altered_image)
+                    .unsqueeze(0)
+                    .to(st.session_state.device)
+                )
 
-    vis_images = [transform(vis_t) for vis_t in image_visualizations]
+                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,
+                )
 
-    if st.session_state.vision_mode == "cropped":
-        resized_img.paste(vis_images[0], (left, top))
-        vis_images = [resized_img]
+                image_visualizations.append(vis_t)
 
-    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]
+        transform = ToPILImage()
 
-    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)
+        vis_images = [transform(vis_t) for vis_t in image_visualizations]
 
-    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)
+        if st.session_state.vision_mode == "cropped":
+            resized_img.paste(vis_images[0], (left, top))
+            vis_images = [resized_img]
 
-    else:
-        st.session_state.activations_image = vis_images[0].resize(scaled_dims)
+        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]
 
-    image_io = BytesIO()
-    st.session_state.activations_image.save(image_io, "PNG")
-    dataurl = "data:image/png;base64," + b64encode(image_io.getvalue()).decode("ascii")
+        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)
 
-    st.html(
-        f"""<div style="display: flex; flex-direction: column; align-items: center;">
-                <img src="{dataurl}" />
-            </div>"""
-    )
+        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)
 
-    info_text.empty()
+        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;">
+                    <img src="{dataurl}" />
+                </div>"""
+        )
 
-    tokenized_text = [tok for tok in tokenized_text if tok != "▁"]
+    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) < 15
+        and len(tokenized_text) < 25
         and st.button(
             "Calculate text importance (may take some time)",
         )
@@ -394,17 +460,31 @@ def visualize_gradcam(viz_image_id):
         progress_bar = st.progress(0.0, text=progress_text)
 
         for t, tok in enumerate(tokenized_text):
-            token = tok.replace("▁", "")
-            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,
-            )
+            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
@@ -429,7 +509,7 @@ def visualize_gradcam(viz_image_id):
 
 
 def format_vision_mode(mode_stub):
-    return f"Vision mode: {mode_stub.capitalize()}"
+    return mode_stub.capitalize()
 
 
 @st.dialog(" ", width="large")
@@ -469,7 +549,7 @@ st.markdown(
     unsafe_allow_html=True,
 )
 
-search_row = st.columns([45, 5, 1, 15, 1, 8, 25], vertical_alignment="center")
+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",
@@ -483,10 +563,10 @@ with search_row[1]:
         "Search", on_click=string_search, use_container_width=True, type="primary"
     )
 with search_row[2]:
-    st.empty()
+    st.markdown("**Vision mode:**")
 with search_row[3]:
     st.selectbox(
-        "Vision mode:",
+        "Vision mode",
         options=["tiled", "stretched", "cropped"],
         key="vision_mode",
         help="How to consider images that aren't square",
@@ -497,56 +577,59 @@ with search_row[3]:
 with search_row[4]:
     st.empty()
 with search_row[5]:
-    st.markdown("**CLIP Model:**")
+    st.markdown("**CLIP model:**")
 with search_row[6]:
-    st.radio(
-        "CLIP Model",
-        options=["M-CLIP (multiple languages)", "J-CLIP (日本語)"],
+    st.selectbox(
+        "CLIP Model:",
+        options=[
+            "M-CLIP (multilingual ViT)",
+            "J-CLIP (日本語 ViT)",
+            "Legacy (multilingual ResNet)",
+        ],
         key="active_model",
         on_change=string_search,
-        horizontal=True,
         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 == "M-CLIP (multiple languages)":
+if st.session_state.active_model == "J-CLIP (日本語 ViT)":
     with canned_searches[1]:
         st.button(
-            "negative space",
+            "間",
             on_click=clip_search,
-            args=["negative space"],
+            args=["間"],
             use_container_width=True,
         )
     with canned_searches[2]:
-        st.button("間", on_click=clip_search, args=["間"], use_container_width=True)
+        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)
+        st.button("山", on_click=clip_search, args=["山"], use_container_width=True)
     with canned_searches[4]:
         st.button(
-            "αρνητικός χώρος",
+            "花に酔えり 羽織着て刀 さす女",
             on_click=clip_search,
-            args=["αρνητικός χώρος"],
+            args=["花に酔えり 羽織着て刀 さす女"],
             use_container_width=True,
         )
 else:
     with canned_searches[1]:
         st.button(
-            "間",
+            "negative space",
             on_click=clip_search,
-            args=["間"],
+            args=["negative space"],
             use_container_width=True,
         )
     with canned_searches[2]:
-        st.button("奥", on_click=clip_search, args=["奥"], use_container_width=True)
+        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)
+        st.button("음각", on_click=clip_search, args=["음각"], use_container_width=True)
     with canned_searches[4]:
         st.button(
-            "花に酔えり 羽織着て刀 さす女",
+            "αρνητικός χώρος",
             on_click=clip_search,
-            args=["花に酔えり 羽織着て刀 さす女"],
+            args=["αρνητικός χώρος"],
             use_container_width=True,
         )
 

requirements.txt

@@ -1,3 +1,4 @@
+clip==1.0
 ftfy==6.2.0
 multilingual_clip==1.0.10
 numpy==1.26