Spaces:
Running
Running
| import streamlit as st | |
| import requests | |
| from PIL import Image | |
| # CONFIGURATION | |
| API_URL = "http://localhost:8000/search" | |
| st.set_page_config(page_title="Face Search Engine", layout="wide") | |
| st.title("AI Face Search Engine") | |
| st.markdown(""" | |
| Upload a photo of a person to find other photos of them in the database. | |
| """) | |
| st.divider() | |
| # Layout: Two columns (Upload vs Results) | |
| col1, col2 = st.columns([1, 2]) | |
| with col1: | |
| st.subheader("1. Upload Image") | |
| uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"]) | |
| if uploaded_file is not None: | |
| # Display the uploaded image | |
| image = Image.open(uploaded_file) | |
| st.image(image, caption="Uploaded Image", use_container_width=True) | |
| # Search Button | |
| if st.button("Search Database", type="primary"): | |
| with st.spinner("Searching Vector Database..."): | |
| try: | |
| # Reset pointer to beginning of file so it can be read again | |
| uploaded_file.seek(0) | |
| # Send POST request to FastAPI | |
| files = {"file": (uploaded_file.name, uploaded_file, uploaded_file.type)} | |
| response = requests.post(API_URL, files=files) | |
| if response.status_code == 200: | |
| data = response.json() | |
| st.session_state.results = data.get("matches", []) | |
| st.session_state.count = data.get("total_matches", 0) | |
| st.success(f"Search Complete! Found {st.session_state.count} matches.") | |
| else: | |
| st.error(f"Error: API returned status {response.status_code}") | |
| except Exception as e: | |
| st.error(f"Connection Error: {e}") | |
| with col2: | |
| st.subheader("2. Search Results") | |
| if "results" in st.session_state and st.session_state.results: | |
| # Create a grid of images | |
| cols = st.columns(3) # 3 images per row | |
| for idx, img_url in enumerate(st.session_state.results): | |
| with cols[idx % 3]: | |
| st.image(img_url, use_container_width=True, caption=f"Match {idx+1}") | |
| else: | |
| st.info("No results to display yet. Upload an image to start.") |