Update app.py

app.py

@@ -2,12 +2,10 @@ 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
-# Custom activation functions
+import neurokit2 as nk
+
+# Custom activation functions required by the model
 def sin_activation(x):
     return tf.math.sin(x)
 
@@ -18,7 +16,7 @@ def cos_activation(x):
 @st.cache_resource
 def load_model():
     return tf.keras.models.load_model(
-        "model.keras",  # Use the .keras format instead of .h5
+        "model.keras",
         custom_objects={
             'sin': sin_activation,
             'cos': cos_activation,
@@ -28,8 +26,8 @@ def load_model():
 
 model = load_model()
 
-# AAMI class map
-class_map = {
+# AAMI class mapping matching training code
+CLASS_MAP = {
     0: "Normal",
     1: "Supraventricular Ectopic (SVEB)",
     2: "Ventricular Ectopic (VEB)",
@@ -37,73 +35,73 @@ class_map = {
     4: "Unknown"
 }
 
-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)
+def process_mitbih_file(dat_file):
+    """Process MIT-BIH .dat file using NeuroKit2"""
+    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
         
-        if end - start == window_size:
-            beat = signal[start:end]
-            beats.append(beat)
-            
-    return np.array(beats)
+        # 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)
+    except Exception as e:
+        st.error(f"File processing error: {str(e)}")
+        return None
 
-st.title("ECG Arrhythmia Classification")
-st.write("Upload MIT-BIH record files (.dat, .hea, .atr)")
+# Streamlit UI
+st.title("ECG Arrhythmia Detection")
+st.write("Upload MIT-BIH .dat file")
 
-uploaded_files = st.file_uploader(
-    "Choose files", 
-    type=["dat", "hea", "atr"],
-    accept_multiple_files=True
+uploaded_file = st.file_uploader(
+    "Select .dat file",
+    type=["dat"],
+    accept_multiple_files=False
 )
 
-if uploaded_files:
-    with tempfile.TemporaryDirectory() as tmpdir:
-        # Save uploaded files
-        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())
-        
-        # Find base record name
-        base_names = {os.path.splitext(f.name)[0] for f in uploaded_files}
-        common_base = list(base_names)[0]  # Get first base name
-        
-        try:
-            # Read record
-            record = wfdb.rdrecord(os.path.join(tmpdir, common_base))
-            annotation = wfdb.rdann(os.path.join(tmpdir, common_base), 'atr')
+if uploaded_file is not None:
+    if st.button("Analyze"):
+        with st.spinner("Processing ECG signal..."):
+            beats = process_mitbih_file(uploaded_file)
             
-            # Process beats
-            beats = extract_beats(record, annotation)
-            if len(beats) == 0:
-                st.error("No valid beats found in the record")
-                st.stop()
+            if beats is not None and len(beats) > 0:
+                # Prepare data for model
+                beats = beats.reshape((-1, 257, 1)).astype(np.float32)
                 
-            # Preprocess and predict
-            beats = beats.reshape((-1, 257, 1)).astype(np.float32)
-            predictions = model.predict(beats)
-            predicted_classes = np.argmax(predictions, axis=1)
-            
-            # Display results
-            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)
-            
-            # Add visualization
-            st.subheader("Sample ECG Beat")
-            fig, ax = plt.subplots()
-            ax.plot(beats[0].flatten())
-            st.pyplot(fig)
-            
-        except Exception as e:
-            st.error(f"Error processing files: {str(e)}")
+                # 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")