Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import pandas as pd | |
| import numpy as np | |
| import matplotlib | |
| matplotlib.use('Agg') | |
| import matplotlib.pyplot as plt | |
| import seaborn as sns | |
| import joblib | |
| import shap | |
| import pandas as pd | |
| def add_features(df: pd.DataFrame) -> pd.DataFrame: | |
| df = df.copy() | |
| df['creatinine_sodium_ratio'] = df['serum_creatinine'] / df['serum_sodium'] | |
| df['is_elderly'] = (df['age'] >= 65).astype(int) | |
| return df | |
| xgb_pipeline = joblib.load('heart_failure_xgb_pipeline.joblib') | |
| cph_model = joblib.load('heart_failure_cph_model.joblib') | |
| kmeans_artifact = joblib.load('heart_failure_kmeans.joblib') | |
| xgb_model = xgb_pipeline.named_steps['xgb'] | |
| xgb_scaler = xgb_pipeline.named_steps['scaler'] | |
| kmeans_model = kmeans_artifact['model'] | |
| kmeans_scaler = kmeans_artifact['scaler'] | |
| kmeans_features = kmeans_artifact['features'] | |
| cluster_profiles = kmeans_artifact['profiles'] | |
| explainer = shap.TreeExplainer(xgb_model) | |
| PHENOTYPE_NAMES = { | |
| 0: "Phenotype A", | |
| 1: "Phenotype B", | |
| 2: "Phenotype C", | |
| } | |
| PHENOTYPE_COLORS = { | |
| 0: "#6366F1", | |
| 1: "#F59E0B", | |
| 2: "#EC4899", | |
| } | |
| def describe_phenotype(cluster_id, profiles): | |
| """Generates a human-readable clinical description for a cluster.""" | |
| row = profiles.loc[cluster_id] | |
| mortality_pct = row['DEATH_EVENT'] * 100 | |
| age_desc = "elderly" if row['age'] >= 65 else "middle-aged" if row['age'] >= 50 else "younger" | |
| ef_desc = "reduced" if row['ejection_fraction'] < 35 else "mildly reduced" if row['ejection_fraction'] < 45 else "preserved" | |
| cr_desc = "elevated" if row['serum_creatinine'] > 1.5 else "borderline" if row['serum_creatinine'] > 1.2 else "normal" | |
| return ( | |
| f"Predominantly {age_desc} patients (avg age {row['age']:.0f}) with " | |
| f"{ef_desc} ejection fraction ({row['ejection_fraction']:.0f}%) and " | |
| f"{cr_desc} serum creatinine ({row['serum_creatinine']:.1f} mg/dL). " | |
| f"Cohort mortality rate: {mortality_pct:.0f}%." | |
| ) | |
| def predict_and_explain( | |
| age, anaemia, cpk, diabetes, ejection_fraction, | |
| high_blood_pressure, platelets, serum_creatinine, | |
| serum_sodium, sex, smoking | |
| ): | |
| anaemia_val = 1 if anaemia == "Yes" else 0 | |
| diabetes_val = 1 if diabetes == "Yes" else 0 | |
| hbp_val = 1 if high_blood_pressure == "Yes" else 0 | |
| sex_val = 1 if sex == "Male" else 0 | |
| smoking_val = 1 if smoking == "Yes" else 0 | |
| patient_dict = { | |
| 'age': float(age), | |
| 'anaemia': int(anaemia_val), | |
| 'creatinine_phosphokinase': int(cpk), | |
| 'diabetes': int(diabetes_val), | |
| 'ejection_fraction': int(ejection_fraction), | |
| 'high_blood_pressure': int(hbp_val), | |
| 'platelets': float(platelets), | |
| 'serum_creatinine': float(serum_creatinine), | |
| 'serum_sodium': int(serum_sodium), | |
| 'sex': int(sex_val), | |
| 'smoking': int(smoking_val) | |
| } | |
| df_patient = pd.DataFrame([patient_dict]) | |
| df_patient = add_features(df_patient) | |
| prob = float(xgb_pipeline.predict_proba(df_patient)[0][1]) | |
| risk_label = "HIGH RISK" if prob >= 0.5 else "STANDARD RISK" | |
| color = "#FF4B4B" if prob >= 0.5 else "#2E7D32" | |
| risk_md = f"<div style='text-align: center; padding: 15px; border-radius: 8px; background-color: {color}20; border: 2px solid {color};'>" | |
| risk_md += f"<p style='margin: 0; font-size: 14px; font-weight: bold; color: {color};'>CLASSIFICATION</p>" | |
| risk_md += f"<h2 style='margin: 5px 0 0 0; font-size: 28px; font-weight: 900; color: {color};'>{risk_label}</h2>" | |
| risk_md += f"<p style='margin: 5px 0 0 0; font-size: 16px; color: {color};'>Calculated Risk Probability: {prob:.2%}</p>" | |
| risk_md += "</div>" | |
| patient_continuous = df_patient[kmeans_features].values | |
| patient_scaled_cluster = kmeans_scaler.transform(patient_continuous) | |
| cluster_id = int(kmeans_model.predict(patient_scaled_cluster)[0]) | |
| pheno_name = PHENOTYPE_NAMES[cluster_id] | |
| pheno_color = PHENOTYPE_COLORS[cluster_id] | |
| pheno_desc = describe_phenotype(cluster_id, cluster_profiles) | |
| pheno_md = f"<div style='padding: 15px; border-radius: 8px; background-color: {pheno_color}15; border: 2px solid {pheno_color};'>" | |
| pheno_md += f"<p style='margin: 0; font-size: 14px; font-weight: bold; color: {pheno_color};'>PATIENT PHENOTYPE</p>" | |
| pheno_md += f"<h2 style='margin: 5px 0 0 0; font-size: 24px; font-weight: 900; color: {pheno_color};'>{pheno_name}</h2>" | |
| pheno_md += f"<p style='margin: 8px 0 0 0; font-size: 13px; color: #D1D5DB; line-height: 1.5;'>{pheno_desc}</p>" | |
| pheno_md += "<div style='margin-top: 12px; display: flex; gap: 8px; flex-wrap: wrap;'>" | |
| profile_row = cluster_profiles.loc[cluster_id] | |
| labels = { | |
| 'age': 'Avg Age', 'ejection_fraction': 'Avg EF%', | |
| 'serum_creatinine': 'Avg Creatinine', 'serum_sodium': 'Avg Sodium', | |
| 'creatinine_phosphokinase': 'Avg CPK', 'platelets': 'Avg Platelets' | |
| } | |
| for feat in kmeans_features: | |
| val = profile_row[feat] | |
| fmt = f"{val:,.0f}" if feat in ('platelets', 'creatinine_phosphokinase') else f"{val:.1f}" | |
| pheno_md += f"<span style='background: {pheno_color}30; color: {pheno_color}; padding: 3px 8px; border-radius: 4px; font-size: 11px; font-weight: 600;'>{labels[feat]}: {fmt}</span>" | |
| pheno_md += "</div></div>" | |
| patient_scaled_xgb = pd.DataFrame(xgb_scaler.transform(df_patient), columns=df_patient.columns) | |
| shap_vals = explainer(patient_scaled_xgb) | |
| plt.close('all') | |
| shap.plots.waterfall(shap_vals[0], show=False) | |
| fig_shap = plt.gcf() | |
| fig_shap.set_size_inches(9, 4.5) | |
| plt.tight_layout() | |
| patient_cph_df = df_patient.copy() | |
| surv = cph_model.predict_survival_function(patient_cph_df) | |
| plt.figure(figsize=(9, 4.5)) | |
| fig_surv = plt.gcf() | |
| plt.plot(surv.index, surv.values, color='dodgerblue', linewidth=2.5) | |
| plt.title("Predicted Patient Survival Probability Over Time", fontsize=12, fontweight='bold', pad=15) | |
| plt.xlabel("Follow-up Duration (Days)", fontsize=10) | |
| plt.ylabel("Survival Probability", fontsize=10) | |
| plt.ylim(0, 1.05) | |
| plt.grid(True, linestyle='--', alpha=0.5) | |
| plt.tight_layout() | |
| return risk_md, pheno_md, fig_surv, fig_shap | |
| app = gr.Blocks() | |
| with app: | |
| gr.HTML(""" | |
| <div style="text-align: center; margin-bottom: 25px;"> | |
| <h1 style="font-weight: 900; font-size: 32px; color: #1E3A8A; margin-bottom: 5px;">Heart Failure Mortality & Survival Predictor</h1> | |
| <p style="font-size: 16px; color: #4B5563; margin-top: 0;">Dual ML Pipeline: XGBoost Risk Classification & Cox Proportional Hazards Survival Curves</p> | |
| </div> | |
| """) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| gr.HTML(""" | |
| <div style="background-color: #F3F4F6; padding: 10px 15px; border-radius: 6px; border-left: 4px solid #1E3A8A; margin-bottom: 15px;"> | |
| <h3 style="margin: 0; font-size: 16px; font-weight: bold; color: #1E3A8A;">Patient Clinical Profile</h3> | |
| </div> | |
| """) | |
| age = gr.Slider(minimum=40, maximum=100, value=65, step=1, label="Age (Years)") | |
| with gr.Row(): | |
| sex = gr.Dropdown(choices=["Female", "Male"], value="Male", label="Sex") | |
| smoking = gr.Dropdown(choices=["No", "Yes"], value="No", label="Smoking Status") | |
| with gr.Row(): | |
| anaemia = gr.Dropdown(choices=["No", "Yes"], value="No", label="Anaemia") | |
| diabetes = gr.Dropdown(choices=["No", "Yes"], value="No", label="Diabetes") | |
| high_blood_pressure = gr.Dropdown(choices=["No", "Yes"], value="No", label="Hypertension") | |
| ejection_fraction = gr.Slider(minimum=10, maximum=80, value=35, step=1, label="Ejection Fraction (%)") | |
| serum_creatinine = gr.Slider(minimum=0.5, maximum=10.0, value=1.2, step=0.1, label="Serum Creatinine (mg/dL)") | |
| serum_sodium = gr.Slider(minimum=110, maximum=150, value=137, step=1, label="Serum Sodium (mEq/L)") | |
| with gr.Row(): | |
| cpk = gr.Number(value=250, label="CPK Enzyme Level (mcg/L)", precision=0) | |
| platelets = gr.Number(value=263000, label="Platelets Count (kiloplatelets/mL)", precision=0) | |
| predict_btn = gr.Button("Evaluate Patient Risk Profile", variant="primary") | |
| with gr.Column(scale=2): | |
| gr.HTML(""" | |
| <div style="background-color: #F3F4F6; padding: 10px 15px; border-radius: 6px; border-left: 4px solid #10B981; margin-bottom: 15px;"> | |
| <h3 style="margin: 0; font-size: 16px; font-weight: bold; color: #065F46;">Diagnostic Results</h3> | |
| </div> | |
| """) | |
| output_risk = gr.HTML(value="<div style='text-align: center; color: #6B7280; padding: 20px;'>Enter patient parameters on the left and click 'Evaluate' to run diagnostic pipeline.</div>") | |
| output_pheno = gr.HTML() | |
| output_surv = gr.Plot(label="Survival Probability Curve") | |
| output_shap = gr.Plot(label="SHAP Risk Attribution") | |
| predict_btn.click( | |
| fn=predict_and_explain, | |
| inputs=[ | |
| age, anaemia, cpk, diabetes, ejection_fraction, | |
| high_blood_pressure, platelets, serum_creatinine, | |
| serum_sodium, sex, smoking | |
| ], | |
| outputs=[output_risk, output_pheno, output_surv, output_shap] | |
| ) | |
| if __name__ == '__main__': | |
| app.launch(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="slate")) | |