Initial HeartWatch AI demo release

- Added app.py: Gradio interface with upload and sample gallery tabs
- Added inference.py: DeepECG inference engine with 4 models
- Added visualization.py: ECG waveform, diagnosis bars, risk gauges
- Added class_names.json: 77 ECG diagnosis class names
- Added sample ECG data from MIT-BIH database
- Added requirements.txt and .gitignore

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

.gitignore

@@ -0,0 +1,59 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Virtual environments
+venv/
+ENV/
+env/
+.venv/
+
+# IDE
+.idea/
+.vscode/
+*.swp
+*.swo
+*~
+
+# Jupyter
+.ipynb_checkpoints/
+
+# Model weights (downloaded at runtime)
+weights/
+*.pt
+*.pth
+*.ckpt
+
+# Local development
+.env
+.env.local
+
+# OS
+.DS_Store
+Thumbs.db
+
+# Logs
+*.log
+logs/
+
+# Temporary files
+*.tmp
+*.temp

README.md

@@ -1,13 +1,43 @@
 ---
 title: HeartWatchAI
-emoji: 🦀
-colorFrom: yellow
-colorTo: indigo
+emoji: ❤️
+colorFrom: red
+colorTo: pink
 sdk: gradio
-sdk_version: 6.5.1
+sdk_version: 5.12.0
 app_file: app.py
 pinned: false
-short_description: AI-based ECG heart Analysis
+short_description: AI-powered 12-Lead ECG Analysis
+hf_oauth: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# HeartWatch AI
+
+AI-powered 12-Lead ECG analysis using deep learning models.
+
+## Features
+
+- **77-Class ECG Diagnosis**: Detect 77 different cardiac conditions
+- **LVEF Prediction**: Predict left ventricular ejection fraction < 40% and < 50%
+- **AFib Risk**: 5-year atrial fibrillation risk prediction
+- **Interactive Visualization**: Clinical 4x3 lead layout with ECG paper grid
+
+## Models
+
+This demo uses EfficientNetV2 models from the DeepECG project:
+
+- `heartwise/EfficientNetV2_77_Classes`
+- `heartwise/EfficientNetV2_LVEF_40`
+- `heartwise/EfficientNetV2_LVEF_50`
+- `heartwise/EfficientNetV2_AFIB_5y`
+
+## Input Format
+
+- NumPy array (.npy file)
+- Shape: (2500, 12) or (12, 2500)
+- 12 standard leads: I, II, III, aVR, aVL, aVF, V1-V6
+- 10 seconds at 250 Hz sampling rate
+
+## Disclaimer
+
+This is a research demonstration tool. Predictions should NOT be used for clinical decision-making. Always consult qualified healthcare professionals for medical advice.

app.py

@@ -0,0 +1,326 @@
+"""
+HeartWatch AI - ECG Analysis Demo
+==================================
+
+A Gradio-based web application for AI-powered ECG analysis using DeepECG models.
+
+Features:
+- 77-class ECG diagnosis
+- LVEF < 40% prediction
+- LVEF < 50% prediction
+- 5-year AFib risk assessment
+- Interactive 12-lead ECG visualization
+"""
+
+import os
+import logging
+import numpy as np
+import gradio as gr
+from pathlib import Path
+
+# Local imports
+from inference import DeepECGInference
+from visualization import (
+    plot_ecg_waveform,
+    plot_diagnosis_bars,
+    plot_risk_gauges,
+    generate_thumbnail
+)
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Global inference engine
+inference_engine = None
+
+
+def load_inference_engine():
+    """Load the inference engine on startup."""
+    global inference_engine
+    if inference_engine is None:
+        logger.info("Loading DeepECG inference engine...")
+        inference_engine = DeepECGInference()
+        inference_engine.load_models()
+        logger.info("Inference engine loaded successfully")
+    return inference_engine
+
+
+def get_sample_ecgs():
+    """Get list of sample ECG files from demo_data directory."""
+    sample_dir = Path(__file__).parent / "demo_data" / "samples"
+    if not sample_dir.exists():
+        return []
+
+    samples = []
+    for npy_file in sorted(sample_dir.glob("*.npy")):
+        samples.append({
+            "path": str(npy_file),
+            "name": npy_file.stem.replace("_", " ").title()
+        })
+    return samples
+
+
+def analyze_ecg(ecg_signal: np.ndarray, filename: str = "Uploaded ECG"):
+    """
+    Analyze an ECG signal and return all visualizations.
+
+    Args:
+        ecg_signal: ECG signal array
+        filename: Name to display
+
+    Returns:
+        Tuple of (ecg_plot, diagnosis_plot, risk_plot, summary_text)
+    """
+    engine = load_inference_engine()
+
+    # Run inference
+    results = engine.predict(ecg_signal)
+
+    # Generate ECG waveform plot
+    ecg_fig = plot_ecg_waveform(ecg_signal, sample_rate=250, title=filename)
+
+    # Generate diagnosis bar chart
+    if "diagnosis_77" in results:
+        probs = results["diagnosis_77"]["probabilities"]
+        class_names = results["diagnosis_77"]["class_names"]
+        diagnosis_dict = dict(zip(class_names, probs))
+        diagnosis_fig = plot_diagnosis_bars(diagnosis_dict, top_n=10)
+    else:
+        diagnosis_fig = None
+
+    # Generate risk gauges
+    lvef_40 = results.get("lvef_40", 0.0)
+    lvef_50 = results.get("lvef_50", 0.0)
+    afib_5y = results.get("afib_5y", 0.0)
+    risk_fig = plot_risk_gauges(lvef_40, lvef_50, afib_5y)
+
+    # Generate summary text
+    inference_time = results.get("inference_time_ms", 0)
+    summary = f"""## Analysis Summary
+
+**Inference Time:** {inference_time:.1f} ms
+
+### Risk Predictions
+- **LVEF < 40%:** {lvef_40*100:.1f}% probability
+- **LVEF < 50%:** {lvef_50*100:.1f}% probability
+- **5-year AFib Risk:** {afib_5y*100:.1f}% probability
+
+### Top Diagnoses
+"""
+    if "diagnosis_77" in results:
+        probs = results["diagnosis_77"]["probabilities"]
+        class_names = results["diagnosis_77"]["class_names"]
+        top_indices = np.argsort(probs)[::-1][:5]
+        for i, idx in enumerate(top_indices, 1):
+            summary += f"{i}. {class_names[idx]}: {probs[idx]*100:.1f}%\n"
+
+    return ecg_fig, diagnosis_fig, risk_fig, summary
+
+
+def analyze_uploaded_file(file):
+    """Handle uploaded .npy file."""
+    if file is None:
+        return None, None, None, "Please upload a .npy file containing ECG data."
+
+    try:
+        ecg_signal = np.load(file.name)
+        filename = Path(file.name).stem
+        return analyze_ecg(ecg_signal, filename)
+    except Exception as e:
+        logger.error(f"Error loading file: {e}")
+        return None, None, None, f"Error loading file: {str(e)}"
+
+
+def analyze_sample(sample_name: str):
+    """Analyze a sample ECG from the gallery."""
+    samples = get_sample_ecgs()
+
+    for sample in samples:
+        if sample["name"] == sample_name:
+            ecg_signal = np.load(sample["path"])
+            return analyze_ecg(ecg_signal, sample["name"])
+
+    return None, None, None, "Sample not found."
+
+
+def create_demo_interface():
+    """Create the Gradio interface."""
+
+    # Custom CSS for styling
+    custom_css = """
+    .gradio-container {
+        font-family: 'Inter', sans-serif;
+    }
+    .main-header {
+        text-align: center;
+        padding: 20px;
+        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+        color: white;
+        border-radius: 10px;
+        margin-bottom: 20px;
+    }
+    .main-header h1 {
+        margin: 0;
+        font-size: 2.5em;
+    }
+    .main-header p {
+        margin: 10px 0 0 0;
+        opacity: 0.9;
+    }
+    """
+
+    with gr.Blocks(css=custom_css, title="HeartWatch AI") as demo:
+        # Header
+        gr.HTML("""
+        <div class="main-header">
+            <h1>HeartWatch AI</h1>
+            <p>AI-Powered 12-Lead ECG Analysis</p>
+        </div>
+        """)
+
+        gr.Markdown("""
+        This demo analyzes 12-lead ECG signals using deep learning models trained on large clinical datasets.
+
+        **Models:**
+        - 77-class ECG diagnosis classifier
+        - LVEF < 40% prediction
+        - LVEF < 50% prediction
+        - 5-year Atrial Fibrillation risk
+
+        **Note:** This is a research demo. Results should not be used for clinical decision-making.
+        """)
+
+        with gr.Tabs():
+            # Tab 1: Upload ECG
+            with gr.TabItem("Upload ECG"):
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        file_input = gr.File(
+                            label="Upload ECG (.npy file)",
+                            file_types=[".npy"],
+                            type="filepath"
+                        )
+                        analyze_btn = gr.Button("Analyze ECG", variant="primary")
+
+                        gr.Markdown("""
+                        **Expected Format:**
+                        - NumPy array shape: (2500, 12) or (12, 2500)
+                        - 12 leads: I, II, III, aVR, aVL, aVF, V1-V6
+                        - 2500 samples (10 seconds at 250 Hz)
+                        """)
+
+                    with gr.Column(scale=2):
+                        summary_output = gr.Markdown(label="Summary")
+
+                with gr.Row():
+                    ecg_plot = gr.Plot(label="12-Lead ECG")
+
+                with gr.Row():
+                    with gr.Column():
+                        diagnosis_plot = gr.Plot(label="Diagnosis Probabilities")
+                    with gr.Column():
+                        risk_plot = gr.Plot(label="Risk Assessment")
+
+                analyze_btn.click(
+                    fn=analyze_uploaded_file,
+                    inputs=[file_input],
+                    outputs=[ecg_plot, diagnosis_plot, risk_plot, summary_output]
+                )
+
+            # Tab 2: Sample Gallery
+            with gr.TabItem("Sample Gallery"):
+                gr.Markdown("### Select a sample ECG to analyze")
+
+                samples = get_sample_ecgs()
+                if samples:
+                    sample_names = [s["name"] for s in samples]
+                    sample_dropdown = gr.Dropdown(
+                        choices=sample_names,
+                        label="Select Sample",
+                        value=sample_names[0] if sample_names else None
+                    )
+                    analyze_sample_btn = gr.Button("Analyze Sample", variant="primary")
+
+                    with gr.Row():
+                        sample_summary = gr.Markdown(label="Summary")
+
+                    with gr.Row():
+                        sample_ecg_plot = gr.Plot(label="12-Lead ECG")
+
+                    with gr.Row():
+                        with gr.Column():
+                            sample_diagnosis_plot = gr.Plot(label="Diagnosis Probabilities")
+                        with gr.Column():
+                            sample_risk_plot = gr.Plot(label="Risk Assessment")
+
+                    analyze_sample_btn.click(
+                        fn=analyze_sample,
+                        inputs=[sample_dropdown],
+                        outputs=[sample_ecg_plot, sample_diagnosis_plot, sample_risk_plot, sample_summary]
+                    )
+                else:
+                    gr.Markdown("*No sample ECGs available. Upload your own in the Upload tab.*")
+
+            # Tab 3: About
+            with gr.TabItem("About"):
+                gr.Markdown("""
+                ## About HeartWatch AI
+
+                HeartWatch AI is a deep learning-based ECG analysis system that can:
+
+                ### Models
+
+                1. **77-Class Diagnosis Model**
+                   - Trained to detect 77 different ECG patterns and conditions
+                   - Based on EfficientNetV2 architecture
+                   - Outputs probability for each condition
+
+                2. **LVEF Prediction Models**
+                   - LVEF < 40%: Identifies patients with reduced ejection fraction
+                   - LVEF < 50%: Identifies patients with moderately reduced ejection fraction
+
+                3. **AFib Risk Model**
+                   - Predicts 5-year risk of developing Atrial Fibrillation
+
+                ### Technical Details
+
+                - **Input:** 12-lead ECG, 10 seconds, 250 Hz sampling rate
+                - **Architecture:** EfficientNetV2 (TorchScript optimized)
+                - **Inference:** CPU-optimized for accessibility
+
+                ### Disclaimer
+
+                This is a research demonstration tool. The predictions provided should NOT be used
+                for clinical decision-making. Always consult qualified healthcare professionals
+                for medical advice and diagnosis.
+
+                ### Contact
+
+                For questions or issues, please visit our GitHub repository.
+                """)
+
+        # Footer
+        gr.Markdown("""
+        ---
+        *Built with Gradio and PyTorch. Models from DeepECG project.*
+        """)
+
+    return demo
+
+
+# Create and launch the demo
+if __name__ == "__main__":
+    # Pre-load the inference engine
+    try:
+        load_inference_engine()
+    except Exception as e:
+        logger.warning(f"Could not pre-load models: {e}")
+
+    # Create and launch demo
+    demo = create_demo_interface()
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False
+    )

class_names.json

@@ -0,0 +1,79 @@
+[
+    "Sinusal",
+    "Regular",
+    "Monomorph",
+    "QS complex in V1-V2-V3",
+    "R complex in V5-V6",
+    "T wave inversion (inferior - II, III, aVF)",
+    "Left bundle branch block",
+    "RaVL > 11 mm",
+    "SV1 + RV5 or RV6 > 35 mm",
+    "T wave inversion (lateral -I, aVL, V5-V6)",
+    "T wave inversion (anterior - V3-V4)",
+    "Left axis deviation",
+    "Left ventricular hypertrophy",
+    "Bradycardia",
+    "Q wave (inferior - II, III, aVF)",
+    "Afib",
+    "Irregularly irregular",
+    "Atrial tachycardia (>= 100 BPM)",
+    "Nonspecific intraventricular conduction delay",
+    "Premature ventricular complex",
+    "Polymorph",
+    "T wave inversion (septal- V1-V2)",
+    "Right bundle branch block",
+    "Ventricular paced",
+    "ST elevation (anterior - V3-V4)",
+    "ST elevation (septal - V1-V2)",
+    "1st degree AV block",
+    "Premature atrial complex",
+    "Atrial flutter",
+    "rSR' in V1-V2",
+    "qRS in V5-V6-I, aVL",
+    "Left anterior fascicular block",
+    "Right axis deviation",
+    "2nd degree AV block - mobitz 1",
+    "ST depression (inferior - II, III, aVF)",
+    "Acute pericarditis",
+    "ST elevation (inferior - II, III, aVF)",
+    "Low voltage",
+    "Regularly irregular",
+    "Junctional rhythm",
+    "Left atrial enlargement",
+    "ST elevation (lateral - I, aVL, V5-V6)",
+    "Atrial paced",
+    "Right ventricular hypertrophy",
+    "Delta wave",
+    "Wolff-Parkinson-White (Pre-excitation syndrome)",
+    "Prolonged QT",
+    "ST depression (anterior - V3-V4)",
+    "QRS complex negative in III",
+    "Q wave (lateral- I, aVL, V5-V6)",
+    "Supraventricular tachycardia",
+    "ST downslopping",
+    "ST depression (lateral - I, avL, V5-V6)",
+    "2nd degree AV block - mobitz 2",
+    "U wave",
+    "R/S ratio in V1-V2 >1",
+    "RV1 + SV6 > 11 mm",
+    "Left posterior fascicular block",
+    "Right atrial enlargement",
+    "ST depression (septal- V1-V2)",
+    "Q wave (septal- V1-V2)",
+    "Q wave (anterior - V3-V4)",
+    "ST upslopping",
+    "Right superior axis",
+    "Ventricular tachycardia",
+    "ST elevation (posterior - V7-V8-V9)",
+    "Ectopic atrial rhythm (< 100 BPM)",
+    "Lead misplacement",
+    "Third Degree AV Block",
+    "Acute MI",
+    "Early repolarization",
+    "Q wave (posterior - V7-V9)",
+    "Bi-atrial enlargement",
+    "LV pacing",
+    "Brugada",
+    "Ventricular Rhythm",
+    "no_qrs"
+]

demo_data/samples/Atrial_Flutter.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8704d5fb100d5a9080640587aef6cb69c725b82b4dbed4ecf72510f67acd92c9
+size 120128

demo_data/samples/Normal_Sinus_Rhythm.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8a1191aa4d9a7dc79b67fcd2eac378b86456d7b1ed035a443819ed88b0c7065f
+size 120128

demo_data/samples/Ventricular_Tachycardia.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a43e7555616d2601daf958140b9770a69522cce0126ed1402f546829646263dc
+size 120128

inference.py

@@ -0,0 +1,414 @@
+"""
+DeepECG Inference Module for HeartWatch AI
+===========================================
+
+This module provides CPU-optimized inference for 4 EfficientNetV2 models:
+- 77-class ECG diagnosis
+- LVEF <= 40% prediction
+- LVEF < 50% prediction
+- 5-year AFib risk prediction
+
+The preprocessing exactly replicates DeepECG's pipeline:
+1. Load signal as (samples, leads) = (2500, 12)
+2. Transpose to (leads, samples) = (12, 2500)
+3. Apply MHI factor scaling: signal *= (1/0.0048)
+4. Apply sigmoid to model logits
+
+Models are downloaded from HuggingFace Hub using HF_TOKEN from environment.
+"""
+
+import os
+import json
+import time
+import logging
+from typing import Dict, Optional, Any, Union
+from pathlib import Path
+
+import numpy as np
+import torch
+from huggingface_hub import snapshot_download
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# CPU optimizations for HuggingFace Spaces (no GPU)
+torch.set_num_threads(2)
+torch.set_flush_denormal(True)
+
+
+class DeepECGInference:
+    """
+    CPU-optimized inference engine for DeepECG EfficientNetV2 models.
+
+    Loads 4 models from HuggingFace Hub:
+    - heartwise/EfficientNetV2_77_Classes: 77-class ECG diagnosis
+    - heartwise/EfficientNetV2_LVEF_40: LVEF <= 40% prediction
+    - heartwise/EfficientNetV2_LVEF_50: LVEF < 50% prediction
+    - heartwise/EfficientNetV2_AFIB_5y: 5-year AFib risk prediction
+
+    Attributes:
+        device: Always CPU for HF Spaces
+        models: Dict containing loaded TorchScript models
+        class_names: List of 77 ECG diagnosis class names
+        mhi_factor: Scaling factor for signal preprocessing (1/0.0048)
+    """
+
+    # Model repository mappings
+    MODEL_REPOS = {
+        "diagnosis_77": "heartwise/EfficientNetV2_77_Classes",
+        "lvef_40": "heartwise/EfficientNetV2_LVEF_40",
+        "lvef_50": "heartwise/EfficientNetV2_LVEF_50",
+        "afib_5y": "heartwise/EfficientNetV2_AFIB_5y",
+    }
+
+    # Expected input specifications
+    EXPECTED_LEADS = 12
+    EXPECTED_SAMPLES = 2500  # 10 seconds at 250 Hz
+    SAMPLING_RATE = 250  # Hz
+
+    # Preprocessing constants from DeepECG
+    MHI_FACTOR = 1 / 0.0048  # ~208.33
+
+    def __init__(self, cache_dir: Optional[str] = None):
+        """
+        Initialize the inference engine.
+
+        Args:
+            cache_dir: Directory to cache downloaded models.
+                      Defaults to ./weights
+        """
+        self.device = torch.device("cpu")
+        self.cache_dir = cache_dir or os.path.join(os.getcwd(), "weights")
+        self.models: Dict[str, torch.jit.ScriptModule] = {}
+        self.class_names: list = []
+        self._load_class_names()
+
+    def _load_class_names(self) -> None:
+        """Load the 77 ECG class names from class_names.json."""
+        class_names_path = os.path.join(
+            os.path.dirname(os.path.abspath(__file__)),
+            "class_names.json"
+        )
+        try:
+            with open(class_names_path, "r") as f:
+                self.class_names = json.load(f)
+            logger.info(f"Loaded {len(self.class_names)} class names")
+        except FileNotFoundError:
+            logger.warning(f"class_names.json not found at {class_names_path}")
+            self.class_names = []
+
+    def _get_hf_token(self) -> Optional[str]:
+        """Get HuggingFace token from environment variable."""
+        token = os.environ.get("HF_TOKEN")
+        if not token:
+            logger.warning("HF_TOKEN environment variable not set")
+        return token
+
+    def _download_model(self, repo_id: str, model_name: str) -> str:
+        """
+        Download model from HuggingFace Hub.
+
+        Args:
+            repo_id: HuggingFace repository ID
+            model_name: Local name for the model
+
+        Returns:
+            Path to the downloaded model directory
+        """
+        local_dir = os.path.join(self.cache_dir, model_name)
+
+        if os.path.exists(local_dir):
+            logger.info(f"Model {model_name} already cached at {local_dir}")
+            return local_dir
+
+        logger.info(f"Downloading {repo_id} to {local_dir}")
+        os.makedirs(local_dir, exist_ok=True)
+
+        hf_token = self._get_hf_token()
+        local_dir = snapshot_download(
+            repo_id=repo_id,
+            local_dir=local_dir,
+            repo_type="model",
+            token=hf_token
+        )
+
+        logger.info(f"Downloaded {repo_id} to {local_dir}")
+        return local_dir
+
+    def _load_model_from_dir(self, model_dir: str) -> torch.jit.ScriptModule:
+        """
+        Load a TorchScript model from a directory.
+
+        Args:
+            model_dir: Directory containing the .pt file
+
+        Returns:
+            Loaded TorchScript model
+
+        Raises:
+            ValueError: If no .pt file is found in the directory
+        """
+        pt_file = next(
+            (f for f in os.listdir(model_dir) if f.endswith('.pt')),
+            None
+        )
+        if not pt_file:
+            raise ValueError(f"No .pt file found in {model_dir}")
+
+        model_path = os.path.join(model_dir, pt_file)
+        model = torch.jit.load(model_path, map_location=self.device)
+        model.eval()
+
+        return model
+
+    def load_models(self) -> None:
+        """
+        Download and load all 4 models from HuggingFace Hub.
+
+        Uses HF_TOKEN from os.environ for authentication.
+        Models are loaded in eval mode on CPU.
+        """
+        logger.info("Loading DeepECG models...")
+
+        for model_key, repo_id in self.MODEL_REPOS.items():
+            try:
+                model_dir = self._download_model(repo_id, model_key)
+                self.models[model_key] = self._load_model_from_dir(model_dir)
+                logger.info(f"Loaded model: {model_key} from {repo_id}")
+            except Exception as e:
+                logger.error(f"Failed to load {model_key}: {e}")
+                raise
+
+        logger.info(f"Successfully loaded {len(self.models)} models")
+
+    def preprocess_ecg(
+        self,
+        ecg_signal: Union[np.ndarray, torch.Tensor]
+    ) -> torch.Tensor:
+        """
+        Preprocess ECG signal to match DeepECG's exact preprocessing.
+
+        The preprocessing pipeline:
+        1. Ensure signal is numpy array with correct shape
+        2. Handle shape: expect (samples, leads) = (2500, 12) or (12, 2500)
+        3. Transpose to (leads, samples) = (12, 2500) if needed
+        4. Convert to float32 tensor
+        5. Add batch dimension: (1, 12, 2500)
+        6. Apply MHI factor scaling: signal *= (1/0.0048)
+
+        Args:
+            ecg_signal: Raw ECG signal, shape (samples, leads) or (leads, samples)
+                       Expected: 12 leads, 2500 samples (10s at 250Hz)
+
+        Returns:
+            Preprocessed tensor ready for model inference, shape (1, 12, 2500)
+
+        Raises:
+            ValueError: If signal shape is invalid
+        """
+        # Convert to numpy if tensor
+        if isinstance(ecg_signal, torch.Tensor):
+            ecg_signal = ecg_signal.numpy()
+
+        # Ensure float32
+        ecg_signal = ecg_signal.astype(np.float32)
+
+        # Handle shape - expect (samples, leads) = (2500, 12) or (12, 2500)
+        if ecg_signal.ndim != 2:
+            raise ValueError(
+                f"Expected 2D signal, got shape {ecg_signal.shape}"
+            )
+
+        # Determine orientation and transpose if needed
+        # If shape is (samples, leads) = (2500, 12), transpose to (12, 2500)
+        # If shape is (12, 2500), it's already correct
+        if ecg_signal.shape[0] == self.EXPECTED_SAMPLES and ecg_signal.shape[1] == self.EXPECTED_LEADS:
+            # Shape is (2500, 12) -> transpose to (12, 2500)
+            ecg_signal = ecg_signal.T
+        elif ecg_signal.shape[0] == self.EXPECTED_LEADS and ecg_signal.shape[1] == self.EXPECTED_SAMPLES:
+            # Shape is already (12, 2500)
+            pass
+        else:
+            # Try to infer orientation
+            if ecg_signal.shape[1] == self.EXPECTED_LEADS:
+                ecg_signal = ecg_signal.T
+            elif ecg_signal.shape[0] != self.EXPECTED_LEADS:
+                raise ValueError(
+                    f"Invalid signal shape {ecg_signal.shape}. "
+                    f"Expected (2500, 12) or (12, 2500)"
+                )
+
+        # Verify final shape
+        if ecg_signal.shape[0] != self.EXPECTED_LEADS:
+            raise ValueError(
+                f"Signal must have {self.EXPECTED_LEADS} leads, "
+                f"got {ecg_signal.shape[0]}"
+            )
+
+        # Convert to tensor and add batch dimension
+        signal_tensor = torch.from_numpy(ecg_signal).float()
+        signal_tensor = signal_tensor.unsqueeze(0)  # (1, 12, samples)
+
+        # Move to device (CPU)
+        signal_tensor = signal_tensor.to(self.device)
+
+        # Apply MHI factor scaling (this is done in model __call__ in DeepECG)
+        signal_tensor = signal_tensor * self.MHI_FACTOR
+
+        return signal_tensor
+
+    def predict(
+        self,
+        ecg_signal: Union[np.ndarray, torch.Tensor]
+    ) -> Dict[str, Any]:
+        """
+        Run inference on an ECG signal using all 4 models.
+
+        Args:
+            ecg_signal: Raw ECG signal, shape (samples, leads) or (leads, samples)
+                       Expected: 12 leads, 2500 samples (10s at 250Hz)
+
+        Returns:
+            Dictionary containing:
+            - diagnosis_77: Dict with 'probabilities' (77 floats) and 'class_names'
+            - lvef_40: Probability of LVEF <= 40%
+            - lvef_50: Probability of LVEF < 50%
+            - afib_5y: Probability of AFib within 5 years
+            - inference_time_ms: Total inference time in milliseconds
+        """
+        if not self.models:
+            raise RuntimeError("Models not loaded. Call load_models() first.")
+
+        start_time = time.time()
+
+        # Preprocess the signal
+        signal_tensor = self.preprocess_ecg(ecg_signal)
+
+        results = {}
+
+        with torch.no_grad():
+            # 77-class diagnosis
+            if "diagnosis_77" in self.models:
+                logits = self.models["diagnosis_77"](signal_tensor)
+                probs = torch.sigmoid(logits)
+                probs_list = probs.squeeze().cpu().numpy().tolist()
+                results["diagnosis_77"] = {
+                    "probabilities": probs_list,
+                    "class_names": self.class_names if self.class_names else None,
+                }
+
+            # LVEF <= 40%
+            if "lvef_40" in self.models:
+                logits = self.models["lvef_40"](signal_tensor)
+                prob = torch.sigmoid(logits)
+                results["lvef_40"] = float(prob.squeeze().cpu().numpy())
+
+            # LVEF < 50%
+            if "lvef_50" in self.models:
+                logits = self.models["lvef_50"](signal_tensor)
+                prob = torch.sigmoid(logits)
+                results["lvef_50"] = float(prob.squeeze().cpu().numpy())
+
+            # 5-year AFib risk
+            if "afib_5y" in self.models:
+                logits = self.models["afib_5y"](signal_tensor)
+                prob = torch.sigmoid(logits)
+                results["afib_5y"] = float(prob.squeeze().cpu().numpy())
+
+        end_time = time.time()
+        results["inference_time_ms"] = (end_time - start_time) * 1000
+
+        return results
+
+    def predict_diagnosis_top_k(
+        self,
+        ecg_signal: Union[np.ndarray, torch.Tensor],
+        k: int = 5
+    ) -> Dict[str, Any]:
+        """
+        Get top-k diagnoses from the 77-class model.
+
+        Args:
+            ecg_signal: Raw ECG signal
+            k: Number of top predictions to return
+
+        Returns:
+            Dictionary with top-k predictions sorted by probability
+        """
+        results = self.predict(ecg_signal)
+
+        if "diagnosis_77" not in results:
+            raise RuntimeError("77-class diagnosis model not loaded")
+
+        probs = results["diagnosis_77"]["probabilities"]
+        class_names = results["diagnosis_77"]["class_names"] or [f"Class_{i}" for i in range(77)]
+
+        # Get top-k indices
+        top_k_indices = np.argsort(probs)[::-1][:k]
+
+        top_k_predictions = [
+            {
+                "class_name": class_names[idx],
+                "probability": probs[idx],
+                "class_index": int(idx)
+            }
+            for idx in top_k_indices
+        ]
+
+        return {
+            "top_k_predictions": top_k_predictions,
+            "inference_time_ms": results["inference_time_ms"]
+        }
+
+
+def get_inference_engine(cache_dir: Optional[str] = None) -> DeepECGInference:
+    """
+    Factory function to create and initialize a DeepECGInference instance.
+
+    Args:
+        cache_dir: Optional directory to cache models
+
+    Returns:
+        Initialized DeepECGInference with models loaded
+    """
+    engine = DeepECGInference(cache_dir=cache_dir)
+    engine.load_models()
+    return engine
+
+
+if __name__ == "__main__":
+    # Example usage / testing
+    print("DeepECG Inference Module")
+    print("=" * 50)
+
+    # Create inference engine
+    engine = DeepECGInference()
+
+    # Load models (requires HF_TOKEN environment variable)
+    try:
+        engine.load_models()
+        print("Models loaded successfully!")
+
+        # Create dummy signal for testing
+        dummy_signal = np.random.randn(2500, 12).astype(np.float32)
+
+        # Run inference
+        results = engine.predict(dummy_signal)
+
+        print(f"\nInference time: {results['inference_time_ms']:.2f} ms")
+        print(f"LVEF <= 40%: {results['lvef_40']:.4f}")
+        print(f"LVEF < 50%: {results['lvef_50']:.4f}")
+        print(f"5-year AFib risk: {results['afib_5y']:.4f}")
+        print(f"77-class diagnosis: {len(results['diagnosis_77']['probabilities'])} classes")
+
+        # Get top-5 diagnoses
+        top_5 = engine.predict_diagnosis_top_k(dummy_signal, k=5)
+        print("\nTop 5 diagnoses:")
+        for pred in top_5["top_k_predictions"]:
+            print(f"  {pred['class_name']}: {pred['probability']:.4f}")
+
+    except Exception as e:
+        print(f"Error: {e}")
+        print("\nMake sure HF_TOKEN environment variable is set:")
+        print("  export HF_TOKEN='your_huggingface_token'")

requirements.txt

@@ -0,0 +1,7 @@
+gradio>=4.0.0
+torch>=2.0.0
+numpy>=1.21.0
+matplotlib>=3.5.0
+Pillow>=9.0.0
+huggingface_hub>=0.16.0
+scipy>=1.7.0

test_inference.py

@@ -0,0 +1,115 @@
+#!/usr/bin/env python3
+"""
+Test script for DeepECG Inference Module
+=========================================
+
+Run this script to verify the inference engine works correctly.
+
+Usage:
+    # Set HF_TOKEN environment variable first
+    export HF_TOKEN='your_huggingface_token'
+
+    # Run the test
+    python test_inference.py
+
+Expected output:
+    - Models download from HuggingFace Hub
+    - Dummy signal inference completes
+    - Results for all 4 models are printed
+"""
+
+import os
+import sys
+import numpy as np
+
+# Ensure HF_TOKEN is set
+if not os.environ.get("HF_TOKEN"):
+    print("ERROR: HF_TOKEN environment variable not set")
+    print("Please run: export HF_TOKEN='your_token'")
+    sys.exit(1)
+
+print("=" * 60)
+print("DeepECG Inference Test")
+print("=" * 60)
+
+# Import the inference module
+try:
+    from inference import DeepECGInference
+    print("[OK] Import successful")
+except ImportError as e:
+    print(f"[FAIL] Import failed: {e}")
+    sys.exit(1)
+
+# Create inference engine
+try:
+    engine = DeepECGInference()
+    print(f"[OK] Engine created with {len(engine.class_names)} class names")
+except Exception as e:
+    print(f"[FAIL] Engine creation failed: {e}")
+    sys.exit(1)
+
+# Load models
+print("\nLoading models from HuggingFace Hub...")
+try:
+    engine.load_models()
+    print(f"[OK] Loaded {len(engine.models)} models")
+    for name in engine.models:
+        print(f"     - {name}")
+except Exception as e:
+    print(f"[FAIL] Model loading failed: {e}")
+    sys.exit(1)
+
+# Test with dummy signal
+print("\nTesting inference with dummy signal...")
+try:
+    # Create dummy 10-second ECG (2500 samples at 250Hz, 12 leads)
+    dummy_signal = np.random.randn(2500, 12).astype(np.float32)
+
+    # Run inference
+    results = engine.predict(dummy_signal)
+
+    print(f"[OK] Inference completed in {results['inference_time_ms']:.2f} ms")
+    print(f"\nResults:")
+    print(f"  - LVEF <= 40%: {results['lvef_40']:.4f}")
+    print(f"  - LVEF < 50%:  {results['lvef_50']:.4f}")
+    print(f"  - 5-year AFib: {results['afib_5y']:.4f}")
+    print(f"  - 77-class diagnosis: {len(results['diagnosis_77']['probabilities'])} probabilities")
+
+except Exception as e:
+    print(f"[FAIL] Inference failed: {e}")
+    import traceback
+    traceback.print_exc()
+    sys.exit(1)
+
+# Test top-k predictions
+print("\nTop 5 diagnoses:")
+try:
+    top_5 = engine.predict_diagnosis_top_k(dummy_signal, k=5)
+    for pred in top_5["top_k_predictions"]:
+        print(f"  {pred['class_name']}: {pred['probability']:.4f}")
+except Exception as e:
+    print(f"[FAIL] Top-k prediction failed: {e}")
+    sys.exit(1)
+
+# Test preprocessing edge cases
+print("\nTesting preprocessing with different input shapes...")
+try:
+    # Test (2500, 12) shape
+    signal_1 = np.random.randn(2500, 12).astype(np.float32)
+    tensor_1 = engine.preprocess_ecg(signal_1)
+    assert tensor_1.shape == (1, 12, 2500), f"Expected (1, 12, 2500), got {tensor_1.shape}"
+    print(f"[OK] Shape (2500, 12) -> {tuple(tensor_1.shape)}")
+
+    # Test (12, 2500) shape
+    signal_2 = np.random.randn(12, 2500).astype(np.float32)
+    tensor_2 = engine.preprocess_ecg(signal_2)
+    assert tensor_2.shape == (1, 12, 2500), f"Expected (1, 12, 2500), got {tensor_2.shape}"
+    print(f"[OK] Shape (12, 2500) -> {tuple(tensor_2.shape)}")
+
+except Exception as e:
+    print(f"[FAIL] Preprocessing test failed: {e}")
+    sys.exit(1)
+
+print("\n" + "=" * 60)
+print("ALL TESTS PASSED!")
+print("=" * 60)

visualization.py

@@ -0,0 +1,433 @@
+"""
+HeartWatch AI Visualization Module
+
+This module provides visualization functions for ECG analysis including:
+- 12-lead ECG waveform plotting with clinical layout
+- Diagnosis probability bar charts
+- Risk assessment gauges
+- ECG thumbnail generation for galleries
+"""
+
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.patches as mpatches
+from matplotlib.patches import Wedge
+from PIL import Image
+import io
+
+
+# Standard 12-lead ECG names in clinical order
+LEAD_NAMES = ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6']
+
+# Clinical layout: 4 columns x 3 rows
+# Col 1: I, II, III | Col 2: aVR, aVL, aVF | Col 3: V1, V2, V3 | Col 4: V4, V5, V6
+LEAD_LAYOUT = [
+    ['I', 'aVR', 'V1', 'V4'],
+    ['II', 'aVL', 'V2', 'V5'],
+    ['III', 'aVF', 'V3', 'V6']
+]
+
+
+def plot_ecg_waveform(ecg_signal: np.ndarray, sample_rate: int = 250,
+                      title: str = "12-Lead ECG") -> plt.Figure:
+    """
+    Plot a 12-lead ECG waveform in clinical layout format.
+
+    Parameters
+    ----------
+    ecg_signal : np.ndarray
+        ECG signal array of shape (12, n_samples) or (n_samples, 12)
+        Each row/column represents one of the 12 standard leads
+    sample_rate : int, optional
+        Sampling rate in Hz, default 250
+    title : str, optional
+        Figure title, default "12-Lead ECG"
+
+    Returns
+    -------
+    plt.Figure
+        Matplotlib figure with 4x3 ECG layout
+    """
+    # Ensure correct shape (12, n_samples)
+    if ecg_signal.shape[0] != 12:
+        if ecg_signal.shape[1] == 12:
+            ecg_signal = ecg_signal.T
+        else:
+            raise ValueError(f"ECG signal must have 12 leads, got shape {ecg_signal.shape}")
+
+    n_samples = ecg_signal.shape[1]
+
+    # 2.5 seconds per column
+    samples_per_col = int(2.5 * sample_rate)
+
+    # Create figure with clinical dimensions
+    fig, axes = plt.subplots(3, 4, figsize=(14, 8))
+    fig.suptitle(title, fontsize=14, fontweight='bold', y=0.98)
+
+    # Create lead index mapping
+    lead_to_idx = {name: i for i, name in enumerate(LEAD_NAMES)}
+
+    for row in range(3):
+        for col in range(4):
+            ax = axes[row, col]
+            lead_name = LEAD_LAYOUT[row][col]
+            lead_idx = lead_to_idx[lead_name]
+
+            # Get signal segment for this column (2.5 sec)
+            start_sample = 0
+            end_sample = min(samples_per_col, n_samples)
+
+            signal_segment = ecg_signal[lead_idx, start_sample:end_sample]
+            time_segment = np.arange(len(signal_segment)) / sample_rate
+
+            # Set up ECG paper grid background (pink/red)
+            ax.set_facecolor('#fff5f5')
+
+            # Major grid (0.5 sec, 0.5 mV equivalent)
+            ax.set_axisbelow(True)
+            ax.grid(True, which='major', color='#ffcccc', linewidth=0.8, linestyle='-')
+            ax.grid(True, which='minor', color='#ffe6e6', linewidth=0.4, linestyle='-')
+
+            # Set tick spacing for major/minor grids
+            ax.set_xticks(np.arange(0, 2.6, 0.5))
+            ax.set_xticks(np.arange(0, 2.6, 0.1), minor=True)
+
+            # Calculate y-limits based on signal range
+            signal_min, signal_max = signal_segment.min(), signal_segment.max()
+            signal_range = signal_max - signal_min
+            if signal_range < 0.1:
+                signal_range = 2.0  # Default range if signal is flat
+            padding = signal_range * 0.1
+            y_min = signal_min - padding
+            y_max = signal_max + padding
+
+            # Set y-ticks for grid
+            y_tick_spacing = signal_range / 4
+            ax.set_yticks(np.arange(y_min, y_max + y_tick_spacing, y_tick_spacing))
+            ax.set_yticks(np.arange(y_min, y_max + y_tick_spacing/5, y_tick_spacing/5), minor=True)
+
+            # Plot ECG waveform
+            ax.plot(time_segment, signal_segment, color='black', linewidth=0.8)
+
+            # Add lead label
+            ax.text(0.02, 0.98, lead_name, transform=ax.transAxes,
+                   fontsize=10, fontweight='bold', verticalalignment='top',
+                   bbox=dict(boxstyle='round,pad=0.2', facecolor='white',
+                            edgecolor='none', alpha=0.7))
+
+            # Set axis limits
+            ax.set_xlim(0, 2.5)
+            ax.set_ylim(y_min, y_max)
+
+            # Remove tick labels for cleaner look (except bottom row and left column)
+            if row < 2:
+                ax.set_xticklabels([])
+            else:
+                ax.set_xlabel('Time (s)', fontsize=8)
+
+            if col > 0:
+                ax.set_yticklabels([])
+            else:
+                ax.set_ylabel('Amplitude (mV)', fontsize=8)
+
+            ax.tick_params(axis='both', which='both', labelsize=6)
+
+    plt.tight_layout(rect=[0, 0, 1, 0.96])
+    return fig
+
+
+def plot_diagnosis_bars(diagnosis_77: dict, top_n: int = 10,
+                        ground_truth: list = None) -> plt.Figure:
+    """
+    Plot horizontal bar chart of diagnosis probabilities.
+
+    Parameters
+    ----------
+    diagnosis_77 : dict
+        Dictionary mapping diagnosis names to probabilities (0-1)
+    top_n : int, optional
+        Number of top diagnoses to display, default 10
+    ground_truth : list, optional
+        List of ground truth diagnosis names to mark with star
+
+    Returns
+    -------
+    plt.Figure
+        Matplotlib figure with horizontal bar chart
+    """
+    if ground_truth is None:
+        ground_truth = []
+
+    # Sort diagnoses by probability (descending)
+    sorted_diagnoses = sorted(diagnosis_77.items(), key=lambda x: x[1], reverse=True)
+    top_diagnoses = sorted_diagnoses[:top_n]
+
+    # Extract names and probabilities
+    names = [d[0] for d in top_diagnoses]
+    probs = [d[1] for d in top_diagnoses]
+
+    # Determine colors based on probability thresholds
+    colors = []
+    for p in probs:
+        if p >= 0.7:
+            colors.append('#2ecc71')  # Green for high confidence
+        elif p >= 0.3:
+            colors.append('#f1c40f')  # Yellow for moderate
+        else:
+            colors.append('#95a5a6')  # Gray for low confidence
+
+    # Create figure
+    fig, ax = plt.subplots(figsize=(8, 6))
+
+    # Create horizontal bar chart
+    y_pos = np.arange(len(names))
+    bars = ax.barh(y_pos, probs, color=colors, edgecolor='black', linewidth=0.5)
+
+    # Add probability labels on bars
+    for i, (bar, prob) in enumerate(zip(bars, probs)):
+        width = bar.get_width()
+        label_x = width + 0.02 if width < 0.85 else width - 0.08
+        label_color = 'black' if width < 0.85 else 'white'
+        ax.text(label_x, bar.get_y() + bar.get_height()/2,
+                f'{prob:.1%}', va='center', fontsize=9, color=label_color)
+
+    # Mark ground truth with star
+    display_names = []
+    for name in names:
+        if name in ground_truth:
+            display_names.append(f'{name} \u2605')  # Unicode star
+        else:
+            display_names.append(name)
+
+    # Set y-axis labels
+    ax.set_yticks(y_pos)
+    ax.set_yticklabels(display_names, fontsize=9)
+
+    # Set axis limits and labels
+    ax.set_xlim(0, 1.0)
+    ax.set_xlabel('Probability', fontsize=11)
+    ax.set_title('Diagnosis Probabilities (Top {})'.format(top_n),
+                 fontsize=12, fontweight='bold', pad=10)
+
+    # Add legend
+    legend_elements = [
+        mpatches.Patch(facecolor='#2ecc71', edgecolor='black', label='High (\u2265 70%)'),
+        mpatches.Patch(facecolor='#f1c40f', edgecolor='black', label='Moderate (30-70%)'),
+        mpatches.Patch(facecolor='#95a5a6', edgecolor='black', label='Low (< 30%)')
+    ]
+    if ground_truth:
+        legend_elements.append(mpatches.Patch(facecolor='white', edgecolor='white',
+                                               label='\u2605 = Ground Truth'))
+    ax.legend(handles=legend_elements, loc='lower right', fontsize=8)
+
+    # Add grid for readability
+    ax.xaxis.grid(True, linestyle='--', alpha=0.7)
+    ax.set_axisbelow(True)
+
+    # Invert y-axis so highest probability is at top
+    ax.invert_yaxis()
+
+    plt.tight_layout()
+    return fig
+
+
+def _draw_gauge(ax, value: float, title: str):
+    """
+    Draw a semicircular gauge on the given axes.
+
+    Parameters
+    ----------
+    ax : matplotlib.axes.Axes
+        Axes to draw on
+    value : float
+        Value between 0 and 1 to display
+    title : str
+        Gauge title
+    """
+    # Clear axes
+    ax.clear()
+    ax.set_xlim(-1.5, 1.5)
+    ax.set_ylim(-0.3, 1.3)
+    ax.set_aspect('equal')
+    ax.axis('off')
+
+    # Create gradient background arc (Green -> Yellow -> Red)
+    n_segments = 100
+    for i in range(n_segments):
+        theta1 = 180 - i * (180 / n_segments)
+        theta2 = 180 - (i + 1) * (180 / n_segments)
+
+        # Calculate color based on position
+        pos = i / n_segments
+        if pos < 0.3:
+            # Green zone
+            color = '#2ecc71'
+        elif pos < 0.6:
+            # Yellow zone (transition from green to yellow)
+            t = (pos - 0.3) / 0.3
+            r = int(46 + t * (241 - 46))
+            g = int(204 + t * (196 - 204))
+            b = int(113 + t * (15 - 113))
+            color = f'#{r:02x}{g:02x}{b:02x}'
+        else:
+            # Red zone (transition from yellow to red)
+            t = (pos - 0.6) / 0.4
+            r = int(241 + t * (231 - 241))
+            g = int(196 - t * 196)
+            b = int(15 - t * 15)
+            color = f'#{r:02x}{g:02x}{b:02x}'
+
+        wedge = Wedge((0, 0), 1.0, theta2, theta1, width=0.3, facecolor=color,
+                      edgecolor='white', linewidth=0.5)
+        ax.add_patch(wedge)
+
+    # Draw needle
+    needle_angle = 180 - value * 180
+    needle_rad = np.radians(needle_angle)
+    needle_length = 0.85
+    needle_x = needle_length * np.cos(needle_rad)
+    needle_y = needle_length * np.sin(needle_rad)
+
+    ax.annotate('', xy=(needle_x, needle_y), xytext=(0, 0),
+                arrowprops=dict(arrowstyle='->', color='#2c3e50', lw=2))
+
+    # Draw center circle
+    center_circle = plt.Circle((0, 0), 0.1, color='#2c3e50', zorder=5)
+    ax.add_patch(center_circle)
+
+    # Add value text
+    ax.text(0, -0.15, f'{value*100:.0f}%', ha='center', va='top',
+            fontsize=14, fontweight='bold', color='#2c3e50')
+
+    # Add title
+    ax.text(0, 1.2, title, ha='center', va='bottom',
+            fontsize=11, fontweight='bold', color='#2c3e50')
+
+    # Add risk labels
+    ax.text(-1.1, -0.05, 'Low', ha='center', va='top', fontsize=8, color='#27ae60')
+    ax.text(0, 1.05, 'Moderate', ha='center', va='bottom', fontsize=8, color='#f39c12')
+    ax.text(1.1, -0.05, 'High', ha='center', va='top', fontsize=8, color='#c0392b')
+
+    # Add threshold markers
+    for pct, label in [(0.3, '30%'), (0.6, '60%')]:
+        angle = 180 - pct * 180
+        rad = np.radians(angle)
+        x_outer = 1.05 * np.cos(rad)
+        y_outer = 1.05 * np.sin(rad)
+        ax.text(x_outer, y_outer, label, ha='center', va='center', fontsize=7, color='#7f8c8d')
+
+
+def plot_risk_gauges(lvef_40: float, lvef_50: float, afib_5y: float) -> plt.Figure:
+    """
+    Plot risk assessment gauges for LVEF and AFib predictions.
+
+    Parameters
+    ----------
+    lvef_40 : float
+        Probability (0-1) of LVEF < 40%
+    lvef_50 : float
+        Probability (0-1) of LVEF < 50%
+    afib_5y : float
+        Probability (0-1) of AFib within 5 years
+
+    Returns
+    -------
+    plt.Figure
+        Matplotlib figure with 3 semicircular gauges
+    """
+    # Clamp values to [0, 1]
+    lvef_40 = np.clip(lvef_40, 0, 1)
+    lvef_50 = np.clip(lvef_50, 0, 1)
+    afib_5y = np.clip(afib_5y, 0, 1)
+
+    # Create figure with 3 subplots
+    fig, axes = plt.subplots(1, 3, figsize=(14, 4))
+    fig.suptitle('Risk Assessment', fontsize=14, fontweight='bold', y=0.98)
+
+    # Draw each gauge
+    _draw_gauge(axes[0], lvef_40, 'LVEF < 40%')
+    _draw_gauge(axes[1], lvef_50, 'LVEF < 50%')
+    _draw_gauge(axes[2], afib_5y, 'AFib (5-year)')
+
+    plt.tight_layout(rect=[0, 0, 1, 0.95])
+    return fig
+
+
+def generate_thumbnail(ecg_signal: np.ndarray, label: str,
+                       sample_rate: int = 250) -> Image.Image:
+    """
+    Generate a thumbnail preview image of Lead II for gallery display.
+
+    Parameters
+    ----------
+    ecg_signal : np.ndarray
+        ECG signal array of shape (12, n_samples) or (n_samples, 12)
+    label : str
+        Label text to display on thumbnail
+    sample_rate : int, optional
+        Sampling rate in Hz, default 250
+
+    Returns
+    -------
+    PIL.Image.Image
+        Thumbnail image approximately 300x150 pixels
+    """
+    # Ensure correct shape (12, n_samples)
+    if ecg_signal.shape[0] != 12:
+        if ecg_signal.shape[1] == 12:
+            ecg_signal = ecg_signal.T
+        else:
+            raise ValueError(f"ECG signal must have 12 leads, got shape {ecg_signal.shape}")
+
+    # Extract Lead II (index 1)
+    lead_ii = ecg_signal[1, :]
+    n_samples = len(lead_ii)
+    time = np.arange(n_samples) / sample_rate
+
+    # Create figure with appropriate DPI for ~300x150 pixel output
+    fig, ax = plt.subplots(figsize=(3, 1.5), dpi=100)
+
+    # Clean, minimal design
+    ax.plot(time, lead_ii, color='#e74c3c', linewidth=1.0)
+
+    # Set background
+    ax.set_facecolor('#fafafa')
+    fig.patch.set_facecolor('#fafafa')
+
+    # Remove axes for clean look
+    ax.set_xticks([])
+    ax.set_yticks([])
+    for spine in ax.spines.values():
+        spine.set_visible(False)
+
+    # Add label
+    ax.text(0.02, 0.98, label, transform=ax.transAxes,
+            fontsize=8, fontweight='bold', verticalalignment='top',
+            color='#2c3e50')
+
+    # Add "Lead II" indicator
+    ax.text(0.98, 0.02, 'Lead II', transform=ax.transAxes,
+            fontsize=6, verticalalignment='bottom', horizontalalignment='right',
+            color='#7f8c8d')
+
+    plt.tight_layout(pad=0.2)
+
+    # Convert to PIL Image
+    buf = io.BytesIO()
+    fig.savefig(buf, format='png', facecolor=fig.get_facecolor(),
+                edgecolor='none', bbox_inches='tight', pad_inches=0.05)
+    plt.close(fig)
+
+    buf.seek(0)
+    img = Image.open(buf)
+
+    # Resize to ensure ~300x150 pixels
+    img = img.resize((300, 150), Image.Resampling.LANCZOS)
+
+    return img
+
+
+if __name__ == '__main__':
+    # Quick test
+    print("Visualization module loaded successfully.")
+    print(f"Available functions: plot_ecg_waveform, plot_diagnosis_bars, plot_risk_gauges, generate_thumbnail")