Update app.py

app.py

@@ -32,16 +32,52 @@ def setup_user():
         st.session_state.neg = 0
         st.session_state.pos = 0
 
+import requests
+
+def check_image_url_accessible(url):
+    try:
+        # Send a HEAD request to save bandwidth
+        response = requests.head(url, allow_redirects=True, timeout=5)
+        # If the HEAD request fails, fallback to GET request
+        if response.status_code != 200:
+            response = requests.get(url, stream=True, timeout=5)
+        
+        # Check the status code
+        if response.status_code == 200:
+            # Verify if it's an image
+            content_type = response.headers.get("Content-Type", "")
+            if "image" in content_type:
+                return True
+            else:
+                return False
+        else:
+            return False
+    except requests.RequestException:
+        return False
+
+
+
 def get_filename():
     if 'preds' in st.session_state:
         p = st.session_state.preds**4
         p /= sum(p)
         choice = np.random.choice(range(len(p)),p=p)
         st.session_state.pred = st.session_state.preds[choice]
-        return embeddings.index[choice]
+        url = get_s3_url(embeddings.index[choice])
+        if check_image_url_accessible(url):
+            st.toast(url)
+            return embeddings.index[choice]
+        else:
+            return get_filename()
     # st.toast('Random for now')
-    return np.random.choice(embeddings.index)
-
+    choice = np.random.choice(embeddings.index)
+    url = get_s3_url(choice)
+    if check_image_url_accessible(url):
+        st.toast(url)
+        return choice
+    else:
+        return get_filename()
+    
 st.title('What does attractive mean to you?')
 st.session_state.name = st.text_input(label='Invent a unique alias (and remember it)')