Add image query

app.py

@@ -310,45 +310,42 @@ def encode_image(image_filename, model):
     return image_embedding
 
 
-st.title("いらすと検索（日本語CLIPゼロショット）")
-description_text = st.empty()
-
-if "model" not in st.session_state:
+@st.cache
+def load_model_and_data():
     description_text.text("日本語CLIPモデル読み込み中... ")
     device = "cuda" if torch.cuda.is_available() else "cpu"
     model = ClipModel("sonoisa/clip-vit-b-32-japanese-v1", device=device)
-    st.session_state.model = model
 
-    print("extract dataset")
     pyminizip.uncompress(
         "clip_zeroshot_irasuto_items_20210224.pq.zip", st.secrets["ZIP_PASSWORD"], None, 1
     )
     
-    print("loading dataset")
     df = pq.read_table("clip_zeroshot_irasuto_items_20210224.parquet",
         columns=["page", "description", "image_url", "sentence_vector", "image_vector"]).to_pandas()
     
     sentence_vectors = np.stack(df["sentence_vector"])
     image_vectors = np.stack(df["image_vector"])
 
-    st.session_state.df = df
-    st.session_state.sentence_vectors = sentence_vectors
-    st.session_state.image_vectors = image_vectors
+    return model, df, sentence_vectors, image_vectors
 
-    print("finished loading model and dataset")
-    
-model = st.session_state.model
-df = st.session_state.df
-sentence_vectors = st.session_state.sentence_vectors
-image_vectors = st.session_state.image_vectors
+
+st.title("いらすと検索（日本語CLIPゼロショット）")
+description_text = st.empty()
+
+model, df, sentence_vectors, image_vectors = load_model_and_data()
 
 description_text.text("日本語CLIPモデル（ゼロショット）を用いて、説明文の意味が近い「いらすとや」画像を検索します。\nキーワードを列挙するよりも、自然な文章を入力した方が精度よく検索できます。\n画像は必ずリンク先の「いらすとや」さんのページを開き、そこからダウンロードしてください。")
 
 def clear_result():
     result_text.text("")
 
+query_type = st.radio(label="クエリ種別", options=("説明文", "画像"))
+
+col1, col2 = st.columns(2)
+
 prev_query = ""
-query_input = st.text_input(label="説明文", value="", on_change=clear_result)
+query_input = col1.text_input(label="説明文", value="", on_change=clear_result, use_column_width=True)
+query_image = col2.file_uploader(label="画像", type=["png", "jpg", "jpeg"], on_change=clear_result, use_column_width=True)
 
 closest_n = st.number_input(label="検索数", min_value=1, value=10, max_value=100)