Spaces:
No application file
No application file
| import os | |
| import io | |
| import json | |
| import faiss | |
| import torch | |
| import pandas as pd | |
| import numpy as np | |
| from PIL import Image | |
| import requests | |
| import gradio as gr | |
| from sentence_transformers import SentenceTransformer | |
| DATA_DIR = os.path.join(os.path.dirname(__file__), "data") | |
| CSV_PATH = os.path.join(DATA_DIR, "products.csv") | |
| EMB_PATH = os.path.join(DATA_DIR, "embeddings.npy") | |
| INDEX_PATH = os.path.join(DATA_DIR, "index.faiss") | |
| MAP_PATH = os.path.join(DATA_DIR, "row_mapping.json") | |
| MODEL_NAME = "clip-ViT-B-32" | |
| # Lazy globals | |
| _df = None | |
| _index = None | |
| _row_map = None | |
| _model = None | |
| def load_resources(): | |
| global _df, _index, _row_map, _model | |
| if _df is None: | |
| _df = pd.read_csv(CSV_PATH) | |
| if _index is None: | |
| if not os.path.exists(INDEX_PATH) or not os.path.exists(EMB_PATH) or not os.path.exists(MAP_PATH): | |
| raise gr.Error("Index not found. Please run build_index.py first (or deploy to Spaces and let it run once).") | |
| _index = faiss.read_index(INDEX_PATH) | |
| with open(MAP_PATH, "r", encoding="utf-8") as f: | |
| _row_map = json.load(f) | |
| if _model is None: | |
| _model = SentenceTransformer(MODEL_NAME) | |
| def fetch_image_from_url(url: str) -> Image.Image: | |
| r = requests.get(url, timeout=15) | |
| r.raise_for_status() | |
| return Image.open(io.BytesIO(r.content)).convert("RGB") | |
| def embed_image(pil_img: Image.Image) -> np.ndarray: | |
| with torch.no_grad(): | |
| emb = _model.encode([pil_img], convert_to_tensor=True) | |
| emb = torch.nn.functional.normalize(emb, p=2, dim=1) | |
| return emb.cpu().numpy().astype(np.float32) | |
| def search(query_img, top_k=12, min_score=0.2, categories=None): | |
| try: | |
| load_resources() | |
| except Exception as e: | |
| raise gr.Error(str(e)) | |
| if query_img is None: | |
| raise gr.Error("Please upload an image or provide an image URL.") | |
| q_emb = embed_image(query_img) # shape (1, d) | |
| scores, idxs = _index.search(q_emb, top_k + 50) # oversample, we will filter by category/score | |
| scores = scores.flatten().tolist() | |
| idxs = idxs.flatten().tolist() | |
| # Build results list | |
| rows = [] | |
| for score, local_idx in zip(scores, idxs): | |
| row_id = _row_map[local_idx] | |
| row = _df.iloc[row_id] | |
| rows.append({ | |
| "id": int(row["id"]), | |
| "name": row["name"], | |
| "category": row["category"], | |
| "price": float(row["price"]), | |
| "image_url": row["image_url"], | |
| "similarity": float(score), | |
| }) | |
| # Filter by category (if any) and score | |
| if categories: | |
| rows = [r for r in rows if r["category"] in categories] | |
| rows = [r for r in rows if r["similarity"] >= min_score] | |
| # Take top_k after filtering | |
| rows = rows[:top_k] | |
| # Prepare gallery & table | |
| gallery = [[r["image_url"], f'{r["name"]} | {r["category"]} | ₹{r["price"]} | sim={r["similarity"]:.3f}'] for r in rows] | |
| table = pd.DataFrame(rows, columns=["id","name","category","price","similarity","image_url"]) | |
| return gallery, table | |
| def infer(upload_image, image_url, top_k, min_score, categories): | |
| # Prefer uploaded image; if absent, try URL | |
| img = None | |
| if upload_image is not None: | |
| img = upload_image | |
| elif image_url and image_url.strip(): | |
| try: | |
| img = fetch_image_from_url(image_url.strip()) | |
| except Exception as e: | |
| raise gr.Error(f"Could not fetch image from URL: {e}") | |
| else: | |
| raise gr.Error("Please upload an image or enter a valid image URL.") | |
| return search(img, top_k=top_k, min_score=min_score, categories=categories) | |
| with gr.Blocks(css="footer {visibility: hidden}") as demo: | |
| gr.Markdown("# 🔎 Visual Product Matcher\nUpload a product photo (or paste an image URL) to find visually similar items.") | |
| with gr.Row(): | |
| with gr.Column(): | |
| upload = gr.Image(type="pil", label="Upload Image") | |
| url = gr.Textbox(label="...or Image URL") | |
| topk = gr.Slider(1, 30, value=12, step=1, label="Results to show") | |
| mins = gr.Slider(0.0, 1.0, value=0.2, step=0.01, label="Minimum similarity") | |
| cats = gr.CheckboxGroup(choices=sorted(list(set(pd.read_csv(CSV_PATH)["category"]))), label="Filter by category (optional)") | |
| btn = gr.Button("Search", variant="primary") | |
| with gr.Column(): | |
| gr.Markdown("### Query preview") | |
| # The query preview is just the uploaded image; URL-only flow won't preview until fetched internally. | |
| preview = gr.Image(label="Uploaded Image", interactive=False) | |
| gallery = gr.Gallery(label="Similar Products").style(grid=[3], height="auto") | |
| table = gr.Dataframe(headers=["id","name","category","price","similarity","image_url"], | |
| label="Results (details)", wrap=True) | |
| upload.change(lambda i: i, inputs=upload, outputs=preview) | |
| btn.click(infer, inputs=[upload, url, topk, mins, cats], outputs=[gallery, table]) | |
| if __name__ == "__main__": | |
| # For local run: `python build_index.py` first, then `python app.py` | |
| demo.launch() | |