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)