|
import streamlit as st |
|
from helper import load_hf_datasets, search, get_file_paths, get_images_from_s3_to_display |
|
import os |
|
|
|
|
|
|
|
AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID") |
|
AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY") |
|
|
|
datasets = ["WayveScenes", "StopSign_test"] |
|
|
|
|
|
bucket_name = "datasets-quasara-io" |
|
|
|
|
|
def main(): |
|
st.title("Semantic Search and Image Display") |
|
|
|
|
|
dataset_name = st.selectbox("Select Dataset", datasets) |
|
if dataset_name == 'WayveScenes': |
|
folder_path = 'WayveScenes/' |
|
else: |
|
folder_path = '' |
|
|
|
query = st.text_input("Enter your search query") |
|
|
|
|
|
limit = st.number_input("Number of results to display", min_value=1, max_value=10, value=10) |
|
|
|
|
|
if st.button("Search"): |
|
|
|
if not query: |
|
st.warning("Please enter a search query.") |
|
else: |
|
|
|
df = load_hf_datasets(dataset_name) |
|
|
|
|
|
results = search(query, df, limit, 0, "cosine", search_in_images=True, search_in_small_objects=False) |
|
|
|
|
|
top_k_paths = get_file_paths(df, results) |
|
|
|
|
|
if top_k_paths: |
|
st.write(f"Displaying top {len(top_k_paths)} results for query '{query}':") |
|
|
|
get_images_from_s3_to_display(bucket_name, top_k_paths,AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, folder_path) |
|
else: |
|
st.write("No results found.") |
|
|
|
if __name__ == "__main__": |
|
main() |
|
|
|
|
|
|
|
|