feature: show prob scores as bar chart

image2text.py

@@ -1,8 +1,10 @@
 import streamlit as st
+import requests
 import numpy as np
 import jax
 import jax.numpy as jnp
 from PIL import Image
+import pandas as pd
 
 from utils import load_model
 
@@ -17,20 +19,21 @@ def app(model_name):
     """
     )
 
-    query = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
-    captions = st.text_input("사용하실 캡션을 쉼표 단위로 구분해서 적어주세요", value="고양이,강아지,느티나무...")
+    query1 = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+    query2 = st.text_input("or a URL to an image...")
+
+    captions = st.text_input(
+        "Enter candidate captions in comma-separated form.", 
+        value="귀여운 고양이,멋있는 강아지,트랜스포머"
+    )
 
     if st.button("질문 (Query)"):
-        if query is None:
-            st.error("Please upload an image query.")
+        if not any([query1, query2]):
+            st.error("Please upload an image or paste an image URL.")
         else:
-            image = Image.open(query)
+            image_data = query1 if query1 is not None else requests.get(query2, stream=True).raw
+            image = Image.open(image_data)
             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(
@@ -38,8 +41,10 @@ def app(model_name):
             )
             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]}")
+            score_dict = {captions[idx]: prob for idx, prob in enumerate(*probs)}
+            df = pd.DataFrame(score_dict.values(), index=score_dict.keys())
+            st.bar_chart(df)
+            # for idx, prob in sorted(enumerate(*probs), key=lambda x: x[1], reverse=True):
+            #     st.text(f"Score: `{prob}`, {captions[idx]}")
 
 

requirements.txt

@@ -5,4 +5,6 @@ transformers
 streamlit
 tqdm
 nmslib
-matplotlib
+matplotlib
+pandas
+requests