NEXAS commited on
Commit
95f3563
Β·
verified Β·
1 Parent(s): 116caaa

Update src/streamlit_app.py

Browse files
Files changed (1) hide show
  1. src/streamlit_app.py +18 -6
src/streamlit_app.py CHANGED
@@ -69,11 +69,17 @@ def query_similar_images(image_file, top_k=5):
69
  tmp.flush()
70
  res = image_collection.query(query_uris=[tmp.name], n_results=top_k)
71
  os.remove(tmp.name)
 
 
 
72
  return res['uris'][0]
73
 
74
  # --- Text-to-Image search ---
75
  def search_images_by_text(text, top_k=5):
76
  res = image_collection.query(query_texts=[text], n_results=top_k)
 
 
 
77
  return res['uris'][0]
78
 
79
  # --- Load CIFAR-10 Demo Dataset (500 images) ---
@@ -141,9 +147,12 @@ if q:
141
  st.image(q, caption="Query Image", width=200)
142
  with st.spinner("Finding similar images..."):
143
  results = query_similar_images(q, top_k=5)
144
- st.subheader("πŸ” Top Matches:")
145
- for u in results:
146
- st.image(u, width=150)
 
 
 
147
 
148
  st.divider()
149
  st.subheader("πŸ“ Text-to-Image Search")
@@ -151,6 +160,9 @@ txt = st.text_input("Describe what you’re looking for (e.g., 'a beach', 'a cat
151
  if txt:
152
  with st.spinner("Finding images by semantic similarity..."):
153
  results = search_images_by_text(txt, top_k=5)
154
- st.subheader("πŸ” Semantic Matches:")
155
- for u in results:
156
- st.image(u, width=150)
 
 
 
 
69
  tmp.flush()
70
  res = image_collection.query(query_uris=[tmp.name], n_results=top_k)
71
  os.remove(tmp.name)
72
+ # Safe check for results
73
+ if not res or 'uris' not in res or not res['uris'] or not res['uris'][0]:
74
+ return []
75
  return res['uris'][0]
76
 
77
  # --- Text-to-Image search ---
78
  def search_images_by_text(text, top_k=5):
79
  res = image_collection.query(query_texts=[text], n_results=top_k)
80
+ # Safe check for results
81
+ if not res or 'uris' not in res or not res['uris'] or not res['uris'][0]:
82
+ return []
83
  return res['uris'][0]
84
 
85
  # --- Load CIFAR-10 Demo Dataset (500 images) ---
 
147
  st.image(q, caption="Query Image", width=200)
148
  with st.spinner("Finding similar images..."):
149
  results = query_similar_images(q, top_k=5)
150
+ if not results:
151
+ st.warning("No similar images found.")
152
+ else:
153
+ st.subheader("πŸ” Top Matches:")
154
+ for u in results:
155
+ st.image(u, width=150)
156
 
157
  st.divider()
158
  st.subheader("πŸ“ Text-to-Image Search")
 
160
  if txt:
161
  with st.spinner("Finding images by semantic similarity..."):
162
  results = search_images_by_text(txt, top_k=5)
163
+ if not results:
164
+ st.warning("No semantic matches found.")
165
+ else:
166
+ st.subheader("πŸ” Semantic Matches:")
167
+ for u in results:
168
+ st.image(u, width=150)