Spaces:
Sleeping
Sleeping
File size: 2,022 Bytes
ab9b7a8 |
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 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 |
import streamlit as st
import torch
import os
from PIL import Image
from utils.search import search_images
device = "cuda" if torch.cuda.is_available() else "cpu"
st.set_page_config(layout="wide")
search_path = None
def show_image(result):
col1, col2, col3 = st.columns(3)
image_folder2 = "./data/images_mr"
for idx, image_name in enumerate(result.ids):
if idx % 3 == 0:
with col1:
file_name = str(image_name) + ".jpg"
image_path = os.path.join(image_folder2, file_name)
st.image(image_path, caption=image_name, width=200)
elif idx % 3 == 1:
with col2:
file_name = str(image_name) + ".jpg"
image_path = os.path.join(image_folder2, file_name)
st.image(image_path, caption=image_name, width=200)
else:
with col3:
file_name = str(image_name) + ".jpg"
image_path = os.path.join(image_folder2, file_name)
st.image(image_path, caption=image_name, width=200)
image_folder1 = "./examples"
image_paths = []
for file_name in os.listdir(image_folder1):
image_paths.append(os.path.join(image_folder1, file_name))
# st.write(image_paths)
with st.sidebar:
if st.sidebar.button("Choose a examples"):
search_path = image_paths[0]
st.image(image_paths[0], caption="example", width=150)
search_term = st.file_uploader(label="Chose a file", type=["jpg", "png"])
if search_term is None:
st.text("Please upload a image!")
else:
image = Image.open(search_term).convert('RGB')
st.image(image, width=300)
if search_term:
button = st.sidebar.button("Search")
if button:
search_path = search_term
st.header("Image Retrieval")
st.write("This is a simple app for image retrieval using Resnet18 and Vector Database")
if search_path is not None:
result = search_images(search_path)
show_image(result)
|