Add multi-file upload support with automatic WFDB pair detection

app.py

@@ -11,6 +11,7 @@ import numpy as np
 import plotly.graph_objects as go
 from huggingface_hub import hf_hub_download
 import tempfile
+import shutil
 from pathlib import Path
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -310,7 +311,7 @@ except Exception as e:
     traceback.print_exc()
 
 def predict_ecg(file_obj):
-    """Main prediction function"""
+    """Main prediction function - handles single or multiple files"""
     
     if model is None:
         return (
@@ -320,11 +321,47 @@ def predict_ecg(file_obj):
         )
     
     try:
-        # Read file
-        if isinstance(file_obj, str):
-            file_path = file_obj
+        # Handle multiple file uploads (list) or single file
+        file_path = None
+        
+        if isinstance(file_obj, list):
+            # Multiple files uploaded
+            if not file_obj:
+                return ("**Error**: No files uploaded", None)
+            
+            # Look for WFDB pairs (.hea + .dat)
+            files = [str(f.name) if hasattr(f, 'name') else str(f) for f in file_obj]
+            hea_files = [f for f in files if f.lower().endswith('.hea')]
+            dat_files = [f for f in files if f.lower().endswith('.dat')]
+            
+            if hea_files and dat_files:
+                # WFDB pair detected - both files present
+                # Copy .dat file next to .hea file for WFDB to work
+                import shutil
+                hea_path = hea_files[0]
+                dat_path = dat_files[0]
+                
+                # Get directory of .hea file
+                hea_dir = Path(hea_path).parent
+                dat_filename = Path(dat_path).name
+                target_dat_path = hea_dir / dat_filename
+                
+                # Copy .dat file to same directory as .hea if not already there
+                if str(target_dat_path) != dat_path:
+                    shutil.copy(dat_path, target_dat_path)
+                
+                file_path = hea_path
+                print(f"WFDB pair detected: {hea_path} + {dat_path}")
+            else:
+                # No WFDB pair, use first file
+                file_path = str(file_obj[0].name) if hasattr(file_obj[0], 'name') else str(file_obj[0])
+                print(f"Multiple files uploaded, using first: {file_path}")
         else:
-            file_path = file_obj.name if hasattr(file_obj, 'name') else str(file_obj)
+            # Single file (backward compatible)
+            if isinstance(file_obj, str):
+                file_path = file_obj
+            else:
+                file_path = file_obj.name if hasattr(file_obj, 'name') else str(file_obj)
         
         # Load ECG using universal loader
         print(f"Loading file: {file_path}")
@@ -494,13 +531,15 @@ with gr.Blocks(
             gr.Markdown("""
             ### Upload Your ECG
             
+            **✅ Multi-file upload supported!** Upload multiple files at once, especially for WFDB pairs.
+            
             **Clinical & Standardized Formats:**
             - `.dcm` – DICOM (medical imaging, PACS systems)
             - `.scp` – SCP-ECG (European interoperability standard)
             - `.xml` – HL7 aECG / FDA XML (clinical trials, regulatory)
             
-            ** Research & PhysioNet Formats:**
-            - `.hea` + `.dat` – WFDB (MIT-BIH, PhysioNet) **Requires BOTH files**
+            **Research & PhysioNet Formats:**
+            - `.hea` + `.dat` – WFDB (MIT-BIH, PhysioNet) **Upload both files together**
             - `.edf` – European Data Format (multi-channel biosignals)
             
             **Generic / Export Formats:**
@@ -509,6 +548,7 @@ with gr.Blocks(
             - `.mat` – MATLAB format
             - `.h5 / .hdf5` – HDF5 (efficient large-scale datasets)
             - `.raw / .bin` – Binary ECG data
+            - `.zip` – Archive with multiple files
             
             **Architecture Auto-Conversion:**
             - Multi-lead (12 leads): Used directly
@@ -519,14 +559,15 @@ with gr.Blocks(
             
             ---
             
-            **WFDB Note:** `.hea` and `.dat` must be in the same directory. If they're separate, please upload them as a ZIP archive or both files together if the interface allows multi-file selection.
+            **💡 WFDB Tip:** Upload both `.hea` and `.dat` files together in one go. The system will automatically detect the pair and process them correctly!
             """)
             
             file_input = gr.File(
-                label="ECG File",
+                label="ECG File(s)",
+                file_count="multiple",
                 file_types=[".csv", ".txt", ".tsv", ".npy", ".hea", ".dat", 
                            ".dcm", ".mat", ".h5", ".hdf5", ".edf", ".xml", 
-                           ".raw", ".bin", ".bat", ".ecg"],
+                           ".raw", ".bin", ".bat", ".ecg", ".zip"],
                 type="filepath"
             )