Spaces:
Runtime error
Runtime error
import streamlit as st | |
st.set_page_config(page_title='Image Search App', layout='wide') | |
import torch | |
from transformers import AutoTokenizer, AutoModel | |
import faiss | |
import numpy as np | |
import wget | |
from PIL import Image | |
from io import BytesIO | |
from sentence_transformers import SentenceTransformer | |
import json | |
from zipfile import ZipFile | |
import zipfile | |
from io import BytesIO | |
from PIL import Image | |
# from huggingface_hub import hf_hub_download | |
# hf_hub_download(repo_id="shivangibithel/Flickr8k", filename="Images.zip") | |
# Load the pre-trained sentence encoder | |
model_name = "sentence-transformers/all-distilroberta-v1" | |
tokenizer = AutoTokenizer.from_pretrained(model_name) | |
model = SentenceTransformer(model_name) | |
# # Load the FAISS index | |
# index_name = 'index.faiss' | |
# index_url = './' | |
# wget.download(index_url, index_name) | |
# index = faiss.read_index(faiss_flickr8k.index) | |
# Define the path to the zip folder containing the images | |
zip_path = "Images.zip" | |
# Open the zip folder | |
zip_file = zipfile.ZipFile(zip_path) | |
# Iterate over the images in the zip folder and display them using Streamlit | |
for image_name in zip_file.namelist(): | |
image_data = zip_file.read(image_name) | |
image = Image.open(io.BytesIO(image_data)) | |
st.image(image, caption=image_name) | |
vectors = np.load("./sbert_text_features.npy") | |
vector_dimension = vectors.shape[1] | |
index = faiss.IndexFlatL2(vector_dimension) | |
faiss.normalize_L2(vectors) | |
index.add(vectors) | |
# Map the image ids to the corresponding image URLs | |
image_map_name = 'captions.json' | |
with open(image_map_name, 'r') as f: | |
caption_dict = json.load(f) | |
image_list = list(caption_dict.keys()) | |
caption_list = list(caption_dict.values()) | |
def search(query, k=5): | |
# Encode the query | |
query_embedding = model.encode(query) | |
query_vector = np.array([query_embedding]) | |
faiss.normalize_L2(query_vector) | |
index.nprobe = index.ntotal | |
# Search for the nearest neighbors in the FAISS index | |
D, I = index.search(query_vector, k) | |
# Map the image ids to the corresponding image URLs | |
image_urls = [] | |
for i in I[0]: | |
text_id = i | |
image_id = str(image_list[i]) | |
image_data = zip_file.read(image_id) | |
image = Image.open(io.BytesIO(image_data)) | |
st.image(image, caption=image_name, width=200) | |
# image_url = "./Images/" + image_id | |
# image_urls.append(image_url) | |
# return image_urls | |
st.title("Image Search App") | |
query = st.text_input("Enter your search query here:") | |
if st.button("Search"): | |
if query: | |
search(query) | |
# Display the images | |
# st.image(image_urls, width=200) | |
if __name__ == '__main__': | |
st.cache(allow_output_mutation=True) | |
run_app() | |