Spaces:
Runtime error
Runtime error
Commit
·
226d235
1
Parent(s):
6b4285a
fixed + improved UI
Browse files
app.py
CHANGED
|
@@ -1,165 +1,138 @@
|
|
|
|
|
| 1 |
import gradio as gr
|
| 2 |
-
import uuid
|
| 3 |
-
import base64
|
| 4 |
-
import io
|
| 5 |
from qdrant_client import QdrantClient
|
| 6 |
-
from qdrant_client.
|
| 7 |
from sentence_transformers import SentenceTransformer
|
| 8 |
from PIL import Image
|
|
|
|
| 9 |
|
| 10 |
-
#
|
| 11 |
-
|
| 12 |
-
# --------------------------
|
| 13 |
-
QDRANT_URL = "https://ff4da494-27b1-413c-ba58-d5ea14932fe1.europe-west3-0.gcp.cloud.qdrant.io:6333" # 🔑 Replace with your cluster URL
|
| 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 |
-
VECTOR_SIZE = 512 # CLIP always outputs 512-dim vectors
|
| 22 |
-
|
| 23 |
-
# Qdrant Client
|
| 24 |
-
qclient = QdrantClient(url=QDRANT_URL, api_key=QDRANT_API_KEY)
|
| 25 |
-
|
| 26 |
-
# Ensure collection exists
|
| 27 |
-
qclient.recreate_collection(
|
| 28 |
-
collection_name=COLLECTION_NAME,
|
| 29 |
-
vectors_config=VectorParams(size=VECTOR_SIZE, distance=Distance.COSINE),
|
| 30 |
-
)
|
| 31 |
-
|
| 32 |
-
# --------------------------
|
| 33 |
-
# Helper Functions
|
| 34 |
-
# --------------------------
|
| 35 |
-
|
| 36 |
-
def image_to_base64(img: Image.Image) -> str:
|
| 37 |
-
buf = io.BytesIO()
|
| 38 |
-
img.save(buf, format="PNG")
|
| 39 |
-
return base64.b64encode(buf.getvalue()).decode("utf-8")
|
| 40 |
|
| 41 |
-
|
| 42 |
-
|
| 43 |
-
|
|
|
|
| 44 |
|
| 45 |
-
|
| 46 |
-
|
| 47 |
|
| 48 |
-
|
| 49 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 50 |
|
| 51 |
-
|
| 52 |
-
|
| 53 |
-
|
|
|
|
|
|
|
|
|
|
| 54 |
|
| 55 |
-
|
| 56 |
-
|
|
|
|
|
|
|
| 57 |
|
| 58 |
-
|
| 59 |
-
|
| 60 |
-
|
| 61 |
-
"finder_phone": finder_phone if finder_phone.strip() else "NA",
|
| 62 |
-
"image_b64": img_b64
|
| 63 |
-
}
|
| 64 |
|
|
|
|
| 65 |
qclient.upsert(
|
| 66 |
collection_name=COLLECTION_NAME,
|
| 67 |
points=[
|
| 68 |
-
PointStruct(
|
| 69 |
-
id=
|
| 70 |
-
vector=
|
| 71 |
-
payload=
|
| 72 |
-
)
|
| 73 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 74 |
)
|
| 75 |
-
return "Item successfully added."
|
| 76 |
|
| 77 |
-
|
| 78 |
-
if not query_text and query_image is None:
|
| 79 |
-
return "Enter text or upload an image to search.", []
|
| 80 |
|
| 81 |
-
|
| 82 |
-
|
| 83 |
-
|
| 84 |
-
query_vector = embed_text(query_text)
|
| 85 |
|
| 86 |
-
|
| 87 |
-
|
| 88 |
-
|
| 89 |
-
|
| 90 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 91 |
|
| 92 |
-
if not
|
| 93 |
-
return "No
|
| 94 |
-
|
| 95 |
-
gallery = []
|
| 96 |
-
output_text = "### Matches Found:\n\n"
|
| 97 |
-
for r in results:
|
| 98 |
-
desc = r.payload.get("description", "No description")
|
| 99 |
-
name = r.payload.get("finder_name", "NA")
|
| 100 |
-
phone = r.payload.get("finder_phone", "NA")
|
| 101 |
-
output_text += f"- **{desc}** (Finder: {name}, Phone: {phone})\n"
|
| 102 |
-
|
| 103 |
-
if "image_b64" in r.payload:
|
| 104 |
-
try:
|
| 105 |
-
img = base64_to_image(r.payload["image_b64"])
|
| 106 |
-
gallery.append(img)
|
| 107 |
-
except Exception:
|
| 108 |
-
pass
|
| 109 |
-
|
| 110 |
-
return output_text, gallery
|
| 111 |
-
|
| 112 |
-
def clear_database():
|
| 113 |
-
qclient.delete_collection(COLLECTION_NAME)
|
| 114 |
-
qclient.recreate_collection(
|
| 115 |
-
collection_name=COLLECTION_NAME,
|
| 116 |
-
vectors_config=VectorParams(size=VECTOR_SIZE, distance=Distance.COSINE),
|
| 117 |
-
)
|
| 118 |
-
return "Database cleared."
|
| 119 |
|
| 120 |
-
|
| 121 |
-
# Gradio UI
|
| 122 |
-
# --------------------------
|
| 123 |
|
|
|
|
| 124 |
with gr.Blocks(theme=gr.themes.Soft()) as demo:
|
| 125 |
gr.Markdown(
|
| 126 |
"""
|
| 127 |
-
# Lost & Found
|
| 128 |
-
Upload
|
| 129 |
"""
|
| 130 |
)
|
| 131 |
|
| 132 |
-
with gr.Tab("Add
|
| 133 |
-
with gr.Row():
|
| 134 |
-
image_in = gr.Image(type="pil", label="Upload Image", height=250)
|
| 135 |
-
desc_in = gr.Textbox(label="Item Description", placeholder="e.g. Silver keychain with tag")
|
| 136 |
with gr.Row():
|
| 137 |
-
|
| 138 |
-
|
| 139 |
-
|
| 140 |
-
|
| 141 |
-
|
| 142 |
-
|
|
|
|
|
|
|
| 143 |
with gr.Row():
|
| 144 |
-
|
| 145 |
-
|
| 146 |
-
|
| 147 |
-
|
| 148 |
-
gallery = gr.Gallery(
|
| 149 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 150 |
)
|
|
|
|
| 151 |
|
| 152 |
-
|
| 153 |
-
clear_btn = gr.Button("Clear Database", variant="stop")
|
| 154 |
-
clear_output = gr.Textbox(label="Status", interactive=False)
|
| 155 |
-
|
| 156 |
-
# Actions
|
| 157 |
-
add_btn.click(add_item, inputs=[image_in, desc_in, finder_name, finder_phone], outputs=add_output)
|
| 158 |
-
search_btn.click(search_items, inputs=[search_text, search_image], outputs=[search_output, gallery])
|
| 159 |
-
clear_btn.click(clear_database, outputs=clear_output)
|
| 160 |
-
|
| 161 |
-
# --------------------------
|
| 162 |
-
# Launch
|
| 163 |
-
# --------------------------
|
| 164 |
if __name__ == "__main__":
|
| 165 |
-
demo.launch(
|
|
|
|
| 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 (.env file for local development)
|
| 10 |
+
load_dotenv()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 11 |
|
| 12 |
+
# Get Qdrant Cloud credentials (set these in your .env or hosting platform)
|
| 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 |
+
# Load sentence transformer model
|
| 18 |
+
model = SentenceTransformer("clip-ViT-B-32")
|
| 19 |
|
| 20 |
+
# Initialize Qdrant client (Cloud if creds exist, else local)
|
| 21 |
+
if QDRANT_URL and QDRANT_API_KEY:
|
| 22 |
+
qclient = QdrantClient(
|
| 23 |
+
url=QDRANT_URL,
|
| 24 |
+
api_key=QDRANT_API_KEY,
|
| 25 |
+
)
|
| 26 |
+
else:
|
| 27 |
+
qclient = QdrantClient("localhost", port=6333)
|
| 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)
|
| 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()
|