v10

app.py

@@ -1,91 +1,80 @@
-import os
-import uuid
 import gradio as gr
-import numpy as np
-from PIL import Image
 from qdrant_client import QdrantClient
-from qdrant_client.http.models import Distance, VectorParams, PointStruct
+from qdrant_client.http import models
 from sentence_transformers import SentenceTransformer
+from PIL import Image
+import uuid
+import os
 
-# =========================
-# CONFIG
-# =========================
-UPLOAD_DIR = "uploads"
+# ------------------------------
+# Setup
+# ------------------------------
+COLLECTION = "lost_and_found"
+UPLOAD_DIR = "uploaded_images"
 os.makedirs(UPLOAD_DIR, exist_ok=True)
 
-COLLECTION = "lostfound"
-
-# Connect to Qdrant (local or remote if deployed)
-qclient = QdrantClient(":memory:")  # use ":memory:" for demo, change for persistent DB
+# Connect to Qdrant (local or remote)
+qclient = QdrantClient(path="qdrant_db")
 
-# Create collection only if missing
-if not qclient.collection_exists(COLLECTION):
-    qclient.create_collection(
-        collection_name=COLLECTION,
-        vectors_config=VectorParams(size=512, distance=Distance.COSINE)
-    )
-
-# Load CLIP model
+# SentenceTransformer model
 model = SentenceTransformer("clip-ViT-B-32")
 
+# Recreate collection
+if qclient.collection_exists(COLLECTION):
+    qclient.delete_collection(COLLECTION)
+
+qclient.create_collection(
+    COLLECTION,
+    vectors_config=models.VectorParams(size=512, distance=models.Distance.COSINE),
+)
 
-# =========================
-# ENCODING FUNCTION
-# =========================
+# ------------------------------
+# Helper functions
+# ------------------------------
 def encode_data(text=None, image=None):
-    if image is not None:
-        if isinstance(image, str):  # path
-            img = Image.open(image).convert("RGB")
-        else:  # PIL.Image
-            img = image.convert("RGB")
-        emb = model.encode(img, convert_to_numpy=True, normalize_embeddings=True)
+    """Encode text or image into vector"""
+    if text and image:
+        text_vec = model.encode([text], convert_to_tensor=False)[0]
+        img_vec = model.encode([image], convert_to_tensor=False)[0]
+        return (text_vec + img_vec) / 2
     elif text:
-        emb = model.encode(text, convert_to_numpy=True, normalize_embeddings=True)
+        return model.encode([text], convert_to_tensor=False)[0]
+    elif image:
+        return model.encode([image], convert_to_tensor=False)[0]
     else:
-        raise ValueError("Need text or image")
-    return emb.astype(np.float32)
+        raise ValueError("No input provided")
 
-
-# =========================
-# ADD ITEM
-# =========================
-def add_item(mode, text, image, name, phone):
+def add_item(mode, text, image, finder_name, finder_phone):
     try:
-        vector = encode_data(text=text if text else None, image=image if image else None)
+        img_path = None
+        if image:
+            img_id = str(uuid.uuid4())
+            img_path = os.path.join(UPLOAD_DIR, f"{img_id}.png")
+            image.save(img_path)
 
-        payload = {
-            "mode": mode,
-            "text": text,
-            "has_image": image is not None,
-        }
-
-        # Save image if uploaded
-        if image is not None:
-            if isinstance(image, str):
-                img = Image.open(image).convert("RGB")
-            else:
-                img = image.convert("RGB")
-            fname = f"{uuid.uuid4().hex}.png"
-            fpath = os.path.join(UPLOAD_DIR, fname)
-            img.save(fpath)
-            payload["image_path"] = fpath
-
-        if mode == "found":
-            payload["finder_name"] = name
-            payload["finder_phone"] = phone
+        vector = encode_data(text=text if text else None, image=image if image else None)
 
         qclient.upsert(
             collection_name=COLLECTION,
-            points=[PointStruct(id=str(uuid.uuid4()), vector=vector.tolist(), payload=payload)]
+            points=[
+                models.PointStruct(
+                    id=str(uuid.uuid4()),
+                    vector=vector.tolist(),
+                    payload={
+                        "mode": mode,
+                        "text": text or "",
+                        "finder_name": finder_name if mode == "found" else "",
+                        "finder_phone": finder_phone if mode == "found" else "",
+                        "image_path": img_path,
+                        "has_image": bool(image),
+                    },
+                )
+            ],
         )
-        return "✅ Item added successfully!"
+        return f"✅ Added successfully as {mode}!"
     except Exception as e:
         return f"❌ Error: {e}"
 
-
-# =========================
-# SEARCH FUNCTION
-# =========================
 def search_items(text, image, max_results, min_score):
     try:
         vector = encode_data(text=text if text else None, image=image if image else None)
@@ -95,15 +84,21 @@ def search_items(text, image, max_results, min_score):
             query_vector=vector.tolist(),
             limit=max_results,
             score_threshold=min_score,
-            with_payload=True
+            with_payload=True,
         )
 
         texts, imgs = [], []
         for r in results:
             p = r.payload
-            desc = f"id:{r.id} | score:{r.score:.3f} | mode:{p.get('mode','')} | text:{p.get('text','')}"
-            if p.get("mode") == "found":
-                desc += f" | finder:{p.get('finder_name','')} ({p.get('finder_phone','')})"
+            desc = (
+                f"id:{r.id} | score:{r.score:.3f} | mode:{p.get('mode','')} | text:{p.get('text','')}"
+            )
+
+            # Always show finder info
+            finder_name = p.get("finder_name", "N/A") or "N/A"
+            finder_phone = p.get("finder_phone", "N/A") or "N/A"
+            desc += f" | finder:{finder_name} ({finder_phone})"
+
             texts.append(desc)
             if p.get("has_image") and "image_path" in p:
                 imgs.append(p["image_path"])
@@ -111,47 +106,63 @@ def search_items(text, image, max_results, min_score):
     except Exception as e:
         return f"❌ Error: {e}", []
 
+def delete_all():
+    try:
+        qclient.delete_collection(COLLECTION)
+        qclient.create_collection(
+            COLLECTION,
+            vectors_config=models.VectorParams(size=512, distance=models.Distance.COSINE),
+        )
+        return "🗑️ All items deleted!"
+    except Exception as e:
+        return f"❌ Error: {e}"
 
-# =========================
-# CLEAR DATABASE
-# =========================
-def clear_all_items():
-    qclient.delete(collection_name=COLLECTION, points_selector={"filter": {}})
-    return "🗑️ All items deleted!"
-
-
-# =========================
-# GRADIO APP
-# =========================
+# ------------------------------
+# Gradio UI
+# ------------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("# 🔍 Lost & Found Search with Images + Text")
+    gr.Markdown("## 🔑 Lost & Found System")
 
-    with gr.Tab("➕ Add Item"):
+    with gr.Tab("Add Item"):
         mode = gr.Radio(["lost", "found"], label="Mode", value="lost")
-        text = gr.Textbox(label="Description")
-        image = gr.Image(type="pil", label="Upload Image")
-        name = gr.Textbox(label="Finder Name (only if found)")
-        phone = gr.Textbox(label="Finder Phone (only if found)")
-        add_btn = gr.Button("Add Item")
-        add_output = gr.Textbox(label="Add result")
-        add_btn.click(add_item, [mode, text, image, name, phone], add_output)
-
-    with gr.Tab("🔎 Search"):
-        s_text = gr.Textbox(label="Search by text (optional)")
-        s_image = gr.Image(type="pil", label="Search by image (optional)")
-        max_results = gr.Slider(1, 10, value=5, step=1, label="Max results")
-        min_score = gr.Slider(0.5, 1.0, value=0.9, step=0.01, label="Min similarity threshold")
+        text_in = gr.Textbox(label="Description")
+        img_in = gr.Image(type="pil", label="Upload Image")
+
+        finder_name = gr.Textbox(label="Finder Name (only if found)")
+        finder_phone = gr.Textbox(label="Finder Phone (only if found)")
+
+        add_btn = gr.Button("Add to Database")
+        add_out = gr.Textbox(label="Status")
+
+        add_btn.click(
+            add_item,
+            inputs=[mode, text_in, img_in, finder_name, finder_phone],
+            outputs=add_out,
+        )
+
+    with gr.Tab("Search"):
+        text_s = gr.Textbox(label="Search by text (optional)")
+        img_s = gr.Image(type="pil", label="Search by image (optional)")
+        max_r = gr.Slider(1, 10, value=5, step=1, label="Max results")
+        min_s = gr.Slider(0.5, 1.0, value=0.8, step=0.01, label="Min similarity threshold")
+
         search_btn = gr.Button("Search")
-        search_text = gr.Textbox(label="Search results (text)")
+        search_out_text = gr.Textbox(label="Search results (text)")
         search_gallery = gr.Gallery(label="Search Results", columns=2, height="auto")
-        search_btn.click(search_items, [s_text, s_image, max_results, min_score], [search_text, search_gallery])
 
-    with gr.Tab("🗑️ Admin"):
-        clear_btn = gr.Button("Clear ALL items")
-        clear_out = gr.Textbox(label="Status")
-        clear_btn.click(clear_all_items, outputs=clear_out)
+        search_btn.click(
+            search_items,
+            inputs=[text_s, img_s, max_r, min_s],
+            outputs=[search_out_text, search_gallery],
+        )
 
+    with gr.Tab("Admin"):
+        del_btn = gr.Button("Delete All Items")
+        del_out = gr.Textbox(label="Status")
+        del_btn.click(delete_all, outputs=del_out)
 
+# ------------------------------
 # Launch
+# ------------------------------
 if __name__ == "__main__":
     demo.launch()