Update app.py

app.py

@@ -14,6 +14,9 @@ app = FastAPI(docs_url="/")
 index = faiss.read_index("./model/db_vit_b_16.index")
 feature_extractor = FeatureExtractor(base_model="vit_b_16")
 
+# Use MPS (Apple Silicon) if available, otherwise fallback to CPU
+if torch.backends.mps.is_built():
+    torch.set_default_device("mps")
 
 
 # Helper function to load image from uploaded file
@@ -21,6 +24,10 @@ def load_image(uploaded_file):
     image = Image.open(BytesIO(uploaded_file)).convert("RGB")
     return image
 
+
+image_dir = "./data/images"
+
+
 @app.post("/search-image/")
 async def search_image(file: UploadFile = File(...)):
     try:
@@ -34,20 +41,19 @@ async def search_image(file: UploadFile = File(...)):
             output = output / output.norm(p=2, dim=1, keepdim=True)
 
             # Perform FAISS search for the top 1 similar image
-            D, I = index.search(output.cpu().numpy(), 1)  # Changed from 5 to 1
-        print("I", I)
-        print("D", D)
-        # Load the list of image filenames (assuming you have image_list)
-        image_list = sorted(os.listdir(os.path.join(DATA_DIR, "images")))
+            D, I = index.search(output.cpu().numpy(), 1)
+        print(D, I)
 
-        # Get the path of the most similar image
-        similar_image_path = os.path.join(DATA_DIR, "images", image_list[int(I[0][0])])
+        image_list = sorted(os.listdir(image_dir))
 
-        # Return the image file itself
-        return FileResponse(similar_image_path, media_type="image/jpeg")
+        data_dir = f"{image_dir}/{image_list[int(I[0][0])]}"
+
+        return FileResponse(data_dir, media_type="image/jpeg")
     except Exception as e:
         return JSONResponse(content={"error": str(e)}, status_code=500)
 
+
 if __name__ == "__main__":
     import uvicorn
+
     uvicorn.run(app, host="0.0.0.0", port=8000)