File size: 2,112 Bytes
e418977 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
import streamlit as st
from PIL import Image
from feature_extractor import FeatureExtractor
from similarity_search import SimilaritySearchEngine
def main():
st.title('Image Similarity Search')
# Upload query image
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])
if uploaded_file is not None:
# Load the uploaded image
query_img = Image.open(uploaded_file)
# Resize and display the query image
query_img_resized = query_img.resize((263, 385))
st.image(query_img_resized, caption='Uploaded Image', use_container_width=False)
# Feature extraction and similarity search
if st.button("Search Similar Images"):
with st.spinner("Analyzing query image..."):
try:
# Initialize feature extractor and search engine
extractor = FeatureExtractor()
search_engine = SimilaritySearchEngine()
# Save the uploaded image temporarily
query_img_path = 'temp_query_image.jpg'
query_img.save(query_img_path)
# Extract features from the query image
query_embedding = extractor.extract_features(query_img_path)
# Perform similarity search
similar_images, distances = search_engine.search_similar_images(query_embedding)
# Display similar images
st.subheader('Similar Images')
cols = st.columns(len(similar_images))
for i, (img_path, dist) in enumerate(zip(similar_images, distances)):
with cols[i]:
similar_img = Image.open(img_path).resize((375, 550))
st.image(similar_img, caption=f'Distance: {dist:.2f}', use_container_width=True)
except Exception as e:
st.error(f"Error during similarity search: {e}")
if __name__ == '__main__':
main()
|