fixed search by text

app.py

@@ -1,79 +1,96 @@
 import os
 import uuid
 import gradio as gr
+import numpy as np
+from PIL import Image
 import qdrant_client
-from qdrant_client import models
+from qdrant_client import QdrantClient
+from qdrant_client.http.models import VectorParams, Distance, PointStruct
 from sentence_transformers import SentenceTransformer
-from PIL import Image
 
 # ===============================
-# Setup
+# Config / Setup
 # ===============================
 UPLOAD_DIR = "uploaded_images"
 os.makedirs(UPLOAD_DIR, exist_ok=True)
 
 COLLECTION = "lost_and_found"
 
-# Qdrant client (in-memory for Hugging Face)
-qclient = qdrant_client.QdrantClient(":memory:")
+# Qdrant client (in-memory for Spaces; replace with actual url/api_key if you use a remote Qdrant)
+qclient = QdrantClient(":memory:")
+
+# SentenceTransformer encoder (CLIP)
 encoder = SentenceTransformer("clip-ViT-B-32")
 
-# Create collection if not exists
+# Create collection if missing (use the model vector size)
+VECTOR_SIZE = encoder.get_sentence_embedding_dimension()
 if not qclient.collection_exists(COLLECTION):
     qclient.create_collection(
         collection_name=COLLECTION,
-        vectors_config=models.VectorParams(size=512, distance=models.Distance.COSINE),
+        vectors_config=VectorParams(size=VECTOR_SIZE, distance=Distance.COSINE),
     )
 
 
 # ===============================
-# Encode Function (Text or Image)
+# Encoding function
+# (image handling MUST remain unchanged per request)
 # ===============================
 def encode_data(text=None, image=None):
-    if isinstance(image, Image.Image):  # Image is already PIL
+    """
+    Returns a vector (numpy array) for either a PIL Image, an image path (str),
+    or text (string). Image-handling kept exactly as requested.
+    """
+    # --- IMAGE branch (unchanged) ---
+    if isinstance(image, Image.Image):
+        # NOTE: per your instruction, do not modify the image encoding logic
         return encoder.encode(image.convert("RGB"))
-    elif isinstance(image, str):  # Path to image
+    if isinstance(image, str):
         return encoder.encode(Image.open(image).convert("RGB"))
-    elif text:
+
+    # --- TEXT branch (safe to adjust) ---
+    if text:
         return encoder.encode([text])[0]
-    else:
-        return None
+
+    return None
 
 
 # ===============================
-# Add Item
+# Add Item (finder uploads a found item)
 # ===============================
 def add_item(text, image, uploader_name, uploader_phone):
     try:
         img_path = None
         vector = None
 
-        if isinstance(image, Image.Image):  # PIL image
+        # If image provided (PIL), save and encode by image (image priority)
+        if isinstance(image, Image.Image):
             img_id = str(uuid.uuid4())
             img_path = os.path.join(UPLOAD_DIR, f"{img_id}.png")
             image.save(img_path)
             vector = encode_data(image=image)
 
+        # If no image but text provided -> encode text
         elif text:
             vector = encode_data(text=text)
 
         if vector is None:
-            return "❌ Please provide an image or text."
+            return "❌ Please provide at least an image or some text."
+
+        # Ensure vector is numpy array
+        vec = np.asarray(vector, dtype=float)
+
+        payload = {
+            "text": text or "",
+            "uploader_name": (uploader_name or "N/A"),
+            "uploader_phone": (uploader_phone or "N/A"),
+            "image_path": img_path,
+            "has_image": bool(img_path),
+        }
 
         qclient.upsert(
             collection_name=COLLECTION,
-            points=[
-                models.PointStruct(
-                    id=str(uuid.uuid4()),
-                    vector=vector.tolist(),
-                    payload={
-                        "text": text or "",
-                        "uploader_name": uploader_name or "N/A",
-                        "uploader_phone": uploader_phone or "N/A",
-                        "image_path": img_path,
-                    },
-                )
-            ],
+            points=[PointStruct(id=str(uuid.uuid4()), vector=vec.tolist(), payload=payload)],
+            wait=True,
         )
         return "✅ Item added to database!"
     except Exception as e:
@@ -81,104 +98,133 @@ def add_item(text, image, uploader_name, uploader_phone):
 
 
 # ===============================
-# Search Function
+# Search (fixed to handle text OR image OR both)
+# - keep image coding intact (per request)
+# - if both text+image supplied, average normalized vectors (cross-modal)
 # ===============================
 def search_items(text, image, max_results, min_score):
     try:
-        vector = None
-        if isinstance(image, Image.Image):  # Search with PIL
-            vector = encode_data(image=image)
-        elif text:
-            vector = encode_data(text=text)
-
-        if vector is None:
-            return "❌ Please provide an image or text.", []
-
-        results = qclient.search(
+        text_vec = None
+        img_vec = None
+
+        # get vectors (do not change image encoding)
+        if isinstance(image, Image.Image):
+            img_vec = encode_data(image=image)
+            img_vec = np.asarray(img_vec, dtype=float)
+        if text and len(text.strip()) > 0:
+            text_vec = encode_data(text=text)
+            text_vec = np.asarray(text_vec, dtype=float)
+
+        # If both provided -> combine (normalize then average)
+        if img_vec is not None and text_vec is not None:
+            # normalize
+            n1 = np.linalg.norm(img_vec) + 1e-12
+            n2 = np.linalg.norm(text_vec) + 1e-12
+            v1 = img_vec / n1
+            v2 = text_vec / n2
+            qvec = v1 + v2
+            qvec = qvec / (np.linalg.norm(qvec) + 1e-12)
+        elif img_vec is not None:
+            qvec = img_vec
+        elif text_vec is not None:
+            qvec = text_vec
+        else:
+            return "❌ Please provide an image or some text to search.", []
+
+        # Run search
+        hits = qclient.search(
             collection_name=COLLECTION,
-            query_vector=vector.tolist(),
-            limit=max_results,
-            score_threshold=min_score,
+            query_vector=qvec.tolist(),
+            limit=int(max_results),
+            score_threshold=float(min_score),
+            with_payload=True,
         )
 
-        if not results:
+        if not hits:
             return "No matches found.", []
 
-        # Format results
-        result_texts, result_imgs = [], []
-        for r in results:
-            payload = r.payload
-            result_texts.append(
-                f"id:{r.id} | score:{r.score:.3f} | "
-                f"text:{payload.get('text','')} | "
-                f"finder:{payload.get('uploader_name','N/A')} "
-                f"({payload.get('uploader_phone','N/A')})"
+        result_texts = []
+        gallery_items = []  # list of image paths (or placeholders)
+
+        for h in hits:
+            payload = h.payload or {}
+            score = getattr(h, "score", None)
+            score_str = f"{float(score):.3f}" if score is not None else "N/A"
+            uploader_name = payload.get("uploader_name", "N/A") or "N/A"
+            uploader_phone = payload.get("uploader_phone", "N/A") or "N/A"
+
+            desc = (
+                f"id:{h.id} | score:{score_str} | text:{payload.get('text','')} "
+                f"| finder:{uploader_name} ({uploader_phone})"
             )
-            if payload.get("image_path") and os.path.exists(payload["image_path"]):
-                result_imgs.append(payload["image_path"])
+            result_texts.append(desc)
+
+            img_path = payload.get("image_path")
+            if img_path and os.path.exists(img_path):
+                gallery_items.append(img_path)
+            else:
+                # append a small placeholder (you can also skip adding)
+                # Gradio can display an empty string but better to put a placeholder image path if desired
+                # We'll skip adding placeholders so gallery only shows real images
+                pass
+
+        return "\n".join(result_texts), gallery_items
 
-        return "\n".join(result_texts), result_imgs
     except Exception as e:
         return f"❌ Error: {e}", []
 
 
 # ===============================
-# Delete All
+# Clear DB
 # ===============================
 def clear_database():
-    qclient.delete_collection(COLLECTION)
-    qclient.create_collection(
-        collection_name=COLLECTION,
-        vectors_config=models.VectorParams(size=512, distance=models.Distance.COSINE),
-    )
-    return "🗑️ Database cleared!"
+    try:
+        if qclient.collection_exists(COLLECTION):
+            qclient.delete_collection(COLLECTION)
+        qclient.create_collection(
+            collection_name=COLLECTION,
+            vectors_config=VectorParams(size=VECTOR_SIZE, distance=Distance.COSINE),
+        )
+        # delete uploaded images
+        for f in os.listdir(UPLOAD_DIR):
+            try:
+                os.remove(os.path.join(UPLOAD_DIR, f))
+            except Exception:
+                pass
+        return "🗑️ Database cleared!"
+    except Exception as e:
+        return f"❌ Error clearing DB: {e}"
 
 
 # ===============================
 # Gradio UI
 # ===============================
 with gr.Blocks() as demo:
-    gr.Markdown("## 🗝️ Lost & Found - Database")
-
-    # --- Add Item Tab ---
-    with gr.Tab("➕ Add Item"):
-        with gr.Row():
-            text_input = gr.Textbox(label="Description (optional)")
-            img_input = gr.Image(type="pil", label="Upload Image")
-        with gr.Row():
-            uploader_name = gr.Textbox(label="Finder Name")
-            uploader_phone = gr.Textbox(label="Finder Phone")
-        add_btn = gr.Button("Add to Database")
-        add_output = gr.Textbox(label="Status")
-
-        add_btn.click(
-            add_item,
-            inputs=[text_input, img_input, uploader_name, uploader_phone],
-            outputs=add_output,
-        )
-
-    # --- Search Tab ---
-    with gr.Tab("🔍 Search"):
-        with gr.Row():
-            search_text = gr.Textbox(label="Search by text (optional)")
-            search_img = gr.Image(type="pil", label="Search by image (optional)")
-        with gr.Row():
-            max_results = gr.Slider(1, 10, value=5, step=1, label="Max results")
-            min_score = gr.Slider(0.5, 1.0, value=0.8, step=0.01, label="Min similarity threshold")
+    gr.Markdown("## 🔎 Lost & Found — Add Found Items (finder) & Search (lost)")
+
+    with gr.Tab("➕ Add Found Item"):
+        text_in = gr.Textbox(label="Description (optional)")
+        img_in = gr.Image(type="pil", label="Upload Image (optional)")
+        uploader_name = gr.Textbox(label="Your name (finder)")
+        uploader_phone = gr.Textbox(label="Your phone (finder)")
+        add_btn = gr.Button("Add to database")
+        add_status = gr.Textbox(label="Status")
+        add_btn.click(add_item, inputs=[text_in, img_in, uploader_name, uploader_phone], outputs=[add_status])
+
+    with gr.Tab("🔍 Search Lost Item"):
+        search_text = gr.Textbox(label="Search by text (optional)")
+        search_img = gr.Image(type="pil", label="Search by image (optional)")
+        max_results = gr.Slider(1, 20, value=5, step=1, label="Max results")
+        min_score = gr.Slider(0.0, 1.0, value=0.75, step=0.01, label="Min similarity score")
         search_btn = gr.Button("Search")
         search_text_out = gr.Textbox(label="Search results (text)")
         search_gallery = gr.Gallery(label="Search Results", columns=2, height="auto")
+        search_btn.click(search_items, inputs=[search_text, search_img, max_results, min_score], outputs=[search_text_out, search_gallery])
 
-        search_btn.click(
-            search_items,
-            inputs=[search_text, search_img, max_results, min_score],
-            outputs=[search_text_out, search_gallery],
-        )
-
-    # --- Admin Tab ---
     with gr.Tab("🗑️ Admin"):
-        clear_btn = gr.Button("Clear Database")
+        clear_btn = gr.Button("Clear database")
         clear_out = gr.Textbox(label="Status")
-        clear_btn.click(clear_database, outputs=clear_out)
+        clear_btn.click(clear_database, outputs=[clear_out])
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)