Update app.py

app.py

@@ -1,168 +1,200 @@
 import gradio as gr
 import numpy as np
-import pandas as pd
 import matplotlib.pyplot as plt
-import os
-
-# Unified Classification Information (English)
-CLASS_INFO = {
-    0: {"name": "Normal Beat (N)",                          "icon": "✅", "color": "#27ae60", "risk": "Low Risk",         "desc": "Normal cardiac rhythm."},
-    1: {"name": "Supraventricular Premature Beat (S)",      "icon": "⚠️", "color": "#e67e22", "risk": "Medium Risk",      "desc": "Arrhythmia originating above the ventricles."},
-    2: {"name": "Premature Ventricular Contraction (V)",    "icon": "🚨", "color": "#e74c3c", "risk": "High Risk",        "desc": "Early ventricular contraction — Critical condition."},
-    3: {"name": "Fusion of Ventricular and Normal (F)",     "icon": "⚡", "color": "#8e44ad", "risk": "Medium-High Risk", "desc": "A beat that fuses normal and ventricular pathways."},
-    4: {"name": "Unclassifiable Beat (Q)",                  "icon": "❓", "color": "#7f8c8d", "risk": "Undetermined",     "desc": "Unclassifiable morphology — Manual review required."},
+from sklearn.preprocessing import MinMaxScaler
+import tensorflow as tf
+from tensorflow.keras import layers, models, regularizers
+
+# ── Class Info ────────────────────────────────────────────────────────────────
+CLASS_MAPPING = {
+    0: "Normal Beat (N)",
+    1: "Supraventricular Premature Beat (S)",
+    2: "Premature Ventricular Contraction (V)",
+    3: "Fusion of Ventricular and Normal Beat (F)",
+    4: "Unclassifiable Beat (Q)"
+}
+CLASS_RISK = {
+    0: "✅ Normal — No action required",
+    1: "⚠️ Moderate — Monitor and consult physician",
+    2: "🔴 High Risk — Immediate medical attention recommended",
+    3: "🔴 High Risk — Immediate medical attention recommended",
+    4: "❓ Unknown — Further evaluation needed"
 }
 
-MODEL_PATH = "cnn_ecg_final.keras"
-model = None
-model_status = ""
+# ── Build & Load Model ────────────────────────────────────────────────────────
+def build_model():
+    model = models.Sequential([
+        layers.Input(shape=(187, 1)),
+        layers.Conv1D(64, 7, padding='same', activation='relu',
+                      kernel_regularizer=regularizers.l2(1e-4)),
+        layers.BatchNormalization(),
+        layers.MaxPooling1D(2),
+        layers.Dropout(0.2),
+        layers.Conv1D(128, 5, padding='same', activation='relu',
+                      kernel_regularizer=regularizers.l2(1e-4)),
+        layers.BatchNormalization(),
+        layers.MaxPooling1D(2),
+        layers.Dropout(0.25),
+        layers.Conv1D(256, 3, padding='same', activation='relu',
+                      kernel_regularizer=regularizers.l2(1e-4)),
+        layers.BatchNormalization(),
+        layers.MaxPooling1D(2),
+        layers.Dropout(0.3),
+        layers.Conv1D(256, 3, padding='same', activation='relu',
+                      kernel_regularizer=regularizers.l2(1e-4)),
+        layers.BatchNormalization(),
+        layers.GlobalAveragePooling1D(),
+        layers.Dropout(0.3),
+        layers.Dense(256, activation='relu',
+                     kernel_regularizer=regularizers.l2(1e-4)),
+        layers.BatchNormalization(),
+        layers.Dropout(0.4),
+        layers.Dense(128, activation='relu',
+                     kernel_regularizer=regularizers.l2(1e-4)),
+        layers.Dropout(0.3),
+        layers.Dense(5, activation='softmax')
+    ], name="CNN_ECG")
+    model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
+    model.load_weights("cnn_weights.weights.h5")
+    return model
+
+model = build_model()
 
-# Auto-load the model from local directory if it exists
-if os.path.exists(MODEL_PATH):
+# ── Predict Function ──────────────────────────────────────────────────────────
+def predict_ecg(text_input):
     try:
-        from tensorflow.keras.models import load_model
-        model = load_model(MODEL_PATH)
-        model_status = "✅ Local Keras Model Loaded Automatically!"
+        # Parse input
+        cleaned = text_input.replace(",", " ").replace("\n", " ").replace("\t", " ")
+        parsed  = [float(x) for x in cleaned.split() if x.strip()]
+        count   = len(parsed)
+
+        if count not in (187, 188):
+            return None, f"❌ Expected 187 or 188 values, got {count}.", "", ""
+
+        features      = np.array(parsed[:187], dtype=np.float32)
+        true_label    = int(parsed[187]) if count == 188 else None
+
+        # Normalize + predict
+        scaler        = MinMaxScaler()
+        signal_scaled = scaler.fit_transform(features.reshape(-1, 1)).reshape(1, 187, 1)
+        probs         = model.predict(signal_scaled, verbose=0)[0]
+        pred_class    = int(np.argmax(probs))
+        confidence    = float(np.max(probs)) * 100
+
+        # ECG Plot
+        fig, axes = plt.subplots(1, 2, figsize=(14, 4))
+        fig.patch.set_facecolor('#0e1117')
+
+        # Signal
+        ax1 = axes[0]
+        ax1.set_facecolor('#1a1a2e')
+        ax1.plot(features, color='#00ff88', linewidth=1.8)
+        ax1.fill_between(range(187), features, alpha=0.15, color='#00ff88')
+        r_idx = int(np.argmax(features))
+        ax1.axvline(r_idx, color='red', linestyle='--', alpha=0.7)
+        ax1.scatter(r_idx, features[r_idx], color='red', s=80, zorder=5)
+        ax1.set_title("ECG Signal", color='white', fontsize=13, fontweight='bold')
+        ax1.set_xlabel("Time Steps", color='#aaa')
+        ax1.set_ylabel("Amplitude", color='#aaa')
+        ax1.tick_params(colors='#aaa')
+        ax1.grid(True, alpha=0.2, color='#444')
+        for spine in ax1.spines.values():
+            spine.set_edgecolor('#333')
+
+        # Probabilities
+        ax2 = axes[1]
+        ax2.set_facecolor('#1a1a2e')
+        colors = ['#2ecc71','#3498db','#e74c3c','#f39c12','#9b59b6']
+        bars   = ax2.barh([CLASS_MAPPING[i] for i in range(5)],
+                          probs * 100, color=colors, alpha=0.85)
+        for bar, val in zip(bars, probs):
+            ax2.text(val * 100 + 0.5, bar.get_y() + bar.get_height() / 2,
+                     f"{val*100:.1f}%", va='center', color='white', fontsize=9)
+        ax2.set_xlabel("Probability (%)", color='#aaa')
+        ax2.set_title("Class Probabilities", color='white', fontsize=13, fontweight='bold')
+        ax2.set_xlim(0, 115)
+        ax2.tick_params(colors='#aaa')
+        ax2.grid(True, alpha=0.2, axis='x', color='#444')
+        for spine in ax2.spines.values():
+            spine.set_edgecolor('#333')
+
+        plt.tight_layout()
+
+        # Result text
+        result_text = f"""
+**🤖 Predicted Class:** {CLASS_MAPPING[pred_class]}
+**📊 Confidence:** {confidence:.2f}%
+**⚕️ Risk:** {CLASS_RISK[pred_class]}
+        """.strip()
+
+        true_text = ""
+        if true_label is not None:
+            match = "✅ Correct!" if pred_class == true_label else "❌ Incorrect"
+            true_text = f"**🏷️ True Label:** {CLASS_MAPPING.get(true_label, 'Unknown')} — {match}"
+
+        stats_text = f"""
+**Signal Stats:**
+- R-peak amplitude: {features.max():.4f} @ timestep {r_idx}
+- Min amplitude: {features.min():.4f}
+- Mean: {features.mean():.4f}
+- Std Dev: {features.std():.4f}
+        """.strip()
+
+        return fig, result_text, true_text, stats_text
+
     except Exception as e:
-        model_status = f"❌ Error loading local model: {str(e)}"
-else:
-    model_status = "🔬 Running in Demo Mode (Model file not found in directory)"
-
-def preprocess(signal):
-    s = (signal - signal.mean()) / (signal.std() + 1e-8)
-    return s.reshape(1, 187, 1).astype("float32")
-
-# Main inference and plotting function
-def process_ecg(input_text, csv_file, row_index):
-    global model
-    features = None
-    true_label_text = "Not Provided"
-
-    # 1. Data Extraction (CSV or Text Input)
-    if csv_file is not None:
-        try:
-            df = pd.read_csv(csv_file.name, header=None)
-            idx = int(row_index) if row_index < len(df) else 0
-            row = df.iloc[idx]
-            if df.shape[1] == 188:
-                features = row.iloc[:-1].values.astype(float)
-                true_cls = int(row.iloc[-1])
-                true_label_text = f"{CLASS_INFO[true_cls]['icon']} {CLASS_INFO[true_cls]['name']}"
-            elif df.shape[1] == 187:
-                features = row.values.astype(float)
-        except Exception as e:
-            return None, f"Error parsing CSV: {str(e)}", "", "", "", model_status
-    elif input_text.strip():
-        try:
-            vals = [float(x) for x in input_text.replace(",", " ").split() if x.strip()]
-            if len(vals) == 187:
-                features = np.array(vals)
-            elif len(vals) == 188:
-                features = np.array(vals[:-1])
-                true_cls = int(vals[-1])
-                true_label_text = f"{CLASS_INFO[true_cls]['icon']} {CLASS_INFO[true_cls]['name']}"
-            else:
-                return None, f"🚨 Input error: {len(vals)} values entered. Exactly 187 or 188 values are required.", "", "", "", model_status
-        except ValueError:
-            return None, "🚨 Error: Text contains non-numeric values.", "", "", "", model_status
-
-    if features is None:
-        return None, "⚠️ Please provide ECG signal data (Text or CSV file).", "", "", "", model_status
-
-    # 2. Model Inference or Demo Mode
-    if model is not None:
-        probs = model.predict(preprocess(features), verbose=0)[0]
-        pred_cls = int(np.argmax(probs))
-        conf = float(probs[pred_cls]) * 100
-        is_demo = False
-    else:
-        # Smart simulated inference for Demo Mode
-        pred_cls = int(abs(features.sum()) % 5)
-        np.random.seed(int(abs(features.sum()) * 100) % 9999)
-        probs = np.random.dirichlet(np.ones(5))
-        probs[pred_cls] = np.random.uniform(0.7, 0.95)
-        probs /= probs.sum()
-        conf = float(probs[pred_cls]) * 100
-        is_demo = True
-
-    pi = CLASS_INFO[pred_cls]
-
-    # 3. Generate ECG Waveform Plot
-    fig, ax = plt.subplots(figsize=(10, 3.5))
-    ax.plot(features, color=pi["color"], lw=2)
-    ax.fill_between(range(187), features, alpha=0.15, color=pi["color"])
-    ax.set_title(f"ECG Waveform — Predicted: {pi['name']} ({conf:.1f}%)", fontsize=11, fontweight="bold", color=pi["color"])
-    ax.set_xlabel("Time Steps")
-    ax.set_ylabel("Amplitude")
-    ax.grid(True, ls="--", alpha=0.4)
-    ax.spines["top"].set_visible(False)
-    ax.spines["right"].set_visible(False)
-    plt.tight_layout()
-    
-    # 4. Prepare UI Output Elements
-    result_html = f"""
-    <div style='background: #f8f9fa; border-left: 8px solid {pi["color"]}; border-radius: 8px; padding: 15px; margin-top: 10px;'>
-        <h2 style='color: {pi["color"]}; margin: 0;'>{pi["icon"]} {pi["name"]}</h2>
-        <p style='font-size: 1.1em; color: #333;'><b>Diagnosis:</b> {pi["desc"]}</p>
-        <p style='margin: 5px 0;'><b>Confidence Level:</b> <span style='font-size: 1.2em; color: {pi["color"]}; font-weight:bold;'>{conf:.2f}%</span></p>
-        <p style='margin: 5px 0;'><b>Risk Level:</b> <span style='font-weight:bold;'>{pi["risk"]}</span> {' (🔬 Demo Mode)' if is_demo else ''}</p>
-    </div>
-    """
-    
-    # Medical Alerts
-    if pred_cls == 2:
-        alert_html = "<div style='background-color: #ffd2d2; color: #d8000c; padding: 10px; border-radius: 5px; font-weight: bold;'>🚨 Alert: Premature Ventricular Contraction detected — Consult a cardiologist immediately.</div>"
-    elif pred_cls in [1, 3]:
-        alert_html = "<div style='background-color: #ffe0b2; color: #e65100; padding: 10px; border-radius: 5px; font-weight: bold;'>⚠️ Warning: Cardiac arrhythmia detected — Medical follow-up is highly recommended.</div>"
-    elif pred_cls == 4:
-        alert_html = "<div style='background-color: #e0f7fa; color: #006064; padding: 10px; border-radius: 5px; font-weight: bold;'>ℹ️ Notice: Unclassifiable heartbeat — Manual review required.</div>"
-    else:
-        alert_html = "<div style='background-color: #d4edda; color: #155724; padding: 10px; border-radius: 5px; font-weight: bold;'>✅ Normal heartbeat detected.</div>"
-
-    label_dict = {CLASS_INFO[i]["name"]: float(probs[i]) for i in range(5)}
-
-    return fig, result_html, alert_html, label_dict, true_label_text, model_status
-
-# ── Gradio Blocks UI Layout ──────────────────────
-with gr.Blocks() as demo:
-    gr.HTML("""
-    <div style='text-align:center; background: linear-gradient(135deg, #1a1a2e, #0f3460); color:white; padding:20px; border-radius:12px;'>
-        <h1 style='margin:0;'>🫀 ECG Heartbeat Classification</h1>
-        <p style='margin:5px 0 0 0; opacity:0.8;'>CNN Model · MIT-BIH Dataset · 96.05% Accuracy</p>
-    </div>
+        return None, f"❌ Error: {str(e)}", "", ""
+
+# ── Gradio UI ─────────────────────────────────────────────────────────────────
+with gr.Blocks(
+    title="ECG Arrhythmia Classification",
+    theme=gr.themes.Base(primary_hue="green")
+) as demo:
+
+    gr.Markdown("""
+    # 🫀 ECG Arrhythmia Classification System
+    **Smart Medical System powered by Deep Learning (1D CNN)**
+    > MIT-BIH Arrhythmia Dataset | Accuracy: **97.25%** | Macro F1: **0.8796**
+    ---
     """)
-    
+
     with gr.Row():
         with gr.Column(scale=1):
-            gr.Markdown("### 🔍 Model Status")
-            status_out = gr.Textbox(value=model_status, label="Current System Status", interactive=False)
-            
-            gr.Markdown("### 📁 Input Data")
-            with gr.Tab("Upload CSV"):
-                csv_input = gr.File(label="Choose ECG CSV File", file_types=[".csv"])
-                row_idx = gr.Number(label="Select Heartbeat Row Index", value=0, precision=0)
-            
-            with gr.Tab("Manual Input"):
-                text_input = gr.Textbox(label="Paste 187 or 188 numbers", placeholder="e.g., 0.5 0.2 0.1 ...", lines=5)
-            
-            btn = gr.Button("🚀 Classify This Heartbeat", variant="primary")
-            
-        with gr.Column(scale=1):
-            gr.Markdown("### 📊 Output Results")
-            plot_output = gr.Plot(label="ECG Waveform Graph")
-            true_label_output = gr.Textbox(label="True Label (From Data)", interactive=False)
-            
-            result_output = gr.HTML()
-            alert_output = gr.HTML()
-            label_output = gr.Label(label="Class Probabilities")
-
-    btn.click(
-        fn=process_ecg,
-        inputs=[text_input, csv_input, row_idx],
-        outputs=[plot_output, result_output, alert_output, label_output, true_label_output, status_out]
+            gr.Markdown("### 📥 Input ECG Signal")
+            text_input = gr.Textbox(
+                label="Paste 187 or 188 ECG values (comma or space or newline separated)",
+                placeholder="0.5, 0.8, 0.3, 1.0, 0.2 ...\n(187 signal values, optionally followed by class label)",
+                lines=8
+            )
+            predict_btn = gr.Button("▶ Run Prediction", variant="primary", size="lg")
+
+            gr.Markdown("""
+            ### 📋 Classes
+            | Label | Class | Risk |
+            |-------|-------|------|
+            | 0 | Normal (N) | ✅ None |
+            | 1 | Supraventricular (S) | ⚠️ Moderate |
+            | 2 | Ventricular (V) | 🔴 High |
+            | 3 | Fusion (F) | 🔴 High |
+            | 4 | Unknown (Q) | ❓ Unknown |
+            """)
+
+        with gr.Column(scale=2):
+            gr.Markdown("### 📈 Visualization & Results")
+            plot_out   = gr.Plot(label="ECG Signal & Probabilities")
+            result_out = gr.Markdown(label="Prediction Result")
+            true_out   = gr.Markdown(label="True Label Check")
+            stats_out  = gr.Markdown(label="Signal Statistics")
+
+    predict_btn.click(
+        fn=predict_ecg,
+        inputs=[text_input],
+        outputs=[plot_out, result_out, true_out, stats_out]
     )
-    
-    gr.HTML("<p style='text-align:center; color:#999; font-size:0.85em; margin-top:20px;'>⚠️ For research and educational purposes only — Does not replace professional medical advice.</p>")
 
-if __name__ == "__main__":
-    demo.launch(theme=gr.themes.Soft())
+    gr.Markdown("""
+    ---
+    *DEPI Final Project 2025 | ECG Arrhythmia Classification | CNN Model*
+    """)
+
+demo.launch()