hackerloi45 commited on
Commit
7a8dd54
·
1 Parent(s): 109869f

fixed+improved UI

Browse files
Files changed (1) hide show
  1. app.py +140 -108
app.py CHANGED
@@ -1,138 +1,170 @@
1
- import os
2
  import gradio as gr
 
 
 
3
  from qdrant_client import QdrantClient
4
- from qdrant_client.http import models
5
  from sentence_transformers import SentenceTransformer
6
  from PIL import Image
7
- from dotenv import load_dotenv
8
 
9
- # Load environment variables from .env
10
- load_dotenv()
11
-
12
- # Get Qdrant Cloud credentials (must be in .env or deployment secrets)
13
- QDRANT_URL = os.getenv("https://ff4da494-27b1-413c-ba58-d5ea14932fe1.europe-west3-0.gcp.cloud.qdrant.io:6333")
14
- QDRANT_API_KEY = os.getenv("eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJhY2Nlc3MiOiJtIn0.jjeB1JgnUSlb1hOOKMdRpVvMrUER57-udT-X1AWXT1E")
15
  COLLECTION_NAME = "lost_and_found"
16
 
17
- # Check creds
18
- if not QDRANT_URL or not QDRANT_API_KEY:
19
- raise RuntimeError(
20
- "❌ Missing Qdrant Cloud credentials. Please set QDRANT_URL and QDRANT_API_KEY in your environment."
21
- )
22
 
23
- # Initialize Qdrant client
24
- qclient = QdrantClient(url=QDRANT_URL, api_key=QDRANT_API_KEY)
25
 
26
- # Load sentence transformer model
27
- model = SentenceTransformer("clip-ViT-B-32")
 
 
 
28
 
29
- # Create collection if it doesn’t exist
30
- if not qclient.collection_exists(COLLECTION_NAME):
31
- qclient.create_collection(
32
- collection_name=COLLECTION_NAME,
33
- vectors_config=models.VectorParams(size=512, distance=models.Distance.COSINE),
34
- )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
35
 
36
- # Add item function
37
- def add_item(image, description):
38
- if image is None or not description.strip():
39
- return "❌ Please provide both an image and a description."
40
 
41
- # Encode description and image
42
- vector_desc = model.encode(description).tolist()
43
- vector_img = model.encode(Image.open(image)).tolist()
 
 
 
44
 
45
- # Upload to Qdrant (store both text & image vectors)
46
  qclient.upsert(
47
  collection_name=COLLECTION_NAME,
48
  points=[
49
- models.PointStruct(
50
- id=None,
51
- vector=vector_desc,
52
- payload={"description": description, "image": image.name},
53
- ),
54
- models.PointStruct(
55
- id=None,
56
- vector=vector_img,
57
- payload={"description": description, "image": image.name},
58
- ),
59
- ],
60
  )
 
61
 
62
- return f"✅ Item added successfully: {description}"
63
-
64
- # Search function (text or image)
65
- def search_items(text_query, image_query):
66
- search_results = []
67
-
68
- if text_query and text_query.strip():
69
- vector = model.encode(text_query).tolist()
70
- hits = qclient.search(
71
- collection_name=COLLECTION_NAME,
72
- query_vector=vector,
73
- limit=5,
74
- )
75
- search_results.extend(hits)
76
-
77
- if image_query is not None:
78
- vector = model.encode(Image.open(image_query)).tolist()
79
- hits = qclient.search(
80
- collection_name=COLLECTION_NAME,
81
- query_vector=vector,
82
- limit=5,
83
- )
84
- search_results.extend(hits)
85
-
86
- # Remove duplicates by description
87
- seen = set()
88
- unique_results = []
89
- for hit in search_results:
90
- desc = hit.payload.get("description")
91
- if desc not in seen:
92
- seen.add(desc)
93
- unique_results.append(
94
- (hit.payload.get("image"), hit.payload.get("description"))
95
- )
96
 
97
- if not unique_results:
98
- return [("not_found.png", "No match found")]
 
 
 
99
 
100
- return unique_results
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
101
 
 
102
  # Gradio UI
103
- with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
 
104
  gr.Markdown(
105
  """
106
- # 🔍 Lost & Found Database
107
- Upload found items or search using text/image.
108
  """
109
  )
110
 
111
- with gr.Tab("Add Items"):
 
112
  with gr.Row():
113
- with gr.Column():
114
- add_img = gr.Image(type="filepath", label="Upload Image")
115
- add_desc = gr.Textbox(label="Description", placeholder="e.g. Silver key with round head")
116
- add_btn = gr.Button("Add Item")
117
- add_output = gr.Textbox(label="Status", interactive=False)
118
- add_btn.click(fn=add_item, inputs=[add_img, add_desc], outputs=add_output)
119
-
120
- with gr.Tab("🔎 Search"):
 
 
121
  with gr.Row():
122
- with gr.Column():
123
- search_text = gr.Textbox(label="Search by Text", placeholder="e.g. key, wallet, phone")
124
- search_img = gr.Image(type="filepath", label="Or Search by Image")
125
- search_btn = gr.Button("Search")
126
- gallery = gr.Gallery(
127
- label="Matched Items",
128
- show_label=True,
129
- elem_id="gallery",
130
- columns=[3],
131
- rows=[2],
132
- height="auto"
133
- )
134
- search_btn.click(fn=search_items, inputs=[search_text, search_img], outputs=gallery)
135
-
136
- # Launch app
 
 
 
 
137
  if __name__ == "__main__":
138
- demo.launch()
 
 
1
  import gradio as gr
2
+ import uuid
3
+ import base64
4
+ import io
5
  from qdrant_client import QdrantClient
6
+ from qdrant_client.models import PointStruct, VectorParams, Distance
7
  from sentence_transformers import SentenceTransformer
8
  from PIL import Image
 
9
 
10
+ # --------------------------
11
+ # Qdrant Cloud Connection
12
+ # --------------------------
13
+ QDRANT_URL = "https://ff4da494-27b1-413c-ba58-d5ea14932fe1.europe-west3-0.gcp.cloud.qdrant.io:6333"
14
+ QDRANT_API_KEY = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJhY2Nlc3MiOiJtIn0.jjeB1JgnUSlb1hOOKMdRpVvMrUER57-udT-X1AWXT1E"
 
15
  COLLECTION_NAME = "lost_and_found"
16
 
17
+ # CLIP model (text + image embeddings)
18
+ MODEL_NAME = "sentence-transformers/clip-ViT-B-32"
19
+ embedder = SentenceTransformer(MODEL_NAME)
 
 
20
 
21
+ # CLIP ViT-B/32 always gives 512-dimensional embeddings
22
+ VECTOR_SIZE = 512
23
 
24
+ # Qdrant Client (Cloud)
25
+ qclient = QdrantClient(
26
+ url=QDRANT_URL,
27
+ api_key=QDRANT_API_KEY
28
+ )
29
 
30
+ # Ensure collection exists
31
+ qclient.recreate_collection(
32
+ collection_name=COLLECTION_NAME,
33
+ vectors_config=VectorParams(size=VECTOR_SIZE, distance=Distance.COSINE),
34
+ )
35
+
36
+ # --------------------------
37
+ # Helper Functions
38
+ # --------------------------
39
+
40
+ def image_to_base64(img: Image.Image) -> str:
41
+ buf = io.BytesIO()
42
+ img.save(buf, format="PNG")
43
+ return base64.b64encode(buf.getvalue()).decode("utf-8")
44
+
45
+ def base64_to_image(b64_str: str) -> Image.Image:
46
+ img_bytes = base64.b64decode(b64_str)
47
+ return Image.open(io.BytesIO(img_bytes))
48
+
49
+ def embed_text(text: str):
50
+ return embedder.encode(text).tolist()
51
+
52
+ def embed_image(img: Image.Image):
53
+ return embedder.encode(img).tolist()
54
+
55
+ def add_item(image, description, finder_name, finder_phone):
56
+ if image is None or description.strip() == "":
57
+ return "Please provide both an image and a description."
58
 
59
+ embedding = embed_image(image)
60
+ img_b64 = image_to_base64(image)
 
 
61
 
62
+ metadata = {
63
+ "description": description,
64
+ "finder_name": finder_name if finder_name.strip() else "NA",
65
+ "finder_phone": finder_phone if finder_phone.strip() else "NA",
66
+ "image_b64": img_b64
67
+ }
68
 
 
69
  qclient.upsert(
70
  collection_name=COLLECTION_NAME,
71
  points=[
72
+ PointStruct(
73
+ id=str(uuid.uuid4()),
74
+ vector=embedding,
75
+ payload=metadata
76
+ )
77
+ ]
 
 
 
 
 
78
  )
79
+ return "Item successfully added!"
80
 
81
+ def search_items(query_text, query_image):
82
+ if not query_text and query_image is None:
83
+ return "Please enter text or upload an image to search.", []
84
+
85
+ if query_image:
86
+ query_vector = embed_image(query_image)
87
+ else:
88
+ query_vector = embed_text(query_text)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
89
 
90
+ results = qclient.search(
91
+ collection_name=COLLECTION_NAME,
92
+ query_vector=query_vector,
93
+ limit=5
94
+ )
95
 
96
+ if not results:
97
+ return "No matches found.", []
98
+
99
+ gallery = []
100
+ output_text = "### Matches Found\n\n"
101
+ for r in results:
102
+ desc = r.payload.get("description", "No description")
103
+ name = r.payload.get("finder_name", "NA")
104
+ phone = r.payload.get("finder_phone", "NA")
105
+ output_text += f"- **{desc}** — Finder: {name}, Phone: {phone}\n"
106
+
107
+ if "image_b64" in r.payload:
108
+ try:
109
+ img = base64_to_image(r.payload["image_b64"])
110
+ gallery.append(img)
111
+ except Exception:
112
+ pass
113
+
114
+ return output_text, gallery
115
+
116
+ def clear_database():
117
+ qclient.delete_collection(COLLECTION_NAME)
118
+ qclient.recreate_collection(
119
+ collection_name=COLLECTION_NAME,
120
+ vectors_config=VectorParams(size=VECTOR_SIZE, distance=Distance.COSINE),
121
+ )
122
+ return "Database cleared successfully."
123
 
124
+ # --------------------------
125
  # Gradio UI
126
+ # --------------------------
127
+
128
+ with gr.Blocks() as demo:
129
  gr.Markdown(
130
  """
131
+ # Lost & Found System
132
+ A simple platform to add and search lost items.
133
  """
134
  )
135
 
136
+ with gr.Tab("Add Item"):
137
+ gr.Markdown("### Add a Found Item")
138
  with gr.Row():
139
+ image_in = gr.Image(type="pil", label="Upload Image")
140
+ desc_in = gr.Textbox(label="Item Description")
141
+ with gr.Row():
142
+ finder_name = gr.Textbox(label="Finder's Name (optional)")
143
+ finder_phone = gr.Textbox(label="Finder's Phone Number (optional)")
144
+ add_btn = gr.Button("Submit Item", variant="primary")
145
+ add_output = gr.Textbox(label="Status", interactive=False)
146
+
147
+ with gr.Tab("Search"):
148
+ gr.Markdown("### Search Lost Items")
149
  with gr.Row():
150
+ search_text = gr.Textbox(label="Search by Text")
151
+ search_image = gr.Image(type="pil", label="Or Upload Image")
152
+ search_btn = gr.Button("Search", variant="primary")
153
+ search_output = gr.Markdown()
154
+ gallery = gr.Gallery(label="Matched Items", show_label=True, elem_id="gallery")
155
+
156
+ with gr.Tab("Admin"):
157
+ gr.Markdown("### Admin Controls")
158
+ clear_btn = gr.Button("Clear Database", variant="stop")
159
+ clear_output = gr.Textbox(label="Status", interactive=False)
160
+
161
+ # Button actions
162
+ add_btn.click(add_item, inputs=[image_in, desc_in, finder_name, finder_phone], outputs=add_output)
163
+ search_btn.click(search_items, inputs=[search_text, search_image], outputs=[search_output, gallery])
164
+ clear_btn.click(clear_database, outputs=clear_output)
165
+
166
+ # --------------------------
167
+ # Launch App
168
+ # --------------------------
169
  if __name__ == "__main__":
170
+ demo.launch(server_name="0.0.0.0", server_port=7860)