image search fixed

app.py

@@ -1,58 +1,63 @@
 import os
 import uuid
 import gradio as gr
+import qdrant_client
+from qdrant_client import models
+from sentence_transformers import SentenceTransformer
 from PIL import Image
-from qdrant_client import QdrantClient, models
-from transformers import CLIPProcessor, CLIPModel
 
-# ==============================
+# ===============================
 # Setup
-# ==============================
+# ===============================
 UPLOAD_DIR = "uploaded_images"
 os.makedirs(UPLOAD_DIR, exist_ok=True)
 
-qclient = QdrantClient(":memory:")
 COLLECTION = "lost_and_found"
 
+# Qdrant client
+qclient = qdrant_client.QdrantClient(":memory:")  # In-memory for Hugging Face
+encoder = SentenceTransformer("clip-ViT-B-32")
+
 # Create collection
 if not qclient.collection_exists(COLLECTION):
     qclient.create_collection(
-        COLLECTION,
+        collection_name=COLLECTION,
         vectors_config=models.VectorParams(size=512, distance=models.Distance.COSINE),
     )
 
-# Load CLIP
-clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
-clip_proc = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
-
-
-# ==============================
-# Encode Function
-# ==============================
+# ===============================
+# Encode Function (Text or Image)
+# ===============================
 def encode_data(text=None, image=None):
-    if text:
-        inputs = clip_proc(text=[text], return_tensors="pt", padding=True)
-        return clip_model.get_text_features(**inputs).detach().numpy()[0]
-
-    if image:
-        inputs = clip_proc(images=image, return_tensors="pt")
-        return clip_model.get_image_features(**inputs).detach().numpy()[0]
-
-    raise ValueError("Need either text or image for encoding")
-
-
-# ==============================
-# Add Item (finder uploads)
-# ==============================
+    if image is not None:
+        return encoder.encode(Image.open(image).convert("RGB"))
+    elif text:
+        return encoder.encode([text])[0]
+    else:
+        return None
+
+# ===============================
+# Add Item
+# ===============================
 def add_item(text, image, uploader_name, uploader_phone):
     try:
         img_path = None
+        vector = None
+
         if image:
             img_id = str(uuid.uuid4())
             img_path = os.path.join(UPLOAD_DIR, f"{img_id}.png")
             image.save(img_path)
 
-        vector = encode_data(text=text if text else None, image=image if image else None)
+            # ✅ Always store image embedding if available
+            vector = encode_data(image=img_path)
+
+        elif text:
+            # ✅ Fallback: text embedding
+            vector = encode_data(text=text)
+
+        if vector is None:
+            return "❌ Please provide an image or text."
 
         qclient.upsert(
             collection_name=COLLECTION,
@@ -65,7 +70,6 @@ def add_item(text, image, uploader_name, uploader_phone):
                         "uploader_name": uploader_name or "N/A",
                         "uploader_phone": uploader_phone or "N/A",
                         "image_path": img_path,
-                        "has_image": bool(image),
                     },
                 )
             ],
@@ -74,93 +78,100 @@ def add_item(text, image, uploader_name, uploader_phone):
     except Exception as e:
         return f"❌ Error: {e}"
 
-
-# ==============================
-# Search Items (for lost things)
-# ==============================
+# ===============================
+# 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)
+        vector = None
+        if image:
+            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(
             collection_name=COLLECTION,
             query_vector=vector.tolist(),
             limit=max_results,
             score_threshold=min_score,
-            with_payload=True,
         )
 
-        texts, imgs = [], []
+        if not results:
+            return "No matches found.", []
+
+        # Format results
+        result_texts, result_imgs = [], []
         for r in results:
-            p = r.payload
-            desc = f"id:{r.id} | score:{r.score:.3f} | text:{p.get('text','')}"
-
-            # Show finder details
-            uploader_name = p.get("uploader_name", "N/A") or "N/A"
-            uploader_phone = p.get("uploader_phone", "N/A") or "N/A"
-            desc += f" | finder:{uploader_name} ({uploader_phone})"
-
-            texts.append(desc)
-            if p.get("has_image") and "image_path" in p:
-                imgs.append(p["image_path"])
-        return "\n".join(texts) if texts else "No matches", imgs
+            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')})"
+            )
+            if payload.get("image_path") and os.path.exists(payload["image_path"]):
+                result_imgs.append(payload["image_path"])
+
+        return "\n".join(result_texts), result_imgs
     except Exception as e:
         return f"❌ Error: {e}", []
 
-
-# ==============================
+# ===============================
 # Delete All
-# ==============================
-def clear_all():
-    qclient.recreate_collection(
-        COLLECTION, vectors_config=models.VectorParams(size=512, distance=models.Distance.COSINE)
+# ===============================
+def clear_database():
+    qclient.delete_collection(COLLECTION)
+    qclient.create_collection(
+        collection_name=COLLECTION,
+        vectors_config=models.VectorParams(size=512, distance=models.Distance.COSINE),
     )
-    return "🗑️ All items cleared."
-
+    return "🗑️ Database cleared!"
 
-# ==============================
+# ===============================
 # Gradio UI
-# ==============================
+# ===============================
 with gr.Blocks() as demo:
-    gr.Markdown("# 🔑 Lost & Found Database")
-
-    with gr.Tab("➕ Add Found Item"):
-        text_in = gr.Textbox(label="Description (optional)")
-        img_in = gr.Image(type="pil", label="Upload Image")
-        uploader_name = gr.Textbox(label="Finder Name")
-        uploader_phone = gr.Textbox(label="Finder Phone")
+    gr.Markdown("## 🗝️ Lost & Found - Database")
+
+    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_out = gr.Textbox(label="Status")
+        add_output = gr.Textbox(label="Status")
 
         add_btn.click(
             add_item,
-            inputs=[text_in, img_in, uploader_name, uploader_phone],
-            outputs=add_out,
+            inputs=[text_input, img_input, uploader_name, uploader_phone],
+            outputs=add_output,
         )
 
-    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, 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")
+    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")
         search_btn = gr.Button("Search")
-        search_out = gr.Textbox(label="Search results (text)")
-        search_gallery = gr.Gallery(label="Search Results")
+        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_out, search_gallery],
+            outputs=[search_text_out, search_gallery],
         )
 
     with gr.Tab("🗑️ Admin"):
-        clear_btn = gr.Button("Clear All Items")
+        clear_btn = gr.Button("Clear Database")
         clear_out = gr.Textbox(label="Status")
-        clear_btn.click(clear_all, outputs=clear_out)
-
+        clear_btn.click(clear_database, outputs=clear_out)
 
-# ==============================
-# Launch
-# ==============================
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()