VPR_final / app.py
Pairavi's picture
Update app.py
ed3c2dc
import streamlit as st
import numpy as np
import similarSearch
import pandas as pd
import cv2
import os
from PIL import Image
df = pd.read_csv('./gallery.csv')
def read_image(uploaded_file):
with open(uploaded_file.name, "wb") as f:
f.write(uploaded_file.getbuffer())
img = cv2.imread(uploaded_file.name, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
if img is None:
raise ValueError('Failed to read {}'.format(uploaded_file.name))
os.remove(uploaded_file.name) # Remove the temporary file
return img
def image_open(image_file):
img = cv2.imread(
image_file, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION
)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
if img is None:
raise ValueError('Failed to read {}'.format(image_file))
return img
# Function to display images from their paths
def display_images(image_ids):
images = []
captions = []
for img_id in image_ids:
image_path = df.loc[df['seller_img_id'] == img_id, 'img_path'].values[0]
print(image_path)
image = image_open(image_path)
images.append(image)
captions.append(img_id)
return images, captions
# Function to find similar image IDs for a newly updated image
def find_similar_images(image):
# Placeholder for your code to find similar image IDs based on the new_image
similar_ids = similarSearch.predict(image) # Replace with the actual logic for finding similar IDs
images, captions = display_images(similar_ids)
length = len(images)
return similar_ids,length
# return similar_ids
def main():
st.set_page_config(
page_title='See it, search it',
page_icon=':mag:',
layout='wide',
initial_sidebar_state='expanded'
)
st.title('See it, search it')
st.markdown(
"""
<style>
.stTitle {
text-align: center;
font-size: 36px;
margin-bottom: 30px;
}
.stRadio > div {
display: flex;
justify-content: center;
align-items: center;
}
.stImageContainer {
display: flex;
justify-content: center;
align-items: center;
margin-bottom: 30px;
}
.stImage {
max-width: 100%;
height: auto;
}
</style>
""",
unsafe_allow_html=True
)
image_columns = st.columns(5)
# Define the desired width and height for the resized images
image_width = 150
image_height = 150
# Calculate the number of columns based on the available width
num_columns = 5
# Calculate the number of rows required
images_per_page = 50
# Check if the image is uploaded by the user
uploaded_image = st.file_uploader("Upload Image", type=['png', 'jpg', 'jpeg'])
if uploaded_image is not None:
# Convert the uploaded image to an array using the read_image function
image = read_image(uploaded_image)
image = Image.fromarray(image)
image = image.resize((224, 224))
# Display the resized image
st.image(image, caption="Query Image", width=224)
# Find similar images
similar_images, length = find_similar_images(image)
# Calculate the total number of pages required
total_pages = (length + images_per_page - 1) // images_per_page
page_number = st.sidebar.number_input("Select Page", min_value=1, max_value=total_pages)
# Calculate the range of images to display for the selected page
start_index = (page_number - 1) * images_per_page
end_index = min(start_index + images_per_page, length)
image_columns = st.columns(num_columns)
index = 0
for img_id in similar_images[start_index:end_index]:
row = df[df['seller_img_id'] == img_id]
if not row.empty:
img_path = row['img_path'].values[0]
product_id = row['product_id'].values[0]
# Construct the full path to the image
full_img_path = img_path
# Open and resize the image to the desired dimensions
image = image_open(full_img_path)
resized_image = cv2.resize(image, (image_width, image_height))
# Determine the column to place the image in
col_index = index % num_columns
with image_columns[col_index]:
st.image(resized_image, caption=f'Product ID: {product_id}\nImage ID: {img_id}', width=image_width)
index += 1
else:
st.write(f"Could not find details for image with ID: {img_id}")
if __name__ == '__main__':
main()