Show matching images

app.py

@@ -1,6 +1,8 @@
 import streamlit as st
 import pandas as pd
 import numpy as np
+import os
+import matplotlib.pyplot as plt
 from transformers import AutoTokenizer, CLIPProcessor, ViTFeatureExtractor
 from medclip.modeling_hybrid_clip import FlaxHybridCLIP
 
@@ -17,10 +19,10 @@ def load_image_embeddings():
     image_files = np.asarray(embeddings_df['files'].tolist())
     return image_files, image_embeds
 
-# def app():
 k = 5
 image_list, image_embeddings = load_image_embeddings()
 model, processor = load_model()
+img_dir = './images'
 
 query = st.text_input("Search:")
 
@@ -34,5 +36,12 @@ if st.button("Search"):
     query_embedding = query_embedding / np.linalg.norm(query_embedding, axis=-1, keepdims=True)
     dot_prod = np.sum(np.multiply(query_embedding, image_embeddings), axis=1)
     matching_images = image_list[dot_prod.argsort()[-k:]]
-    st.write(f"matching images: {matching_images}")
+
+    # st.write(f"matching images: {matching_images}")
+    #show images
+    
+    for img_path in matching_images:
+        img = plt.imread(os.path.join(img_dir, img_path))
+        st.write(img_path)
+        st.image(img)
         

requirements.txt

@@ -5,5 +5,6 @@ streamlit==0.84.1
 torch==1.9.0
 torchvision==0.10.0
 pandas==1.3.0
+matplotlib>=3.4.0
 transformers==4.8.2
 watchdog==2.1.3