app.py

# main.py — FastAPI backend
# Endpoints:
#   GET  /                  -> serve static/index.html
#   GET  /health            -> health check
#   GET  /benchmark         -> cached benchmark results
#   GET  /models/info       -> metadata about loaded models
#   POST /predict           -> TF model prediction
#   POST /predict/onnx      -> ONNX FP32 prediction
#   POST /predict/dynamic   -> Dynamic INT8 prediction
#   POST /predict/static    -> Static INT8 prediction
#   POST /predict/gradcam   -> TF prediction + Grad-CAM overlay image

import os
import io
import json
import time
import base64
import tempfile
import numpy as np
import cv2
import matplotlib
matplotlib.use("Agg")
from PIL import Image

from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel

from src.utils import load_config, get_logger
from predict import BrainTumorPredictor

logger = get_logger("api")
cfg = load_config("config.yaml")

# ---------------------------------------------------------------------------
# App setup
# ---------------------------------------------------------------------------

app = FastAPI(
    title       = "Brain Tumor Classification API",
    description = "MRI brain tumor classification with TF, ONNX FP32, Dynamic INT8, Static INT8 + Grad-CAM",
    version     = "1.0.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins     = ["*"],
    allow_credentials = True,
    allow_methods     = ["*"],
    allow_headers     = ["*"],
)

BASE_DIR   = os.path.dirname(os.path.abspath(__file__))
STATIC_DIR = os.path.join(BASE_DIR, "static")

if os.path.exists(STATIC_DIR):
    app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
else:
    logger.warning(f"Static directory not found: {STATIC_DIR}")

# ---------------------------------------------------------------------------
# Load predictors at startup
# ---------------------------------------------------------------------------

predictors:  dict = {}
load_errors: dict = {}   # stores error messages for failed backends

@app.on_event("startup")
async def load_models():
    global predictors, load_errors
    logger.info("Loading models ...")

    # TF backend — always try first
    try:
        predictors["tensorflow"] = BrainTumorPredictor(cfg, backend="tensorflow")
        logger.info("[OK] TF model loaded.")
    except Exception as e:
        load_errors["tensorflow"] = str(e)
        logger.error(f"[FAIL] TF model failed to load: {e}")

    # ONNX backends — optional, warn if not found
    for backend in ["onnx_fp32", "onnx_dynamic", "onnx_static"]:
        try:
            predictors[backend] = BrainTumorPredictor(cfg, backend=backend)
            logger.info(f"[OK] {backend} loaded.")
        except FileNotFoundError as e:
            load_errors[backend] = str(e)
            logger.warning(f"[SKIP] {backend} — file not found. Run export_onnx.py first.")
        except Exception as e:
            load_errors[backend] = str(e)
            logger.error(f"[FAIL] {backend} failed: {e}")

    logger.info(f"Loaded backends : {list(predictors.keys())}")
    if load_errors:
        logger.warning(f"Failed backends : {list(load_errors.keys())}")


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def save_upload_temp(file: UploadFile) -> str:
    suffix = os.path.splitext(file.filename)[-1] or ".jpg"
    tmp    = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
    tmp.write(file.file.read())
    tmp.flush()
    return tmp.name


def numpy_to_b64(img_array: np.ndarray) -> str:
    img_pil = Image.fromarray(img_array.astype(np.uint8))
    buf     = io.BytesIO()
    img_pil.save(buf, format="PNG")
    return base64.b64encode(buf.getvalue()).decode("utf-8")


def predict_with_backend(backend: str, tmp_path: str) -> dict:
    if backend not in predictors:
        # Give a specific error message explaining what to do
        if backend in load_errors:
            detail = (
                f"Backend '{backend}' failed to load: {load_errors[backend]}. "
                f"Run export_onnx.py first to generate ONNX models."
            )
        else:
            detail = (
                f"Backend '{backend}' is not available. "
                f"Run export_onnx.py to generate ONNX models."
            )
        raise HTTPException(status_code=503, detail=detail)

    t0     = time.time()
    result = predictors[backend].predict(tmp_path)
    result["latency_ms"] = round((time.time() - t0) * 1000, 2)
    return result


# ---------------------------------------------------------------------------
# Pydantic response models
# ---------------------------------------------------------------------------

class PredictionResponse(BaseModel):
    predicted_class   : str
    confidence        : float
    all_probabilities : dict
    backend           : str
    latency_ms        : float


class GradCamResponse(BaseModel):
    predicted_class   : str
    confidence        : float
    all_probabilities : dict
    backend           : str
    latency_ms        : float
    gradcam_b64       : str
    heatmap_b64       : str


# ---------------------------------------------------------------------------
# Routes
# ---------------------------------------------------------------------------

@app.get("/")
async def serve_index():
    index_path = os.path.join(STATIC_DIR, "index.html")
    if not os.path.exists(index_path):
        raise HTTPException(status_code=404, detail="static/index.html not found.")
    return FileResponse(index_path)


@app.get("/health")
async def health():
    return {
        "status"          : "ok",
        "loaded_backends" : list(predictors.keys()),
        "failed_backends" : load_errors,   
    }


@app.get("/models/info")
async def models_info():
    meta_path = os.path.join(cfg["models"]["save_dir"], "model_metadata.json")
    if not os.path.exists(meta_path):
        raise HTTPException(status_code=404, detail="model_metadata.json not found.")
    with open(meta_path) as f:
        return json.load(f)


@app.get("/benchmark")
async def get_benchmark():
    bench_path = os.path.join(cfg["models"]["onnx_dir"], "benchmark_results.json")
    if not os.path.exists(bench_path):
        raise HTTPException(
            status_code=404,
            detail="benchmark_results.json not found. Run export_onnx.py first.",
        )
    with open(bench_path) as f:
        return json.load(f)


@app.post("/predict", response_model=PredictionResponse)
async def predict_tf(file: UploadFile = File(...)):
    tmp = save_upload_temp(file)
    try:
        return predict_with_backend("tensorflow", tmp)
    finally:
        os.unlink(tmp)


@app.post("/predict/onnx", response_model=PredictionResponse)
async def predict_onnx_fp32(file: UploadFile = File(...)):
    tmp = save_upload_temp(file)
    try:
        return predict_with_backend("onnx_fp32", tmp)
    finally:
        os.unlink(tmp)


@app.post("/predict/dynamic", response_model=PredictionResponse)
async def predict_dynamic(file: UploadFile = File(...)):
    tmp = save_upload_temp(file)
    try:
        return predict_with_backend("onnx_dynamic", tmp)
    finally:
        os.unlink(tmp)


@app.post("/predict/static", response_model=PredictionResponse)
async def predict_static(file: UploadFile = File(...)):
    tmp = save_upload_temp(file)
    try:
        return predict_with_backend("onnx_static", tmp)
    finally:
        os.unlink(tmp)


@app.post("/predict/gradcam", response_model=GradCamResponse)
async def predict_gradcam(file: UploadFile = File(...)):
    if "tensorflow" not in predictors:
        raise HTTPException(status_code=503, detail="TF model not loaded.")

    tmp = save_upload_temp(file)
    try:
        t0      = time.time()
        result  = predictors["tensorflow"].predict_with_gradcam(tmp)
        latency = round((time.time() - t0) * 1000, 2)

        gradcam_b64 = numpy_to_b64(result["gradcam_overlay"])
        heatmap_b64 = numpy_to_b64(
            cv2.applyColorMap(
                np.uint8(255 * result["heatmap"]), cv2.COLORMAP_JET
            )[:, :, ::-1]
        )

        return {
            "predicted_class"  : result["predicted_class"],
            "confidence"       : result["confidence"],
            "all_probabilities": result["all_probabilities"],
            "backend"          : result["backend"],
            "latency_ms"       : latency,
            "gradcam_b64"      : gradcam_b64,
            "heatmap_b64"      : heatmap_b64,
        }
    finally:
        os.unlink(tmp)


# ---------------------------------------------------------------------------
# Run
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import uvicorn
    api_cfg = cfg["api"]
    uvicorn.run(
        "main:app",
        host   = api_cfg["host"],
        port   = api_cfg["port"],
        reload = api_cfg["reload"],
    )