Add support for .hea (WFDB) and .bat file formats

app.py

@@ -95,16 +95,66 @@ def predict_ecg(file_obj):
         # Load ECG data - handle multiple formats
         print(f"Loading file: {file_path}")
         
-        try:
-            # Try space-separated (UCR/ArrowHead format)
-            ecg = np.genfromtxt(file_path, delimiter=None)
-        except:
+        # Detect format by extension
+        if file_path.endswith('.hea'):
+            # WFDB (PhysioNet) format
+            print("Detected WFDB (.hea) format...")
             try:
-                # Try comma-separated
-                ecg = np.loadtxt(file_path, delimiter=',')
+                import wfdb
+                # Load the record (without extension)
+                record_path = file_path.replace('.hea', '')
+                record = wfdb.rdrecord(record_path)
+                ecg = record.p_signal  # Get signal data
+                print(f"WFDB loaded: {ecg.shape}")
+            except ImportError:
+                return (
+                    "**Error**: WFDB library not installed\n"
+                    "Python-wfdb package required for .hea files",
+                    None
+                )
+            except Exception as e:
+                return (f"**WFDB Error**: {str(e)}", None)
+        
+        elif file_path.endswith('.bat'):
+            # BAT format (binary or text batch format)
+            print("Detected BAT format...")
+            try:
+                # Try reading as binary NumPy array first
+                ecg = np.fromfile(file_path, dtype=np.float32)
+                # Reshape if needed - assume 128 samples per lead or similar
+                if len(ecg) % 12 == 0:
+                    ecg = ecg.reshape(12, -1)
+                else:
+                    # Single lead, will be replicated later
+                    ecg = ecg.reshape(1, -1)
+                print(f"BAT loaded: {ecg.shape}")
+            except:
+                try:
+                    # Try text format
+                    ecg = np.loadtxt(file_path)
+                    if ecg.ndim == 1:
+                        ecg = ecg.reshape(1, -1)
+                except Exception as e:
+                    return (f"**BAT Error**: Could not parse file - {str(e)}", None)
+        
+        else:
+            # Try standard text formats (CSV, space-separated, tab-separated, .npy)
+            try:
+                if file_path.endswith('.npy'):
+                    ecg = np.load(file_path)
+                else:
+                    # Try space-separated (UCR/ArrowHead format)
+                    ecg = np.genfromtxt(file_path, delimiter=None)
             except:
-                # Try tab-separated
-                ecg = np.loadtxt(file_path, delimiter='\t')
+                try:
+                    # Try comma-separated
+                    ecg = np.loadtxt(file_path, delimiter=',')
+                except:
+                    try:
+                        # Try tab-separated
+                        ecg = np.loadtxt(file_path, delimiter='\t')
+                    except Exception as e:
+                        return (f"**File Error**: Could not parse file format - {str(e)}", None)
         
         print(f"Loaded shape: {ecg.shape}")
         
@@ -270,15 +320,19 @@ with gr.Blocks(
             ### Upload Your ECG
             
             **Supported Formats:**
-            - CSV / TSV / TXT
+            - CSV / TSV / TXT (space, comma, or tab-separated)
             - NumPy .npy file
+            - WFDB .hea (PhysioNet format)
+            - BAT (binary or batch ECG files)
+            - UCR/ArrowHead format (UCR5000, etc.)
             
             **Requirements:**
-            - **Shape:** 12 leads × N samples
-            - **Sampling Rate:** Any (will be normalized)
-            - **Format:** Raw ECG values (not images)
+            - Multi-lead: (N leads, M samples)
+            - Single-lead: Will be replicated to 12 leads
+            - Sampling Rate: Any (will be normalized)
+            - Will auto-pad/trim to 5000 samples
             
-            **Example Structure:**
+            **Example Structure (CSV/TXT):**
             ```
             lead_I, lead_II, ..., lead_aVF
             0.123, 0.456, ..., 0.789
@@ -288,7 +342,7 @@ with gr.Blocks(
             
             file_input = gr.File(
                 label="ECG File",
-                file_types=[".csv", ".txt", ".tsv", ".npy"],
+                file_types=[".csv", ".txt", ".tsv", ".npy", ".hea", ".bat"],
                 type="filepath"
             )
             

requirements.txt

@@ -6,3 +6,4 @@ numpy>=1.24.0
 scipy>=1.10.0
 plotly>=5.17.0
 huggingface-hub>=0.19.0
+wfdb>=4.0.0