Upload app.py

app.py

@@ -10,6 +10,7 @@ from src.modules import FeatureExtractor
 from src.config import DATA_DIR
 from fastapi.middleware.cors import CORSMiddleware
 import zipfile
+
 app = FastAPI(docs_url="/")
 origins = ["*"]
 
@@ -45,33 +46,46 @@ def image_to_base64(image: Image.Image) -> str:
     buffered = BytesIO()
     image.save(buffered, format="JPEG")
     return base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+
+# Helper function to convert PIL image to base64 string
+def image_to_base64(image: Image.Image) -> str:
+    buffered = BytesIO()
+    image.save(buffered, format="JPEG")
+    return base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+
 def unzip_folder(zip_file_path, extract_to_path):
     # Check if the zip file exists
     if not os.path.exists(zip_file_path):
         raise FileNotFoundError(f"Zip file not found: {zip_file_path}")
-    
+
     # Unzip the folder
-    with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
+    with zipfile.ZipFile(zip_file_path, "r") as zip_ref:
         zip_ref.extractall(extract_to_path)
         print(f"Extracted all files to: {extract_to_path}")
 
+
 # Example usage
-zip_file = './images.zip'   # Replace with your zip file path
-extract_path = './data'  # Replace with the directory you want to extract to
+zip_file = "./images.zip"  # Replace with your zip file path
+extract_path = "./data"  # Replace with the directory you want to extract to
 unzip_folder(zip_file, extract_path)
 
 
 image_dir = "./data/images"
 from pydantic import BaseModel, Field
+
+
 class ImageSearchBody(BaseModel):
     base64_image: str = Field(..., title="Base64 Image String")
 
+
 @app.post("/search-image/")
 async def search_image(body: ImageSearchBody):
     try:
         # Convert the Base64 string to an image
         image = base64_to_image(body.base64_image)
-    
+
         # Extract features from the image
         with torch.no_grad():
             output = feature_extractor.extract_features(image)
@@ -85,13 +99,23 @@ async def search_image(body: ImageSearchBody):
         image_list = sorted(os.listdir(image_dir))
 
         # Load the matched image from the directory
+
+        image_list = sorted(os.listdir(image_dir))
+        image_name = image_list[int(I[0][0])]
         matched_image_path = f"{image_dir}/{image_list[int(I[0][0])]}"
         matched_image = Image.open(matched_image_path)
 
         # Convert the matched image to Base64 string
         matched_image_base64 = image_to_base64(matched_image)
 
-        return JSONResponse(content={"image_base64": matched_image_base64})
+        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)