import gradio as gr import numpy as np import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt from sklearn.preprocessing import MinMaxScaler import os import spaces os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' 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", 3: "🔴 High Risk — Immediate medical attention", 4: "❓ Unknown — Further evaluation needed" } COLORS = ['#2ecc71','#3498db','#e74c3c','#f39c12','#9b59b6'] # ── Load model lazily inside GPU function ───────────────────────────────────── model = None @spaces.GPU def predict_ecg(text_input): global model if model is None: import tensorflow as tf from tensorflow.keras import layers, models, regularizers m = 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") m.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy']) m.load_weights("cnn_weights.weights.h5") model = m try: 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 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 fig, axes = plt.subplots(1, 2, figsize=(14, 4)) fig.patch.set_facecolor('#0e1117') 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 sp in ax1.spines.values(): sp.set_edgecolor('#333') ax2 = axes[1] ax2.set_facecolor('#1a1a2e') 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 sp in ax2.spines.values(): sp.set_edgecolor('#333') plt.tight_layout() result = f"**🤖 Predicted:** {CLASS_MAPPING[pred_class]}\n\n**📊 Confidence:** {confidence:.2f}%\n\n**⚕️ Risk:** {CLASS_RISK[pred_class]}" true_out = "" if true_label is not None: match = "✅ Correct!" if pred_class == true_label else "❌ Incorrect" true_out = f"**🏷️ True Label:** {CLASS_MAPPING.get(true_label,'Unknown')} — {match}" stats = f"**Signal Stats:** R-peak={features.max():.4f} @ t={r_idx} | Min={features.min():.4f} | Mean={features.mean():.4f} | Std={features.std():.4f}" return fig, result, true_out, stats except Exception as e: return None, f"❌ Error: {str(e)}", "", "" # ── UI ──────────────────────────────────────────────────────────────────────── with gr.Blocks(title="ECG Classification") as demo: gr.Markdown(""" # 🫀 ECG Arrhythmia Classification **1D CNN | MIT-BIH Dataset | Accuracy: 97.25% | Macro F1: 0.8796** --- """) with gr.Row(): with gr.Column(scale=1): text_input = gr.Textbox( label="Paste 187 or 188 ECG values", lines=8, placeholder="0.5, 0.8, 0.3, 1.0 ..." ) predict_btn = gr.Button("▶ Run Prediction", variant="primary") gr.Markdown(""" | Label | Class | Risk | |-------|-------|------| | 0 | Normal (N) | ✅ | | 1 | Supraventricular (S) | ⚠️ | | 2 | Ventricular (V) | 🔴 | | 3 | Fusion (F) | 🔴 | | 4 | Unknown (Q) | ❓ | """) with gr.Column(scale=2): plot_out = gr.Plot(label="ECG Signal & Probabilities") result_out = gr.Markdown() true_out = gr.Markdown() stats_out = gr.Markdown() predict_btn.click( predict_ecg, inputs=[text_input], outputs=[plot_out, result_out, true_out, stats_out] ) gr.Markdown("*DEPI Final Project 2025*") demo.launch(ssr_mode=False)