Upload 2 files

app.py

@@ -27,7 +27,7 @@ from CLIP_Explainability.vit_cam import (
 
 from pytorch_grad_cam.grad_cam import GradCAM
 
-RUN_LITE = False  # Load vision model for CAM viz explainability for M-CLIP only
+RUN_LITE = True  # Load vision model for CAM viz explainability for M-CLIP only
 
 MAX_IMG_WIDTH = 500
 MAX_IMG_HEIGHT = 800
@@ -58,7 +58,10 @@ def encode_search_query(search_query, model_type):
             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"
+                search_query,
+                padding=True,
+                return_tensors="pt",
+                device=st.session_state.device,
             )
             text_encoded = st.session_state.ja_model.get_text_features(**t_text)
             text_encoded /= text_encoded.norm(dim=-1, keepdim=True)
@@ -67,7 +70,7 @@ def encode_search_query(search_query, model_type):
             text_encoded /= text_encoded.norm(dim=-1, keepdim=True)
 
     # Retrieve the feature vector
-    return text_encoded
+    return text_encoded.to(st.session_state.device)
 
 
 def clip_search(search_query):
@@ -153,7 +156,9 @@ def load_image_features():
 def init():
     st.session_state.current_page = 1
 
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+    # device = "cuda" if torch.cuda.is_available() else "cpu"
+    device = "cpu"
+
     st.session_state.device = device
 
     # Load the open CLIP models
@@ -168,7 +173,7 @@ def init():
 
         st.session_state.ml_model = (
             pt_multilingual_clip.MultilingualCLIP.from_pretrained(ml_model_name)
-        )
+        ).to(device)
         st.session_state.ml_tokenizer = AutoTokenizer.from_pretrained(ml_model_name)
 
         ja_model_name = "hakuhodo-tech/japanese-clip-vit-h-14-bert-wider"
@@ -193,7 +198,7 @@ def init():
 
         st.session_state.rn_model = legacy_multilingual_clip.load_model(
             "M-BERT-Base-69"
-        )
+        ).to(device)
         st.session_state.rn_tokenizer = BertTokenizer.from_pretrained(
             "bert-base-multilingual-cased"
         )
@@ -210,7 +215,6 @@ def init():
     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..."):
@@ -221,233 +225,271 @@ 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
+def get_overlay_vis(image, img_dim, image_model):
+    orig_img_dims = image.size
 
-    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()
+    ##### If the features are based on tiled image slices
+    tile_behavior = None
 
-    st.markdown(
-        f"**Query text:** {st.session_state.search_field_value} | **Image relevance:** {round(st.session_state.search_image_scores[viz_image_id], 3)}"
-    )
+    if st.session_state.vision_mode == "tiled":
+        scaled_dims = [img_dim, img_dim]
 
-    with st.spinner("Calculating..."):
-        # info_text = st.text("Calculating activation regions...")
+        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)
 
-        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")
+        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))
 
-        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)
+        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
 
-            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))]
 
-            # Crop the center of the image
-            image_tiles = [resized_img.crop((left, top, right, bottom))]
+    image_visualizations = []
+    image_features = []
+    image_similarities = []
 
-        image_visualizations = []
+    if st.session_state.active_model == "M-CLIP (multilingual ViT)":
+        text_features = st.session_state.ml_model.forward(
+            st.session_state.search_field_value, st.session_state.ml_tokenizer
+        )
+
+        if st.session_state.device == "cpu":
+            text_features = text_features.float().to(st.session_state.device)
+        else:
+            text_features = text_features.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
+        for altered_image in image_tiles:
+            p_image = (
+                st.session_state.ml_image_preprocess(altered_image)
+                .unsqueeze(0)
+                .to(st.session_state.device)
             )
 
-            text_features = st.session_state.ml_model.forward(
-                st.session_state.search_field_value, st.session_state.ml_tokenizer
+            vis_t, img_feats, similarity = interpret_vit_overlapped(
+                p_image.type(image_model.dtype),
+                text_features.type(image_model.dtype),
+                image_model.visual,
+                st.session_state.device,
+                img_dim=img_dim,
             )
 
-            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
+            image_visualizations.append(vis_t)
+            image_features.append(img_feats)
+            image_similarities.append(similarity.item())
+
+    elif st.session_state.active_model == "J-CLIP (日本語 ViT)":
+        t_text = st.session_state.ja_tokenizer(
+            st.session_state.search_field_value,
+            return_tensors="pt",
+            device=st.session_state.device,
+        )
+
+        text_features = st.session_state.ja_model.get_text_features(**t_text)
+
+        if st.session_state.device == "cpu":
+            text_features = text_features.float().to(st.session_state.device)
+        else:
+            text_features = text_features.to(st.session_state.device)
+
+        for altered_image in image_tiles:
+            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"
+            vis_t, img_feats, similarity = interpret_vit_overlapped(
+                p_image.type(image_model.dtype),
+                text_features.type(image_model.dtype),
+                image_model.visual,
+                st.session_state.device,
+                img_dim=img_dim,
             )
-            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
+
+            image_visualizations.append(vis_t)
+            image_features.append(img_feats)
+            image_similarities.append(similarity.item())
+
+    else:  # st.session_state.active_model == Legacy
+        text_features = st.session_state.rn_model(st.session_state.search_field_value)
+
+        if st.session_state.device == "cpu":
+            text_features = text_features.float().to(st.session_state.device)
+        else:
+            text_features = text_features.to(st.session_state.device)
+
+        for altered_image in image_tiles:
+            p_image = (
+                st.session_state.rn_image_preprocess(altered_image)
+                .unsqueeze(0)
+                .to(st.session_state.device)
             )
 
-            text_features = st.session_state.rn_model(
-                st.session_state.search_field_value
+            vis_t = interpret_rn_overlapped(
+                p_image.type(image_model.dtype),
+                text_features.type(image_model.dtype),
+                image_model.visual,
+                GradCAM,
+                st.session_state.device,
+                img_dim=img_dim,
             )
 
-            image_model = st.session_state.rn_image_model
+            text_features_norm = text_features.norm(dim=-1, keepdim=True)
+            text_features_new = text_features / text_features_norm
 
-            for altered_image in image_tiles:
-                p_image = (
-                    st.session_state.rn_image_preprocess(altered_image)
-                    .unsqueeze(0)
-                    .to(st.session_state.device)
-                )
+            image_feats = image_model.encode_image(p_image.type(image_model.dtype))
+            image_feats_norm = image_feats.norm(dim=-1, keepdim=True)
+            image_feats_new = image_feats / image_feats_norm
 
-                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,
-                )
+            similarity = image_feats_new[0].dot(text_features_new[0])
 
-                image_visualizations.append(vis_t)
+            image_visualizations.append(vis_t)
+            image_features.append(p_image)
+            image_similarities.append(similarity.item())
 
-        transform = ToPILImage()
+    transform = ToPILImage()
 
-        vis_images = [transform(vis_t) for vis_t in image_visualizations]
+    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 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 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)
+    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))
+        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)
+    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))
+        activations_image = vis_image.resize(scaled_dims)
 
-        else:
-            st.session_state.activations_image = vis_images[0].resize(scaled_dims)
+    else:
+        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"
+    return activations_image, image_features, np.mean(image_similarities)
+
+
+def visualize_gradcam(image):
+    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()
+
+    if st.session_state.active_model == "M-CLIP (multilingual ViT)":
+        img_dim = 240
+        image_model = st.session_state.ml_image_model
+        # Sometimes used for token importance viz
+        tokenized_text = st.session_state.ml_tokenizer.tokenize(
+            st.session_state.search_field_value
+        )
+    elif st.session_state.active_model == "Legacy (multilingual ResNet)":
+        img_dim = 288
+        image_model = st.session_state.rn_image_model
+        # Sometimes used for token importance viz
+        tokenized_text = st.session_state.rn_tokenizer.tokenize(
+            st.session_state.search_field_value
+        )
+    else:  # J-CLIP
+        img_dim = 224
+        image_model = st.session_state.ja_image_model
+        # Sometimes used for token importance viz
+        tokenized_text = st.session_state.ja_tokenizer.tokenize(
+            st.session_state.search_field_value
         )
 
-        st.html(
-            f"""<div style="display: flex; flex-direction: column; align-items: center;">
-                    <img src="{dataurl}" />
-                </div>"""
+    with st.spinner("Calculating..."):
+        # info_text = st.text("Calculating activation regions...")
+
+        activations_image, image_features, similarity_score = get_overlay_vis(
+            image, img_dim, image_model
         )
 
-    tokenized_text = [tok.replace("▁", "") for tok in tokenized_text if tok != "▁"]
+        st.markdown(
+            f"**Query text:** {st.session_state.search_field_value} | **Approx. image relevance:** {round(similarity_score.item(), 3)}"
+        )
+
+        st.image(activations_image)
+
+        # image_io = BytesIO()
+        # 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.replace("▁", "").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"]
+        tok
+        for tok in tokenized_text
+        if tok
+        not in ["s", "ed", "a", "the", "an", "ing", "て", "に", "の", "は", "と", "た"]
     ]
 
     if (
@@ -457,8 +499,7 @@ def visualize_gradcam(viz_image_id):
             "Calculate text importance (may take some time)",
         )
     ):
-        search_tokens = []
-        token_scores = []
+        scores_per_token = {}
 
         progress_text = f"Processing {len(tokenized_text)} text tokens"
         progress_bar = st.progress(0.0, text=progress_text)
@@ -466,34 +507,37 @@ def visualize_gradcam(viz_image_id):
         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)
+            for img_feats in image_features:
+                if st.session_state.active_model == "Legacy (multilingual ResNet)":
+                    word_rel = rn_perword_relevance(
+                        img_feats,
+                        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(
+                        img_feats,
+                        st.session_state.search_field_value,
+                        image_model,
+                        tokenize,
+                        st.session_state.device,
+                        token,
+                        img_dim=img_dim,
+                    )
+                avg_score = np.mean(word_rel)
+                if avg_score == 0 or np.isnan(avg_score):
+                    continue
+
+                if token not in scores_per_token:
+                    scores_per_token[token] = [1 / avg_score]
+                else:
+                    scores_per_token[token].append(1 / avg_score)
 
             progress_bar.progress(
                 (t + 1) / len(tokenized_text),
@@ -501,24 +545,48 @@ def visualize_gradcam(viz_image_id):
             )
         progress_bar.empty()
 
-        normed_scores = torch.softmax(torch.tensor(token_scores), dim=0)
+        avg_scores_per_token = [
+            np.mean(scores_per_token[tok]) for tok in list(scores_per_token.keys())
+        ]
+
+        normed_scores = torch.softmax(torch.tensor(avg_scores_per_token), 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}
+            {"token": list(scores_per_token.keys()), "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(image):
+    visualize_gradcam(image)
 
 
-@st.dialog(" ", width="large")
-def image_modal(vis_image_id):
-    visualize_gradcam(vis_image_id)
+def vis_known_image(vis_image_id):
+    image_url = st.session_state.images_info.loc[vis_image_id]["image_url"]
+    image_response = requests.get(image_url)
+    image = Image.open(BytesIO(image_response.content), formats=["JPEG", "GIF", "PNG"])
+    image = image.convert("RGB")
+
+    image_modal(image)
+
+
+def vis_uploaded_image():
+    uploaded_file = st.session_state.uploaded_image
+    if uploaded_file is not None:
+        # To read file as bytes:
+        bytes_data = uploaded_file.getvalue()
+        image = Image.open(BytesIO(bytes_data), formats=["JPEG", "GIF", "PNG"])
+        image = image.convert("RGB")
+
+        image_modal(image)
+
+
+def format_vision_mode(mode_stub):
+    return mode_stub.capitalize()
 
 
 st.title("Explore Japanese visual aesthetics with CLIP models")
@@ -637,7 +705,7 @@ else:
             use_container_width=True,
         )
 
-controls = st.columns([35, 5, 35, 5, 20], gap="large", vertical_alignment="center")
+controls = st.columns([25, 25, 20, 35], gap="large", vertical_alignment="center")
 with controls[0]:
     im_per_pg = st.columns([30, 70], vertical_alignment="center")
     with im_per_pg[0]:
@@ -647,8 +715,6 @@ with controls[0]:
             "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:**")
@@ -657,9 +723,7 @@ with controls[2]:
             "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]:
+with controls[2]:
     pager = st.columns([40, 60], vertical_alignment="center")
     with pager[0]:
         st.markdown(f"Page **{st.session_state.current_page}** of **{num_batches}** ")
@@ -672,6 +736,14 @@ with controls[4]:
             label_visibility="collapsed",
             key="current_page",
         )
+with controls[3]:
+    st.file_uploader(
+        "Upload an image",
+        type=["jpg", "jpeg", "gif", "png"],
+        key="uploaded_image",
+        label_visibility="collapsed",
+        on_change=vis_uploaded_image,
+    )
 
 
 if len(st.session_state.search_image_ids) == 0:
@@ -708,7 +780,7 @@ for image_id in batch:
         if not RUN_LITE or st.session_state.active_model == "M-CLIP (multilingual ViT)":
             st.button(
                 "Explain this",
-                on_click=image_modal,
+                on_click=vis_known_image,
                 args=[image_id],
                 use_container_width=True,
                 key=image_id,

requirements.txt

@@ -1,5 +1,6 @@
 clip @ git+https://github.com/openai/CLIP.git
 ftfy==6.2.0
+matplotlib==3.8.1
 multilingual_clip==1.0.10
 numpy==1.26
 opencv-python==4.10.0.84
@@ -7,7 +8,6 @@ pandas==2.1.2
 pillow==10.1.0
 requests==2.31.0
 sentencepiece==0.2.0
-streamlit
 torch==2.4.0
 torchvision==0.19.0
 transformers==4.35.0