Update app.py

app.py

@@ -1,8 +1,10 @@
-from fastapi import FastAPI, UploadFile, File, HTTPException
+import os
+from fastapi import FastAPI, UploadFile, File, HTTPException, Header
 from fastapi.responses import JSONResponse
 from transformers import pipeline
 from PIL import Image
 import torch
+import numpy as np
 import io
 
 # Set up the device for the model
@@ -15,6 +17,9 @@ feature_extraction_pipeline = pipeline(
     device=0 if torch.cuda.is_available() else -1
 )
 
+# Read SECRET_KEY from the environment
+SECRET_KEY = os.getenv("SECRET_KEY", "default-secret")
+
 # Initialize FastAPI
 app = FastAPI()
 
@@ -24,8 +29,12 @@ def read_root():
 
 # Endpoint to extract image features
 @app.post("/extract-features/")
-async def extract_features(file: UploadFile = File(...)):
+async def extract_features(file: UploadFile = File(...), secret_key: str = Header(None)):
     try:
+        # Verify the SECRET_KEY
+        if secret_key != SECRET_KEY:
+            raise HTTPException(status_code=403, detail="Invalid SECRET_KEY")
+        
         # Validate file format
         if file.content_type not in ["image/jpeg", "image/png", "image/jpg"]:
             raise HTTPException(status_code=400, detail=f"Unsupported file format. Upload a JPEG or PNG image. Received {file.content_type}")
@@ -43,6 +52,8 @@ async def extract_features(file: UploadFile = File(...)):
         # Return the embedding vector
         return JSONResponse(content={"features": cls_embedding.tolist()})
     
+    except HTTPException:
+        raise  # Reraise HTTPExceptions for proper status codes
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))