README.md

metadata

title: SBM Stratify
emoji: 🧠
colorFrom: blue
colorTo: indigo
sdk: docker
app_port: 7860
pinned: false

SBM Stratify Training Pipeline

A streamlined, robust, and highly configurable machine learning pipeline for medical tabular data. It handles both classification and regression automatically, supports complex temporal/predefined splitting strategies, trains Scikit-Learn and PyTorch models, and exports publication-ready scientific plots.

🌐 Web Application (publishable site)

The clinician-facing product is a single, deployable Flask website.

python doctor_app.py

• Public site → http://localhost:5001/ — a professional website for clinicians (overview, predicted outcomes, methodology, and the interactive Risk Calculator).
• Administrator panel → http://localhost:5001/admin/ — the model training studio (grid-search best-parameter training, benchmarks, plots). It is mounted as a Flask Blueprint inside the same app and is intentionally not linked from the clinician navigation — only via discreet footer links — so doctors are not exposed to raw model internals. It is hidden, not access-controlled; anyone with the URL can reach it.

The admin studio can also be run on its own for local development:

python flask_app.py      # admin-only app on http://localhost:5000/admin/

Deploying online

Use a production WSGI server pointing at the single unified app object doctor_app:app:

# Linux / macOS
gunicorn -w 2 -b 0.0.0.0:8000 doctor_app:app

# Windows
waitress-serve --port=8000 doctor_app:app

This serves both the public site (/) and the administrator panel (/admin/) from one process. To restrict the admin panel in production, place an authentication layer (reverse-proxy basic-auth, or an app middleware) in front of the /admin/ path.

🎨 Credits

The brain glyph in the SBM Stratify logo (static/sbm-stratify-logo.svg) is derived from the "brain" icon by Lorc via game-icons.net, licensed under CC BY 3.0.

📂 Project Structure

.
├── train.py                  # Main training execution script
├── preprocessing.py          # Custom scikit-learn transformers (dates, multilabel)
├── utils/
│   ├── logger.py             # Standardized terminal logging
│   └── vis.py                # Publication-ready plotting utilities (matplotlib/seaborn)
├── nn/
│   ├── torch_mlp.py          # PyTorch Multi-Layer Perceptron
│   └── torch_ft_transformer.py # PyTorch FT-Transformer
├── experiments/
│   └── train.sh              # Bash script for easy experiment configuration
├── data_config.json          # Maps the dataset features and target
└── parameters.json           # Defines hyperparameters for the models

⚙️ Configuration

1. data_config.json

Defines your dataset. Group your features appropriately so the pipeline knows how to scale and encode them.

{
    "input_file": "data/SBM1212.xlsx",
    "input_features": ["Age", "Sex", "Pre-Op KPS", "Radio_Tumor side"],
    "cols_string": ["Sex", "Radio_Tumor side"],
    "cols_date": [],
    "cols_multi": []
}

2. parameters.json

Define the hyperparameters for any model you wish to use (hgb, rf, lr, ridge, svc, torch_mlp, torch_ft_transformer).

🚀 Usage

You can run the script directly via python:

python train.py --target "Severe_complication" --split_strategy temporal --date_column "Date of surgery"

Or use the provided bash script for easier experiment management:

cd experiments
./train.sh

Train All Targets With Grid-Search Best Parameters

Run the temporal grid search first. This writes one best_parameters.json file per target under gridsearch/preoperative/.

bash experiments/grid_search.sh

Then train every available model for every configured target using those saved best parameters:

bash experiments/train_from_gridsearch.sh

This script now enables --feature_importance by default, so each trained model also writes:

• feature_importance.csv
• feature_importance.png
• feature_importance.pdf

The outputs are written under outputs/preoperative_from_gridsearch/<target>/<model>/.

🧠 Inference (Loading Saved Weights)

The script automatically saves the entire trained pipeline (imputers, scalers, encoders, and the model itself) as pipeline.joblib. To use this model on new, unseen patients later:

import joblib
import pandas as pd

# Load the saved pipeline
pipeline = joblib.load("benchmark_output/rf/pipeline.joblib")

# Load new patient data (must contain the same features defined in data_config.json)
new_patients = pd.read_csv("new_patients.csv")

# Predict directly! The pipeline handles all preprocessing internally.
predictions = pipeline.predict(new_patients)
probabilities = pipeline.predict_proba(new_patients)

---

📊 How to Read the Generated Plots

When training finishes, the output folder will contain a metrics.json file and several high-resolution (300 DPI) plots tailored for scientific publication.

1. Confusion Matrix (confusion_matrix.png)

• What it shows: A grid comparing the Actual patient outcomes (True Label) against the Predicted outcomes by the model.
• How to read it: * Diagonal cells (top-left to bottom-right) represent correct predictions (True Positives and True Negatives).
  • Off-diagonal cells represent errors (False Positives and False Negatives). In clinical settings, predicting a complication when there isn't one (False Positive) is usually preferred over missing a fatal complication (False Negative).

2. ROC Curve (roc_curve.png)

• What it shows: The trade-off between the True Positive Rate (Sensitivity) and the False Positive Rate (1 - Specificity) across different probability thresholds.
• How to read it: * The dashed diagonal line represents random guessing (AUC = 0.50).
  • The closer the solid curve gets to the top-left corner, the better the model is at distinguishing between the two classes.
  • AUC (Area Under the Curve): A value of 1.0 means perfect separation. A value > 0.80 is generally considered excellent for clinical models.

3. Precision-Recall Curve (pr_curve.png)

• What it shows: The trade-off between Precision (Positive Predictive Value) and Recall (Sensitivity).
• How to read it: This plot is highly recommended over the ROC curve when your dataset is imbalanced (e.g., only 5% of patients have the complication). A model that stays close to the top-right corner is highly effective at finding the rare minority class without throwing too many false alarms.

4. Actual vs Predicted Plot (actual_vs_predicted.png)

• What it shows: Used only for regression tasks (e.g., predicting "Days of hospitalization"). It plots the model's prediction on the Y-axis against the actual truth on the X-axis.
• How to read it: * The red dashed line represents perfect prediction ($y = x$).
  • Points clustered tightly along this line indicate high accuracy.
  • If points fan out heavily at higher values, the model is struggling to predict extreme/high outcomes.