Update app.py

app.py

@@ -1,160 +1,272 @@
 import os
+from contextlib import asynccontextmanager
+
+# ─── 1. MANDATORY BACKEND CONFIG (Must be at the very top) ───────────────────
 os.environ["KERAS_BACKEND"] = "jax"
-os.environ["JAX_PLATFORM_NAME"] = "cpu"  # Fast startup, no GPU needed
+os.environ["JAX_PLATFORM_NAME"] = "cpu"
+os.environ["HF_HOME"] = "/tmp/huggingface"  # HF Spaces writable directory
 
 import base64
 import logging
 import numpy as np
 import cv2
 import requests
-from typing import List, Dict, Any
+from typing import List, Dict, Any, Optional
 
 import keras
 from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.responses import JSONResponse
 from huggingface_hub import hf_hub_download
+from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
 
-# ─── Logging & Backend ────────────────────────────────────────────────────────
-logging.basicConfig(level=logging.INFO)
+# ─── 2. LOGGING & CONFIG ─────────────────────────────────────────────────────
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-
-# ─── Config (use Space Secrets for API_KEY) ───────────────────────────────────
-API_KEY = os.getenv("ROBOFLOW_API_KEY")  # Set in Space Settings → Secrets
-WORKFLOW_URL = "https://serverless.roboflow.com/my-workspace-iy9ll/find-teeth"  # From your SDK
+API_KEY = os.getenv("ROBOFLOW_API_KEY")
+WORKFLOW_URL = "https://serverless.roboflow.com/my-workspace-iy9ll/find-teeth"
 IMG_SIZE = 256
-THRESHOLD = 0.05  # mean+2*std from healthy images [file:1]
-
+THRESHOLD = 0.05
 MODEL_REPO = "sairamv1/TeethCariesDetection_CNNAutoencoder"
-MODEL_FILENAME = "model.keras"  # Change if your file is named differently
+MODEL_FILENAME = "model.keras"
 
-# ─── Load Model at Startup ────────────────────────────────────────────────────
 MODEL = None
-try:
-    logger.info("🚀 Downloading model...")
-    local_model_path = hf_hub_download(
+
+# ─── 3. RETRY-ENABLED MODEL DOWNLOAD ─────────────────────────────────────────
+@retry(
+    stop=stop_after_attempt(3),
+    wait=wait_exponential(multiplier=1, min=2, max=10),
+    retry=retry_if_exception_type((requests.RequestException, OSError))
+)
+def download_model() -> str:
+    """Download model from HuggingFace Hub with retry logic."""
+    logger.info(f"🚀 Downloading model from {MODEL_REPO}...")
+    local_path = hf_hub_download(
         repo_id=MODEL_REPO,
         filename=MODEL_FILENAME,
-        cache_dir="/tmp"
+        cache_dir="/tmp/huggingface"
     )
-    logger.info(f"📥 Model at: {local_model_path}")
+    logger.info(f"✅ Model downloaded to {local_path}")
+    return local_path
+
+# ─── 4. LIFESPAN EVENT HANDLER ───────────────────────────────────────────────
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Manage model loading at startup and cleanup at shutdown."""
+    global MODEL
+    try:
+        local_model_path = download_model()
+        logger.info("🔄 Loading Keras model with JAX backend...")
+        MODEL = keras.saving.load_model(local_model_path)
+        logger.info(f"✅ Model loaded successfully! Input shape: {MODEL.input_shape}")
+    except Exception as e:
+        logger.error(f"❌ Model load failed: {e}")
+        MODEL = None
+    
+    yield  # Application runs here
     
-    logger.info("🔄 Loading Keras model...")
-    MODEL = keras.saving.load_model(local_model_path)
-    logger.info("✅ Model loaded successfully!")
-except Exception as e:
-    logger.error(f"❌ Model error: {e}")
-    MODEL = None
+    # Cleanup (if needed)
+    logger.info("🛑 Shutting down...")
 
-app = FastAPI(title="Teeth Caries Detection API")
+app = FastAPI(
+    title="Teeth Caries Detection API",
+    description="Dental anomaly detection using CNN Autoencoder",
+    version="1.0.0",
+    lifespan=lifespan
+)
 
-def first(seq: Any) -> Any:
-    return seq[0] if seq else None
+# ─── 5. ROBOFLOW API WITH RETRY ──────────────────────────────────────────────
+@retry(
+    stop=stop_after_attempt(2),
+    wait=wait_exponential(multiplier=1, min=1, max=5),
+    retry=retry_if_exception_type(requests.RequestException)
+)
+def call_roboflow_api(image_base64: str) -> Dict:
+    """Call Roboflow workflow with retry logic."""
+    response = requests.post(
+        WORKFLOW_URL,
+        params={"api_key": API_KEY},
+        json={"inputs": {"image": {"type": "base64", "value": image_base64}}},
+        timeout=30
+    )
+    response.raise_for_status()
+    return response.json()
 
+# ─── 6. ROUTES ───────────────────────────────────────────────────────────────
 @app.get("/")
 async def root():
+    """Health check endpoint."""
     return {
-        "status": "alive ✅",
-        "model_loaded": MODEL is not None,
-        "endpoints": ["/analyze"],
-        "tooth_size": f"{IMG_SIZE}x{IMG_SIZE}",
-        "threshold": THRESHOLD
+        "status": "active",
+        "model_ready": MODEL is not None,
+        "backend": os.environ.get("KERAS_BACKEND"),
+        "api_configured": API_KEY is not None
     }
 
+@app.get("/health")
+async def health():
+    """Detailed health check."""
+    if MODEL is None:
+        return JSONResponse(
+            status_code=503,
+            content={"status": "unhealthy", "reason": "Model not loaded"}
+        )
+    return {"status": "healthy", "model_input_shape": str(MODEL.input_shape)}
+
 @app.post("/analyze")
 async def analyze(file: UploadFile = File(...)):
-    if MODEL is None:
-        raise HTTPException(status_code=500, detail="Model failed to load at startup")
+    """
+    Analyze dental image for caries detection.
     
-    # 1. Decode image
-    contents = await file.read()
-    nparr = np.frombuffer(contents, np.uint8)
-    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
-    if img is None:
-        raise HTTPException(status_code=400, detail="Invalid image")
+    Returns anomaly scores for each detected tooth.
+    """
+    if MODEL is None:
+        raise HTTPException(status_code=503, detail="Model not initialized. Please retry later.")
     
-    img_h, img_w = img.shape[:2]
+    if not API_KEY:
+        raise HTTPException(status_code=500, detail="ROBOFLOW_API_KEY not configured")
 
-    # 2. Roboflow workflow
-    image_base64 = base64.b64encode(contents).decode("utf-8")
+    # ─── Read and validate image ─────────────────────────────────────────────
     try:
-        response = requests.post(
-            WORKFLOW_URL,
-            params={"api_key": API_KEY},
-            json={"inputs": {"image": {"type": "base64", "value": image_base64}}},
-            timeout=60
-        )
-        response.raise_for_status()
+        contents = await file.read()
+        if len(contents) == 0:
+            raise HTTPException(status_code=400, detail="Empty file uploaded")
+        
+        nparr = np.frombuffer(contents, np.uint8)
+        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+        
+        if img is None or img.size == 0:
+            raise HTTPException(status_code=400, detail="Invalid image format or corrupted file")
+        
+        img_h, img_w = img.shape[:2]
+        if img_h < 50 or img_w < 50:
+            raise HTTPException(status_code=400, detail="Image too small (min 50x50 pixels)")
+            
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Image decode error: {e}")
+        raise HTTPException(status_code=400, detail=f"Image processing failed: {str(e)}")
+
+    # ─── Call Roboflow detection ─────────────────────────────────────────────
+    try:
+        image_base64 = base64.b64encode(contents).decode("utf-8")
+        result_json = call_roboflow_api(image_base64)
+    except requests.Timeout:
+        raise HTTPException(status_code=504, detail="Detection service timeout")
+    except requests.RequestException as e:
+        logger.error(f"Roboflow API error: {e}")
+        raise HTTPException(status_code=502, detail="Detection service unavailable")
     except Exception as e:
-        raise HTTPException(status_code=502, detail=f"Roboflow error: {str(e)}")
+        logger.error(f"Unexpected error: {e}")
+        raise HTTPException(status_code=500, detail="Internal processing error")
 
-    # 3. Parse predictions (flexible for different response formats)
-    result_json = response.json()
-    predictions_list = []
+    # ─── Parse detections ────────────────────────────────────────────────────
     try:
         if isinstance(result_json, list):
-            outputs = first(result_json).get("outputs", {})
+            output_data = result_json[0].get("outputs", {})
         else:
-            outputs = result_json.get("outputs", [{}])[0]
-        pred_blob = outputs.get("predictions", {})
-        predictions_list = pred_blob.get("predictions", []) if isinstance(pred_blob, dict) else pred_blob
-    except:
-        predictions_list = []
+            output_data = result_json.get("outputs", [{}])[0]
+        
+        predictions = output_data.get("predictions", [])
+        if isinstance(predictions, dict):
+            predictions = predictions.get("predictions", [])
+    except Exception as e:
+        logger.warning(f"Response parsing issue: {e}")
+        predictions = []
 
-    if not predictions_list:
-        return {"status": "success", "tooth_count": 0, "results": [], "message": "No teeth detected"}
+    if not predictions:
+        return {
+            "status": "success",
+            "tooth_count": 0,
+            "results": [],
+            "message": "No teeth detected in image"
+        }
 
-    # 4. Crop + preprocess teeth
+    # ─── Process crops ───────────────────────────────────────────────────────
     batch_images = []
     metadata = []
-    for i, det in enumerate(predictions_list[:5]):  # Max 5 teeth
-        xc, yc, w, h = det["x"], det["y"], det["width"], det["height"]
-        x1, y1 = max(0, int(xc - w/2)), max(0, int(yc - h/2))
-        x2, y2 = min(img_w, int(xc + w/2)), min(img_h, int(yc + h/2))
-        
-        crop = img[y1:y2, x1:x2]
-        if crop.size == 0: continue
-
-        # Grayscale + CLAHE enhancement
-        gray = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY)
-        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
-        gray = clahe.apply(gray)
-
-        # Resize with padding (preserves aspect ratio)
-        ch, cw = gray.shape
-        scale = IMG_SIZE / max(ch, cw)
-        new_w, new_h = max(1, int(cw * scale)), max(1, int(ch * scale))
-        resized = cv2.resize(gray, (new_w, new_h), interpolation=cv2.INTER_AREA)
-        
-        canvas = np.zeros((IMG_SIZE, IMG_SIZE), dtype=np.uint8)
-        y0, x0 = (IMG_SIZE - new_h)//2, (IMG_SIZE - new_w)//2
-        canvas[y0:y0+new_h, x0:x0+new_w] = resized
-        
-        # Normalize to [0,1] + channel dim
-        final_img = (canvas.astype(np.float32) / 255.0)[..., np.newaxis]
-        batch_images.append(final_img)
-        metadata.append({"id": i, "box": [x1,y1,x2,y2]})
+
+    for i, det in enumerate(predictions[:12]):  # Limit for performance
+        try:
+            xc, yc, w, h = det["x"], det["y"], det["width"], det["height"]
+            x1 = max(0, int(xc - w/2))
+            y1 = max(0, int(yc - h/2))
+            x2 = min(img_w, int(xc + w/2))
+            y2 = min(img_h, int(yc + h/2))
+            
+            if x2 <= x1 or y2 <= y1:
+                continue
+            
+            crop = img[y1:y2, x1:x2]
+            if crop.size == 0:
+                continue
+
+            # Preprocessing pipeline
+            gray = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY)
+            clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
+            gray = clahe.apply(gray)
+            
+            # Aspect-ratio preserving resize with padding
+            h_c, w_c = gray.shape
+            scale = IMG_SIZE / max(h_c, w_c)
+            nw, nh = int(w_c * scale), int(h_c * scale)
+            resized = cv2.resize(gray, (nw, nh), interpolation=cv2.INTER_AREA)
+            
+            canvas = np.zeros((IMG_SIZE, IMG_SIZE), dtype=np.uint8)
+            y_offset = (IMG_SIZE - nh) // 2
+            x_offset = (IMG_SIZE - nw) // 2
+            canvas[y_offset:y_offset + nh, x_offset:x_offset + nw] = resized
+            
+            final_img = (canvas.astype(np.float32) / 255.0)[..., np.newaxis]
+            batch_images.append(final_img)
+            metadata.append({"id": i, "box": [x1, y1, x2, y2]})
+            
+        except Exception as e:
+            logger.warning(f"Skipping detection {i} due to error: {e}")
+            continue
 
     if not batch_images:
-        return {"status": "success", "tooth_count": 0, "results": []}
+        return {
+            "status": "success",
+            "tooth_count": 0,
+            "results": [],
+            "message": "No valid tooth crops extracted"
+        }
 
-    # 5. Autoencoder inference + anomaly scoring
-    batch_array = np.array(batch_images)
-    reconstructions = MODEL.predict(batch_array, verbose=0)
-    errors = np.mean(np.abs(batch_array - reconstructions), axis=(1,2,3))
+    # ─── Model inference ─────────────────────────────────────────────────────
+    try:
+        batch_array = np.array(batch_images)
+        reconstructions = MODEL.predict(batch_array, verbose=0)
+        errors = np.mean(np.abs(batch_array - reconstructions), axis=(1, 2, 3))
+    except Exception as e:
+        logger.error(f"Inference error: {e}")
+        raise HTTPException(status_code=500, detail="Model inference failed")
 
+    # ─── Format results ──────────────────────────────────────────────────────
     results = []
     for idx, score in enumerate(errors):
+        score_val = float(score)
         results.append({
             "tooth_index": metadata[idx]["id"],
-            "score": float(score),
-            "is_anomaly": score > THRESHOLD,
-            "confidence": float(1.0 - score),  # Lower error = higher confidence healthy
-            "bounding_box": metadata[idx]["box"]
+            "anomaly_score": round(score_val, 4),
+            "is_caries_detected": score_val > THRESHOLD,
+            "bounding_box": metadata[idx]["box"],
+            "confidence": "high" if score_val > THRESHOLD * 1.5 else "medium" if score_val > THRESHOLD else "low"
         })
 
     return {
         "status": "success",
         "tooth_count": len(results),
-        "results": results,
-        "threshold_used": THRESHOLD
+        "detection_threshold": THRESHOLD,
+        "results": results
     }
+
+# ─── 7. ERROR HANDLER ────────────────────────────────────────────────────────
+@app.exception_handler(Exception)
+async def global_exception_handler(request, exc):
+    logger.error(f"Unhandled exception: {exc}", exc_info=True)
+    return JSONResponse(
+        status_code=500,
+        content={"detail": "Internal server error occurred"}
+    )