Add text embedding, allowing input to compare with both text and image

app.py

@@ -5,7 +5,7 @@ import streamlit as st
 
 
 
-#st.set_page_config(layout="wide")
+st.set_page_config(layout="wide")
 
 st.sidebar.title("Multi-task Vision–Language AI for Pathology")
 

data/twitter.asset

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:73aa98497b501cef03980d0ff0be5b3a02ff88d377bf5513e4eca8dab0870153
-size 145886932
+oid sha256:8804057c2b910dd56a2cde6f02d317fed9dacc51e6e0ace5fa57effdf06f8c34
+size 266592030

helper.py

@@ -0,0 +1,65 @@
+import streamlit as st
+import pandas as pd
+from plip_support import embed_text
+import numpy as np
+from PIL import Image
+import requests
+import tokenizers
+import os
+from io import BytesIO
+import pickle
+import base64
+
+import torch
+from transformers import (
+    VisionTextDualEncoderModel,
+    AutoFeatureExtractor,
+    AutoTokenizer,
+    CLIPModel,
+    AutoProcessor
+)
+import streamlit.components.v1 as components
+from st_clickable_images import clickable_images #pip install st-clickable-images
+
+
+@st.cache(
+    hash_funcs={
+        torch.nn.parameter.Parameter: lambda _: None,
+        tokenizers.Tokenizer: lambda _: None,
+        tokenizers.AddedToken: lambda _: None
+    }
+)
+def load_path_clip():
+    model = CLIPModel.from_pretrained("vinid/plip")
+    processor = AutoProcessor.from_pretrained("vinid/plip")
+    return model, processor
+
+@st.cache
+def init():
+    with open('data/twitter.asset', 'rb') as f:
+        data = pickle.load(f)
+    meta = data['meta'].reset_index(drop=True)
+    image_embedding = data['image_embedding']
+    text_embedding = data['text_embedding']
+    print(meta.shape, image_embedding.shape)
+    validation_subset_index = meta['source'].values == 'Val_Tweets'
+    return meta, image_embedding, text_embedding, validation_subset_index
+
+def embed_images(model, images, processor):
+    inputs = processor(images=images)
+    pixel_values = torch.tensor(np.array(inputs["pixel_values"]))
+
+    with torch.no_grad():
+        embeddings = model.get_image_features(pixel_values=pixel_values)
+    return embeddings
+
+def embed_texts(model, texts, processor):
+    inputs = processor(text=texts, padding="longest")
+    input_ids = torch.tensor(inputs["input_ids"])
+    attention_mask = torch.tensor(inputs["attention_mask"])
+
+    with torch.no_grad():
+        embeddings = model.get_text_features(
+            input_ids=input_ids, attention_mask=attention_mask
+        )
+    return embeddings

image2image.py

@@ -1,6 +1,5 @@
 import streamlit as st
 import pandas as pd
-from plip_support import embed_text
 import numpy as np
 from PIL import Image
 import requests
@@ -21,50 +20,38 @@ from transformers import (
 import streamlit.components.v1 as components
 from st_clickable_images import clickable_images #pip install st-clickable-images
 
-
-def embed_images(model, images, processor):
-    inputs = processor(images=images)
-    pixel_values = torch.tensor(np.array(inputs["pixel_values"]))
-
-    with torch.no_grad():
-        embeddings = model.get_image_features(pixel_values=pixel_values)
-    return embeddings
-
-@st.cache
-def load_embeddings(embeddings_path):
-    print("loading embeddings")
-    return np.load(embeddings_path)
-
-@st.cache(
-    hash_funcs={
-        torch.nn.parameter.Parameter: lambda _: None,
-        tokenizers.Tokenizer: lambda _: None,
-        tokenizers.AddedToken: lambda _: None
-    }
-)
-def load_path_clip():
-    model = CLIPModel.from_pretrained("vinid/plip")
-    processor = AutoProcessor.from_pretrained("vinid/plip")
-    return model, processor
-
-def init():
-    with open('data/twitter.asset', 'rb') as f:
-        data = pickle.load(f)
-    meta = data['meta'].reset_index(drop=True)
-    image_embedding = data['embedding']
-    print(meta.shape, image_embedding.shape)
-    validation_subset_index = meta['source'].values == 'Val_Tweets'
-    return meta, image_embedding, validation_subset_index
+from helper import load_path_clip, init, embed_images
 
 
 def app():
     st.title('Image to Image Retrieval')
     st.markdown('#### A pathology image search engine that correlate images with images.')
 
-    meta, image_embedding, validation_subset_index = init()
+    meta, image_embedding, text_embedding, validation_subset_index = init()
     model, processor = load_path_clip()
 
-    st.markdown('Click following examples:')
+
+    col1, col2 = st.columns(2)
+    with col1:
+        data_options = ["All twitter data (2006-03-21 — 2023-01-15)",
+                        "Twitter validation data (2022-11-16 — 2023-01-15)"]
+        st.radio(
+            "Choose dataset for image retrieval 👉",
+            key="datapool",
+            options=data_options,
+        )
+    with col2:
+        retrieval_options = ["Image only",
+                            "Text and image (beta)",
+                             ]
+        st.radio(
+            "Similarity calcuation 👉",
+            key="calculation_option",
+            options=retrieval_options,
+        )
+
+
+    st.markdown('Try out following examples:')
     example_path = 'data/example_images'
     list_of_examples = [os.path.join(example_path, v) for v in os.listdir(example_path)]
     example_imgs = []
@@ -86,18 +73,9 @@ def app():
 
 
 
-    data_options = ["All twitter data (2006-03-21 — 2023-01-15)",
-                    "Twitter validation data (2022-11-16 — 2023-01-15)"]
-    st.radio(
-        "Or choose dataset for image retrieval 👉",
-        key="datapool",
-        options=data_options,
-    )
-
 
 
-
-    col1, col2 = st.columns(2)
+    col1, col2, _ = st.columns(3)
     with col1:
         query = st.file_uploader("Choose a file to upload")
 
@@ -113,49 +91,77 @@ def app():
         with col2:
             st.image(image, caption='Your upload')
 
-        single_image = embed_images(model, [image], processor)[0].detach().cpu().numpy()
+        input_image = embed_images(model, [image], processor)[0].detach().cpu().numpy()
 
-        single_image = single_image/np.linalg.norm(single_image)
+        input_image = input_image/np.linalg.norm(input_image)
         
         # Sort IDs by cosine-similarity from high to low
-        similarity_scores = single_image.dot(image_embedding.T)
 
+        if st.session_state.calculation_option == retrieval_options[0]: # Image only
+            similarity_scores = input_image.dot(image_embedding.T)
+        else: # Text and Image
+            similarity_scores_i = input_image.dot(image_embedding.T)
+            similarity_scores_t = input_image.dot(text_embedding.T)
+            similarity_scores_i = similarity_scores_i/np.max(similarity_scores_i)
+            similarity_scores_t = similarity_scores_t/np.max(similarity_scores_t)
+            similarity_scores = (similarity_scores_i + similarity_scores_t)/2
 
+
+        ############################################################
+        # Get top results
+        ############################################################
         topn = 5
-        if st.session_state.datapool == data_options[0]:
-            #Use all twitter data
-            id_sorted = np.argsort(similarity_scores)[::-1]
-            best_ids = id_sorted[:topn]
-            best_scores = similarity_scores[best_ids]
-            target_weblinks = meta["weblink"].values[best_ids]
-        else:
-            #Use validation twitter data
-            similarity_scores = similarity_scores[validation_subset_index]
-            # Sort IDs by cosine-similarity from high to low
-            id_sorted = np.argsort(similarity_scores)[::-1]
-            best_ids = id_sorted[:topn]
-            best_scores = similarity_scores[best_ids]
-            target_weblinks = meta["weblink"].values[validation_subset_index][best_ids]
-        #TODO: Avoid duplicated ID
+        df = pd.DataFrame(np.c_[np.arange(len(meta)), similarity_scores, meta['weblink'].values], columns = ['idx', 'score', 'twitterlink'])
+        if st.session_state.datapool == data_options[1]: #Use val twitter data
+            df = df.loc[validation_subset_index,:]
+        df = df.sort_values('score', ascending=False)
+        df = df.drop_duplicates(subset=['twitterlink'])
+        best_id_topk = df['idx'].values[:topn]
+        target_scores = df['score'].values[:topn]
+        target_weblinks = df['twitterlink'].values[:topn]
+
+
 
+        ############################################################
+        # Display results
+        ############################################################
+        
+        st.markdown('#### Top 5 results:')
         topk_options = ['1st', '2nd', '3rd', '4th', '5th']
-        st.radio(
-            "Choose the most similar  👉",
-            key="top_k",
-            options=topk_options,
-            horizontal=True
-        )
-        topn_txt = st.session_state.top_k
-        topn_value = int(st.session_state.top_k[0])-1
-        st.caption(f'The {topn_txt} relevant image (similarity = {best_scores[topn_value]:.4f})')
-        components.html('''
-            <blockquote class="twitter-tweet">
-                <a href="%s"></a>
-            </blockquote>
-            <script async src="https://platform.twitter.com/widgets.js" charset="utf-8">
-            </script>
-            ''' % target_weblinks[topn_value],
-        height=800)
+        tab = {}
+        tab[0], tab[1], tab[2] = st.columns(3)
+        for i in [0,1,2]:
+            with tab[i]:
+                topn_value = i
+                topn_txt = topk_options[i]
+                st.caption(f'The {topn_txt} relevant image (similarity = {target_scores[topn_value]:.4f})')
+                components.html('''
+                    <blockquote class="twitter-tweet">
+                        <a href="%s"></a>
+                    </blockquote>
+                    <script async src="https://platform.twitter.com/widgets.js" charset="utf-8">
+                    </script>
+                    ''' % target_weblinks[topn_value],
+                height=800)
+
+        tab[3], tab[4], tab[5] = st.columns(3)
+        for i in [3,4]:
+            with tab[i]:
+                topn_value = i
+                topn_txt = topk_options[i]
+                st.caption(f'The {topn_txt} relevant image (similarity = {target_scores[topn_value]:.4f})')
+                components.html('''
+                    <blockquote class="twitter-tweet">
+                        <a href="%s"></a>
+                    </blockquote>
+                    <script async src="https://platform.twitter.com/widgets.js" charset="utf-8">
+                    </script>
+                    ''' % target_weblinks[topn_value],
+                height=800)
+
+
+
+
 
 
 

plip_support.py

@@ -1,9 +0,0 @@
-import clip
-import torch
-
-
-
-
-def embed_text(plip, text, device="cpu"):
-    idx = clip.tokenize([text], truncate=True).to(device)
-    return plip.encode_text(idx).detach().cpu().numpy()[0]

text2image.py

@@ -1,6 +1,5 @@
 import streamlit as st
 import pandas as pd
-from plip_support import embed_text
 import numpy as np
 from PIL import Image
 import requests
@@ -17,43 +16,9 @@ from transformers import (
 )
 import streamlit.components.v1 as components
 
+from helper import load_path_clip, init, embed_texts
 
-def embed_texts(model, texts, processor):
-    inputs = processor(text=texts, padding="longest")
-    input_ids = torch.tensor(inputs["input_ids"])
-    attention_mask = torch.tensor(inputs["attention_mask"])
 
-    with torch.no_grad():
-        embeddings = model.get_text_features(
-            input_ids=input_ids, attention_mask=attention_mask
-        )
-    return embeddings
-
-@st.cache
-def load_embeddings(embeddings_path):
-    print("loading embeddings")
-    return np.load(embeddings_path)
-
-@st.cache(
-    hash_funcs={
-        torch.nn.parameter.Parameter: lambda _: None,
-        tokenizers.Tokenizer: lambda _: None,
-        tokenizers.AddedToken: lambda _: None
-    }
-)
-def load_path_clip():
-    model = CLIPModel.from_pretrained("vinid/plip")
-    processor = AutoProcessor.from_pretrained("vinid/plip")
-    return model, processor
-
-def init():
-    with open('data/twitter.asset', 'rb') as f:
-        data = pickle.load(f)
-    meta = data['meta'].reset_index(drop=True)
-    image_embedding = data['embedding']
-    print(meta.shape, image_embedding.shape)
-    validation_subset_index = meta['source'].values == 'Val_Tweets'
-    return meta, image_embedding, validation_subset_index
 
 def app():
 
@@ -61,16 +26,29 @@ def app():
     st.markdown('#### A pathology image search engine that correlate texts directly with images.')
     st.caption('Note: The searching query matches images only. The twitter text does not used for searching.')
     
-    meta, image_embedding, validation_subset_index = init()
+    meta, image_embedding, text_embedding, validation_subset_index = init()
     model, processor = load_path_clip()
 
-    data_options = ["All twitter data (2006-03-21 — 2023-01-15)",
-                    "Twitter validation data (2022-11-16 — 2023-01-15)"]
-    st.radio(
-        "Choose dataset for image retrieval 👉",
-        key="datapool",
-        options=data_options,
-    )
+
+    col1, col2 = st.columns(2)
+    with col1:
+        data_options = ["All twitter data (2006-03-21 — 2023-01-15)",
+                        "Twitter validation data (2022-11-16 — 2023-01-15)"]
+        st.radio(
+            "Choose dataset for image retrieval 👉",
+            key="datapool",
+            options=data_options,
+        )
+    with col2:
+        retrieval_options = ["Image only",
+                            "text and image (beta)",
+                             ]
+        st.radio(
+            "Similarity calcuation Mapping input with 👉",
+            key="calculation_option",
+            options=retrieval_options,
+        )
+
         
 
     col1, col2 = st.columns(2)
@@ -106,46 +84,74 @@ def app():
     else:
         query = query_2
 
-    text_embedding = embed_texts(model, [query], processor)[0].detach().cpu().numpy()
-    text_embedding = text_embedding/np.linalg.norm(text_embedding)
+
+    input_text = embed_texts(model, [query], processor)[0].detach().cpu().numpy()
+    input_text = input_text/np.linalg.norm(input_text)
+    
     
-    similarity_scores = text_embedding.dot(image_embedding.T)
+    if st.session_state.calculation_option == retrieval_options[0]: # Image only
+        similarity_scores = input_text.dot(image_embedding.T)
+    else: # Text and Image
+        similarity_scores_i = input_text.dot(image_embedding.T)
+        similarity_scores_t = input_text.dot(text_embedding.T)
+        similarity_scores_i = similarity_scores_i/np.max(similarity_scores_i)
+        similarity_scores_t = similarity_scores_t/np.max(similarity_scores_t)
+        similarity_scores = (similarity_scores_i + similarity_scores_t)/2
 
-    topn = 5
-    if st.session_state.datapool == data_options[0]:
-        #Use all twitter data
-        id_sorted = np.argsort(similarity_scores)[::-1]
-        best_ids = id_sorted[:topn]
-        best_scores = similarity_scores[best_ids]
-        target_weblinks = meta["weblink"].values[best_ids]
-    else:
-        #Use validation twitter data
-        similarity_scores = similarity_scores[validation_subset_index]
-        # Sort IDs by cosine-similarity from high to low
-        id_sorted = np.argsort(similarity_scores)[::-1]
-        best_ids = id_sorted[:topn]
-        best_scores = similarity_scores[best_ids]
-        target_weblinks = meta["weblink"].values[validation_subset_index][best_ids]
-    #TODO: Avoid duplicated ID
 
+
+
+    ############################################################
+    # Get top results
+    ############################################################
+    topn = 5
+    df = pd.DataFrame(np.c_[np.arange(len(meta)), similarity_scores, meta['weblink'].values], columns = ['idx', 'score', 'twitterlink'])
+    if st.session_state.datapool == data_options[1]: #Use val twitter data
+        df = df.loc[validation_subset_index,:]
+    df = df.sort_values('score', ascending=False)
+    df = df.drop_duplicates(subset=['twitterlink'])
+    best_id_topk = df['idx'].values[:topn]
+    target_scores = df['score'].values[:topn]
+    target_weblinks = df['twitterlink'].values[:topn]
+
+
+    ############################################################
+    # Display results
+    ############################################################
+    
+    st.markdown('Your input query: %s' % query)
+    st.markdown('#### Top 5 results:')
     topk_options = ['1st', '2nd', '3rd', '4th', '5th']
-    st.radio(
-        "Choose the most similar  👉",
-        key="top_k",
-        options=topk_options,
-        horizontal=True
-    )
-    topn_txt = st.session_state.top_k
-    topn_value = int(st.session_state.top_k[0])-1
-    st.caption(f'The {topn_txt} relevant image (similarity = {best_scores[topn_value]:.4f})')
-    components.html('''
-        <blockquote class="twitter-tweet">
-            <a href="%s"></a>
-        </blockquote>
-        <script async src="https://platform.twitter.com/widgets.js" charset="utf-8">
-        </script>
-        ''' % target_weblinks[topn_value],
-    height=800)
+    tab = {}
+    tab[0], tab[1], tab[2] = st.columns(3)
+    for i in [0,1,2]:
+        with tab[i]:
+            topn_value = i
+            topn_txt = topk_options[i]
+            st.caption(f'The {topn_txt} relevant image (similarity = {target_scores[topn_value]:.4f})')
+            components.html('''
+                <blockquote class="twitter-tweet">
+                    <a href="%s"></a>
+                </blockquote>
+                <script async src="https://platform.twitter.com/widgets.js" charset="utf-8">
+                </script>
+                ''' % target_weblinks[topn_value],
+            height=800)
+
+    tab[3], tab[4], tab[5] = st.columns(3)
+    for i in [3,4]:
+        with tab[i]:
+            topn_value = i
+            topn_txt = topk_options[i]
+            st.caption(f'The {topn_txt} relevant image (similarity = {target_scores[topn_value]:.4f})')
+            components.html('''
+                <blockquote class="twitter-tweet">
+                    <a href="%s"></a>
+                </blockquote>
+                <script async src="https://platform.twitter.com/widgets.js" charset="utf-8">
+                </script>
+                ''' % target_weblinks[topn_value],
+            height=800)