Spaces:
Runtime error
Runtime error
| import numpy as np | |
| import gradio as gr | |
| from sentence_transformers import util as st_util | |
| import pandas as pd | |
| import os | |
| from utils import load_models, get_image_embedding, img_folder, model_name_to_ids, data_path, model_names | |
| def search(input_img, num_outputs): | |
| results = [] | |
| for model_name in model_names: | |
| query_embedding = get_image_embedding(model_name, input_img) | |
| top_results = st_util.semantic_search(query_embedding, | |
| np.vstack(list(corpus_embeddings[model_name + '-embedding'])), | |
| top_k=int(num_outputs))[0] | |
| results.append([os.path.join(img_folder, | |
| corpus_embeddings.iloc[hit['corpus_id']]['name']) for hit in top_results]) | |
| return results | |
| load_models() | |
| corpus_embeddings = pd.read_parquet( | |
| os.path.join(data_path, 'metadata/patagonia_losGatos_embeddings.pq')) | |
| # Create the Gradio interface | |
| iface = gr.Interface( | |
| fn=search, | |
| inputs=[gr.Image(type="pil"), | |
| gr.inputs.Number(label="Number of results", default=3)], | |
| outputs=[gr.Gallery(label=model_name, type='filepath') for model_name in model_names], | |
| title="Search Similar Images", | |
| description="Upload an image and find similar images", | |
| ) | |
| # Launch the Gradio interface | |
| iface.launch(debug=True) | |