add selective db / feat / lang

app.py

@@ -3,7 +3,7 @@ import numpy as np
 import base64
 from io import BytesIO
 from multilingual_clip import pt_multilingual_clip
-from transformers import CLIPTokenizerFast, AutoTokenizer
+from transformers import CLIPTokenizerFast, AutoTokenizer, CLIPModel
 import torch
 import logging
 from os import environ
@@ -12,30 +12,22 @@ environ['TOKENIZERS_PARALLELISM'] = 'true'
 
 
 db_name_map = {
-    "Unsplash Photos 25K": "mqdb_demo.unsplash_25k_clip_indexer",
-    "RSICD: Remote Sensing Images 11K": "mqdb_demo.rsicd_clip_b_32",
+    "Unsplash Photos 25K": lambda feat: f"mqdb_demo.unsplash_25k_{feat}_indexer",
+    "RSICD: Remote Sensing Images 11K": lambda feat: f"mqdb_demo.rsicd_{feat}_b_32",
 }
+feat_name_map = {
+    'Vanilla CLIP': "clip",
+    'CLIP finetuned on RSICD': "cliprsicd"
+}
+
 
 DB_NAME = "mqdb_demo.unsplash_25k_clip_indexer"
-MODEL_ID = 'M-CLIP/XLM-Roberta-Large-Vit-B-32'
 DIMS = 512
 # Ignore some bad links (broken in the dataset already)
 BAD_IDS = {'9_9hzZVjV8s', 'RDs0THr4lGs', 'vigsqYux_-8',
            'rsJtMXn3p_c', 'AcG-unN00gw', 'r1R_0ZNUcx0'}
 
 
-@st.experimental_singleton(show_spinner=False)
-def init_clip():
-    """ Initialize CLIP Model
-
-    Returns:
-        Tokenizer: CLIPTokenizerFast (which convert words into embeddings)
-    """
-    clip = pt_multilingual_clip.MultilingualCLIP.from_pretrained(MODEL_ID)
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-    return tokenizer, clip
-
-
 @st.experimental_singleton(show_spinner=False)
 def init_db():
     """ Initialize the Database Connection
@@ -82,15 +74,15 @@ def query(xq, top_k=10):
                 # Using PREWHERE allows you to do column filter before vector search
                 xc = st.session_state.index.fetch(f"SELECT id, url, vector,\
                         distance('topK={top_k}')(vector, {xq_s}) AS dist\
-                        FROM {db_name_map[st.session_state.db_name_ref]} \
+                        FROM {db_name_map[st.session_state.db_name_ref](feat_name_map[st.session_state.feat_name])} \
                         PREWHERE id NOT IN ({exclude_list})")
             else:
                 xc = st.session_state.index.fetch(f"SELECT id, url, vector,\
                         distance('topK={top_k}')(vector, {xq_s}) AS dist\
-                        FROM {db_name_map[st.session_state.db_name_ref]}")
+                        FROM {db_name_map[st.session_state.db_name_ref](feat_name_map[st.session_state.feat_name])}")
             real_xc = st.session_state.index.fetch(f"SELECT id, url, vector,\
                         distance('topK={top_k}')(vector, {xq_s}) AS dist\
-                        FROM {db_name_map[st.session_state.db_name_ref]}")
+                        FROM {db_name_map[st.session_state.db_name_ref](feat_name_map[st.session_state.feat_name])}")
             top_k = real_xc
             xc = [xi for xi in xc if xi['id'] not in st.session_state.meta or
                   st.session_state.meta[xi['id']] < 1]
@@ -166,38 +158,6 @@ class NormalizingLayer(torch.nn.Module):
         return x / torch.norm(x, dim=-1, keepdim=True)
 
 
-def prompt2vec(prompt: str):
-    """ Convert prompt into a computational vector
-
-    Args:
-        prompt (str): Text to be tokenized
-
-    Returns:
-        xq: vector from the tokenizer, representing the original prompt
-    """
-    # inputs = tokenizer(prompt, return_tensors='pt')
-    # out = clip.get_text_features(**inputs)
-    out = clip.forward(prompt, tokenizer)
-    xq = out.squeeze(0).cpu().detach().numpy().tolist()
-    return xq
-
-
-def pil_to_bytes(img):
-    """ Convert a Pillow image into base64
-
-    Args:
-        img (PIL.Image): Pillow (PIL) Image
-
-    Returns:
-        img_bin: image in base64 format
-    """
-    with BytesIO() as buf:
-        img.save(buf, format='jpeg')
-        img_bin = buf.getvalue()
-        img_bin = base64.b64encode(img_bin).decode('utf-8')
-    return img_bin
-
-
 def card(i, url):
     return f'<img id="img{i}" src="{url}" width="200px;">'
 
@@ -286,6 +246,63 @@ def delete_element(element):
     del element
 
 
+@st.experimental_singleton(show_spinner=False)
+def init_clip_mlang():
+    """ Initialize CLIP Model
+
+    Returns:
+        Tokenizer: CLIPTokenizerFast (which convert words into embeddings)
+    """
+    MODEL_ID = 'M-CLIP/XLM-Roberta-Large-Vit-B-32'
+    clip = pt_multilingual_clip.MultilingualCLIP.from_pretrained(MODEL_ID)
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+    return tokenizer, clip
+
+@st.experimental_singleton(show_spinner=False)
+def init_clip_vanilla():
+    """ Initialize CLIP Model
+
+    Returns:
+        Tokenizer: CLIPTokenizerFast (which convert words into embeddings)
+    """
+    MODEL_ID = "openai/clip-vit-base-patch32"
+    tokenizer = CLIPTokenizerFast.from_pretrained(MODEL_ID)
+    clip = CLIPModel.from_pretrained(MODEL_ID)
+    return tokenizer, clip
+
+
+@st.experimental_singleton(show_spinner=False)
+def init_clip_rsicd():
+    """ Initialize CLIP Model
+
+    Returns:
+        Tokenizer: CLIPTokenizerFast (which convert words into embeddings)
+    """
+    MODEL_ID = "flax-community/clip-rsicd"
+    tokenizer = CLIPTokenizerFast.from_pretrained(MODEL_ID)
+    clip = CLIPModel.from_pretrained(MODEL_ID)
+    return tokenizer, clip
+
+
+def prompt2vec_mlang(prompt: str, tokenizer, clip):
+    """ Convert prompt into a computational vector
+
+    Args:
+        prompt (str): Text to be tokenized
+
+    Returns:
+        xq: vector from the tokenizer, representing the original prompt
+    """
+    out = clip.forward(prompt, tokenizer)
+    xq = out.squeeze(0).cpu().detach().numpy().tolist()
+    return xq
+
+def prompt2vec_vanilla(prompt: str, tokenizer, clip):
+    inputs = tokenizer(prompt, return_tensors='pt')
+    out = clip.get_text_features(**inputs)
+    xq = out.squeeze(0).cpu().detach().numpy().tolist()
+    return xq 
+
 st.markdown("""
 <link
   rel="stylesheet"
@@ -323,13 +340,23 @@ messages = [
     """
 ]
 
+text_model_map = {
+    'Multi Lingual': {'Vanilla CLIP': [prompt2vec_mlang, ]},
+    'English': {'Vanilla CLIP': [prompt2vec_vanilla, ],
+                'CLIP finetuned on RSICD': [prompt2vec_vanilla, ],
+                }
+    }
+
+
 with st.spinner("Connecting DB..."):
     st.session_state.meta, st.session_state.index = init_db()
 
 with st.spinner("Loading Models..."):
     # Initialize CLIP model
     if 'xq' not in st.session_state:
-        tokenizer, clip = init_clip()
+        text_model_map['Multi Lingual']['Vanilla CLIP'].append(init_clip_mlang())
+        text_model_map['English']['Vanilla CLIP'].append(init_clip_vanilla())
+        text_model_map['English']['CLIP finetuned on RSICD'].append(init_clip_rsicd())
         st.session_state.query_num = 0
 
 if 'xq' not in st.session_state:
@@ -347,8 +374,15 @@ if 'xq' not in st.session_state:
         start = [st.empty(), st.empty(), st.empty(), st.empty(),
                  st.empty(), st.empty(), st.empty()]
         start[0].info(msg)
-        st.session_state.db_name_ref = start[1].selectbox(
-            "Select Database:", list(db_name_map.keys()))
+        start_col = start[1].columns(3)
+        st.session_state.db_name_ref = start_col[0].selectbox("Select Database:", list(db_name_map.keys()))
+        st.session_state.lang = start_col[1].selectbox("Select Language:", list(text_model_map.keys()))
+        st.session_state.feat_name = start_col[2].selectbox("Select Image Feature:", 
+                                                            list(text_model_map[st.session_state.lang].keys()))
+        if st.session_state.db_name_ref == "RSICD: Remote Sensing Images 11K":
+            st.warning('If you are searching for Remote Sensing Images, \
+                        try to use prompt "An aerial photograph of <your-real-query>" \
+                        to obtain best search experience!')
         prompt = start[2].text_input(
             "Prompt:", value="", placeholder="Examples: playing corgi, 女人举着雨伞, mouette volant au-dessus de la mer, ガラスの花瓶の花 ...")
         if len(prompt) > 0:
@@ -388,7 +422,8 @@ if 'xq' not in st.session_state:
         else:
             print(f"Input prompt is {prompt}")
             # Tokenize the vectors
-            xq = prompt2vec(prompt)
+            p2v_func, args = text_model_map[st.session_state.lang][st.session_state.feat_name]
+            xq = p2v_func(prompt, *args)
         st.session_state.xq = xq
         st.session_state.orig_xq = xq
         _ = [elem.empty() for elem in start]