Update app.py

app.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, UploadFile, File
+from fastapi import FastAPI, UploadFile, File, HTTPException, Header
 from fastapi.responses import JSONResponse
 from PIL import Image
 import numpy as np
@@ -7,22 +7,26 @@ import io
 
 app = FastAPI()
 
-# Load the model
+# Load your model
 model = tf.keras.models.load_model("2.keras")
 CLASS_NAMES = ['Fungi', 'Healthy', 'Nematode', 'Pest', 'Phytopthora', 'Virus']
 
+# Define your API key (keep it secret in prod)
+API_KEY = "mysecretkey"
+
 @app.post("/predict")
-async def predict(file: UploadFile = File(...)):
+async def predict(file: UploadFile = File(...), x_api_key: str = Header(None)):
+    if x_api_key != API_KEY:
+        raise HTTPException(status_code=401, detail="Invalid or missing API Key")
+
     try:
         contents = await file.read()
 
-        # Load and process image
+        # Process the image
         image = Image.open(io.BytesIO(contents)).convert("RGB")
         image = image.resize((224, 224))
-
-        # Don't normalize — match tf.data image_dataset_from_directory
         img_array = np.array(image).astype("float32")
-        img_array = np.expand_dims(img_array, axis=0)  # (1, 224, 224, 3)
+        img_array = np.expand_dims(img_array, axis=0)
 
         # Predict
         prediction = model.predict(img_array)