Update app.py

app.py

@@ -45,16 +45,38 @@ with tab_caption:
         else:
             st.error("Error: " + resp.text)
 
+# with tab_text2img:
+#     if text_input and st.button("Search Images"):
+#         data = {"model_name": model_name, "query": text_input, "top_k": top_k}
+#         resp = requests.post(f"{API_BASE}/search/text2img", data=data)
+#         if resp.status_code == 200:
+#             results = resp.json()
+#             for res in results:
+#                 st.image(res["image_path"], caption=f"Score: {res['score']:.3f}")
+#         else:
+#             st.error("Error: " + resp.text)
+
 with tab_text2img:
     if text_input and st.button("Search Images"):
         data = {"model_name": model_name, "query": text_input, "top_k": top_k}
         resp = requests.post(f"{API_BASE}/search/text2img", data=data)
         if resp.status_code == 200:
             results = resp.json()
-            for res in results:
-                st.image(res["image_path"], caption=f"Score: {res['score']:.3f}")
+            if results:
+                cols = st.columns(3)
+                for idx, res in enumerate(results):
+                    with cols[idx % 3]:
+                        try:
+                            st.image(res["image_path"], 
+                                   caption=f"Score: {res['score']:.3f}", 
+                                   use_column_width=True)
+                        except Exception as e:
+                            st.warning(f"Could not load: {res['image_path']}")
+                            st.write(f"Score: {res['score']:.3f}")
+            else:
+                st.info("No matching images in demo set.")
         else:
-            st.error("Error: " + resp.text)
+            st.error(f"Backend error: {resp.status_code} - {resp.text}")
 
 with tab_img2text:
     if image_file and st.button("Retrieve Text"):
@@ -66,17 +88,59 @@ with tab_img2text:
         else:
             st.error("Error: " + resp.text)
 
+# with tab_img2img:
+#     if image_file and st.button("Retrieve Similar Images"):
+#         files = {"file": image_file.getvalue()}
+#         data = {"model_name": model_name, "top_k": top_k}
+#         resp = requests.post(f"{API_BASE}/search/img2img", files=files, data=data)
+#         if resp.status_code == 200:
+#             results = resp.json()
+#             for res in results:
+#                 st.image(res["image_path"], caption=f"Score: {res['score']:.3f}")
+#         else:
+#             st.error("Error: " + resp.text)
+
 with tab_img2img:
     if image_file and st.button("Retrieve Similar Images"):
         files = {"file": image_file.getvalue()}
         data = {"model_name": model_name, "top_k": top_k}
         resp = requests.post(f"{API_BASE}/search/img2img", files=files, data=data)
+        
         if resp.status_code == 200:
             results = resp.json()
-            for res in results:
-                st.image(res["image_path"], caption=f"Score: {res['score']:.3f}")
+            if results:
+                cols = st.columns(3)
+                for idx, res in enumerate(results):
+                    with cols[idx % 3]:
+                        try:
+                            st.image(
+                                res["image_path"],
+                                caption=f"Score: {res['score']:.3f}",
+                                use_column_width=True
+                            )
+                        except Exception as e:
+                            st.warning(f"Could not load image: {res['image_path']}")
+                            st.write(f"Score: {res['score']:.3f}")
+            else:
+                st.info("No similar images found in the demo set.")
         else:
-            st.error("Error: " + resp.text)
+            st.error(f"Backend error: {resp.status_code} - {resp.text}")
 
 with tab_text2text:
-    st.info("Text → Text not implemented yet. Add to model interface if needed.")
+    text_input_tt = st.text_input("Enter text to find similar captions", 
+                                 placeholder="A child playing with water in the garden")
+    
+    if text_input_tt and st.button("Search Similar Captions"):
+        data = {"model_name": model_name, "query": text_input_tt, "top_k": top_k}
+        resp = requests.post(f"{API_BASE}/search/text2text", data=data)
+        
+        if resp.status_code == 200:
+            results = resp.json()
+            if results:
+                st.subheader("Top similar captions:")
+                for idx, res in enumerate(results, 1):
+                    st.markdown(f"**{idx}.** {res['caption']}  \nScore: `{res['score']:.4f}`")
+            else:
+                st.info("No similar captions found.")
+        else:
+            st.error(f"Error: {resp.status_code} - {resp.text}")