backend/server.py

"""
FastAPI server for Diseased ROI Extraction.
Routes, middleware, and static file serving.
"""

import sys
from pathlib import Path
from contextlib import asynccontextmanager

from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, FileResponse, Response
from fastapi.staticfiles import StaticFiles
from starlette.middleware.base import BaseHTTPMiddleware
from fastapi.middleware.gzip import GZipMiddleware

from model import load_model, warmup
from inference import decode_image, validate_image, run_prediction, DEFAULT_THRESHOLD

BASE = Path(__file__).parent.parent
FRONTEND = BASE / "frontend"

compiled_model = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Load and warm up model before accepting requests."""
    global compiled_model
    compiled_model = load_model()
    try:
        warmup(compiled_model)
    except Exception as e:
        print(f"[server] FATAL: Warm-up failed — {e}", file=sys.stderr)
        sys.exit(1)
    print("[server] Ready to accept requests.")
    yield


app = FastAPI(lifespan=lifespan)

# Add GZip compression for all responses
app.add_middleware(GZipMiddleware, minimum_size=1000)

# Security headers + cache control middleware
class SecurityHeaders(BaseHTTPMiddleware):
    async def dispatch(self, request, call_next):
        response = await call_next(request)
        response.headers["X-Content-Type-Options"] = "nosniff"
        response.headers["Referrer-Policy"] = "strict-origin-when-cross-origin"
        # Cache static assets aggressively, but never cache HTML or API
        path = request.url.path
        if path.endswith((".css", ".js", ".jpg", ".png", ".ico", ".woff2")):
            response.headers["Cache-Control"] = "public, max-age=86400, immutable"
        elif path == "/" or path.endswith(".html"):
            response.headers["Cache-Control"] = "no-cache"
        return response


app.add_middleware(SecurityHeaders)


@app.get("/")
def index():
    return FileResponse(FRONTEND / "index.html")


@app.get("/health")
def health():
    return {"status": "healthy", "model_loaded": compiled_model is not None}


@app.post("/predict")
async def predict(request: Request):
    """Run disease segmentation on uploaded image."""
    if compiled_model is None:
        return JSONResponse(
            {"error": "Service is starting up, please try again shortly"},
            status_code=503,
        )

    # Parse JSON body
    try:
        body = await request.json()
    except Exception:
        return JSONResponse({"error": "Invalid JSON body"}, status_code=400)

    image_data = body.get("image")
    if not image_data:
        return JSONResponse({"error": "No image provided"}, status_code=400)

    # Decode image
    try:
        img = decode_image(image_data)
    except Exception:
        return JSONResponse({"error": "Invalid image data. Accepted formats: JPEG, PNG"}, status_code=400)

    # Validate
    err = validate_image(img)
    if err:
        return JSONResponse({"error": err}, status_code=400)

    # Parse options
    threshold = float(body.get("conf", DEFAULT_THRESHOLD))
    if threshold < 0.0 or threshold > 1.0:
        return JSONResponse(
            {"error": "Confidence threshold must be between 0.0 and 1.0"},
            status_code=422,
        )

    options = {
        "show_boxes": body.get("show_boxes", True),
        "show_masks": body.get("show_masks", True),
        "show_confidence": body.get("show_confidence", False),
        "show_heatmap": body.get("show_heatmap", False),
        "color": body.get("color", "orange"),
    }

    # Run prediction
    try:
        result = run_prediction(compiled_model, img, threshold, options)
        return result
    except Exception as e:
        print(f"[server] Inference error: {e}")
        return JSONResponse({"error": "Inference failed. Please try a different image."}, status_code=500)


# Serve frontend static files
if FRONTEND.exists():
    app.mount("/", StaticFiles(directory=str(FRONTEND), html=True), name="frontend")
else:
    print("[server] WARNING: frontend/ directory not found.")


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)