Update app.py

app.py

@@ -2,10 +2,13 @@ import streamlit as st
 import numpy as np
 import pandas as pd
 import tensorflow as tf
+import wfdb
+import tempfile
+import os
+from scipy.signal import resample
 import matplotlib.pyplot as plt
-import neurokit2 as nk
 
-# Custom activation functions required by the model
+# Custom activation functions
 def sin_activation(x):
     return tf.math.sin(x)
 
@@ -26,8 +29,8 @@ def load_model():
 
 model = load_model()
 
-# AAMI class mapping matching training code
-CLASS_MAP = {
+# AAMI class map
+class_map = {
     0: "Normal",
     1: "Supraventricular Ectopic (SVEB)",
     2: "Ventricular Ectopic (VEB)",
@@ -35,73 +38,82 @@ CLASS_MAP = {
     4: "Unknown"
 }
 
-def process_mitbih_file(dat_file):
-    """Process MIT-BIH .dat file using NeuroKit2"""
+def extract_beats(record, annotation, window_size=257):
+    beats = []
+    r_locs = annotation.sample
+    signal = record.p_signal[:, 0]  # Using first channel
+
+    for r in r_locs:
+        start = max(0, r - window_size//2)
+        end = min(len(signal), r + window_size//2 + 1)
+
+        if end - start == window_size:
+            beat = signal[start:end]
+            beats.append(beat)
+
+    return np.array(beats)
+
+st.title("ECG Arrhythmia Classification")
+st.write("Upload MIT-BIH record files (.dat, .hea, .atr) or load record 108")
+
+record_loaded = False
+record = None
+annotation = None
+
+# Load Record 108 Button
+if st.button("Load Record 108"):
     try:
-        # Read raw signal data (assuming format 16, 360Hz, gain=200)
-        signal = np.frombuffer(dat_file.getbuffer(), dtype=np.int16).astype(np.float32)
-        signal /= 200.0  # Convert to mV using standard gain
-        
-        # NeuroKit2 processing with assumed 360Hz sampling rate
-        ecg_signals, info = nk.ecg_process(signal, sampling_rate=360)
-        r_peaks = info["ECG_R_Peaks"]
-        
-        # Extract beats with same parameters as training
-        window_size = 257
-        beats = []
-        for r in r_peaks:
-            start = max(0, r - window_size//2)
-            end = start + window_size
-            if end <= len(signal):
-                beat = signal[start:end]
-                beats.append(beat)
-                
-        return np.array(beats)
+        base_name = "108"
+        record = wfdb.rdrecord(base_name)
+        annotation = wfdb.rdann(base_name, 'atr')
+        record_loaded = True
     except Exception as e:
-        st.error(f"File processing error: {str(e)}")
-        return None
-
-# Streamlit UI
-st.title("ECG Arrhythmia Detection")
-st.write("Upload MIT-BIH .dat file")
+        st.error(f"Error loading Record 108: {str(e)}")
 
-uploaded_file = st.file_uploader(
-    "Select .dat file",
-    type=["dat"],
-    accept_multiple_files=False
+# File uploader
+uploaded_files = st.file_uploader(
+    "Or upload your own files",
+    type=["dat", "hea", "atr"],
+    accept_multiple_files=True
 )
 
-if uploaded_file is not None:
-    if st.button("Analyze"):
-        with st.spinner("Processing ECG signal..."):
-            beats = process_mitbih_file(uploaded_file)
-            
-            if beats is not None and len(beats) > 0:
-                # Prepare data for model
-                beats = beats.reshape((-1, 257, 1)).astype(np.float32)
-                
-                # Make predictions
-                predictions = model.predict(beats)
-                pred_classes = np.argmax(predictions, axis=1)
-                
-                # Show results
-                st.subheader("Analysis Results")
-                results = pd.DataFrame({
-                    "Beat Index": range(len(beats)),
-                    "Predicted Class": [CLASS_MAP[c] for c in pred_classes],
-                    "Confidence": [f"{np.max(p):.1%}" for p in predictions]
-                })
-                st.dataframe(results)
-                
-                # Visualizations
-                st.subheader("ECG Signal")
-                fig, ax = plt.subplots(1, 2, figsize=(15, 4))
-                ax[0].plot(beats[0].flatten())
-                ax[0].set_title("Sample ECG Beat")
-                class_dist = results["Predicted Class"].value_counts()
-                ax[1].bar(class_dist.index, class_dist.values)
-                ax[1].set_title("Class Distribution")
-                ax[1].tick_params(axis='x', rotation=45)
-                st.pyplot(fig)
-            else:
-                st.error("Failed to extract valid beats from the signal")
+if uploaded_files and not record_loaded:
+    with tempfile.TemporaryDirectory() as tmpdir:
+        for f in uploaded_files:
+            file_path = os.path.join(tmpdir, f.name)
+            with open(file_path, "wb") as f_out:
+                f_out.write(f.getbuffer())
+
+        base_names = {os.path.splitext(f.name)[0] for f in uploaded_files}
+        common_base = list(base_names)[0]
+
+        try:
+            record = wfdb.rdrecord(os.path.join(tmpdir, common_base))
+            annotation = wfdb.rdann(os.path.join(tmpdir, common_base), 'atr')
+            record_loaded = True
+        except Exception as e:
+            st.error(f"Error reading uploaded files: {str(e)}")
+
+# Run processing if record is loaded
+if record_loaded and record is not None and annotation is not None:
+    beats = extract_beats(record, annotation)
+    if len(beats) == 0:
+        st.error("No valid beats found in the record")
+        st.stop()
+
+    beats = beats.reshape((-1, 257, 1)).astype(np.float32)
+    predictions = model.predict(beats)
+    predicted_classes = np.argmax(predictions, axis=1)
+
+    st.subheader("Classification Results")
+    results = pd.DataFrame({
+        "Beat Index": range(len(beats)),
+        "Predicted Class": [class_map[c] for c in predicted_classes],
+        "Confidence": np.max(predictions, axis=1)
+    })
+    st.dataframe(results)
+
+    st.subheader("Sample ECG Beat")
+    fig, ax = plt.subplots()
+    ax.plot(beats[0].flatten())
+    st.pyplot(fig)