IVF Success Predictor

A two-stage statistical model that estimates IVF success probabilities for a given patient profile.

Stage 1 — a survival model (Random Survival Forest or Cox PH) predicts the probability of accumulating k ≥ 1…10 transferable (euploid/mosaic) embryos by retrieval cycle c.

Stage 2 — a dynamic-programming implantation model converts the embryo distribution into clinical pregnancy probabilities for 1–5 pregnancies, with the per-attempt success rate varying by age.

Both stages are exposed through a small Flask web app with a Chart.js frontend, and also available as a standalone command-line script.

---

Setup

python -m venv venv
source venv/bin/activate          # Windows: venv\Scripts\activate
pip install -r requirements.txt

---

Running the web app

The trained model files (model_rsf.pkl/model_cox.pkl and model_results.pkl) are not included in this repository — they must be generated from your own dataset (see Training below).

model_results.pkl (the clinical pregnancy / implantation model) is optional. There's no script in this repo to train it — it needs a separate per-patient dataset with a binary pregnant column — so if you only have the embryo accumulation model, that's fine: the app detects that model_results.pkl is missing and automatically serves embryo-accumulation predictions only, with the clinical pregnancy tab disabled in the UI.

Once you have at least the embryo model file:

python app.py

Open http://localhost:8000 in your browser. To use the Cox model instead of RSF:

MODEL_TYPE=cox python app.py

---

Command-line predictions

To get predictions without starting the web server, edit the patient values at the top of predict.py and run:

python predict.py

---

Training on your own data

IVF success rates vary between clinics due to differences in patient populations and lab protocols. Training on your own data produces predictions calibrated to your patient population.

Dataset format

Prepare a CSV where each row represents one embryo retrieved in one cycle for one patient. Required columns:

Column	Description
pid	Unique patient identifier
opuid	Retrieval procedure ID (used to order cycles per patient)
ploidy_encoding	Embryo ploidy: 1 = euploid, 0 = mosaic, -1 = aneuploid
age_of_oocyte_provider	Age of the egg provider (years)
BMI_of_pt	Patient BMI (kg/m²)
AFC_oocyte_provider	Antral follicle count
AMH_oocyte_provider_pmoll	AMH level (pmol/L)
SPERM_PROVIDER_AGE	Sperm provider age (years)

If your column names differ, update the columns section in train_config.yaml.

You will also need a separate dataset with one row per patient and a binary pregnant column (1/0) to retrain the implantation probability model (model_results.pkl).

Training

Fill in train_config.yaml with your data file path and any column name overrides, then run:

source venv/bin/activate
python train_rsf.py

This overwrites model_rsf.pkl and saves evaluation metrics to metrics_rsf.json. Restart the Flask app to pick up the new model.

To train the Cox model instead:

python train_cox.py --data your_data.csv

Configuration

train_config.yaml covers all training parameters so you don't need to pass CLI flags:

data_file: your_data.csv
output_model: model_rsf.pkl
test_size: 0.20
random_seed: 42
targets: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

columns:
  patient_id: pid
  age: age_of_oocyte_provider
  # ... etc

CLI flags always override the config file if both are provided.

---

Project structure

app.py              Flask web app (loads models, serves predictions)
model.py            IVFSurvivalAnalysis (Cox) and IVFRSFAnalysis (RSF) classes
train_rsf.py        Train the RSF embryo accumulation model
train_cox.py        Train the Cox PH embryo accumulation model
train_config.yaml   Training parameters (edit before running train_rsf.py)
predict.py          Standalone CLI predictions without the web server
templates/
  index.html        Single-page frontend
requirements.txt

---

Disclaimer

This tool does not provide medical advice, diagnosis, or treatment. Estimates are for informational purposes only and may not reflect actual outcomes. Estimates are less reliable at the extremes of age, BMI, AMH, and AFC. Please consult a doctor or healthcare provider for personalised advice.