app.py

import gradio as gr
import joblib
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import json
from catboost import CatBoostRegressor, Pool
from huggingface_hub import hf_hub_download
import spaces                                       # GPU

REPO_ID = "Elegbede/stunting-catboost"

def _dl(fname):
    return hf_hub_download(repo_id=REPO_ID, filename=fname)

model = CatBoostRegressor()
model.load_model(_dl("catboost_final.cbm"))

features     = json.load(open(_dl("features.json")))
cat_features = json.load(open(_dl("cat_features.json")))
breslow      = joblib.load(_dl("breslow_cat.pkl"))
risk_centre  = float(np.load(_dl("risk_centre.npy"))[0])

wealth_map  = {"poorest": 0, "poorer": 1, "middle": 2, "richer": 3, "richest": 4}
mat_edu_map = {"no education": 0, "primary": 1, "secondary": 2, "higher": 3}


def create_inputs():
    return [
        gr.Number(label="Birth Order", value=1, minimum=1, maximum=20),
        gr.Number(label="Maternal Age (years)", value=25, minimum=12, maximum=49),
        gr.Dropdown(
            choices=[("Poorest", "poorest"), ("Poorer", "poorer"), ("Middle", "middle"),
                     ("Richer", "richer"), ("Richest", "richest")],
            label="Wealth Index", value="middle"),
        gr.Dropdown(
            choices=[("No Education", "no education"), ("Primary", "primary"),
                     ("Secondary", "secondary"), ("Higher", "higher")],
            label="Maternal Education", value="primary"),
        gr.Dropdown(choices=[("No", 0), ("Yes", 1)], label="Multiple Birth", value=0),
        gr.Dropdown(choices=["male", "female"], label="Child Sex", value="male"),
        gr.Dropdown(
            choices=["islam", "catholic", "other christian", "traditionalist", "other"],
            label="Religion", value="islam"),
        gr.Dropdown(choices=["Improved", "Unimproved"],
                    label="Water Source", value="Improved"),
        gr.Dropdown(choices=["Improved", "Unimproved"],
                    label="Sanitation", value="Improved"),
        gr.Dropdown(choices=["no education", "primary", "secondary", "higher"],
                    label="Partner Education", value="primary"),
        gr.Dropdown(
            choices=["North West", "North East", "North Central",
                     "South West", "South East", "South South"],
            label="Region", value="North West"),
        gr.Dropdown(choices=["rural", "urban"], label="Residence", value="rural"),
    ]


@spaces.GPU                                          # GPU
def predict_stunting(birth_order, mat_age, wealth_label, mat_edu_label,
                     multiple_birth_binary, child_sex, religion,
                     water_improved, sanitation_improved,
                     partner_edu, region, residence):

    row = pd.DataFrame([{
        "birth_order":           float(birth_order),
        "mat_age":               float(mat_age),
        "wealth_ordinal":        float(wealth_map[wealth_label]),
        "mat_edu_ordinal":       float(mat_edu_map[mat_edu_label]),
        "multiple_birth_binary": float(multiple_birth_binary),
        "child_sex":             str(child_sex),
        "religion":              str(religion),
        "water_improved":        str(water_improved),
        "sanitation_improved":   str(sanitation_improved),
        "partner_edu":           str(partner_edu),
        "region":                str(region),
        "residence":             str(residence),
    }])[features]

    pool = Pool(row, cat_features=cat_features)
    log_hazard = float(model.predict(pool, prediction_type="RawFormulaVal")[0])
    log_hazard -= risk_centre

    times = np.array([6, 12, 18, 24, 36, 48])
    surv_fn = breslow.get_survival_function(np.array([log_hazard]))[0]
    domain_min, domain_max = surv_fn.domain
    surv_probs = surv_fn(np.clip(times, domain_min, domain_max))
    stunting_risks = 1 - surv_probs

    fig, ax = plt.subplots(figsize=(8, 4))
    ax.plot(times, stunting_risks, marker="o", linestyle="-",
            color="steelblue", linewidth=2, markersize=8)
    ax.axhline(y=0.20, color="red", linestyle="--", label="20% Risk Threshold")
    ax.set_xlabel("Age (months)")
    ax.set_ylabel("Stunting Risk")
    ax.set_title("Stunting Risk by Age")
    ax.set_ylim(0, 1)
    ax.grid(True, linestyle=":", alpha=0.7)
    ax.legend()
    ax.set_xticks(times)

    THRESH = 0.20
    age_lines = "\n".join(f"- **{int(a)} months:** {r:.1%} estimated risk"
                          for a, r in zip(times, stunting_risks))
    crossing = next((int(a) for a, r in zip(times, stunting_risks) if r >= THRESH), None)
    threshold_text = (
        f"Predicted stunting risk reaches the {THRESH:.0%} threshold by about "
        f"**{crossing} months**."
        if crossing else
        f"Predicted stunting risk does not reach {THRESH:.0%} within 48 months."
    )

    flagged = []
    if wealth_map[wealth_label] <= 1:        flagged.append("lower household wealth quintile")
    if mat_edu_map[mat_edu_label] == 0:      flagged.append("mother has no formal education")
    if water_improved == "Unimproved":       flagged.append("unimproved water source")
    if sanitation_improved == "Unimproved":  flagged.append("unimproved sanitation")
    if multiple_birth_binary == 1:           flagged.append("multiple birth")
    if mat_age < 18 or mat_age > 40:         flagged.append("maternal age outside 18-40")
    if birth_order >= 5:                     flagged.append("birth order 5th or later")

    factors_text = ("Factors flagged as elevated-risk in this input:\n"
                    + "\n".join(f"- {f}" for f in flagged)) if flagged else \
                   "No major elevated-risk factors flagged."

    summary_md = f"""
## Stunting Risk Prediction (CatBoost Cox model)

### Predicted risk over time
{age_lines}

{threshold_text}

### Factors noted
{factors_text}

---
⚠️ **Research prototype only** — external validation not performed.
Do not use for clinical or policy decisions.
"""
    return summary_md, fig


demo = gr.Interface(
    fn=predict_stunting,
    inputs=create_inputs(),
    outputs=[gr.Markdown(label="Risk Summary"), gr.Plot(label="Risk Curve")],
    title="Stunting Risk Predictor — CatBoost",
    description="""
    Predicts stunting risk at ages 6, 12, 18, 24, 36, and 48 months using a
    **CatBoost Cox model** trained on 42,981 Nigerian children.

    ⚠️ **Research prototype** — external validation not performed.
    """,
    examples=[
        [1, 25, "middle",  "primary",       0, "male",   "islam", "Improved",   "Improved",   "primary",      "North West", "rural"],
        [4, 18, "poorest", "no education",  0, "female", "catholic", "Unimproved", "Unimproved", "no education", "South East", "rural"],
        [2, 35, "richest", "higher",        1, "male",   "other christian", "Improved", "Improved", "secondary", "South West", "urban"],
    ],
    flagging_mode="never",
)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)