Spaces:
Running
Running
File size: 3,386 Bytes
4bb166c c17d554 4bb166c c17d554 4bb166c 7026488 eb80b2b aff9ecd 7026488 aff9ecd 7026488 aff9ecd 4bb166c c17d554 4bb166c |
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 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 |
import os
import numpy as np
import boto3
import streamlit as st
import faiss
import pandas as pd
from PIL import Image
from model_prediction import Ranker
from io import BytesIO
import cv2
@st.cache
def load_model():
return Ranker()
def load_faiss_index():
return faiss.read_index('embeddings.index')
def load_labels():
return pd.read_csv("labels.csv")
class ModelLoader:
model = None
index = None
labels = None
@classmethod
def get_model(cls):
if cls.model is None:
cls.model = load_model()
return cls.model
@classmethod
def get_index(cls):
if cls.index is None:
cls.index = load_faiss_index()
return cls.index
@classmethod
def get_labels(cls):
if cls.labels is None:
cls.labels = load_labels()
return cls.labels
target_size = (224, 224)
st.set_page_config(page_title="Product Retrieval App")
st.title("Product Retrieval App")
st.markdown("""The Product Retrieval App is a demonstration of a computer vision model created by <a href="https://intelliarts.com/">Intelliarts</a>. It can analyze and interpret visual data , i.e., shapes, colors, and textures from uploaded digital images. The data is then used to conduct a search on the web. The output of the computer vision model is a set of images that are predicted to be most similar to the input image.
To use the Product Retrieval App, you need to:
1. Select an image that depicts the item of interest. Acceptable formats are JPG, JPEG, and PNG.
2. Upload the image by either dragging and dropping the file into the search field or selecting a file from your computer using the “browse files” button.
3. Scroll to the bottom of the page to review the output results.""",
unsafe_allow_html=True)
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
loading_text = st.empty()
s3 = boto3.client(
's3',
aws_access_key_id='AKIAUUWYLZEQYT6ESW4Q',
aws_secret_access_key='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',
region_name='eu-west-1'
)
bucket_name = "product-retrieval"
if uploaded_file is not None:
image = Image.open(uploaded_file)
image = np.asarray(image)
if len(image.shape) > 2 and image.shape[2] == 4:
image = cv2.cvtColor(image, cv2.COLOR_BGRA2BGR)
image = Image.fromarray(image)
st.image(image, caption="Uploaded image", use_column_width=True)
loading_text.text("Loading predictions...")
model = ModelLoader.get_model()
index = ModelLoader.get_index()
labels = ModelLoader.get_labels()
image_embedding = model.predict(image)
distances, indices = index.search(image_embedding, 12)
predicted_images = labels["path"][indices[0]].to_list()
loading_text.empty()
col1, col2, col3, col4 = st.columns(4)
for i, img_path in enumerate(predicted_images):
response = s3.get_object(Bucket=bucket_name, Key=img_path.split("/")[-1])
image_data = response['Body'].read()
img = Image.open(BytesIO(image_data)).resize(target_size)
if i % 4 == 0:
column = col1
elif i % 4 == 1:
column = col2
elif i % 4 == 2:
column = col3
else:
column = col4
with column:
st.image(img, caption=f"Predicted image {i+1}", use_column_width=True)
|