Update app.py

app.py

@@ -1,105 +1,225 @@
-# ══════════════════════════════════════════════════════════════════
-# STANDALONE ECG PIPELINE TEST  —  paste into a fresh Colab cell
-# Reconnect runtime first (Runtime > Reconnect), then run this whole cell.
-# It loads your V3.1 ensemble from Drive and lets you upload an ECG.
-# ══════════════════════════════════════════════════════════════════
-!pip install -q timm albumentations grad-cam opencv-python-headless 2>/dev/null
+"""
+ECG -> ECHO Screening (V3.1)
+Predicts ejection fraction (regression) and significant RWMA (binary screen)
+from a 12-lead ECG image. Feasibility / research demo -- NOT for clinical use.
 
-from google.colab import drive, files
-drive.mount('/content/drive')
+Model: 5-fold ensemble, EfficientNet-B3 backbone (PTB-XL pretrained),
+       multi-task heads (EF regression + binary RWMA).
+"""
 
-import numpy as np, cv2, torch, torch.nn as nn, timm
+import os
+import numpy as np
+import cv2
+from PIL import Image
+import streamlit as st
+import torch
+import torch.nn as nn
+import timm
 import albumentations as A
 from albumentations.pytorch import ToTensorV2
-import matplotlib.pyplot as plt
 
-MODEL_PATH = '/content/drive/MyDrive/ecg_echo_ai/outputs_clinical_v3_binary/ensemble_clinical_v3_binary.pth'
-BACKBONE, IMG_SIZE, EDGE_CROP = 'tf_efficientnet_b3.ns_jft_in1k', 384, 0.05
-DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+# Config -- must match training exactly
+MODEL_PATH = "ensemble_clinical_v3_binary.pth"
+BACKBONE   = "tf_efficientnet_b3.ns_jft_in1k"
+IMG_SIZE   = 384
+EDGE_CROP  = 0.05
+DEVICE     = torch.device("cpu")
+
+st.set_page_config(page_title="ECG -> ECHO Screening", page_icon=":anatomical_heart:", layout="wide")
+
 
 class ECGNetV3(nn.Module):
     def __init__(self, dropout=0.45):
         super().__init__()
-        self.backbone = timm.create_model(BACKBONE, pretrained=False, num_classes=0,
-                                          global_pool='avg', drop_rate=0.2)
-        f = self.backbone.num_features
+        self.backbone = timm.create_model(
+            BACKBONE, pretrained=False, num_classes=0,
+            global_pool="avg", drop_rate=0.2,
+        )
+        feat = self.backbone.num_features
         self.neck = nn.Sequential(
-            nn.Linear(f,512), nn.LayerNorm(512), nn.GELU(), nn.Dropout(dropout),
-            nn.Linear(512,256), nn.LayerNorm(256), nn.GELU(), nn.Dropout(dropout*0.7),
-            nn.Linear(256,128), nn.LayerNorm(128), nn.GELU(), nn.Dropout(dropout*0.5))
-        self.ef_reg = nn.Linear(128,1); self.rwma_cls = nn.Linear(128,2)
+            nn.Linear(feat, 512), nn.LayerNorm(512), nn.GELU(), nn.Dropout(dropout),
+            nn.Linear(512, 256), nn.LayerNorm(256), nn.GELU(), nn.Dropout(dropout * 0.7),
+            nn.Linear(256, 128), nn.LayerNorm(128), nn.GELU(), nn.Dropout(dropout * 0.5),
+        )
+        self.ef_reg   = nn.Linear(128, 1)
+        self.rwma_cls = nn.Linear(128, 2)
+
     def forward(self, x):
         z = self.neck(self.backbone(x))
-        return {'ef_norm': self.ef_reg(z).squeeze(-1), 'rwma_logits': self.rwma_cls(z)}
+        return {"ef_norm": self.ef_reg(z).squeeze(-1),
+                "rwma_logits": self.rwma_cls(z)}
+
 
-def preprocess(img_bgr):
+def preprocess_image(img_bgr):
     if EDGE_CROP > 0:
-        H, W = img_bgr.shape[:2]; c = EDGE_CROP
+        H, W = img_bgr.shape[:2]
+        c = EDGE_CROP
         img_bgr = img_bgr[int(H*c):int(H*(1-c)), int(W*c):int(W*(1-c))]
-    gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
-    enh  = cv2.createCLAHE(clipLimit=2.5, tileGridSize=(8,8)).apply(gray)
-    sharp = cv2.filter2D(enh, -1, np.array([[0,-1,0],[-1,5,-1],[0,-1,0]]))
+    gray  = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
+    clahe = cv2.createCLAHE(clipLimit=2.5, tileGridSize=(8, 8))
+    enh   = clahe.apply(gray)
+    kernel = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]])
+    sharp = cv2.filter2D(enh, -1, kernel)
     sharp = cv2.normalize(sharp, None, 0, 255, cv2.NORM_MINMAX)
     return cv2.cvtColor(sharp, cv2.COLOR_GRAY2RGB)
 
-val_tf = A.Compose([A.Resize(IMG_SIZE, IMG_SIZE),
-                    A.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]),
-                    ToTensorV2()])
-
-# Load ensemble
-ckpt = torch.load(MODEL_PATH, map_location=DEVICE, weights_only=False)
-models = []
-for s in ckpt['fold_models']:
-    m = ECGNetV3().to(DEVICE); m.load_state_dict(s); m.eval(); models.append(m)
-print(f'Loaded {len(models)}-fold ensemble. CV: '
-      f"EF MAE={ckpt['overall']['ef_mae']:.1f}%, RWMA AUROC={ckpt['overall']['rwma_auroc']:.3f}")
-
-# Upload an ECG
-print('\nUpload an ECG image:')
-up = files.upload()
-fname = list(up.keys())[0]
-img_bgr = cv2.imdecode(np.frombuffer(up[fname], np.uint8), cv2.IMREAD_COLOR)
-proc = preprocess(img_bgr)
-
-# Predict (5-fold ensemble)
-THRESHOLD = 0.45
-tensor = val_tf(image=proc)['image'].unsqueeze(0).to(DEVICE)
-ef_vals, sig_probs = [], []
-with torch.no_grad():
-    for m in models:
-        out = m(tensor)
-        ef_vals.append(float(out['ef_norm'].cpu())*100)
-        sig_probs.append(float(torch.softmax(out['rwma_logits'],-1)[0,1].cpu()))
-ef_mean, ef_std = np.mean(ef_vals), np.std(ef_vals); sig_p = np.mean(sig_probs)
-sev = ('Normal' if ef_mean>=55 else 'Mildly reduced' if ef_mean>=45
-       else 'Moderately reduced' if ef_mean>=35 else 'Severely reduced')
-
-print('\n' + '='*50)
-print(f'  EJECTION FRACTION : {ef_mean:.1f}%  (95% CI {ef_mean-1.96*ef_std:.1f}-{ef_mean+1.96*ef_std:.1f})')
-print(f'  EF severity       : {sev}')
-print(f'  RWMA significant  : {sig_p:.0%} probability  ->  '
-      f'{"SIGNIFICANT (refer for echo)" if sig_p>=THRESHOLD else "Non-significant"}')
-print('='*50)
-
-# Grad-CAM
-from pytorch_grad_cam import GradCAM
-from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
-class WEF(nn.Module):
-    def __init__(s,m): super().__init__(); s.m=m
-    def forward(s,x): return s.m(x)['ef_norm'].unsqueeze(1)
-class WRW(nn.Module):
-    def __init__(s,m): super().__init__(); s.m=m
-    def forward(s,x): return s.m(x)['rwma_logits'][:,1:2]
-def cam(wrap):
-    t = val_tf(image=proc)['image'].unsqueeze(0).to(DEVICE)
-    with GradCAM(model=wrap, target_layers=[models[0].backbone.blocks[-1]]) as g:
-        h = g(input_tensor=t, targets=[ClassifierOutputTarget(0)])[0]
-    h = cv2.resize(h, (proc.shape[1], proc.shape[0]))
-    heat = cv2.cvtColor(cv2.applyColorMap(np.uint8(255*h), cv2.COLORMAP_JET), cv2.COLOR_BGR2RGB)
-    return cv2.addWeighted(proc.astype(np.uint8), 0.55, heat, 0.45, 0)
-
-fig, ax = plt.subplots(1, 3, figsize=(22, 6))
-ax[0].imshow(proc); ax[0].set_title('Preprocessed input'); ax[0].axis('off')
-ax[1].imshow(cam(WEF(models[0]))); ax[1].set_title('EF attention'); ax[1].axis('off')
-ax[2].imshow(cam(WRW(models[0]))); ax[2].set_title('RWMA attention'); ax[2].axis('off')
-plt.tight_layout(); plt.show()
+
+val_tf = A.Compose([
+    A.Resize(IMG_SIZE, IMG_SIZE),
+    A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ToTensorV2(),
+])
+
+
+@st.cache_resource(show_spinner=False)
+def load_models():
+    ckpt = torch.load(MODEL_PATH, map_location="cpu", weights_only=False)
+    states = ckpt["fold_models"]
+    models = []
+    for s in states:
+        m = ECGNetV3().to(DEVICE)
+        m.load_state_dict(s)
+        m.eval()
+        models.append(m)
+    return models, ckpt.get("overall", {})
+
+
+def predict(img_rgb, models, sig_threshold=0.5):
+    tensor = val_tf(image=img_rgb)["image"].unsqueeze(0).to(DEVICE)
+    ef_vals, sig_probs = [], []
+    with torch.no_grad():
+        for m in models:
+            out = m(tensor)
+            ef_vals.append(float(out["ef_norm"].cpu()) * 100)
+            sig_probs.append(float(torch.softmax(out["rwma_logits"], -1)[0, 1].cpu()))
+    ef_mean, ef_std = float(np.mean(ef_vals)), float(np.std(ef_vals))
+    sig_p = float(np.mean(sig_probs))
+
+    if   ef_mean >= 50: ef_sev = "Normal"
+    elif ef_mean >= 40: ef_sev = "Mildly reduced"
+    elif ef_mean >= 30: ef_sev = "Moderately reduced"
+    else:               ef_sev = "Severely reduced"
+
+    return {
+        "ef_value": round(ef_mean, 1),
+        "ef_low":   round(max(0, ef_mean - 1.96 * ef_std), 1),
+        "ef_high":  round(min(100, ef_mean + 1.96 * ef_std), 1),
+        "ef_sev":   ef_sev,
+        "sig_prob": round(sig_p, 3),
+        "sig_flag": sig_p >= sig_threshold,
+    }
+
+
+def grad_cam(model, img_rgb, mode="ef"):
+    from pytorch_grad_cam import GradCAM
+    from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
+
+    class WrapEF(nn.Module):
+        def __init__(s, m): super().__init__(); s.m = m
+        def forward(s, x): return s.m(x)["ef_norm"].unsqueeze(1)
+
+    class WrapRWMA(nn.Module):
+        def __init__(s, m): super().__init__(); s.m = m
+        def forward(s, x): return s.m(x)["rwma_logits"][:, 1:2]
+
+    wrapper = WrapEF(model) if mode == "ef" else WrapRWMA(model)
+    tensor = val_tf(image=img_rgb)["image"].unsqueeze(0).to(DEVICE)
+    with GradCAM(model=wrapper, target_layers=[model.backbone.blocks[-1]]) as cam:
+        g = cam(input_tensor=tensor, targets=[ClassifierOutputTarget(0)])[0]
+    H, W = img_rgb.shape[:2]
+    g = cv2.resize(g, (W, H))
+    heat = cv2.applyColorMap(np.uint8(255 * g), cv2.COLORMAP_JET)
+    heat = cv2.cvtColor(heat, cv2.COLOR_BGR2RGB)
+    return cv2.addWeighted(img_rgb.astype(np.uint8), 0.55, heat, 0.45, 0)
+
+
+st.markdown(
+    "<h1 style='margin-bottom:0'>ECG -> ECHO Screening</h1>"
+    "<p style='color:#666;margin-top:4px'>Estimates ejection fraction and screens for "
+    "significant wall-motion abnormality from a 12-lead ECG image.</p>",
+    unsafe_allow_html=True,
+)
+
+st.warning(
+    "**Research / feasibility demo -- NOT a medical device.** "
+    "Trained on 500 ECG-echo pairs from a single center. Outputs are not validated "
+    "for clinical decisions and must never replace echocardiography or physician judgment."
+)
+
+with st.sidebar:
+    st.header("Settings")
+    threshold = st.slider(
+        "RWMA referral threshold", 0.20, 0.70, 0.45, 0.05,
+        help="Lower = more sensitive (catches more significant cases, more false referrals).",
+    )
+    st.caption("A 12-lead ECG image (phone photo or scan) works best.")
+
+try:
+    models, overall = load_models()
+    model_ok = True
+except Exception as e:
+    model_ok = False
+    st.error(f"Could not load model file `{MODEL_PATH}`. "
+             f"Make sure it is uploaded to this Space.\n\n{e}")
+
+if model_ok and overall:
+    with st.expander("Model performance (cross-validated, n=500)"):
+        c1, c2, c3, c4 = st.columns(4)
+        c1.metric("EF MAE", f"{overall.get('ef_mae', float('nan')):.1f}%")
+        c2.metric("EF within +/-10%", f"{overall.get('ef_within_10', float('nan')):.0f}%")
+        c3.metric("RWMA AUROC", f"{overall.get('rwma_auroc', float('nan')):.2f}")
+        c4.metric("RWMA recall", f"{overall.get('rwma_recall', float('nan')):.2f}")
+
+uploaded = st.file_uploader("Upload a 12-lead ECG image", type=["jpg", "jpeg", "png"])
+
+if uploaded and model_ok:
+    file_bytes = np.frombuffer(uploaded.read(), np.uint8)
+    img_bgr = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
+    if img_bgr is None:
+        st.error("Could not read that image. Try a different file.")
+    else:
+        proc = preprocess_image(img_bgr)
+        with st.spinner("Running 5-model ensemble..."):
+            result = predict(proc, models, sig_threshold=threshold)
+
+        st.subheader("Results")
+        col1, col2 = st.columns(2)
+
+        with col1:
+            ef = result["ef_value"]
+            ef_color = "#2e7d32" if ef >= 50 else "#f9a825" if ef >= 40 else "#e65100" if ef >= 30 else "#c62828"
+            st.markdown(
+                f"<div style='border:1px solid #ddd;border-radius:12px;padding:18px'>"
+                f"<div style='color:#888;font-size:14px'>EJECTION FRACTION</div>"
+                f"<div style='font-size:42px;font-weight:700;color:{ef_color}'>{ef}%</div>"
+                f"<div style='color:#555'>95% CI: {result['ef_low']}-{result['ef_high']}%</div>"
+                f"<div style='margin-top:6px;font-weight:600;color:{ef_color}'>{result['ef_sev']}</div>"
+                f"</div>", unsafe_allow_html=True,
+            )
+
+        with col2:
+            flag = result["sig_flag"]
+            box  = "#c62828" if flag else "#2e7d32"
+            label = "SIGNIFICANT -- consider echo referral" if flag else "Non-significant"
+            st.markdown(
+                f"<div style='border:1px solid #ddd;border-radius:12px;padding:18px'>"
+                f"<div style='color:#888;font-size:14px'>WALL-MOTION ABNORMALITY</div>"
+                f"<div style='font-size:26px;font-weight:700;color:{box};margin-top:6px'>{label}</div>"
+                f"<div style='color:#555;margin-top:8px'>Significant probability: "
+                f"{result['sig_prob']:.0%} (threshold {threshold:.0%})</div>"
+                f"</div>", unsafe_allow_html=True,
+            )
+
+        st.divider()
+        st.subheader("Where the model is looking (Grad-CAM)")
+        st.caption("Heatmaps should fall on the ECG waveforms, not borders or text.")
+        t1, t2, t3 = st.tabs(["Preprocessed input", "EF attention", "RWMA attention"])
+        with t1:
+            st.image(proc, use_column_width=True)
+        with t2:
+            st.image(grad_cam(models[0], proc, "ef"), use_column_width=True)
+        with t3:
+            st.image(grad_cam(models[0], proc, "rwma"), use_column_width=True)
+
+elif not uploaded:
+    st.info("Upload an ECG image to run the pipeline.")