Update app.py

app.py

@@ -45,7 +45,7 @@ if not os.path.exists(index_path):
     raise FileNotFoundError(f"Index file not found: {index_path}")
 
 try:
-    index = None
+    index = faiss.read_index(index_path)
 except RuntimeError as e:
     raise RuntimeError(f"Error reading FAISS index: {e}")
 feature_extractor = FeatureExtractor(base_model="vit_b_16")
@@ -104,32 +104,32 @@ class ImageSearchBody(BaseModel):
 
 @app.post("/search-image/")
 async def search_image(body: ImageSearchBody):
-    # try:
-    image = base64_to_image(body.base64_image)
-    with torch.no_grad():
-        output = feature_extractor.extract_features(image)
-        output = output.view(output.size(0), -1)
-        output = output / output.norm(p=2, dim=1, keepdim=True)
-        D, I = index.search(output.cpu().numpy(), 1)
-    print(D, I)
-    image_list = sorted([f for f in os.listdir(extract_path) if is_image_file(f)])
-    print(image_list)
-    image_name = image_list[int(I[0][0])]
-    matched_image_path = f"{extract_path}/{image_list[int(I[0][0])]}"
-    matched_image = Image.open(matched_image_path)
-    matched_image_base64 = image_to_base64(matched_image)
-
-    return JSONResponse(
-        content={
-            "image_base64": matched_image_base64,
-            "image_name": image_name,
-            "similarity_score": float(D[0][0]),
-        },
-        status_code=200,
-    )
-
-    # except Exception as e:
-    #     return JSONResponse(content={"error": str(e)}, status_code=500)
+    try:
+        image = base64_to_image(body.base64_image)
+        with torch.no_grad():
+            output = feature_extractor.extract_features(image)
+            output = output.view(output.size(0), -1)
+            output = output / output.norm(p=2, dim=1, keepdim=True)
+            D, I = index.search(output.cpu().numpy(), 1)
+        print(D, I)
+        image_list = sorted([f for f in os.listdir(extract_path) if is_image_file(f)])
+        print(image_list)
+        image_name = image_list[int(I[0][0])]
+        matched_image_path = f"{extract_path}/{image_list[int(I[0][0])]}"
+        matched_image = Image.open(matched_image_path)
+        matched_image_base64 = image_to_base64(matched_image)
+    
+        return JSONResponse(
+            content={
+                "image_base64": matched_image_base64,
+                "image_name": image_name,
+                "similarity_score": float(D[0][0]),
+            },
+            status_code=200,
+        )
+
+    except Exception as e:
+        return JSONResponse(content={"error": str(e)}, status_code=500)
 
 
 from src.firebase.firebase_provider import process_images