File size: 8,195 Bytes
0930d33 df464b1 0930d33 01eeb3f 0930d33 01eeb3f 0930d33 01eeb3f 0930d33 01eeb3f df464b1 01eeb3f 0930d33 01eeb3f 60ac74d 01eeb3f 0930d33 01eeb3f 0930d33 01eeb3f 0930d33 01eeb3f df464b1 01eeb3f 0930d33 df464b1 01eeb3f 0930d33 df464b1 01eeb3f 0930d33 01eeb3f df464b1 01eeb3f |
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 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 |
import os
import sys
src_directory = os.path.abspath(os.path.join(os.path.dirname(__file__), "../..", "src"))
sys.path.append(src_directory)
import streamlit as st
from utils import logger
from database_pinecone import create_database, query_database
from model.clip_model import ClipModel
from data import data_set
clip_model = ClipModel()
logger = logger.get_logger()
PAGE_TITLE = "Look-a-Like: Image Finder"
PAGE_LAYOUT = "wide"
SIDEBAR_TITLE = "Find Similar Images"
PHOTO_ID_KEY = "photo_id"
IMAGE_URL_KEY = "photo_image_url"
PINECONE_INDEX = create_database.get_index()
def setup_page():
st.set_page_config(page_title=PAGE_TITLE, layout=PAGE_LAYOUT)
st.markdown(f"""
<h1 style='color:darkblue; text-align:center; font-size:32px; margin-top:-10px;'>
<i>{PAGE_TITLE} ππΈ</i>
</h1>
""", unsafe_allow_html=True)
# st.toast("β¨ Welcome to Look-a-Like: The Ultimate Image Finder! Start searching now. π")
logger.info(f"Page successfully configured with title: {PAGE_TITLE}")
def search_tab():
st.markdown("<hr>", unsafe_allow_html=True) # To add a Horizontal line below title
if "search_query" not in st.session_state:
st.session_state.search_query = ""
if "uploaded_image" not in st.session_state:
st.session_state.uploaded_image = None
with st.container():
col1, col2 = st.columns([7, 5], gap="small")
with col1:
search_query = st.text_input(
label="π Search for Images",
placeholder="Type keywords (e.g., 'sunset beach', 'city skyline')",
value=st.session_state.search_query
)
if search_query.strip() and search_query != st.session_state.search_query:
st.session_state.search_query = search_query.strip()
st.session_state.uploaded_image = None
with col2:
uploaded_image = st.file_uploader(
label="π€ Upload an Image",
type=["png", "jpg", "jpeg"],
help="Upload an image to find visually similar results."
)
if uploaded_image is not None and uploaded_image != st.session_state.uploaded_image:
st.session_state.uploaded_image = uploaded_image
st.session_state.search_query = ""
# with col3:
# st.markdown("<br>", unsafe_allow_html=True)
# if st.button(label="ποΈ Clear", help="Clear search input and uploaded image"):
# st.session_state.search_query = ""
# st.session_state.uploaded_image = None
# st.session_state.clear()
with st.container():
if st.session_state.search_query:
get_images_by_text(st.session_state.search_query)
st.session_state.search_query = ""
if st.session_state.uploaded_image:
st.image(st.session_state.uploaded_image, caption="Uploaded Image", use_container_width=True)
get_images_by_image(st.session_state.uploaded_image)
st.session_state.uploaded_image = None
def get_user_selection(options):
selected_option = st.sidebar.selectbox("Select the option", options)
return selected_option
def get_search_image_input():
uploaded_image = st.file_uploader("Upload the image to get similar images", type=['png', 'jpeg'])
return uploaded_image
def get_search_text_input():
user_search = st.sidebar.text_input("Enter the text to search")
return user_search
def display_images(response):
logger.info("Loading the images to dispay")
if response:
cols = st.columns(2)
for i, result in enumerate(response.matches):
with cols[i % 2]:
st.image(result.metadata["url"], width=500)
st.write(f"Image score : {result.score}")
logger.info("Displayed the images successfully")
def write_message(message):
st.write(message)
def get_images_by_text(query):
embedding = clip_model.get_text_embedding(query)
response = query_database.fetch_data(embedding)
message = f"π Showing search results for {query}"
write_message(message)
images = display_images(response)
def get_images_by_image(query):
embedding = clip_model.get_uploaded_image_embedding(query)
response = query_database.fetch_data(embedding)
message = f"π Showing search results of relevant images"
write_message(message)
images = display_images(response)
def load_data():
st.sidebar.header("π Data Loading Parameters")
start_index = st.sidebar.number_input("Select start index", min_value=0, value=0)
end_index = st.sidebar.number_input("Select end index", min_value=0, value=100)
if start_index > end_index:
st.sidebar.error("β οΈ Start index must be earlier than the end index.")
return
if "load_clicked" not in st.session_state:
st.session_state.load_clicked = False
try:
st.sidebar.info(f"Click the button to load data from index **{start_index} to {end_index}**.")
if st.sidebar.button("π Upsert Data", disabled=st.session_state.load_clicked, help="Click to insert data into the database"):
st.session_state.load_clicked = True
with st.spinner("β³ Upserting data... Please wait"):
df = data_set.get_df(start_index, end_index)
if df.empty:
st.warning("β οΈ No data found in the selected range.")
return
success_message = st.empty()
progress_bar = st.progress(0)
start = 0
end = len(df)
for i, data in df.iterrows():
create_database.process_and_upsert_data(PINECONE_INDEX, data, IMAGE_URL_KEY, PHOTO_ID_KEY)
success_message.success(f"Row {i + 1} (ID: {data.get('photo_id', 'unknown')}) added successfully!")
logger.info(f"Row {i + 1} (ID: {data.get('photo_id', 'unknown')}) upserted successfully.")
start = start + 1
progress = int((start) / end * 100)
progress_bar.progress(progress)
progress_bar.empty()
success_message.success("All data loaded and added to the database successfully!")
st.session_state.load_clicked = False
st.rerun()
except Exception as e:
st.error(f"Error loading data: {e}")
logger.error(f"Error loading data: {e}")
st.session_state.load_clicked = False
def get_or_greet_user_name():
if 'user_name' not in st.session_state:
st.session_state.user_name = None
logger.info("User_name not found in session_state, setting to None.")
if st.session_state.user_name is None:
logger.info("User_name is None, requesting user input.")
user_name = st.text_input("Please let me know your name:",
placeholder="Enter your name buddy")
if user_name:
st.session_state.user_name = user_name
logger.info(f"User entered name: {user_name}. Setting session_state.user_name.")
st.rerun()
else:
logger.info(f"User already entered a name: {st.session_state.user_name}. Displaying greeting.")
return st.toast(f"Hello {st.session_state.user_name}! Happy Searching")
def about_us():
if st.button("About us"):
st.write("""
This app allows you to search for images in two powerful ways:
1. **Text-based Query**: You can simply type a description or keyword, and we will fetch the most relevant images from our database.
2. **Image-based Query**: Alternatively, you can upload an image, and we'll search for similar images based on your input image.
Whether you're looking for images based on a specific text query or searching using an image, our app makes it easy to find exactly what you're looking for. Simply enter your query and get results instantly!
Explore and discover the images you need. Enjoy the search experience! π
""") |