feature: add image2text feature

image2text.py

@@ -1,5 +1,6 @@
 import streamlit as st
 import numpy as np
+import jax
 import jax.numpy as jnp
 from PIL import Image
 
@@ -24,17 +25,21 @@ def app(model_name):
             st.error("Please upload an image query.")
         else:
             image = Image.open(query)
-            pixel_values = processor(
-                text=[""], images=image, return_tensors="jax", padding=True
-            ).pixel_values
-            pixel_values = jnp.transpose(pixel_values, axes=[0, 2, 3, 1])
-            vec = np.asarray(model.get_image_features(pixel_values))
-            # ids, dists = index.knnQuery(vec, k=10)
-            # result_files = map(lambda id: files[id], ids)
-            # result_imgs, result_captions = [], []
-            # for file, dist in zip(result_files, dists):
-            #     result_imgs.append(plt.imread(os.path.join(images_directory, file)))
-            #     result_captions.append("{:s} (유사도: {:.3f})".format(file, 1.0 - dist))
-
+            st.image(image)
+            # pixel_values = processor(
+            #     text=[""], images=image, return_tensors="jax", padding=True
+            # ).pixel_values
+            # pixel_values = jnp.transpose(pixel_values, axes=[0, 2, 3, 1])
+            # vec = np.asarray(model.get_image_features(pixel_values))
+            captions = captions.split(",")
+            inputs = processor(text=captions, images=image, return_tensors="jax", padding=True)
+            inputs["pixel_values"] = jnp.transpose(
+                inputs["pixel_values"], axes=[0, 2, 3, 1]
+            )
+            outputs = model(**inputs)
+            probs = jax.nn.softmax(outputs.logits_per_image, axis=1)
+
+            for idx, prob in sorted(enumerate(*probs), key=lambda x: x[1], reverse=True):
+                st.text(f"Score: `{prob}`, {captions[idx]}")