Add Gradio app for ECG classification

app.py

@@ -0,0 +1,278 @@
+#!/usr/bin/env python3
+"""
+Gradio interface for ECG classification
+Deploy to Hugging Face Spaces
+"""
+
+import gradio as gr
+import torch
+import numpy as np
+import plotly.graph_objects as go
+from transformers import AutoModel, AutoConfig
+import tempfile
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Constants
+MODEL_ID = "Tumo505/SSL-ECG-CLASSIFICATION"
+CLASS_LABELS = ["NORM", "MI", "STTC", "HYP", "CD"]
+CLASS_COLORS = {
+    "NORM": "#90EE90",
+    "MI": "#FF6B6B",
+    "STTC": "#FFD93D",
+    "HYP": "#6C5CE7",
+    "CD": "#A29BFE"
+}
+
+# Load model
+model = None
+try:
+    print("Loading model from Hub...")
+    model = AutoModel.from_pretrained(MODEL_ID, trust_remote_code=True)
+    model.to(device)
+    model.eval()
+    print("Model loaded successfully")
+except Exception as e:
+    print(f"Error loading model: {e}")
+
+def predict_ecg(file_obj):
+    """Main prediction function"""
+    
+    if model is None:
+        return (
+            "**Model Loading Error**\n"
+            "The model failed to load. Please try again or contact support.",
+            None,
+            None
+        )
+    
+    try:
+        # Read file
+        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 data
+        if file_path.endswith(('.csv', '.txt')):
+            ecg = np.loadtxt(file_path, delimiter=',')
+        elif file_path.endswith('.npy'):
+            ecg = np.load(file_path)
+        else:
+            ecg = np.genfromtxt(file_path)
+        
+        # Validation
+        if ecg.ndim != 2:
+            return (
+                "**Invalid Format**\n"
+                f"Expected 2D array, got shape {{ecg.shape}}\n"
+                "Expected: (12 leads, N samples)",
+                None,
+                None
+            )
+        
+        # Handle transposition
+        if ecg.shape[0] != 12:
+            if ecg.shape[1] == 12:
+                ecg = ecg.T
+            else:
+                return (
+                    "**Invalid Dimensions**\n"
+                    f"Got shape {{ecg.shape}}, expected (12, N)\n"
+                    "Ensure file has 12 leads (rows) × N samples (columns)",
+                    None,
+                    None
+                )
+        
+        # Resize to 5000 samples
+        if ecg.shape[1] < 5000:
+            ecg = np.pad(ecg, ((0, 0), (0, 5000 - ecg.shape[1])), mode='edge')
+        else:
+            ecg = ecg[:, :5000]
+        
+        # Normalize each lead independently
+        ecg = (ecg - ecg.mean(axis=1, keepdims=True)) / (ecg.std(axis=1, keepdims=True) + 1e-8)
+        
+        # Convert to tensor
+        x = torch.tensor(ecg, dtype=torch.float32).unsqueeze(0).to(device)
+        
+        # Predict
+        with torch.no_grad():
+            output = model(x)
+            logits = output["logits"][0].cpu().numpy()
+            probs = torch.softmax(torch.tensor(logits), dim=0).numpy()
+        
+        # Get prediction
+        pred_idx = int(np.argmax(probs))
+        pred_class = CLASS_LABELS[pred_idx]
+        confidence = float(probs[pred_idx])
+        
+        # Create visualization
+        fig = go.Figure()
+        
+        fig.add_trace(go.Bar(
+            y=CLASS_LABELS,
+            x=probs,
+            orientation='h',
+            marker=dict(
+                color=[CLASS_COLORS.get(c, '#87CEEB') for c in CLASS_LABELS],
+                line=dict(
+                    color=['#000000' if i == pred_idx else '#CCCCCC' for i in range(5)],
+                    width=[3 if i == pred_idx else 1 for i in range(5)]
+                )
+            ),
+            text=[f'{p:.1%}' for p in probs],
+            textposition='auto',
+            hovertemplate='<b>%{y}</b><br>Probability: %{x:.2%}<extra></extra>'
+        ))
+        
+        fig.update_layout(
+            title=dict(
+                text=f"ECG Classification Results<br><sub>Prediction: <b>{pred_class}</b> ({confidence:.1%})</sub>",
+                x=0.5,
+                xanchor='center'
+            ),
+            xaxis_title="Model Confidence",
+            yaxis_title="Diagnostic Class",
+            height=450,
+            showlegend=False,
+            font=dict(size=12),
+            plot_bgcolor='rgba(240,240,240,0.5)'
+        )
+        
+        # Format output text
+        output_md = f"""
+## Prediction Complete
+
+### Primary Diagnosis: **{pred_class}**
+### Confidence: **{confidence:.1%}**
+
+---
+
+### All Class Probabilities:
+
+| Class | Probability |
+|-------|-------------|
+| {CLASS_LABELS[0]} | {probs[0]:.2%} |
+| {CLASS_LABELS[1]} | {probs[1]:.2%} |
+| {CLASS_LABELS[2]} | {probs[2]:.2%} |
+| {CLASS_LABELS[3]} | {probs[3]:.2%} |
+| {CLASS_LABELS[4]} | {probs[4]:.2%} |
+
+---
+
+**Model Information:**
+- Framework: SimCLR SSL
+- Training Data: PTB-XL (10% labeled)
+- Test AUROC: 0.8717
+- Input: 12-lead ECG @ 100 Hz
+
+**Disclaimer:** This is a research model for demonstration only. Not validated for clinical use.
+"""
+        
+        return output_md, fig, None
+        
+    except FileNotFoundError:
+        return "**File Error:** Could not read uploaded file", None, None
+    except Exception as e:
+        import traceback
+        error_msg = f"**Error:** {{str(e)}}\n\nDebug: {{traceback.format_exc()}}"
+        return error_msg, None, None
+
+
+# Create interface
+with gr.Blocks(
+    title="ECG Classification with Self-Supervised Learning",
+    theme=gr.themes.Soft(primary_hue="emerald")
+) as demo:
+    
+    gr.Markdown("""
+    # ECG Classification with Self-Supervised Learning
+    
+    **Test ECG cardiovascular disease classification** using a SimCLR pre-trained model fine-tuned on the PTB-XL dataset.
+    
+    **Model Performance:** AUROC 0.8717 | Accuracy 0.8234 | 10% labeled data
+    
+    ---
+    """)
+    
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("""
+            ### Upload Your ECG
+            
+            **Supported Formats:**
+            - CSV / TSV / TXT
+            - NumPy .npy file
+            
+            **Requirements:**
+            - **Shape:** 12 leads × N samples
+            - **Sampling Rate:** Any (will be normalized)
+            - **Format:** Raw ECG values (not images)
+            
+            **Example Structure:**
+            ```
+            lead_I, lead_II, ..., lead_aVF
+            0.123, 0.456, ..., 0.789
+            ...
+            ```
+            """)
+            
+            file_input = gr.File(
+                label="ECG File",
+                file_types=[".csv", ".txt", ".tsv", ".npy"],
+                type="filepath"
+            )
+            
+            submit_btn = gr.Button("Classify ECG", variant="primary", size="lg")
+        
+        with gr.Column():
+            gr.Markdown("""
+            ### Results
+            
+            Predictions appear here after classification.
+            """)
+            
+            output_text = gr.Markdown(
+                "Upload an ECG file to see predictions",
+                label="Classification Results"
+            )
+    
+    with gr.Row():
+        chart_output = gr.Plot(label="Probability Distribution")
+    
+    # Connect button
+    submit_btn.click(
+        fn=predict_ecg,
+        inputs=[file_input],
+        outputs=[output_text, chart_output, None]
+    )
+    
+    # Info section
+    gr.Markdown("""
+    ---
+    
+    ### About This Model
+    
+    **Architecture:** 1D CNN with SimCLR self-supervised pre-training
+    
+    **Training:**
+    - Pre-training: SimCLR on 17.5K unlabeled PTB-XL ECGs
+    - Fine-tuning: Supervised on 1.7K labeled ECGs (10%)
+    
+    **Classes Predicted:**
+    - NORM: Normal ECG
+    - MI: Myocardial Infarction
+    - STTC: ST/T Changes
+    - HYP: Hypertrophy
+    - CD: Conduction Disturbances
+    
+    **Research Only** - Not validated for clinical use
+    
+    [View Model Card](https://huggingface.co/{MODEL_ID})  
+    [GitHub Repository](https://github.com/Tumo505/SSL-for-ECG-classification)
+    """)
+
+
+if __name__ == "__main__":
+    demo.launch(share=False)

requirements.txt

@@ -0,0 +1,8 @@
+torch>=2.0.0
+transformers>=4.36.0
+safetensors>=0.4.0
+gradio>=4.0
+numpy>=1.24.0
+scipy>=1.10.0
+plotly>=5.17.0
+huggingface-hub>=0.19.0