Update app.py

app.py

@@ -1,444 +1,105 @@
-import os
-import tempfile
-import hashlib
-from datetime import datetime
+# ══════════════════════════════════════════════════════════════════
+# 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
 
-import streamlit as st
-import numpy as np
-import torch
-import torch.nn as nn
-import cv2
-from PIL import Image
-import timm
+from google.colab import drive, files
+drive.mount('/content/drive')
+
+import numpy as np, cv2, torch, torch.nn as nn, timm
 import albumentations as A
 from albumentations.pytorch import ToTensorV2
+import matplotlib.pyplot as plt
 
-from pytorch_grad_cam import GradCAM
-from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
-
-
-# ═════════════════════════════════════════════════════════════════════════════
-#  PAGE CONFIG
-# ═════════════════════════════════════════════════════════════════════════════
-st.set_page_config(
-    page_title='ECG → ECHO Screening AI',
-    page_icon='🫀',
-    layout='wide',
-    initial_sidebar_state='collapsed',
-)
-
-st.markdown("""
-<style>
-    #MainMenu {visibility: hidden;}
-    footer {visibility: hidden;}
-    header {visibility: hidden;}
-    html, body, [class*="css"] { font-family: 'Inter', system-ui, sans-serif; }
-    .hero {
-        background: linear-gradient(135deg, #0f172a 0%, #1e3a8a 50%, #3b82f6 100%);
-        color: white; padding: 32px; border-radius: 14px; margin-bottom: 24px;
-        box-shadow: 0 4px 20px rgba(30, 58, 138, 0.3); text-align: center;
-    }
-    .hero h1 { margin: 0; font-size: 2.4em; color: white;}
-    .hero p { margin: 8px 0 0 0; opacity: 0.95; font-size: 1.1em;}
-    .hero small { opacity: 0.7; font-size: 0.85em;}
-    .risk-high {
-        background: linear-gradient(135deg, #dc2626, #b91c1c); color: white;
-        padding: 16px; border-radius: 10px; text-align: center;
-        font-size: 1.4em; font-weight: 600; margin: 10px 0;
-    }
-    .risk-low {
-        background: linear-gradient(135deg, #16a34a, #15803d); color: white;
-        padding: 16px; border-radius: 10px; text-align: center;
-        font-size: 1.4em; font-weight: 600; margin: 10px 0;
-    }
-    .report-card {
-        background: white; border: 1px solid #e5e7eb; border-radius: 12px;
-        padding: 20px; box-shadow: 0 1px 3px rgba(0,0,0,0.05); margin-bottom: 14px;
-    }
-    .report-card h3 { color: #1e3a8a; margin-top: 0;}
-    .disclaimer {
-        background: #fef3c7; border-left: 5px solid #f59e0b; padding: 16px;
-        margin-top: 24px; border-radius: 8px; color: #78350f;
-    }
-    .methodology {
-        background: #f8fafc; border-left: 4px solid #3b82f6;
-        padding: 12px; border-radius: 6px; font-size: 0.9em; margin: 14px 0;
-    }
-    .heatmap-section {
-        background: linear-gradient(135deg, #fef3c7 0%, #fed7aa 100%);
-        padding: 20px; border-radius: 12px; margin: 20px 0;
-        border: 1px solid #fcd34d;
-    }
-    .heatmap-section h2 { color: #78350f; margin-top: 0;}
-    .stButton > button {
-        background: linear-gradient(135deg, #1e3a8a 0%, #3b82f6 100%);
-        color: white; font-weight: 600; border: none;
-        padding: 0.6em 2em; border-radius: 8px; font-size: 1.05em;
-    }
-    .stButton > button:hover {
-        background: linear-gradient(135deg, #1e40af 0%, #2563eb 100%);
-        transform: translateY(-1px); box-shadow: 0 4px 12px rgba(59, 130, 246, 0.4);
-    }
-</style>
-""", unsafe_allow_html=True)
-
-
-# ═════════════════════════════════════════════════════════════════════════════
-#  CONFIG + MODEL + PREPROCESSING
-# ═════════════════════════════════════════════════════════════════════════════
-IMG_SIZE      = 384
-BACKBONE      = 'tf_efficientnet_b3_ns'
-MEAN          = [0.485, 0.456, 0.406]
-STD           = [0.229, 0.224, 0.225]
-DEVICE        = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-N_TTA         = 4 if DEVICE.type == 'cuda' else 1
-ENSEMBLE_PATH = 'ensemble_clinical_v2.pth'
+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')
 
-
-class ECGNetV2(nn.Module):
+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,
-        )
-        feat = self.backbone.num_features
+        self.backbone = timm.create_model(BACKBONE, pretrained=False, num_classes=0,
+                                          global_pool='avg', drop_rate=0.2)
+        f = self.backbone.num_features
         self.neck = nn.Sequential(
-            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_cls   = nn.Linear(128, 1)
-        self.rwma_cls = nn.Linear(128, 1)
-
+            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)
     def forward(self, x):
-        z = self.backbone(x)
-        z = self.neck(z)
-        return {
-            'ef_logit':   self.ef_cls(z).squeeze(-1),
-            'rwma_logit': self.rwma_cls(z).squeeze(-1),
-        }
-
-
-def preprocess_ecg(img_path):
-    img = cv2.imread(str(img_path))
-    if img is None:
-        img = cv2.cvtColor(np.array(Image.open(str(img_path)).convert('RGB')),
-                           cv2.COLOR_RGB2BGR)
-    gray   = cv2.cvtColor(img, 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)
+        z = self.neck(self.backbone(x))
+        return {'ef_norm': self.ef_reg(z).squeeze(-1), 'rwma_logits': self.rwma_cls(z)}
+
+def preprocess(img_bgr):
+    if EDGE_CROP > 0:
+        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]]))
+    sharp = cv2.normalize(sharp, None, 0, 255, cv2.NORM_MINMAX)
     return cv2.cvtColor(sharp, cv2.COLOR_GRAY2RGB)
 
-
-tta_aug = A.Compose([
-    A.Resize(IMG_SIZE, IMG_SIZE),
-    A.Normalize(mean=MEAN, std=STD),
-    ToTensorV2(),
-])
-
-
-# ═════════════════════════════════════════════════════════════════════════════
-#  GRAD-CAM
-# ═════════════════════════════════════════════════════════════════════════════
-class GradCAMWrapper(nn.Module):
-    """Wraps ECGNetV2 to expose a single classifier head for Grad-CAM."""
-    def __init__(self, model, task='ef'):
-        super().__init__()
-        self.model = model
-        self.key   = 'ef_logit' if task == 'ef' else 'rwma_logit'
-
-    def forward(self, x):
-        out = self.model(x)
-        # Return shape [N, 1] so ClassifierOutputTarget(0) works
-        logit = out[self.key]
-        if logit.dim() == 1:
-            logit = logit.unsqueeze(1)
-        return logit
-
-
-def generate_gradcam_overlay(model, img_array, task='ef', alpha=0.45):
-    """
-    Generate Grad-CAM heatmap and overlay on the preprocessed ECG image.
-
-    Args:
-        model: ECGNetV2 instance (single model, not ensemble)
-        img_array: preprocessed ECG image, shape (H, W, 3), uint8
-        task: 'ef' or 'rwma'
-        alpha: blend factor for overlay (0=original only, 1=heatmap only)
-
-    Returns:
-        Overlay image (H, W, 3) uint8 with heatmap blended on original
-    """
-    wrapper       = GradCAMWrapper(model, task).to(DEVICE).eval()
-    # conv_head is the 1×1 expansion conv just before global avg pool —
-    # spatially richer than blocks[-1] and gradient-friendly for GradCAM
-    target_layers = [model.backbone.conv_head]
-
-    # Prepare model input
-    tensor = tta_aug(image=img_array)['image'].unsqueeze(0).to(DEVICE)
-
-    # Generate CAM
-    with GradCAM(model=wrapper, target_layers=target_layers) as cam:
-        grayscale_cam = cam(
-            input_tensor=tensor,
-            targets=[ClassifierOutputTarget(0)],
-        )[0]    # shape (H_cam, W_cam), values in [0, 1]
-
-    # Resize heatmap to original image dimensions
-    H, W = img_array.shape[:2]
-    heatmap_resized = cv2.resize(grayscale_cam, (W, H))
-
-    # Apply JET colormap (blue=low attention, red=high attention)
-    heatmap_uint8   = np.uint8(255 * heatmap_resized)
-    heatmap_colored = cv2.applyColorMap(heatmap_uint8, cv2.COLORMAP_JET)
-    heatmap_colored = cv2.cvtColor(heatmap_colored, cv2.COLOR_BGR2RGB)
-
-    # Blend with original
-    overlay = cv2.addWeighted(
-        img_array.astype(np.uint8), 1 - alpha,
-        heatmap_colored, alpha, 0
-    )
-    return overlay
-
-
-# ═════════════════════════════════════════════════════════════════════════════
-#  ENSEMBLE LOADING + PREDICTION
-# ═════════════════════════════════════════════════════════════════════════════
-@st.cache_resource(show_spinner='Loading AI ensemble (first time only — ~60 seconds)...')
-def load_ensemble():
-    if not os.path.exists(ENSEMBLE_PATH):
-        st.error(f'Model file {ENSEMBLE_PATH} not found. Upload it to the Space root.')
-        st.stop()
-    ckpt = torch.load(ENSEMBLE_PATH, map_location=DEVICE, weights_only=False)
-    models = []
-    for sd in ckpt['fold_models']:
-        m = ECGNetV2(dropout=0.45).to(DEVICE)
-        m.load_state_dict(sd)
-        m.eval()
-        models.append(m)
-    return {
-        'ensemble': models,
-        'ef_thr':   float(ckpt['config']['ef_threshold']),
-        'rwma_thr': float(ckpt['config']['rwma_threshold']),
-        'mean_auc': float(np.mean(ckpt['fold_auc'])),
-    }
-
-
-def predict_from_array(img_array, ensemble):
-    """Run ensemble prediction on a preprocessed image array."""
-    ef_probs, rw_probs = [], []
-    for m in ensemble:
-        for _ in range(N_TTA):
-            t = tta_aug(image=img_array)['image'].unsqueeze(0).to(DEVICE)
-            with torch.no_grad():
-                out = m(t)
-            ef_probs.append(float(torch.sigmoid(out['ef_logit']).cpu()))
-            rw_probs.append(float(torch.sigmoid(out['rwma_logit']).cpu()))
-    return float(np.mean(ef_probs)), float(np.mean(rw_probs))
-
-
-# ═════════════════════════════════════════════════════════════════════════════
-#  UI
-# ═════════════════════════════════════════════════════════════════════════════
-st.markdown("""
-<div class="hero">
-    <h1>🫀 ECG → ECHO Clinical Screening AI</h1>
-    <p>Detect LV Dysfunction & Regional Wall Motion Abnormalities from a 12-lead ECG</p>
-    <small>Now with AI Attention Heatmaps • Research Prototype</small>
-</div>
-""", unsafe_allow_html=True)
-
-data = load_ensemble()
-ensemble = data['ensemble']
-EF_THR   = data['ef_thr']
-RWMA_THR = data['rwma_thr']
-MEAN_AUC = data['mean_auc']
-
-col_left, col_right = st.columns([1, 2], gap='large')
-
-with col_left:
-    st.markdown('### 📤 Upload ECG')
-    uploaded_file = st.file_uploader(
-        'Upload a 12-lead ECG image',
-        type=['jpg', 'jpeg', 'png'],
-        label_visibility='collapsed',
-    )
-
-    with st.expander('👤 Patient Information (optional)'):
-        patient_id  = st.text_input('Patient ID', placeholder='e.g. P-001')
-        col_a, col_b = st.columns(2)
-        with col_a:
-            patient_age = st.number_input('Age', min_value=0, max_value=120, value=None)
-        with col_b:
-            patient_sex = st.radio('Sex', ['Male', 'Female'], horizontal=True, index=None)
-
-    analyze_btn = st.button('🔬 Analyze ECG', type='primary', use_container_width=True)
-
-    if uploaded_file:
-        st.image(uploaded_file, caption='Uploaded ECG', use_column_width=True)
-
-
-with col_right:
-    if analyze_btn and uploaded_file is not None:
-        with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as tmp:
-            tmp.write(uploaded_file.getvalue())
-            tmp_path = tmp.name
-
-        try:
-            # Step 1: Preprocess
-            img_preprocessed = preprocess_ecg(tmp_path)
-
-            # Step 2: Run inference
-            with st.spinner(f'🧠 Running inference ({len(ensemble) * N_TTA} votes)...'):
-                ef_prob, rw_prob = predict_from_array(img_preprocessed, ensemble)
-
-            # Step 3: Patient info bar
-            info_parts = []
-            if patient_id:  info_parts.append(f'**Patient:** {patient_id}')
-            if patient_age: info_parts.append(f'**Age:** {patient_age}')
-            if patient_sex: info_parts.append(f'**Sex:** {patient_sex}')
-            info_parts.append(f"**Date:** {datetime.now().strftime('%d-%b-%Y %H:%M')}")
-
-            st.markdown('## 📋 Clinical Screening Report')
-            st.markdown(' &nbsp;|&nbsp; '.join(info_parts))
-
-            # Step 4: Risk cards
-            ef_pct, rw_pct = ef_prob * 100, rw_prob * 100
-            ef_thr_pct     = EF_THR * 100
-            rw_thr_pct     = RWMA_THR * 100
-            ef_high        = ef_prob >= EF_THR
-            rw_high        = rw_prob >= RWMA_THR
-
-            rc1, rc2 = st.columns(2)
-            with rc1:
-                st.markdown('<div class="report-card"><h3>1. Left Ventricular Function</h3>',
-                            unsafe_allow_html=True)
-                if ef_high:
-                    st.markdown('<div class="risk-high">⚠️ HIGH RISK</div>', unsafe_allow_html=True)
-                    action = '**Recommendation:** Urgent cardiology referral'
-                else:
-                    st.markdown('<div class="risk-low">✓ LOW RISK</div>', unsafe_allow_html=True)
-                    action = '**Recommendation:** Routine follow-up'
-                st.markdown(f"""
-                - **Abnormal probability:** {ef_pct:.1f}%
-                - **Decision threshold:** {ef_thr_pct:.1f}%
-                - {action}
-                """)
-                st.progress(ef_prob, text=f'{ef_pct:.1f}%')
-                st.markdown('</div>', unsafe_allow_html=True)
-
-            with rc2:
-                st.markdown('<div class="report-card"><h3>2. Regional Wall Motion</h3>',
-                            unsafe_allow_html=True)
-                if rw_high:
-                    st.markdown('<div class="risk-high">⚠️ HIGH RISK</div>', unsafe_allow_html=True)
-                    action = '**Recommendation:** Consider coronary angiography'
-                else:
-                    st.markdown('<div class="risk-low">✓ LOW RISK</div>', unsafe_allow_html=True)
-                    action = '**Recommendation:** No RWMA detected'
-                st.markdown(f"""
-                - **RWMA probability:** {rw_pct:.1f}%
-                - **Decision threshold:** {rw_thr_pct:.1f}%
-                - {action}
-                """)
-                st.progress(rw_prob, text=f'{rw_pct:.1f}%')
-                st.markdown('</div>', unsafe_allow_html=True)
-
-            # Step 5: Grad-CAM heatmaps
-            st.markdown("""
-            <div class="heatmap-section">
-                <h2>🔥 AI Attention Heatmaps</h2>
-                <p>The colored regions show <b>where the AI focused</b> when making its prediction.
-                Red/yellow areas had the strongest influence, blue areas were less important.
-                This helps clinicians verify the AI is looking at clinically meaningful regions.</p>
-            </div>
-            """, unsafe_allow_html=True)
-
-            try:
-                with st.spinner('🔥 Generating attention maps (~10-15 sec)...'):
-                    # Use first model from ensemble — Grad-CAM is for visualization, single model is fine
-                    overlay_ef   = generate_gradcam_overlay(ensemble[0], img_preprocessed, task='ef')
-                    overlay_rwma = generate_gradcam_overlay(ensemble[0], img_preprocessed, task='rwma')
-
-                tab_ef, tab_rwma, tab_both = st.tabs([
-                    '🫀 EF Attention',
-                    '📊 RWMA Attention',
-                    '👀 Side-by-Side',
-                ])
-
-                with tab_ef:
-                    st.image(overlay_ef,
-                             caption='Regions influencing Ejection Fraction prediction',
-                             use_column_width=True)
-                    st.caption('Bright/warm areas = stronger influence on the EF prediction')
-
-                with tab_rwma:
-                    st.image(overlay_rwma,
-                             caption='Regions influencing Regional Wall Motion prediction',
-                             use_column_width=True)
-                    st.caption('Bright/warm areas = stronger influence on the RWMA prediction')
-
-                with tab_both:
-                    sub_a, sub_b = st.columns(2)
-                    with sub_a:
-                        st.markdown('**EF Attention**')
-                        st.image(overlay_ef, use_column_width=True)
-                    with sub_b:
-                        st.markdown('**RWMA Attention**')
-                        st.image(overlay_rwma, use_column_width=True)
-            except Exception as e:
-                st.warning(f'Grad-CAM generation failed: {e}. Predictions are still valid.')
-
-            # Step 6: Methodology badge
-            st.markdown(f"""
-            <div class="methodology">
-                <b>🧠 Model:</b> 5-fold ensemble × {N_TTA} TTA = {len(ensemble) * N_TTA} votes
-                &nbsp;|&nbsp; <b>Backbone:</b> EfficientNet-B3 (NoisyStudent)
-                &nbsp;|&nbsp; <b>Validation AUC:</b> {MEAN_AUC:.3f}
-                &nbsp;|&nbsp; <b>Device:</b> {DEVICE.type.upper()}
-                &nbsp;|&nbsp; <b>Explainability:</b> Grad-CAM (conv_head layer)
-            </div>
-            """, unsafe_allow_html=True)
-
-        finally:
-            try: os.unlink(tmp_path)
-            except: pass
-
-    elif analyze_btn and uploaded_file is None:
-        st.warning('⚠️ Please upload an ECG image first.')
-
-    else:
-        st.markdown("""
-        <div style='padding: 60px 20px; text-align: center; color: #94a3b8;
-                    background: #f8fafc; border: 2px dashed #cbd5e1; border-radius: 12px;'>
-            <h2 style='color: #94a3b8;'>📋 Awaiting ECG Analysis</h2>
-            <p>Upload an ECG image and click <b>Analyze</b> to begin</p>
-            <p style='font-size: 0.85em; margin-top: 20px;'>
-              <b>NEW:</b> Now includes Grad-CAM attention heatmaps showing
-              which regions of the ECG the AI used for its predictions.
-            </p>
-        </div>
-        """, unsafe_allow_html=True)
-
-
-# Disclaimer
-st.markdown("""
-<div class="disclaimer">
-    <b>⚠️ Important Medical Disclaimer</b><br>
-    This is an AI <b>research prototype</b> for screening purposes only. It is <b>NOT</b>
-    a substitute for clinical evaluation by a qualified cardiologist. All predictions must
-    be reviewed by a healthcare professional. This system is <b>not approved for clinical
-    diagnostic use</b>.
-</div>
-""", unsafe_allow_html=True)
-
-st.markdown("""
-<div style='text-align: center; margin-top: 24px; color: #6b7280; font-size: 0.85em;'>
-    Built with PyTorch • EfficientNet-B3 • 5-fold Multilabel Stratified K-Fold ensemble
-    • Grad-CAM for explainability
-</div>
-""", unsafe_allow_html=True)
+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()