deepecgkit 0.1.0__py3-none-any.whl

deepecgkit/__init__.py

@@ -0,0 +1,42 @@
+"""
+DeepECG-Kit: A comprehensive Python package for ECG analysis using deep learning.
+
+This package provides tools for:
+- Loading and preprocessing ECG data
+- Training deep learning models on ECG signals
+- Evaluating model performance
+- Traditional ECG analysis methods
+- Visualization and reporting
+
+Basic usage:
+    import deepecgkit as ecg
+
+
+    data = ecg.load_ecg_data("path/to/data")
+
+
+    preprocessed = ecg.preprocess(data)
+
+
+    model = ecg.models.KanResWideX(input_channels=12, output_size=1)
+    trainer = ecg.train_model(model, preprocessed)
+
+
+    results = ecg.evaluate_model(model, test_data)
+"""
+
+from deepecgkit.datasets import ECGDataModule
+from deepecgkit.models import KanResWideX
+from deepecgkit.training import ECGLitModel
+from deepecgkit.utils import read_csv
+
+__version__ = "0.1.0"
+__author__ = "DeepECG Kit Contributors"
+__license__ = "MIT"
+
+__all__ = [
+    "ECGDataModule",
+    "ECGLitModel",
+    "KanResWideX",
+    "read_csv",
+]

deepecgkit/__main__.py

@@ -0,0 +1,6 @@
+"""Entry point for python -m deepecgkit."""
+
+from deepecgkit.cli import main
+
+if __name__ == "__main__":
+    raise SystemExit(main())

deepecgkit/cli/__init__.py

@@ -0,0 +1,61 @@
+"""
+Command-line interface for deepecg-kit.
+
+Usage:
+    deepecg train --model <model> --dataset <dataset> [options]
+    deepecg evaluate --checkpoint <path> --dataset <dataset> [options]
+    deepecg predict --checkpoint <path> --input <file> [options]
+    deepecg resume --checkpoint <path> [options]
+    deepecg list-models
+    deepecg list-datasets
+    deepecg info --model <model>
+
+Datasets are automatically downloaded when first used if not already present.
+Use --force-download to re-download an existing dataset.
+"""
+
+from deepecgkit.cli.commands import (
+    evaluate,
+    list_datasets,
+    list_models,
+    predict,
+    resume,
+    show_info,
+    train,
+)
+from deepecgkit.cli.config import load_config
+from deepecgkit.cli.logger import CLILogger
+from deepecgkit.cli.main import main
+from deepecgkit.cli.registry import (
+    DATASET_INPUT_CHANNELS,
+    DATASET_NAMES,
+    DATASET_NUM_CLASSES,
+    MODEL_DESCRIPTIONS,
+    MODEL_NAMES,
+    MODEL_WEIGHTS,
+    _get_dataset_registry,
+    _get_model_registry,
+)
+from deepecgkit.training import ECGLitModel
+
+__all__ = [
+    "DATASET_INPUT_CHANNELS",
+    "DATASET_NAMES",
+    "DATASET_NUM_CLASSES",
+    "MODEL_DESCRIPTIONS",
+    "MODEL_NAMES",
+    "MODEL_WEIGHTS",
+    "CLILogger",
+    "ECGLitModel",
+    "_get_dataset_registry",
+    "_get_model_registry",
+    "evaluate",
+    "list_datasets",
+    "list_models",
+    "load_config",
+    "main",
+    "predict",
+    "resume",
+    "show_info",
+    "train",
+]

deepecgkit/cli/commands/__init__.py

@@ -0,0 +1,16 @@
+"""CLI command functions."""
+
+from deepecgkit.cli.commands.evaluate import evaluate
+from deepecgkit.cli.commands.info import list_datasets, list_models, show_info
+from deepecgkit.cli.commands.predict import predict
+from deepecgkit.cli.commands.train import resume, train
+
+__all__ = [
+    "evaluate",
+    "list_datasets",
+    "list_models",
+    "predict",
+    "resume",
+    "show_info",
+    "train",
+]

deepecgkit/cli/commands/evaluate.py

@@ -0,0 +1,83 @@
+"""Evaluate command."""
+
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+import pytorch_lightning as pl
+
+from deepecgkit.cli.logger import CLILogger
+from deepecgkit.cli.registry import DATASET_NAMES, _get_dataset_registry
+from deepecgkit.datasets import ECGDataModule
+from deepecgkit.training import ECGLitModel
+
+
+def evaluate(
+    checkpoint: str,
+    dataset_name: str,
+    data_dir: Optional[str] = None,
+    batch_size: int = 32,
+    num_workers: int = 4,
+    accelerator: str = "auto",
+    devices: int = 1,
+    force_download: bool = False,
+    split: str = "test",
+    logger: Optional[CLILogger] = None,
+) -> int:
+    """Evaluate a trained model on a dataset."""
+    logger = logger or CLILogger()
+
+    checkpoint_path = Path(checkpoint)
+    if not checkpoint_path.exists():
+        logger.error(f"Checkpoint not found: {checkpoint}")
+        return 1
+
+    dataset_registry = _get_dataset_registry()
+
+    if dataset_name not in dataset_registry:
+        logger.error(f"Unknown dataset: {dataset_name}")
+        logger.info(f"Available datasets: {', '.join(DATASET_NAMES)}")
+        return 1
+
+    try:
+        logger.info(f"Loading model from {checkpoint}...")
+        lit_model = ECGLitModel.load_from_checkpoint(checkpoint)
+
+        dataset_class = dataset_registry[dataset_name]
+        dataset_kwargs: Dict[str, Any] = {"force_download": force_download}
+        if data_dir:
+            dataset_kwargs["data_dir"] = data_dir
+
+        logger.info(f"Creating data module with {dataset_name} dataset...")
+        data_module = ECGDataModule(
+            dataset_class=dataset_class,
+            batch_size=batch_size,
+            num_workers=num_workers,
+            download=True,
+            dataset_kwargs=dataset_kwargs,
+        )
+
+        trainer = pl.Trainer(
+            accelerator=accelerator,
+            devices=devices,
+            enable_progress_bar=not logger.quiet,
+        )
+
+        logger.info(f"Running evaluation on {split} split...")
+        if split == "test":
+            results = trainer.test(lit_model, data_module)
+        elif split == "val":
+            results = trainer.validate(lit_model, data_module)
+        else:
+            logger.error(f"Invalid split: {split}. Use 'test' or 'val'.")
+            return 1
+
+        if results:
+            logger.info("\nEvaluation Results:")
+            logger.info("-" * 40)
+            for key, value in results[0].items():
+                logger.info(f"  {key}: {value:.4f}")
+
+        return 0
+    except Exception as e:
+        logger.error(f"Evaluation failed: {e}")
+        return 1

deepecgkit/cli/commands/info.py

@@ -0,0 +1,73 @@
+"""Info and listing commands."""
+
+from typing import Optional
+
+from deepecgkit.cli.logger import CLILogger
+from deepecgkit.cli.registry import (
+    DATASET_NAMES,
+    DATASET_NUM_CLASSES,
+    MODEL_DESCRIPTIONS,
+    MODEL_NAMES,
+    MODEL_WEIGHTS,
+    _get_model_registry,
+)
+
+
+def list_models() -> None:
+    """List all available models."""
+    print("Available models:")
+    print("-" * 60)
+    for name in sorted(MODEL_NAMES):
+        desc = MODEL_DESCRIPTIONS.get(name, "")
+        weights = MODEL_WEIGHTS.get(name, [])
+        print(f"  {name}")
+        if desc:
+            print(f"    Description: {desc}")
+        if weights:
+            print(f"    Pretrained weights: {', '.join(weights)}")
+        print()
+
+
+def list_datasets() -> None:
+    """List all available datasets."""
+    print("Available datasets:")
+    print("-" * 60)
+    for name in sorted(DATASET_NAMES):
+        num_classes = DATASET_NUM_CLASSES[name]
+        print(f"  {name}")
+        print(f"    Classes: {num_classes}")
+        print()
+
+
+def show_info(model_name: str, logger: Optional[CLILogger] = None) -> int:
+    """Show detailed information about a model."""
+    logger = logger or CLILogger()
+
+    if model_name not in MODEL_NAMES:
+        logger.error(f"Unknown model: {model_name}")
+        logger.info(f"Available models: {', '.join(MODEL_NAMES)}")
+        return 1
+
+    try:
+        model_registry = _get_model_registry()
+        model_class = model_registry[model_name]
+        model = model_class(input_channels=1, output_size=4)
+
+        total_params = sum(p.numel() for p in model.parameters())
+        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+        print(f"Model: {model_name}")
+        print("=" * 60)
+        print(f"Description: {MODEL_DESCRIPTIONS.get(model_name, 'N/A')}")
+        print(f"Total parameters: {total_params:,}")
+        print(f"Trainable parameters: {trainable_params:,}")
+        print(f"Pretrained weights: {', '.join(MODEL_WEIGHTS.get(model_name, ['None']))}")
+        print()
+        print("Architecture:")
+        print("-" * 60)
+        print(model)
+
+        return 0
+    except Exception as e:
+        logger.error(f"Failed to load model info: {e}")
+        return 1

deepecgkit/cli/commands/predict.py

@@ -0,0 +1,94 @@
+"""Predict command."""
+
+import json
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+import pandas as pd
+import torch
+import wfdb
+
+import deepecgkit.cli as _cli
+from deepecgkit.cli.logger import CLILogger
+
+
+def predict(
+    checkpoint: str,
+    input_path: str,
+    output_path: Optional[str] = None,
+    batch_size: int = 1,
+    accelerator: str = "auto",
+    logger: Optional[CLILogger] = None,
+) -> int:
+    """Run inference on ECG data."""
+    logger = logger or CLILogger()
+
+    checkpoint_file = Path(checkpoint)
+    if not checkpoint_file.exists():
+        logger.error(f"Checkpoint not found: {checkpoint}")
+        return 1
+
+    input_file = Path(input_path)
+    if not input_file.exists():
+        logger.error(f"Input file not found: {input_path}")
+        return 1
+
+    try:
+        logger.info(f"Loading model from {checkpoint}...")
+        lit_model = _cli.ECGLitModel.load_from_checkpoint(checkpoint)
+        lit_model.eval()
+
+        logger.info(f"Loading input data from {input_path}...")
+        if input_file.suffix == ".npy":
+            data = np.load(input_path)
+        elif input_file.suffix == ".csv":
+            data = pd.read_csv(input_path).values
+        elif input_file.suffix in (".dat", ".hea"):
+            record = wfdb.rdrecord(str(input_file.with_suffix("")))
+            data = record.p_signal
+        else:
+            logger.error(f"Unsupported input format: {input_file.suffix}")
+            logger.info("Supported formats: .npy, .csv, .dat/.hea (WFDB)")
+            return 1
+
+        if data.ndim == 1:
+            data = data.reshape(1, 1, -1)
+        elif data.ndim == 2:
+            data = data.reshape(data.shape[0], 1, -1)
+
+        tensor = torch.tensor(data, dtype=torch.float32)
+
+        logger.info("Running inference...")
+        device = "cuda" if accelerator == "gpu" and torch.cuda.is_available() else "cpu"
+        lit_model = lit_model.to(device)
+        tensor = tensor.to(device)
+
+        with torch.no_grad():
+            outputs = lit_model(tensor)
+            if hasattr(outputs, "logits"):
+                outputs = outputs.logits
+            probabilities = torch.softmax(outputs, dim=-1)
+            predictions = torch.argmax(probabilities, dim=-1)
+
+        results = {
+            "predictions": predictions.cpu().numpy().tolist(),
+            "probabilities": probabilities.cpu().numpy().tolist(),
+        }
+
+        if output_path:
+            output_file = Path(output_path)
+            output_file.parent.mkdir(parents=True, exist_ok=True)
+            with open(output_file, "w") as f:
+                json.dump(results, f, indent=2)
+            logger.info(f"Results saved to {output_path}")
+        else:
+            logger.info("\nPrediction Results:")
+            logger.info("-" * 40)
+            for i, (pred, prob) in enumerate(zip(results["predictions"], results["probabilities"])):
+                logger.info(f"  Sample {i}: class={pred}, confidence={max(prob):.4f}")
+
+        return 0
+    except Exception as e:
+        logger.error(f"Prediction failed: {e}")
+        return 1