Update app.py

app.py

@@ -1,48 +1,54 @@
 from fastapi import FastAPI, File, UploadFile, HTTPException, Header, Depends
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import JSONResponse
 import numpy as np
 import tensorflow as tf
 import cv2
 import base64
+import os
+import logging
+from huggingface_hub import hf_hub_download
 
-# ---- Add this for API key protection ----
-API_KEY = "your-secret-api-key"  # 🔐 Replace with your actual key
+# ---------- CONFIG ----------
+API_KEY = "your-secret-api-key"  # Replace this with your key
+IMG_SIZE = 256
+CLASS_COLORS = {0: (0, 0, 0), 1: (0, 255, 0), 2: (0, 0, 255)}
 
-def verify_api_key(x_api_key: str = Header(...)):
-    if x_api_key != API_KEY:
-        raise HTTPException(status_code=403, detail="Invalid API Key")
-# -----------------------------------------
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
+# ---------- API SETUP ----------
 app = FastAPI()
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # In production: allow only trusted domains
+    allow_origins=["*"],
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
-from tensorflow.keras.models import load_model
 
-# Available backend options are: "jax", "torch", "tensorflow".
-import os
-os.environ["HF_HOME"] = "/tmp/huggingface"  # ✅ Place this before hf_hub_download
-from huggingface_hub import hf_hub_download
+def verify_api_key(x_api_key: str = Header(...)):
+    if x_api_key != API_KEY:
+        raise HTTPException(status_code=403, detail="Invalid API Key")
 
-# Make sure HF_HOME is set before using huggingface_hub
-os.environ["HF_HOME"] = "/tmp/huggingface"
+# ---------- LOAD MODEL ----------
+try:
+    os.environ["HF_HOME"] = "/tmp/huggingface"  # Prevent permission issues on Spaces
 
-model_path = hf_hub_download(
-    repo_id="rishab1090/potato",
-    filename="unet_model.keras"
-)
-model = tf.keras.models.load_model(model_path)
+    model_path = hf_hub_download(
+        repo_id="rishab1090/potato",
+        filename="unet_model.keras",  # 👈 Use the exact filename
+        cache_dir="/tmp/hf_cache"     # 👈 Helps avoid read-only FS errors
+    )
 
+    model = tf.keras.models.load_model(model_path)
+    logger.info("✅ Model loaded successfully from .keras file.")
 
-IMG_SIZE = 256
-CLASS_COLORS = {0: (0, 0, 0), 1: (0, 255, 0), 2: (0, 0, 255)}
+except Exception as e:
+    logger.error(f"❌ Failed to load model: {e}")
+    raise RuntimeError(f"Model load failed: {e}")
 
+# ---------- UTILS ----------
 def decode_mask_to_overlay(image_bgr, mask):
     overlay = image_bgr.copy()
     for class_id, color in CLASS_COLORS.items():
@@ -55,16 +61,20 @@ def image_to_base64(img: np.ndarray) -> str:
     _, buffer = cv2.imencode('.png', img)
     return base64.b64encode(buffer).decode("utf-8")
 
-
+# ---------- PREDICTION ROUTE ----------
 @app.post("/predict_severity")
 async def predict_severity(
     file: UploadFile = File(...),
-    x_api_key: str = Depends(verify_api_key)  # 🔐 Require API key
+    x_api_key: str = Depends(verify_api_key)
 ):
     try:
         contents = await file.read()
         file_bytes = np.frombuffer(contents, np.uint8)
         img_bgr = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
+
+        if img_bgr is None:
+            raise ValueError("Invalid image file")
+
         img_resized = cv2.resize(img_bgr, (IMG_SIZE, IMG_SIZE))
         img_norm = img_resized.astype(np.float32) / 255.0
         img_input = np.expand_dims(img_norm, axis=0)
@@ -72,9 +82,6 @@ async def predict_severity(
         prediction = model.predict(img_input)[0]
         mask = np.argmax(prediction, axis=-1).astype(np.uint8)
 
-        center_pixel = prediction[IMG_SIZE // 2, IMG_SIZE // 2]
-        print(f"Center pixel confidence: {center_pixel}")
-
         unique, counts = np.unique(mask, return_counts=True)
         class_counts = {int(k): int(v) for k, v in zip(unique, counts)}
 
@@ -92,4 +99,5 @@ async def predict_severity(
         }
 
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        logger.error(f"Error during prediction: {e}")
+        raise HTTPException(status_code=500, detail=str(e))