Yashashvibhardwaj commited on
Commit
253d7c7
·
verified ·
1 Parent(s): 1204afc

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +23 -14
main.py CHANGED
@@ -7,6 +7,7 @@ from fastapi.middleware.cors import CORSMiddleware
7
  from sentence_transformers import SentenceTransformer
8
  from PIL import Image
9
  import io
 
10
 
11
  # Fix caching permissions for Hugging Face
12
  os.environ["HF_HOME"] = "./cache"
@@ -15,19 +16,21 @@ os.environ["SENTENCE_TRANSFORMERS_HOME"] = "./cache"
15
 
16
  app = FastAPI()
17
 
18
- # Enable CORS (so frontend on Netlify can call backend on HF)
19
  app.add_middleware(
20
  CORSMiddleware,
21
- allow_origins=["*"], # for now allow all, can restrict to Netlify domain
22
  allow_credentials=True,
23
  allow_methods=["*"],
24
  allow_headers=["*"],
25
  )
26
 
27
- # Load product metadata
28
- with open("id_mapping.json", "r", encoding="utf-8") as f:
29
  products = json.load(f)
30
 
 
 
31
  # Load FAISS index
32
  index = faiss.read_index("products.index")
33
 
@@ -50,17 +53,23 @@ def search_text(query: str = Form(...), top_k: int = 5, min_score: float = 0.0):
50
  distances, indices = index.search(query_emb, top_k)
51
 
52
  results = []
53
- for score, idx in zip(distances[0], indices[0]):
54
- if score >= min_score: # filter by threshold
55
- item = products[idx]
56
- item["score"] = float(score)
 
57
  results.append(item)
58
 
59
  return {"matches": results}
60
 
61
 
62
- @app.post("/match") # 👈 Renamed to match frontend
63
- async def search_image(file: UploadFile = File(None), image_url: str = Form(None), top_k: int = 5, min_score: float = 0.0):
 
 
 
 
 
64
  """
65
  Search products using image query (upload or URL).
66
  """
@@ -68,7 +77,6 @@ async def search_image(file: UploadFile = File(None), image_url: str = Form(None
68
  image_bytes = await file.read()
69
  image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
70
  elif image_url:
71
- import requests
72
  response = requests.get(image_url)
73
  image = Image.open(io.BytesIO(response.content)).convert("RGB")
74
  else:
@@ -78,10 +86,11 @@ async def search_image(file: UploadFile = File(None), image_url: str = Form(None
78
  distances, indices = index.search(image_emb, top_k)
79
 
80
  results = []
81
- for score, idx in zip(distances[0], indices[0]):
 
82
  if score >= min_score:
83
- item = products[idx]
84
- item["score"] = float(score)
85
  results.append(item)
86
 
87
  return {"matches": results}
 
7
  from sentence_transformers import SentenceTransformer
8
  from PIL import Image
9
  import io
10
+ import requests
11
 
12
  # Fix caching permissions for Hugging Face
13
  os.environ["HF_HOME"] = "./cache"
 
16
 
17
  app = FastAPI()
18
 
19
+ # Enable CORS (for frontend)
20
  app.add_middleware(
21
  CORSMiddleware,
22
+ allow_origins=["*"], # you can restrict to your Netlify domain later
23
  allow_credentials=True,
24
  allow_methods=["*"],
25
  allow_headers=["*"],
26
  )
27
 
28
+ # Load products directly from products.json
29
+ with open("products.json", "r", encoding="utf-8") as f:
30
  products = json.load(f)
31
 
32
+ print(f"📦 Loaded {len(products)} products")
33
+
34
  # Load FAISS index
35
  index = faiss.read_index("products.index")
36
 
 
53
  distances, indices = index.search(query_emb, top_k)
54
 
55
  results = []
56
+ for dist, idx in zip(distances[0], indices[0]):
57
+ score = float(1 - dist) # convert distance to similarity (optional)
58
+ if score >= min_score:
59
+ item = products[idx].copy()
60
+ item["score"] = score
61
  results.append(item)
62
 
63
  return {"matches": results}
64
 
65
 
66
+ @app.post("/match") # image search
67
+ async def search_image(
68
+ file: UploadFile = File(None),
69
+ image_url: str = Form(None),
70
+ top_k: int = 5,
71
+ min_score: float = 0.0
72
+ ):
73
  """
74
  Search products using image query (upload or URL).
75
  """
 
77
  image_bytes = await file.read()
78
  image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
79
  elif image_url:
 
80
  response = requests.get(image_url)
81
  image = Image.open(io.BytesIO(response.content)).convert("RGB")
82
  else:
 
86
  distances, indices = index.search(image_emb, top_k)
87
 
88
  results = []
89
+ for dist, idx in zip(distances[0], indices[0]):
90
+ score = float(1 - dist) # convert distance to similarity
91
  if score >= min_score:
92
+ item = products[idx].copy()
93
+ item["score"] = score
94
  results.append(item)
95
 
96
  return {"matches": results}