text improvements, CC is the default dataset, check if URLS can be loaded before displaying, small fixes

image2text.py

@@ -21,7 +21,7 @@ def app():
 
         🤌 Italian mode on! 🤌
         
-        For example, try to write "gatto" (cat) in the space for label1 and "dog" (cane) in the space for label2 and the run
+        For example, try to write "gatto" (cat) in the space for label1 and "cane" (dog) in the space for label2 and the run
         "classify"!
 
         """
@@ -62,7 +62,7 @@ def app():
                     text_embeds.extend(text_encoder(c, model, tokenizer))
 
                 text_embeds = jnp.array(text_embeds)
-                image_raw = requests.get(image_url, stream=True).raw
+                image_raw = requests.get(image_url, stream=True,).raw
 
                 image = Image.open(image_raw).convert("RGB")
                 transform = get_image_transform(model.config.vision_config.image_size)
@@ -81,5 +81,9 @@ def app():
 
                 with col2:
                     st.image(image)
-
         gc.collect()
+
+    elif image_url:
+        image_raw = requests.get(image_url, stream=True,).raw
+        image = Image.open(image_raw).convert("RGB")
+        st.image(image)

text2image.py

@@ -4,6 +4,8 @@ import requests
 import zipfile
 import natsort
 import gc
+from PIL import Image
+from PIL import UnidentifiedImageError
 
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 from stqdm import stqdm
@@ -96,6 +98,13 @@ def get_image_transform(image_size):
     )
 
 
+headers = {
+    #'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/71.0.3578.98 Safari/537.36',
+    'User-Agent': 'Googlebot-Image/1.0',  # Pretend to be googlebot
+    'X-Forwarded-For': '64.18.15.200'
+}
+
+
 def app():
 
     st.title("From Text to Image")
@@ -104,9 +113,9 @@ def app():
     
         ### 👋 Ciao!
 
-        Here you can search for images in the Unsplash 25k Photos dataset and the Conceptual Caption dataset.
-        You will see most queries make sense. When you see errors, there might be two possibilities: the model is answering
-        in a wrong way or the image you are looking for and the model is giving you the best answer it can get.
+        Here you can search for ~150.000 images in the Conceptual Captions dataset (CC) or in the Unsplash 25k Photos dataset.
+        Even though we did not train on any of these images you will see most queries make sense. When you see errors, there might be two possibilities: 
+        the model is answering in a wrong way or the image you are looking for are not in the dataset and the model is giving you the best answer it can get.
         
         
         
@@ -143,7 +152,7 @@ def app():
     with col1:
         query = st.text_input("... or insert an Italian query text")
     with col2:
-        dataset_name = st.selectbox("IR dataset", ["Unsplash", "CC"])
+        dataset_name = st.selectbox("IR dataset", ["CC", "Unsplash"])
 
     query = suggestions[sugg_idx] if sugg_idx > -1 else query if query else ""
 
@@ -167,11 +176,21 @@ def app():
             else:
                 raise ValueError()
 
+            N = 3
+
             image_paths = utils.find_image(
-                query, model, dataset, tokenizer, image_features, 1, dataset_name
+                query, model, dataset, tokenizer, image_features, N, dataset_name
             )
 
-        st.image(image_paths)
+        for i, image_url in enumerate(image_paths):
+            try:
+                image_raw = requests.get(image_url, stream=True, allow_redirects=True, headers=headers).raw
+                image = Image.open(image_raw).convert("RGB")
+                st.image(image)
+                break
+            except (UnidentifiedImageError) as e:
+                if i == N - 1:
+                    st.text(f'Tried to show  {N} different image URLS but none of them were reachabele.\nMaybe try a different query?')
 
         gc.collect()