:construction: update app

requirements.txt

@@ -6,4 +6,4 @@ Pillow~=9.0.1
 jax
 jaxlib
 flax
-Jinja2
+Jinja2==3.0.1

text_to_image.py

@@ -34,17 +34,17 @@ def load_captions(caption_file):
     return image2caption
 
 
-def get_image(text):
+def get_image(text, number):
     model = FlaxCLIPModel.from_pretrained("flax-community/clip-rsicd-v2")
     processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
     filename, index = load_index("./vectors/test-bs128x8-lr5e-6-adam-ckpt-1.tsv")
     image2caption = load_captions("./images/test-captions.json")
 
-    inputs = processor(text=[text], image=None, return_tensors="jax", padding=True)
+    inputs = processor(text=[text], images=None, return_tensors="jax", padding=True)
 
     vector = model.get_text_features(**inputs)
     vector = np.asarray(vector)
-    ids, distances = index.knnQuery(vector, k=10)
+    ids, distances = index.knnQuery(vector, k=number)
     result_filenames = [filename[id] for id in ids]
     for rank, (result_filename, score) in enumerate(zip(result_filenames, distances)):
         caption = "{:s} (score: {:.3f})".format(result_filename, 1.0 - score)
@@ -56,8 +56,8 @@ def get_image(text):
         for caption in image2caption[result_filename]:
             caption_text.append("* {:s}".format(caption))
             col3.markdown("".join(caption_text))
-            st.markdown("---")
-            suggest_idx = -1
+        st.markdown("---")
+        suggest_idx = -1
 
 
 def app():
@@ -65,7 +65,8 @@ def app():
     st.text("You want search an image with given text.")
 
     text = st.text_input("Enter text: ")
+    number = st.number_input("Enter number of images result: ", min_value=1, max_value=10)
 
     if st.button("Search"):
-        get_image(text)
+        get_image(text, number)