import io import os from typing import List import pickle import numpy as np import pandas as pd import torch import torch.nn as nn from fastapi import FastAPI, File, UploadFile, HTTPException from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel from safetensors.torch import load_file # ----------------------------- # Config # ----------------------------- FEATURES = [ "ESS_TOTAL", "MCATOT", "GDS_TOTAL", "MCAALTTM", "MCACUBE", "MCASER7", "MCAABSTR", "GDSSATIS", "GDSHAPPY", "GDSENRGY", "ESS1", "ESS2" ] TARGETS = [ "DATSCAN_CAUDATE_R", "DATSCAN_CAUDATE_L", "DATSCAN_PUTAMEN_R", "DATSCAN_PUTAMEN_L" ] MODEL_PATHS = { "state": "model/flake_transformer.safetensors", "sx": "model/scaler_x.pkl", "sy": "model/scaler_y.pkl", } # ----------------------------- # App # ----------------------------- app = FastAPI(title="PD Biomarker Predictor", version="1.0.0") app.add_middleware( CORSMiddleware, allow_origins=["*"], # Tighten for production allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # ----------------------------- # Model definition # ----------------------------- class flakeParkinsonTransformer(nn.Module): def __init__(self, num_features: int, output_dim: int, embed_dim=64, nhead=4, num_layers=3): super().__init__() self.feature_embeds = nn.ModuleList([nn.Linear(1, embed_dim) for _ in range(num_features)]) encoder_layer = nn.TransformerEncoderLayer( d_model=embed_dim, nhead=nhead, dim_feedforward=256, dropout=0.1 ) self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers) self.pool = nn.AdaptiveAvgPool1d(1) self.fc = nn.Sequential( nn.Linear(embed_dim, 128), nn.ReLU(), nn.Dropout(0.1), nn.Linear(128, output_dim), ) def forward(self, x): embeds = [self.feature_embeds[i](x[:, i].unsqueeze(1)) for i in range(x.shape[1])] x = torch.stack(embeds, dim=1) x = x.permute(1, 0, 2) x = self.transformer(x) x = x.permute(1, 2, 0) x = self.pool(x).squeeze(2) x = self.fc(x) return x # ----------------------------- # Load model and scalers # ----------------------------- def load_model_and_scalers(): # Load scalers using pickle try: with open(MODEL_PATHS["sx"], 'rb') as f: scaler_x = pickle.load(f) print("scaler_x loaded successfully.") except Exception as e: raise RuntimeError(f"Error loading scaler_x: {e}") try: with open(MODEL_PATHS["sy"], 'rb') as f: scaler_y = pickle.load(f) print("scaler_y loaded successfully.") except Exception as e: raise RuntimeError(f"Error loading scaler_y: {e}") # Load model using safetensors try: state_dict = load_file(MODEL_PATHS["state"], device="cpu") model = flakeParkinsonTransformer(num_features=len(FEATURES), output_dim=len(TARGETS)) model.load_state_dict(state_dict) model.eval() print("Model loaded and set to evaluation mode.") except Exception as e: raise RuntimeError(f"Error loading model state dictionary: {e}") return model, scaler_x, scaler_y MODEL, SCALER_X, SCALER_Y = load_model_and_scalers() # ----------------------------- # Utilities # ----------------------------- def try_compute_ess_total(df: pd.DataFrame) -> pd.DataFrame: if "ESS_TOTAL" not in df.columns: ess_cols = [f"ESS{i}" for i in range(1, 9) if f"ESS{i}" in df.columns] if len(ess_cols) == 8: df["ESS_TOTAL"] = df[ess_cols].sum(axis=1) return df def try_compute_gds_total(df: pd.DataFrame) -> pd.DataFrame: gds_items = [ "GDSSATIS", "GDSDROPD", "GDSEMPTY", "GDSBORED", "GDSGSPIR", "GDSAFRAD", "GDSHAPPY", "GDSHLPLS", "GDSHOME", "GDSMEMRY", "GDSALIVE", "GDSWRTLS", "GDSENRGY", "GDSHOPLS", "GDSBETER" ] if "GDS_TOTAL" not in df.columns: present = [c for c in gds_items if c in df.columns] if len(present) == 15: df["GDS_TOTAL"] = df[present].sum(axis=1) return df def standardize_and_predict(df_features: pd.DataFrame): X = df_features[FEATURES].values X = SCALER_X.transform(X) X_t = torch.tensor(X, dtype=torch.float32) with torch.no_grad(): preds = MODEL(X_t).numpy() preds = SCALER_Y.inverse_transform(preds) return preds def merge_four_frames(ess: pd.DataFrame, moca: pd.DataFrame, gds: pd.DataFrame, dat: pd.DataFrame) -> pd.DataFrame: ess = try_compute_ess_total(ess) gds = try_compute_gds_total(gds) for df, name in [(ess, "ESS"), (moca, "MoCA"), (gds, "GDS"), (dat, "DaTSCAN")]: if not all(col in df.columns for col in ["PATNO", "EVENT_ID"]): raise HTTPException(status_code=400, detail=f"{name} CSV missing PATNO or EVENT_ID columns") df = ess.merge(moca, on=["PATNO", "EVENT_ID"], how="inner", suffixes=("_ess", "_moca")) df = df.merge(gds, on=["PATNO", "EVENT_ID"], how="inner") df = df.merge(dat, on=["PATNO", "EVENT_ID"], how="inner") missing = [f for f in FEATURES if f not in df.columns] if missing: raise HTTPException(status_code=400, detail=f"Merged CSVs missing required features: {missing}") return df def detect_file_kind(name: str) -> str: l = name.lower() if "datscan" in l or ("dat" in l and "scan" in l): return "datscan" if "moca" in l: return "moca" if "gds" in l: return "gds" if "ess" in l: return "ess" return "unknown" # ----------------------------- # Schemas # ----------------------------- class PatientData(BaseModel): ESS_TOTAL: float MCATOT: float GDS_TOTAL: float MCAALTTM: float MCACUBE: float MCASER7: float MCAABSTR: float GDSSATIS: float GDSHAPPY: float GDSENRGY: float ESS1: float ESS2: float # ----------------------------- # Endpoints # ----------------------------- @app.post("/predict") def predict_json(data: PatientData): df = pd.DataFrame([data.dict()]) preds = standardize_and_predict(df) out = preds[0].tolist() return { "predicted_biomarkers": { TARGETS[0]: out[0], TARGETS[1]: out[1], TARGETS[2]: out[2], TARGETS[3]: out[3], }, "source": "json", "rows": 1 } @app.post("/predict/files") async def predict_files(files: List[UploadFile] = File(...)): if len(files) < 4: raise HTTPException(status_code=400, detail="Please upload four CSV files: ESS, MoCA, GDS, DaTSCAN.") buckets = {"ess": None, "moca": None, "gds": None, "datscan": None} fallback = [] for f in files: content = await f.read() try: df = pd.read_csv(io.BytesIO(content)) except Exception: raise HTTPException(status_code=400, detail=f"Could not parse CSV: {f.filename}") kind = detect_file_kind(f.filename) if kind in buckets and buckets[kind] is None: buckets[kind] = df else: fallback.append((kind, df)) if buckets["ess"] is None: candidates = [df for kind, df in fallback if "ESS1" in df.columns or "ESS_TOTAL" in df.columns] if candidates: buckets["ess"] = candidates[0] if buckets["moca"] is None: candidates = [df for kind, df in fallback if "MCATOT" in df.columns] if candidates: buckets["moca"] = candidates[0] if buckets["gds"] is None: candidates = [df for kind, df in fallback if "GDS_TOTAL" in df.columns or "GDSSATIS" in df.columns] if candidates: buckets["gds"] = candidates[0] if buckets["datscan"] is None: candidates = [df for kind, df in fallback if any(c.startswith("DATSCAN_") for c in df.columns)] if candidates: buckets["datscan"] = candidates[0] if any(v is None for v in buckets.values()): raise HTTPException(status_code=400, detail="Could not identify all four CSVs (ESS, MoCA, GDS, DaTSCAN) by filename/columns.") merged = merge_four_frames(buckets["ess"], buckets["moca"], buckets["gds"], buckets["datscan"]) preds = standardize_and_predict(merged) # Create a list of predictions for each patient results = [] for idx, pred in enumerate(preds): result = { "PATNO": int(merged.iloc[idx]["PATNO"]), "EVENT_ID": str(merged.iloc[idx]["EVENT_ID"]), "predicted_biomarkers": { TARGETS[0]: float(pred[0]), TARGETS[1]: float(pred[1]), TARGETS[2]: float(pred[2]), TARGETS[3]: float(pred[3]), } } results.append(result) return { "predictions": results, "source": "files", "merged_rows": int(merged.shape[0]) }