tharu22 commited on
Commit
a87efec
Β·
1 Parent(s): 05da1e8
Files changed (2) hide show
  1. app.py +116 -0
  2. requirements.txt +7 -0
app.py ADDED
@@ -0,0 +1,116 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from pinecone import Pinecone
3
+ import os
4
+ from PIL import Image
5
+ import requests
6
+ from transformers import AutoProcessor, CLIPModel
7
+ import numpy as np
8
+ import torch
9
+
10
+ # βœ… Set Page Config (Must be the first Streamlit command)
11
+ st.set_page_config(page_title="Image Search App", layout="wide", initial_sidebar_state="expanded")
12
+
13
+ # βœ… Initialize Pinecone
14
+ pc = Pinecone(api_key="pcsk_6r4DPn_4P9LckhZak3PhebvSebnEBKQZuzYFeJL2X93LtLxZVBxyJ93inBAktefa8usvJC") # Replace with your Pinecone API key
15
+ index_name = "unsplash-index"
16
+ unsplash_index = pc.Index(index_name)
17
+
18
+ # βœ… Load CLIP Model & Processor
19
+ @st.cache_resource
20
+ def load_clip_model():
21
+ model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
22
+ processor = AutoProcessor.from_pretrained("openai/clip-vit-base-patch32")
23
+ return model, processor
24
+
25
+ model, processor = load_clip_model()
26
+
27
+ # βœ… Sidebar for Inputs
28
+ st.sidebar.title("πŸ” Search Options")
29
+ top_k = st.sidebar.slider("πŸ”’ Number of Similar Images", 1, 20, 10)
30
+
31
+ # πŸ“Œ **Option 1: Text-to-Image Search**
32
+ st.sidebar.subheader("πŸ“ Search by Text")
33
+ search_query = st.sidebar.text_input("Enter a description (e.g., 'a cute cat', 'a red car')")
34
+ text_search_btn = st.sidebar.button("πŸ” Search by Text")
35
+
36
+ # πŸ“Œ **Option 2: Image-to-Image Search**
37
+ st.sidebar.subheader("πŸ–ΌοΈ Search by Image")
38
+ uploaded_file = st.sidebar.file_uploader("Upload an image...", type=["jpg", "png", "jpeg"])
39
+ image_search_btn = st.sidebar.button("πŸ” Search by Image")
40
+
41
+ # βœ… Function to Generate Embedding from Text
42
+ def get_text_embedding(text):
43
+ inputs = processor(text=[text], return_tensors="pt", padding=True, truncation=True)
44
+ with torch.no_grad():
45
+ text_features = model.get_text_features(**inputs)
46
+ return text_features.detach().cpu().numpy().flatten().tolist()
47
+
48
+ # βœ… Function to Generate Embedding from Image
49
+ def get_image_embedding(image):
50
+ inputs = processor(images=image, return_tensors="pt")
51
+ with torch.no_grad():
52
+ image_features = model.get_image_features(**inputs)
53
+ return image_features.detach().cpu().numpy().flatten().tolist()
54
+
55
+ # βœ… Function to Query Pinecone and Fetch Similar Images
56
+ def search_similar_images(embedding, top_k=10):
57
+ results = unsplash_index.query(
58
+ vector=embedding,
59
+ top_k=top_k,
60
+ include_metadata=True,
61
+ namespace="image-search-dataset"
62
+ )
63
+ return results.get("matches", [])
64
+
65
+ # βœ… Streamlit UI
66
+ st.title("πŸ” Image & Text Search with CLIP & Pinecone")
67
+
68
+ # πŸ“Œ **Process Text-to-Image Search**
69
+ if search_query and text_search_btn:
70
+ with st.spinner("Generating embedding..."):
71
+ embedding = get_text_embedding(search_query)
72
+ with st.spinner("Searching for similar images..."):
73
+ matches = search_similar_images(embedding, top_k=top_k)
74
+
75
+ st.subheader("πŸ”Ž Top Similar Images")
76
+ if matches:
77
+ cols = st.columns(3) # Arrange images in 3 columns
78
+ for i, match in enumerate(matches):
79
+ cosine_distance = 1 - match.get("score", 0) # Convert similarity score to cosine distance
80
+ photo_id = match.get("id", "Unknown ID")
81
+ url = match.get("metadata", {}).get("url", None)
82
+
83
+ with cols[i % 3]: # Alternate images in columns
84
+ st.write(f"πŸ“· **Photo ID**: {photo_id} | πŸ“ **Cosine Distance**: {cosine_distance:.4f}")
85
+ if url:
86
+ st.image(url, caption=f"Photo ID: {photo_id}", use_container_width=True)
87
+ else:
88
+ st.warning(f"⚠️ Image URL not found for Photo ID: {photo_id}")
89
+ else:
90
+ st.warning("⚠️ No similar images found!")
91
+
92
+ # πŸ“Œ **Process Image-to-Image Search**
93
+ if uploaded_file and image_search_btn:
94
+ image = Image.open(uploaded_file).convert("RGB")
95
+ st.image(image, caption="Uploaded Image", use_container_width=True)
96
+
97
+ with st.spinner("Generating embedding..."):
98
+ embedding = get_image_embedding(image)
99
+ with st.spinner("Searching for similar images..."):
100
+ matches = search_similar_images(embedding, top_k=top_k)
101
+
102
+ st.subheader("πŸ”Ž Top Similar Images")
103
+ if matches:
104
+ cols = st.columns(3) # Arrange in 3 columns
105
+ for i, match in enumerate(matches):
106
+ cosine_distance = 1 - match.get("score", 0) # Convert similarity score to cosine distance
107
+ photo_id = match.get("id", "Unknown ID")
108
+ url = match.get("metadata", {}).get("url", None)
109
+
110
+ with cols[i % 3]: # Alternate images in columns
111
+ st.write(f"πŸ“· **Photo ID**: {photo_id} | πŸ“ **Cosine Distance**: {cosine_distance:.4f}")
112
+ if url:
113
+ st.image(url, caption=f"Photo ID: {photo_id}", use_container_width=True)
114
+ else:
115
+ st.warning(f"⚠️ Image URL not found for Photo ID: {photo_id}")
116
+
requirements.txt ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ streamlit
2
+ requests
3
+ torch
4
+ transformers
5
+ pinecone
6
+ numpy
7
+ Pillow