ecgen-eval 0.1.0__py3-none-any.whl

ecgen_eval/__init__.py

@@ -0,0 +1,11 @@
+"""
+ECGEN-Eval: Evaluation metrics and interactive reports for synthetic ECG quality assessment.
+"""
+
+__version__ = "0.1.0"
+__author__ = "Vajira Thambawita"
+
+from ecgen_eval.data.loader import load_ecg_folder, load_ecg_npy, load_ecg_wfdb
+from ecgen_eval.metrics import AVAILABLE_METRICS
+
+__all__ = ["load_ecg_folder", "load_ecg_npy", "load_ecg_wfdb", "AVAILABLE_METRICS"]

ecgen_eval/cli.py

@@ -0,0 +1,277 @@
+"""
+ecgen-eval CLI — entry point.
+
+Commands
+--------
+  ecgen-eval eval         Run evaluation metrics and generate HTML report.
+  ecgen-eval visualize    Generate ECG sample comparison plots only.
+  ecgen-eval list-metrics List available metrics.
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+from typing import List, Optional, Tuple
+
+import click
+
+from ecgen_eval import __version__
+from ecgen_eval.data.loader import load_ecg_folder
+from ecgen_eval.metrics import AVAILABLE_METRICS
+from ecgen_eval.report.html_report import generate_report
+
+
+# ------------------------------------------------------------------ #
+# CLI group                                                            #
+# ------------------------------------------------------------------ #
+
+@click.group()
+@click.version_option(__version__, prog_name="ecgen-eval")
+def main():
+    """ECGEN-Eval: Evaluation metrics for synthetic ECG quality assessment."""
+    pass
+
+
+# ------------------------------------------------------------------ #
+# eval command                                                         #
+# ------------------------------------------------------------------ #
+
+@main.command()
+@click.option(
+    "--real", "-r",
+    required=True,
+    type=click.Path(exists=True),
+    help="Path to folder containing real ECG CSV files.",
+)
+@click.option(
+    "--synthetic", "-s",
+    multiple=True,
+    required=True,
+    type=click.Path(exists=True),
+    help="Path to folder(s) containing synthetic ECG CSV files. Repeatable.",
+)
+@click.option(
+    "--metrics", "-m",
+    multiple=True,
+    default=list(AVAILABLE_METRICS.keys()),
+    show_default=True,
+    type=click.Choice(list(AVAILABLE_METRICS.keys()), case_sensitive=False),
+    help="Metrics to compute. Repeatable. Defaults to all available metrics.",
+)
+@click.option(
+    "--output", "-o",
+    default="ecgen_eval_report.html",
+    show_default=True,
+    help="Output HTML report path.",
+)
+@click.option(
+    "--fs",
+    default=500.0,
+    show_default=True,
+    help="Sampling frequency in Hz.",
+)
+@click.option(
+    "--n-samples",
+    default=3,
+    show_default=True,
+    help="Number of ECG samples to show in gallery.",
+)
+@click.option(
+    "--leads",
+    default=None,
+    help="Comma-separated list of lead names to show (e.g. 'I,II,V1'). Default: all.",
+)
+@click.option(
+    "--compare-mode",
+    default="overlay",
+    show_default=True,
+    type=click.Choice(["overlay", "sidebyside", "grid"], case_sensitive=False),
+    help="ECG comparison mode.",
+)
+@click.option(
+    "--max-samples",
+    default=None,
+    type=int,
+    help="Cap number of loaded recordings per folder (useful for quick tests).",
+)
+@click.option(
+    "--offline",
+    is_flag=True,
+    default=False,
+    help="Embed Plotly.js in HTML (larger file, no internet needed to view).",
+)
+@click.option(
+    "--real-label",
+    default=None,
+    help="Display label for real dataset. Defaults to folder name.",
+)
+@click.option(
+    "--synth-labels",
+    default=None,
+    help="Comma-separated labels for synthetic datasets. Defaults to folder names.",
+)
+def eval(
+    real,
+    synthetic,
+    metrics,
+    output,
+    fs,
+    n_samples,
+    leads,
+    compare_mode,
+    max_samples,
+    offline,
+    real_label,
+    synth_labels,
+):
+    """Run evaluation metrics and generate an interactive HTML report."""
+
+    click.echo(f"\n{'='*60}")
+    click.echo(f"  ECGEN-Eval v{__version__}")
+    click.echo(f"{'='*60}\n")
+
+    # Parse optional labels
+    synth_label_list = (
+        [s.strip() for s in synth_labels.split(",")] if synth_labels else []
+    )
+    leads_list = [l.strip() for l in leads.split(",")] if leads else None
+
+    # -------------------------------------------------------------- #
+    # Load datasets                                                    #
+    # -------------------------------------------------------------- #
+    click.echo("Loading datasets...")
+    real_ds = load_ecg_folder(
+        real, fs=fs, label=real_label, max_samples=max_samples
+    )
+    click.echo(f"  Real:    {real_ds}")
+
+    synth_datasets = []
+    for i, s_path in enumerate(synthetic):
+        lbl = synth_label_list[i] if i < len(synth_label_list) else None
+        ds = load_ecg_folder(s_path, fs=fs, label=lbl, max_samples=max_samples)
+        synth_datasets.append(ds)
+        click.echo(f"  Synth {i+1}: {ds}")
+
+    # -------------------------------------------------------------- #
+    # Compute metrics                                                  #
+    # -------------------------------------------------------------- #
+    click.echo("\nComputing metrics...")
+    metric_results: dict = {}
+
+    for metric_key in metrics:
+        metric_cls = AVAILABLE_METRICS[metric_key.lower()]
+        metric_obj = metric_cls()
+        click.echo(f"  [{metric_key.upper()}]")
+        by_synth = {}
+        for synth_ds in synth_datasets:
+            click.echo(f"    Computing against {synth_ds.label} ...", nl=False)
+            result = metric_obj.compute(
+                real=real_ds.data,
+                synth=synth_ds.data,
+                lead_names=real_ds.lead_names,
+            )
+            by_synth[synth_ds.label] = result
+            click.echo(f" score = {result.score:.6f}")
+        metric_results[metric_key] = by_synth
+
+    # -------------------------------------------------------------- #
+    # Generate report                                                  #
+    # -------------------------------------------------------------- #
+    click.echo(f"\nGenerating report → {output}")
+    out_path = generate_report(
+        real=real_ds,
+        synthetics=synth_datasets,
+        metric_results=metric_results,
+        output_path=output,
+        n_ecg_samples=n_samples,
+        leads=leads_list,
+        compare_mode=compare_mode,
+        plotly_cdn=not offline,
+    )
+    click.echo(f"\nDone!  Report saved to: {out_path.resolve()}")
+    click.echo("Open in any modern browser.\n")
+
+
+# ------------------------------------------------------------------ #
+# visualize command                                                    #
+# ------------------------------------------------------------------ #
+
+@main.command()
+@click.option("--real", "-r", required=True, type=click.Path(exists=True))
+@click.option("--synthetic", "-s", multiple=True, required=True, type=click.Path(exists=True))
+@click.option("--output", "-o", default="ecg_comparison.html", show_default=True)
+@click.option("--fs", default=500.0, show_default=True)
+@click.option("--n-samples", default=3, show_default=True)
+@click.option("--leads", default=None, help="Comma-separated lead names.")
+@click.option(
+    "--compare-mode",
+    default="overlay",
+    type=click.Choice(["overlay", "sidebyside", "grid"]),
+)
+@click.option("--max-samples", default=None, type=int)
+@click.option("--offline", is_flag=True, default=False)
+def visualize(real, synthetic, output, fs, n_samples, leads, compare_mode, max_samples, offline):
+    """Generate ECG sample comparison plots (no metrics)."""
+    from ecgen_eval.visualization.ecg_waveform import plot_ecg_comparison
+
+    leads_list = [l.strip() for l in leads.split(",")] if leads else None
+
+    real_ds = load_ecg_folder(real, fs=fs, max_samples=max_samples)
+    synth_datasets = [
+        load_ecg_folder(s, fs=fs, max_samples=max_samples) for s in synthetic
+    ]
+
+    n_samples = min(n_samples, real_ds.n_samples)
+    plots_html_parts = []
+    include_js = "cdn" if not offline else True
+
+    for i in range(n_samples):
+        fig = plot_ecg_comparison(
+            real=real_ds,
+            synthetics=synth_datasets,
+            sample_index=i,
+            leads=leads_list,
+            mode=compare_mode,
+            add_calibration_pulse=True,
+        )
+        plots_html_parts.append(
+            f"<h3>Sample {i}</h3>"
+            + fig.to_html(
+                full_html=False,
+                include_plotlyjs=include_js if i == 0 else False,
+                config={"responsive": True},
+            )
+        )
+
+    cdn_tag = (
+        '<script src="https://cdn.plot.ly/plotly-2.27.0.min.js"></script>'
+        if not offline else ""
+    )
+    body = "\n".join(plots_html_parts)
+    html = f"""<!DOCTYPE html>
+<html><head><meta charset="UTF-8"/>
+<title>ECG Comparison</title>{cdn_tag}
+<style>body{{font-family:sans-serif;max-width:1400px;margin:auto;padding:24px}}</style>
+</head><body><h1>ECG Comparison ({compare_mode})</h1>{body}</body></html>"""
+
+    Path(output).write_text(html, encoding="utf-8")
+    click.echo(f"Saved: {Path(output).resolve()}")
+
+
+# ------------------------------------------------------------------ #
+# list-metrics command                                                 #
+# ------------------------------------------------------------------ #
+
+@main.command(name="list-metrics")
+def list_metrics():
+    """List all available evaluation metrics."""
+    click.echo(f"\nAvailable metrics in ecgen-eval v{__version__}:\n")
+    for key, cls in AVAILABLE_METRICS.items():
+        obj = cls()
+        click.echo(f"  {key:<12} {obj.description[:80]}")
+        click.echo(f"  {'':12} Ref: {obj.reference[:80]}...\n")
+
+
+if __name__ == "__main__":
+    main()

ecgen_eval/data/__init__.py

@@ -0,0 +1,4 @@
+from ecgen_eval.data.loader import load_ecg_folder, load_ecg_npy, load_ecg_wfdb
+from ecgen_eval.data.dataset import ECGDataset
+
+__all__ = ["load_ecg_folder", "load_ecg_npy", "load_ecg_wfdb", "ECGDataset"]

ecgen_eval/data/dataset.py

@@ -0,0 +1,103 @@
+"""ECGDataset: in-memory container for a collection of ECG recordings."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import List, Optional
+
+import numpy as np
+
+
+# Standard lead names for 12-lead and 8-lead ECGs
+LEADS_12 = ["I", "II", "III", "aVR", "aVL", "aVF", "V1", "V2", "V3", "V4", "V5", "V6"]
+LEADS_8 = ["I", "II", "V1", "V2", "V3", "V4", "V5", "V6"]
+
+
+@dataclass
+class ECGDataset:
+    """
+    Container for a collection of ECG recordings loaded from a folder.
+
+    Attributes
+    ----------
+    data : np.ndarray
+        Shape (n_samples, n_leads, n_timepoints).
+    lead_names : List[str]
+        Ordered list of lead names.
+    fs : float
+        Sampling frequency in Hz.
+    source_path : Path
+        Folder the data was loaded from.
+    label : str
+        Human-readable label used in plots/reports (defaults to folder name).
+    file_names : List[str]
+        Original filenames in load order.
+    """
+
+    data: np.ndarray                    # (N, L, T)
+    lead_names: List[str]
+    fs: float
+    source_path: Path
+    label: str
+    file_names: List[str] = field(default_factory=list)
+
+    # ------------------------------------------------------------------ #
+    # Convenience properties                                               #
+    # ------------------------------------------------------------------ #
+
+    @property
+    def n_samples(self) -> int:
+        return self.data.shape[0]
+
+    @property
+    def n_leads(self) -> int:
+        return self.data.shape[1]
+
+    @property
+    def n_timepoints(self) -> int:
+        return self.data.shape[2]
+
+    @property
+    def duration_sec(self) -> float:
+        return self.n_timepoints / self.fs
+
+    @property
+    def time_axis(self) -> np.ndarray:
+        """Time axis in seconds."""
+        return np.arange(self.n_timepoints) / self.fs
+
+    # ------------------------------------------------------------------ #
+    # Subsetting helpers                                                   #
+    # ------------------------------------------------------------------ #
+
+    def select_leads(self, leads: List[str]) -> "ECGDataset":
+        """Return a new ECGDataset restricted to the requested leads."""
+        indices = [self.lead_names.index(l) for l in leads]
+        return ECGDataset(
+            data=self.data[:, indices, :],
+            lead_names=leads,
+            fs=self.fs,
+            source_path=self.source_path,
+            label=self.label,
+            file_names=self.file_names,
+        )
+
+    def sample(self, n: int, seed: Optional[int] = 42) -> "ECGDataset":
+        """Return a new ECGDataset with n randomly selected recordings."""
+        rng = np.random.default_rng(seed)
+        idx = rng.choice(self.n_samples, size=min(n, self.n_samples), replace=False)
+        return ECGDataset(
+            data=self.data[idx],
+            lead_names=self.lead_names,
+            fs=self.fs,
+            source_path=self.source_path,
+            label=self.label,
+            file_names=[self.file_names[i] for i in idx] if self.file_names else [],
+        )
+
+    def __repr__(self) -> str:
+        return (
+            f"ECGDataset(label={self.label!r}, n_samples={self.n_samples}, "
+            f"leads={self.n_leads}, timepoints={self.n_timepoints}, fs={self.fs} Hz)"
+        )