hackerloi45 commited on
Commit
22ef1d5
Β·
1 Parent(s): a06f639
Files changed (1) hide show
  1. app.py +100 -192
app.py CHANGED
@@ -1,176 +1,85 @@
1
- # app.py
2
- import os
3
  import uuid
4
- import io
5
- import base64
 
6
  from PIL import Image
7
- import gradio as gr
8
  import numpy as np
9
 
10
- # CLIP via Sentence-Transformers
11
- from sentence_transformers import SentenceTransformer
12
-
13
- # Gemini (Google) client
14
- from google import genai
15
-
16
- # Qdrant client & helpers
17
- from qdrant_client import QdrantClient
18
- from qdrant_client.http.models import VectorParams, Distance, PointStruct
19
-
20
- # -------------------------
21
- # CONFIG (reads env vars)
22
- # -------------------------
23
- GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "").strip()
24
- QDRANT_URL = os.environ.get("QDRANT_URL", "").strip()
25
- QDRANT_API_KEY = os.environ.get("QDRANT_API_KEY", "").strip()
26
-
27
- # -------------------------
28
- # Initialize clients/models
29
- # -------------------------
30
- print("Loading CLIP model (this may take 20-60s the first time)...")
31
- MODEL_ID = "sentence-transformers/clip-ViT-B-32-multilingual-v1"
32
- clip_model = SentenceTransformer(MODEL_ID)
33
-
34
- genai_client = genai.Client(api_key=GEMINI_API_KEY) if GEMINI_API_KEY else None
35
-
36
- if not QDRANT_URL:
37
- raise RuntimeError("Please set QDRANT_URL environment variable")
38
-
39
- qclient = QdrantClient(url=QDRANT_URL, api_key=QDRANT_API_KEY)
40
- COLLECTION = "lost_found_items"
41
- VECTOR_SIZE = clip_model.get_sentence_embedding_dimension()
42
 
43
- # Create collection if missing
44
- try:
45
- if not qclient.collection_exists(COLLECTION):
46
- qclient.create_collection(
 
 
 
 
 
 
 
 
 
47
  collection_name=COLLECTION,
48
- vectors_config=VectorParams(size=VECTOR_SIZE, distance=Distance.COSINE),
49
  )
50
- except Exception as e:
51
- print("Error initializing Qdrant collection:", e)
52
-
53
- # -------------------------
54
- # Helpers
55
- # -------------------------
56
- def embed_text(text: str):
57
- return clip_model.encode(text, convert_to_numpy=True)
58
-
59
- def embed_image_pil(pil_img: Image.Image):
60
- pil_img = pil_img.convert("RGB")
61
- return clip_model.encode(pil_img, convert_to_numpy=True)
62
 
63
- def gen_tags_from_image_file(image_bytes: io.BytesIO) -> str:
64
- if genai_client is None:
65
- return ""
66
  try:
67
- file_obj = genai_client.files.upload(file=image_bytes)
68
- prompt_text = (
69
- "Give 4 short tags (comma-separated) describing this item in the image. "
70
- "Tags should be short single words or two-word phrases (e.g. 'black backpack', 'water bottle'). "
71
- "Respond only with tags, no extra explanation."
72
  )
73
- response = genai_client.models.generate_content(
74
- model="gemini-2.5-flash",
75
- contents=[prompt_text, file_obj],
 
76
  )
77
- return response.text.strip()
78
- except Exception:
79
- return ""
80
-
81
- # -------------------------
82
- # App logic: add item
83
- # -------------------------
84
- def add_item(mode: str, uploaded_image, text_description: str, finder_name: str, finder_phone: str):
85
- item_id = str(uuid.uuid4())
86
- payload = {"mode": mode, "text": text_description}
87
-
88
- # Found item extra details
89
- if mode == "found":
90
- payload["finder_name"] = finder_name
91
- payload["finder_phone"] = finder_phone
92
-
93
- if uploaded_image is not None:
94
- img_bytes = io.BytesIO()
95
- uploaded_image.convert("RGB").save(img_bytes, format="PNG")
96
- img_bytes.seek(0)
97
-
98
- vec = embed_image_pil(uploaded_image).tolist()
99
- payload["has_image"] = True
100
-
101
- payload["tags"] = gen_tags_from_image_file(img_bytes)
102
- payload["image_b64"] = base64.b64encode(img_bytes.getvalue()).decode("utf-8")
103
- else:
104
- vec = embed_text(text_description).tolist()
105
- payload["has_image"] = False
106
- if genai_client:
107
- try:
108
- resp = genai_client.models.generate_content(
109
- model="gemini-2.5-flash",
110
- contents=f"Give 4 short, comma-separated tags for this item described as: {text_description}. Reply only with tags."
111
- )
112
- payload["tags"] = resp.text.strip()
113
- except Exception:
114
- payload["tags"] = ""
115
- else:
116
- payload["tags"] = ""
117
-
118
- try:
119
- point = PointStruct(id=item_id, vector=vec, payload=payload)
120
- qclient.upsert(collection_name=COLLECTION, points=[point], wait=True)
121
  except Exception as e:
122
- return f"Error saving to Qdrant: {e}"
123
 
124
- return f"βœ… Saved item id: {item_id}\nTags: {payload.get('tags','')}"
125
-
126
- # -------------------------
127
- # App logic: search
128
- # -------------------------
129
- def search_items(query_image, query_text, limit: int = 5, min_score: float = 0.90):
130
- if query_image is not None:
131
- qvec = embed_image_pil(query_image).tolist()
132
- elif query_text and len(query_text.strip()) > 0:
133
- qvec = embed_text(query_text).tolist()
134
- else:
135
- return "⚠️ Please provide a query image or some query text.", []
136
-
137
- try:
138
- hits = qclient.search(collection_name=COLLECTION, query_vector=qvec, limit=limit)
139
- except Exception as e:
140
- return f"❌ Error querying Qdrant: {e}", []
141
-
142
- if not hits:
143
- return "No results found.", []
144
-
145
- results_text = []
146
- results_imgs = []
147
- for h in hits:
148
- score = getattr(h, "score", None)
149
- if score is None or score < min_score:
150
- continue
151
-
152
- payload = h.payload or {}
153
- text_entry = (
154
- f"id:{h.id} | score:{score:.4f} | mode:{payload.get('mode','')} | tags:{payload.get('tags','')} "
155
- f"| text:{payload.get('text','')} | finder:{payload.get('finder_name','-')} | phone:{payload.get('finder_phone','-')}"
156
- )
157
- results_text.append(text_entry)
158
-
159
- if payload.get("has_image") and "image_b64" in payload:
160
- try:
161
- img = Image.open(io.BytesIO(base64.b64decode(payload["image_b64"])))
162
- results_imgs.append(img)
163
- except Exception:
164
- pass
165
-
166
- if not results_text:
167
- return f"No results above similarity threshold {min_score}", []
168
-
169
- return "\n\n".join(results_text), results_imgs
170
-
171
- # -------------------------
172
- # App logic: clear images
173
- # -------------------------
174
  def clear_all_images():
175
  try:
176
  qclient.delete(
@@ -179,38 +88,37 @@ def clear_all_images():
179
  "filter": {"must": [{"key": "has_image", "match": {"value": True}}]}
180
  }
181
  )
182
- return "πŸ—‘οΈ All items with images have been cleared!"
183
  except Exception as e:
184
- return f"❌ Error while clearing images: {e}"
185
 
186
- # -------------------------
187
  # Gradio UI
188
- # -------------------------
189
- with gr.Blocks(title="Lost & Found β€” Simple Helper") as demo:
190
- gr.Markdown("## Lost & Found Helper (image/text search) β€” upload items, then search by image or text.")
191
- with gr.Row():
192
- with gr.Column():
193
- mode = gr.Radio(choices=["lost", "found"], value="lost", label="Add as")
194
- upload_img = gr.Image(type="pil", label="Item photo (optional)")
195
- text_desc = gr.Textbox(lines=2, placeholder="Short description (e.g. 'black backpack with blue zipper')", label="Description (optional)")
196
- finder_name = gr.Textbox(label="Finder Name (only if found)", placeholder="e.g. John Doe")
197
- finder_phone = gr.Textbox(label="Finder Phone (only if found)", placeholder="e.g. +1234567890")
198
- add_btn = gr.Button("Add item")
199
- add_out = gr.Textbox(label="Add result", interactive=False)
200
- clear_btn = gr.Button("Clear All Images")
201
- clear_out = gr.Textbox(label="Clear Result", interactive=False)
202
- with gr.Column():
203
- gr.Markdown("### Search")
204
- query_img = gr.Image(type="pil", label="Search by image (optional)")
205
- query_text = gr.Textbox(lines=2, label="Search by text (optional)")
206
- score_slider = gr.Slider(0.5, 1.0, value=0.90, step=0.01, label="Min similarity threshold")
207
- search_btn = gr.Button("Search")
208
- search_out = gr.Textbox(label="Search results (text)", interactive=False)
209
- gallery = gr.Gallery(label="Search Results", show_label=True, elem_id="gallery", columns=2, height="auto")
210
-
211
- add_btn.click(add_item, inputs=[mode, upload_img, text_desc, finder_name, finder_phone], outputs=[add_out])
212
- search_btn.click(search_items, inputs=[query_img, query_text, gr.Number(value=5, visible=False), score_slider], outputs=[search_out, gallery])
213
- clear_btn.click(clear_all_images, outputs=[clear_out])
214
-
215
- if __name__ == "__main__":
216
- demo.launch(server_name="0.0.0.0", server_port=7860)
 
1
+ import gradio as gr
 
2
  import uuid
3
+ from qdrant_client import QdrantClient
4
+ from qdrant_client.models import PointStruct
5
+ from sentence_transformers import SentenceTransformer
6
  from PIL import Image
 
7
  import numpy as np
8
 
9
+ # Connect to Qdrant
10
+ COLLECTION = "lost_and_found"
11
+ qclient = QdrantClient(":memory:") # use in-memory for demo, replace with host/port for persistence
12
+
13
+ # Load CLIP model
14
+ model = SentenceTransformer("sentence-transformers/clip-ViT-B-32-multilingual-v1")
15
+
16
+ # Ensure collection exists
17
+ qclient.recreate_collection(
18
+ collection_name=COLLECTION,
19
+ vectors_config={"size": 512, "distance": "Cosine"}
20
+ )
21
+
22
+ # Encode helper
23
+ def encode_data(text=None, image=None):
24
+ if image is not None:
25
+ img = Image.open(image).convert("RGB")
26
+ emb = model.encode(img, convert_to_numpy=True, normalize_embeddings=True)
27
+ elif text:
28
+ emb = model.encode(text, convert_to_numpy=True, normalize_embeddings=True)
29
+ else:
30
+ raise ValueError("Need text or image")
31
+ return emb.astype(np.float32)
 
 
 
 
 
 
 
 
 
32
 
33
+ # Add item
34
+ def add_item(mode, text, image, name, phone):
35
+ try:
36
+ vector = encode_data(text=text if text else None, image=image if image else None)
37
+ payload = {
38
+ "mode": mode,
39
+ "text": text,
40
+ "has_image": image is not None,
41
+ }
42
+ if mode == "found":
43
+ payload["finder_name"] = name
44
+ payload["finder_phone"] = phone
45
+ qclient.upsert(
46
  collection_name=COLLECTION,
47
+ points=[PointStruct(id=str(uuid.uuid4()), vector=vector.tolist(), payload=payload)]
48
  )
49
+ return "βœ… Item added successfully!"
50
+ except Exception as e:
51
+ return f"❌ Error: {e}"
 
 
 
 
 
 
 
 
 
52
 
53
+ # Search items
54
+ def search_items(query_image, query_text, limit, min_score):
 
55
  try:
56
+ query_vector = encode_data(
57
+ text=query_text if query_text else None,
58
+ image=query_image if query_image else None
 
 
59
  )
60
+ results = qclient.search(
61
+ collection_name=COLLECTION,
62
+ query_vector=query_vector.tolist(),
63
+ limit=limit,
64
  )
65
+ out_texts, out_imgs = [], []
66
+ for r in results:
67
+ if r.score < min_score:
68
+ continue
69
+ pl = r.payload
70
+ info = f"id:{r.id} | score:{r.score:.4f} | mode:{pl.get('mode','')}"
71
+ if pl.get("text"):
72
+ info += f" | text:{pl['text']}"
73
+ if pl.get("mode") == "found":
74
+ info += f" | found by: {pl.get('finder_name','?')} ({pl.get('finder_phone','?')})"
75
+ out_texts.append(info)
76
+ if pl.get("has_image"):
77
+ out_imgs.append(query_image) # just echo search image (or store actual image paths if needed)
78
+ return "\n".join(out_texts) if out_texts else "No matches.", out_imgs
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
79
  except Exception as e:
80
+ return f"❌ Error: {e}", []
81
 
82
+ # Clear all images
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
83
  def clear_all_images():
84
  try:
85
  qclient.delete(
 
88
  "filter": {"must": [{"key": "has_image", "match": {"value": True}}]}
89
  }
90
  )
91
+ return "πŸ—‘οΈ All image items cleared!"
92
  except Exception as e:
93
+ return f"❌ Error clearing images: {e}"
94
 
 
95
  # Gradio UI
96
+ with gr.Blocks() as demo:
97
+ gr.Markdown("# πŸ”Ž Lost & Found System")
98
+
99
+ with gr.Tab("βž• Add Item"):
100
+ mode = gr.Radio(["lost", "found"], label="Mode")
101
+ text = gr.Textbox(label="Describe the item (optional)")
102
+ img = gr.Image(type="filepath", label="Upload image (optional)")
103
+ name = gr.Textbox(label="Finder's Name (only if found)", placeholder="John Doe")
104
+ phone = gr.Textbox(label="Finder's Phone (only if found)", placeholder="+1234567890")
105
+ add_btn = gr.Button("Add Item")
106
+ add_out = gr.Textbox(label="Add result")
107
+ add_btn.click(add_item, inputs=[mode, text, img, name, phone], outputs=add_out)
108
+
109
+ with gr.Tab("πŸ” Search"):
110
+ query_text = gr.Textbox(label="Search by text (optional)")
111
+ query_img = gr.Image(type="filepath", label="Search by image (optional)")
112
+ max_results = gr.Slider(1, 10, value=5, step=1, label="Max results")
113
+ score_slider = gr.Slider(0.5, 1.0, value=0.9, step=0.01, label="Min similarity threshold")
114
+ search_btn = gr.Button("Search")
115
+ search_out = gr.Textbox(label="Search results (text)")
116
+ gallery = gr.Gallery(label="Search Results", show_label=True, elem_id="gallery", columns=2, height="auto")
117
+ search_btn.click(search_items, inputs=[query_img, query_text, max_results, score_slider], outputs=[search_out, gallery])
118
+
119
+ with gr.Tab("πŸ—‘οΈ Admin"):
120
+ clear_btn = gr.Button("Clear All Images")
121
+ clear_out = gr.Textbox(label="Clear Result")
122
+ clear_btn.click(clear_all_images, outputs=clear_out)
123
+
124
+ demo.launch()