Add documentation

app.py

@@ -3,7 +3,7 @@ import pandas as pd
 import numpy as np
 import os
 import matplotlib.pyplot as plt
-from transformers import AutoTokenizer, CLIPProcessor, ViTFeatureExtractor
+from transformers import AutoTokenizer, CLIPProcessor
 from medclip.modeling_hybrid_clip import FlaxHybridCLIP
 
 @st.cache(allow_output_mutation=True)
@@ -14,7 +14,7 @@ def load_model():
 
 @st.cache(allow_output_mutation=True)
 def load_image_embeddings():
-    embeddings_df = pd.read_pickle('image_embeddings.pkl')
+    embeddings_df = pd.read_pickle('feature_store/image_embeddings.pkl')
     image_embeds = np.stack(embeddings_df['image_embedding'])
     image_files = np.asarray(embeddings_df['files'].tolist())
     return image_files, image_embeds
@@ -24,7 +24,11 @@ image_list, image_embeddings = load_image_embeddings()
 model, processor = load_model()
 img_dir = './images'
 
-query = st.text_input("Search:")
+st.title("MedCLIP 🩺📎")
+st.markdown("Search for medical images in natural language.")
+st.markdown("""This demo uses a CLIP model finetuned on the
+ [Radiology Objects in COntext (ROCO) dataset](https://github.com/razorx89/roco-dataset).""")
+query = st.text_input("Enter your query here:")
 
 if st.button("Search"):
     st.write(f"Searching our image database for {query}...")
@@ -36,8 +40,6 @@ 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}")
     #show images
     
     for img_path in matching_images:

tools/create_embeddings.py

@@ -0,0 +1,39 @@
+import os
+import jax
+
+from transformers import AutoTokenizer, CLIPProcessor
+from configuration_hybrid_clip import HybridCLIPConfig
+from modeling_hybrid_clip import FlaxHybridCLIP
+from PIL import Image
+
+import matplotlib.pyplot as plt
+import torch
+import torchvision
+from torchvision.transforms.functional import InterpolationMode
+from torchvision.transforms import Resize, Normalize, ConvertImageDtype, ToTensor
+import numpy as np
+import pandas as pd
+
+
+def main():
+    model = FlaxHybridCLIP.from_pretrained("flax-community/medclip-roco")
+    vision_model_name = "openai/clip-vit-base-patch32"
+    img_dir = "/Users/kaumad/Documents/coding/hf-flax/demo/medclip-roco/images"
+
+    processor = CLIPProcessor.from_pretrained(vision_model_name)
+
+    img_list = os.listdir(img_dir)
+    embeddings = []
+
+    for idx, img_path in enumerate(img_list):
+        if idx % 10 == 0:
+            print(f"{idx} images processed")
+        img = Image.open(os.path.join(img_dir, img_path)).convert('RGB')
+        inputs = processor(images=img, return_tensors="jax", padding=True)
+        inputs['pixel_values'] = inputs['pixel_values'].transpose(0, 2, 3, 1)
+        img_vec = model.get_image_features(**inputs)
+        img_vec = np.array(img_vec).reshape(-1).tolist()
+        embeddings.append(img_vec)
+        
+if __name__=='__main__':
+    main()