server.py

# =============================================================================
# server.py — FastAPI Backend
# =============================================================================
# This connects the beautiful HTML frontend to your trained ThermoCAN model.
#
# HOW TO RUN:
#   pip install -r requirements.txt
#   python server.py
#
# Then open browser: http://localhost:8000
# =============================================================================

import os
import sys
import io
import base64
import yaml
import torch
import numpy as np
import pandas as pd
from pathlib import Path
from PIL import Image
from fastapi import FastAPI, File, UploadFile, Form, HTTPException
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from fastapi import Request
from contextlib import asynccontextmanager
import uvicorn

sys.path.insert(0, str(Path(__file__).parent))
from model_loader import load_model
from predictor   import predict, generate_gradcam, tiff_float_to_uint8

# ── Configuration ─────────────────────────────────────────────────────────────
BASE_DIR        = Path(__file__).parent.absolute()
CONFIG_PATH     = BASE_DIR / "configs" / "config.yaml"
CHECKPOINT_PATH = BASE_DIR / "outputs" / "checkpoints" / "best.pth"

# ── Load config ───────────────────────────────────────────────────────────────
with open(CONFIG_PATH) as f:
    CONFIG = yaml.safe_load(f)

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Global state
MODEL = None
ENCODER = None
NUM_VIEWS = 3
MODEL_LOADED = False

@asynccontextmanager
async def lifespan(app: FastAPI):
    global MODEL, ENCODER, NUM_VIEWS, MODEL_LOADED
    print("Device: " + str(DEVICE))
    print("Loading ThermoCAN model...")
    try:
        # Use str(CHECKPOINT_PATH) for os.path operations inside load_model
        MODEL, ENCODER, NUM_VIEWS = load_model(str(CHECKPOINT_PATH), CONFIG, DEVICE)
        MODEL_LOADED = True
        print("Model loaded and ready!")
    except Exception as e:
        print(f"Model load error: {e}")
        MODEL_LOADED = False
    yield
    # Cleanup if needed

# Which views to use
d            = CONFIG["dataset"]
use_all      = d.get("use_all_five_views", False)
primary      = d.get("primary_views", [0, 1, 3])
WANTED_VIEWS = [0, 1, 2, 3, 4] if use_all else primary
IMAGE_SIZE   = d.get("image_size", 224)

# ── FastAPI app ───────────────────────────────────────────────────────────────
app = FastAPI(title="ThermoCAN", version="1.0", lifespan=lifespan)

# Serve HTML templates
templates = Jinja2Templates(directory="templates")


@app.get("/", response_class=HTMLResponse)
async def home(request: Request):
    """Serve the main frontend page."""
    return templates.TemplateResponse(request=request, name="index.html")


@app.get("/favicon.ico")
async def favicon():
    """Handle favicon requests."""
    return JSONResponse(status_code=204, content={})


@app.get("/health")
async def health():
    """Health check — tells you if model is loaded."""
    return {
        "status":       "ok" if MODEL_LOADED else "model_not_loaded",
        "device":       str(DEVICE),
        "num_views":    NUM_VIEWS,
        "wanted_views": WANTED_VIEWS,
        "clinical_dim": ENCODER.feature_dim if ENCODER else 0,
    }


@app.post("/predict")
async def predict_endpoint(
    # Image uploads — view_0 through view_4
    view_0: UploadFile = File(None),
    view_1: UploadFile = File(None),
    view_2: UploadFile = File(None),
    view_3: UploadFile = File(None),
    view_4: UploadFile = File(None),
    # Clinical fields
    age:             str = Form("45"),
    body_temp:       str = Form("36.6"),
    biopsy:          str = Form("No"),
    mammography:     str = Form("No"),
    radiotherapy:    str = Form("No"),
    mastectomy:      str = Form("No"),
    eating_habits:   str = Form("Low fat"),
    protocol_smoked: str = Form("No"),
    marital_status:  str = Form("Married"),
    use_tta:         str = Form("true"),
):
    """
    Main prediction endpoint.
    Receives images + clinical data, returns prediction + Grad-CAM.
    """
    if not MODEL_LOADED:
        raise HTTPException(
            status_code=503,
            detail="Model not loaded. Check CHECKPOINT_PATH in server.py."
        )

    # ── Collect uploaded images ────────────────────────────────────────────
    uploads = {0: view_0, 1: view_1, 2: view_2, 3: view_3, 4: view_4}
    images_dict = {}

    for view_id, upload in uploads.items():
        if upload is not None and upload.filename:
            try:
                raw   = await upload.read()
                img   = Image.open(io.BytesIO(raw))
                # Convert mode F (TIFF float) or any mode to RGB PIL image
                if img.mode == "F":
                    arr = tiff_float_to_uint8(img)
                    img = Image.fromarray(arr)
                elif img.mode != "RGB":
                    img = img.convert("RGB")
                images_dict[view_id] = img
            except Exception as e:
                print(f"Failed to read view_{view_id}: {e}")

    if not images_dict:
        raise HTTPException(
            status_code=400,
            detail="No valid images uploaded. Please upload at least the Frontal view."
        )

    if 0 not in images_dict:
        raise HTTPException(
            status_code=400,
            detail="Frontal view (view_0) is required but was not uploaded."
        )

    # ── Encode clinical features ───────────────────────────────────────────
    clinical_tensor = None
    if ENCODER is not None and ENCODER.feature_dim > 0:
        try:
            row = {
                "age_current":               float(age),
                "age_at_visit":              float(age),
                "body_temperature":          float(body_temp),
                "menarche":                  12.0,
                "biopsy":                    biopsy,
                "mammography":               mammography,
                "radiotherapy":              radiotherapy,
                "use_of_hormone_replacement": "No",
                "plastic_surgery":           "No",
                "prosthesis":                "No",
                "marital_status":            marital_status,
                "eating_habits":             eating_habits,
                "protocol_smoked":           protocol_smoked,
                "protocol_drank_coffee":     "No",
                "mastectomy":                mastectomy,
            }
            df              = pd.DataFrame([row])
            arr             = ENCODER.transform(df)
            clinical_tensor = torch.tensor(arr[0], dtype=torch.float32)
        except Exception as e:
            print(f"Clinical encoding failed: {e} — using image only")
            clinical_tensor = None

    # ── Run prediction ─────────────────────────────────────────────────────
    use_tta_bool = use_tta.lower() in ("true", "1", "yes")

    try:
        probability, label_str, confidence = predict(
            MODEL, images_dict, clinical_tensor, DEVICE,
            NUM_VIEWS, WANTED_VIEWS,
            use_tta=use_tta_bool,
            image_size=IMAGE_SIZE,
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Prediction error: {str(e)}")

    # ── Generate Grad-CAM ──────────────────────────────────────────────────
    gradcam_b64 = None
    try:
        gradcam_pil = generate_gradcam(
            MODEL, images_dict, WANTED_VIEWS, DEVICE, IMAGE_SIZE
        )
        buf = io.BytesIO()
        gradcam_pil.save(buf, format="PNG")
        gradcam_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
    except Exception as e:
        print(f"  ⚠️ Grad-CAM failed: {e}")

    # ── Clean label for frontend ───────────────────────────────────────────
    # Remove emoji from label string for JSON
    clean_label = "MALIGNANT" if probability >= 0.5 else "BENIGN"

    return JSONResponse({
        "probability":    round(float(probability), 4),
        "label":          clean_label,
        "confidence":     confidence,
        "gradcam_base64": gradcam_b64,
        "views_used":     len(images_dict),
        "tta_used":       use_tta_bool,
        "clinical_used":  clinical_tensor is not None,
        "device":         str(DEVICE),
    })


# ── Run server ────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    print("=" * 55)
    print("ThermoCAN -- Starting Web Server")
    print("=" * 55)
    print(f"   Open browser: http://localhost:8000")
    print(f"   Press Ctrl+C to stop")
    print("=" * 55)

    uvicorn.run(
        app,
        host="0.0.0.0",
        port=7860,
        reload=False,
        log_level="info",
    )