deepecgkit 0.1.0__tar.gz → 0.2.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/PKG-INFO +9 -5
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/README.md +1 -1
- deepecgkit-0.2.0/deepecgkit/__init__.py +17 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/cli/__init__.py +2 -2
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/cli/commands/evaluate.py +20 -14
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/cli/commands/predict.py +19 -5
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/cli/commands/train.py +21 -73
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/cli/main.py +30 -3
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/datasets/modules.py +6 -5
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/evaluation/evaluator.py +7 -13
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/evaluation/visualization.py +23 -28
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/transformer_ecg.py +1 -4
- deepecgkit-0.2.0/deepecgkit/training/__init__.py +3 -0
- deepecgkit-0.2.0/deepecgkit/training/train.py +435 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit.egg-info/PKG-INFO +9 -5
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit.egg-info/requires.txt +8 -3
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/pyproject.toml +15 -4
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/tests/test_cli.py +71 -13
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/tests/test_deepecgkit.py +28 -21
- deepecgkit-0.2.0/tests/test_training.py +329 -0
- deepecgkit-0.1.0/deepecgkit/__init__.py +0 -42
- deepecgkit-0.1.0/deepecgkit/training/__init__.py +0 -9
- deepecgkit-0.1.0/deepecgkit/training/train.py +0 -171
- deepecgkit-0.1.0/tests/test_training.py +0 -457
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/LICENSE +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/__main__.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/cli/commands/__init__.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/cli/commands/info.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/cli/config.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/cli/evaluation.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/cli/logger.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/cli/registry.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/datasets/__init__.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/datasets/af_classification.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/datasets/base.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/datasets/ltafdb.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/datasets/mitbih_afdb.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/datasets/preprocessing.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/datasets/ptbxl.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/datasets/splitting.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/datasets/unified_af.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/evaluation/__init__.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/evaluation/metrics.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/__init__.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/af_classifier.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/convnext_v2_1d.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/crnn.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/deep_res_cnn.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/ecg_dualnet.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/fcn_wang.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/gru_model.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/inception_time.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/kanres_wide_x.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/kanres_x.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/lstm_model.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/mamba1d.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/medformer.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/resnet1d.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/resnet1d_wang.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/se_resnet1d.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/simple_cnn.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/tcn.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/xresnet1d.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/xresnet1d_benchmark.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/registry.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/utils/__init__.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/utils/download.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/utils/weights.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit.egg-info/SOURCES.txt +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit.egg-info/dependency_links.txt +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit.egg-info/entry_points.txt +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit.egg-info/top_level.txt +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/setup.cfg +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/tests/test_datasets.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/tests/test_evaluation.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/tests/test_models.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/tests/test_registry.py +0 -0
- {deepecgkit-0.1.0 → deepecgkit-0.2.0}/tests/test_utils.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: deepecgkit
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.2.0
|
|
4
4
|
Summary: A deep learning toolkit for ECG analysis
|
|
5
5
|
Author-email: Steven <steven@simula.no>
|
|
6
6
|
Classifier: Programming Language :: Python :: 3
|
|
@@ -19,20 +19,24 @@ Requires-Dist: wfdb>=3.4.0
|
|
|
19
19
|
Requires-Dist: pyEDFlib>=0.1.0
|
|
20
20
|
Requires-Dist: PyWavelets>=1.1.0
|
|
21
21
|
Requires-Dist: seaborn>=0.11.0
|
|
22
|
-
Requires-Dist:
|
|
23
|
-
Requires-Dist:
|
|
24
|
-
Requires-Dist: tensorboardx>=2.6.2.2
|
|
22
|
+
Requires-Dist: tqdm>=4.60.0
|
|
23
|
+
Requires-Dist: pyyaml>=6.0.3
|
|
25
24
|
Provides-Extra: dev
|
|
26
25
|
Requires-Dist: pytest-cov>=3.0.0; extra == "dev"
|
|
27
26
|
Requires-Dist: black>=22.0.0; extra == "dev"
|
|
28
27
|
Requires-Dist: isort>=5.0.0; extra == "dev"
|
|
29
28
|
Requires-Dist: mypy>=0.900; extra == "dev"
|
|
30
29
|
Requires-Dist: ruff>=0.11.13; extra == "dev"
|
|
30
|
+
Provides-Extra: docs
|
|
31
|
+
Requires-Dist: mkdocs>=1.6.0; extra == "docs"
|
|
32
|
+
Requires-Dist: mkdocs-material>=9.5.0; extra == "docs"
|
|
33
|
+
Requires-Dist: mkdocstrings[python]>=0.25.0; extra == "docs"
|
|
34
|
+
Requires-Dist: mkdocs-autorefs>=1.0.0; extra == "docs"
|
|
31
35
|
Dynamic: license-file
|
|
32
36
|
|
|
33
37
|
# DeepECG-Kit
|
|
34
38
|
|
|
35
|
-
A PyTorch
|
|
39
|
+
A pure-PyTorch deep learning library for ECG analysis and arrhythmia classification. Provides a CLI for training, evaluation, and inference across multiple model architectures and ECG datasets.
|
|
36
40
|
|
|
37
41
|
Install with `pip install deepecgkit`, or for local development clone the repo and run `pip install -e ".[dev]"`.
|
|
38
42
|
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
# DeepECG-Kit
|
|
2
2
|
|
|
3
|
-
A PyTorch
|
|
3
|
+
A pure-PyTorch deep learning library for ECG analysis and arrhythmia classification. Provides a CLI for training, evaluation, and inference across multiple model architectures and ECG datasets.
|
|
4
4
|
|
|
5
5
|
Install with `pip install deepecgkit`, or for local development clone the repo and run `pip install -e ".[dev]"`.
|
|
6
6
|
|
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
"""DeepECG-Kit: Deep learning library for ECG analysis and arrhythmia classification."""
|
|
2
|
+
|
|
3
|
+
from deepecgkit.datasets import ECGDataModule
|
|
4
|
+
from deepecgkit.models import KanResWideX
|
|
5
|
+
from deepecgkit.training import ECGTrainer
|
|
6
|
+
from deepecgkit.utils import read_csv
|
|
7
|
+
|
|
8
|
+
__version__ = "0.1.0"
|
|
9
|
+
__author__ = "DeepECG Kit Contributors"
|
|
10
|
+
__license__ = "MIT"
|
|
11
|
+
|
|
12
|
+
__all__ = [
|
|
13
|
+
"ECGDataModule",
|
|
14
|
+
"ECGTrainer",
|
|
15
|
+
"KanResWideX",
|
|
16
|
+
"read_csv",
|
|
17
|
+
]
|
|
@@ -36,7 +36,7 @@ from deepecgkit.cli.registry import (
|
|
|
36
36
|
_get_dataset_registry,
|
|
37
37
|
_get_model_registry,
|
|
38
38
|
)
|
|
39
|
-
from deepecgkit.training import
|
|
39
|
+
from deepecgkit.training import ECGTrainer
|
|
40
40
|
|
|
41
41
|
__all__ = [
|
|
42
42
|
"DATASET_INPUT_CHANNELS",
|
|
@@ -46,7 +46,7 @@ __all__ = [
|
|
|
46
46
|
"MODEL_NAMES",
|
|
47
47
|
"MODEL_WEIGHTS",
|
|
48
48
|
"CLILogger",
|
|
49
|
-
"
|
|
49
|
+
"ECGTrainer",
|
|
50
50
|
"_get_dataset_registry",
|
|
51
51
|
"_get_model_registry",
|
|
52
52
|
"evaluate",
|
|
@@ -3,17 +3,21 @@
|
|
|
3
3
|
from pathlib import Path
|
|
4
4
|
from typing import Any, Dict, Optional
|
|
5
5
|
|
|
6
|
-
import pytorch_lightning as pl
|
|
7
|
-
|
|
8
6
|
from deepecgkit.cli.logger import CLILogger
|
|
9
|
-
from deepecgkit.cli.registry import
|
|
7
|
+
from deepecgkit.cli.registry import (
|
|
8
|
+
DATASET_NAMES,
|
|
9
|
+
MODEL_NAMES,
|
|
10
|
+
_get_dataset_registry,
|
|
11
|
+
_get_model_registry,
|
|
12
|
+
)
|
|
10
13
|
from deepecgkit.datasets import ECGDataModule
|
|
11
|
-
from deepecgkit.training import
|
|
14
|
+
from deepecgkit.training import ECGTrainer
|
|
12
15
|
|
|
13
16
|
|
|
14
17
|
def evaluate(
|
|
15
18
|
checkpoint: str,
|
|
16
19
|
dataset_name: str,
|
|
20
|
+
model_name: str,
|
|
17
21
|
data_dir: Optional[str] = None,
|
|
18
22
|
batch_size: int = 32,
|
|
19
23
|
num_workers: int = 4,
|
|
@@ -32,15 +36,23 @@ def evaluate(
|
|
|
32
36
|
return 1
|
|
33
37
|
|
|
34
38
|
dataset_registry = _get_dataset_registry()
|
|
39
|
+
model_registry = _get_model_registry()
|
|
35
40
|
|
|
36
41
|
if dataset_name not in dataset_registry:
|
|
37
42
|
logger.error(f"Unknown dataset: {dataset_name}")
|
|
38
43
|
logger.info(f"Available datasets: {', '.join(DATASET_NAMES)}")
|
|
39
44
|
return 1
|
|
40
45
|
|
|
46
|
+
if model_name not in model_registry:
|
|
47
|
+
logger.error(f"Unknown model: {model_name}")
|
|
48
|
+
logger.info(f"Available models: {', '.join(MODEL_NAMES)}")
|
|
49
|
+
return 1
|
|
50
|
+
|
|
41
51
|
try:
|
|
42
52
|
logger.info(f"Loading model from {checkpoint}...")
|
|
43
|
-
|
|
53
|
+
model_class = model_registry[model_name]
|
|
54
|
+
model = model_class(input_channels=1, output_size=4)
|
|
55
|
+
trainer = ECGTrainer.load_checkpoint(checkpoint, model=model, device=accelerator)
|
|
44
56
|
|
|
45
57
|
dataset_class = dataset_registry[dataset_name]
|
|
46
58
|
dataset_kwargs: Dict[str, Any] = {"force_download": force_download}
|
|
@@ -56,17 +68,11 @@ def evaluate(
|
|
|
56
68
|
dataset_kwargs=dataset_kwargs,
|
|
57
69
|
)
|
|
58
70
|
|
|
59
|
-
trainer = pl.Trainer(
|
|
60
|
-
accelerator=accelerator,
|
|
61
|
-
devices=devices,
|
|
62
|
-
enable_progress_bar=not logger.quiet,
|
|
63
|
-
)
|
|
64
|
-
|
|
65
71
|
logger.info(f"Running evaluation on {split} split...")
|
|
66
72
|
if split == "test":
|
|
67
|
-
results = trainer.test(
|
|
73
|
+
results = trainer.test(data_module)
|
|
68
74
|
elif split == "val":
|
|
69
|
-
results = trainer.validate(
|
|
75
|
+
results = trainer.validate(data_module)
|
|
70
76
|
else:
|
|
71
77
|
logger.error(f"Invalid split: {split}. Use 'test' or 'val'.")
|
|
72
78
|
return 1
|
|
@@ -74,7 +80,7 @@ def evaluate(
|
|
|
74
80
|
if results:
|
|
75
81
|
logger.info("\nEvaluation Results:")
|
|
76
82
|
logger.info("-" * 40)
|
|
77
|
-
for key, value in results
|
|
83
|
+
for key, value in results.items():
|
|
78
84
|
logger.info(f" {key}: {value:.4f}")
|
|
79
85
|
|
|
80
86
|
return 0
|
|
@@ -9,8 +9,8 @@ import pandas as pd
|
|
|
9
9
|
import torch
|
|
10
10
|
import wfdb
|
|
11
11
|
|
|
12
|
-
import deepecgkit.cli as _cli
|
|
13
12
|
from deepecgkit.cli.logger import CLILogger
|
|
13
|
+
from deepecgkit.cli.registry import _get_model_registry
|
|
14
14
|
|
|
15
15
|
|
|
16
16
|
def predict(
|
|
@@ -19,6 +19,7 @@ def predict(
|
|
|
19
19
|
output_path: Optional[str] = None,
|
|
20
20
|
batch_size: int = 1,
|
|
21
21
|
accelerator: str = "auto",
|
|
22
|
+
model_name: Optional[str] = None,
|
|
22
23
|
logger: Optional[CLILogger] = None,
|
|
23
24
|
) -> int:
|
|
24
25
|
"""Run inference on ECG data."""
|
|
@@ -36,8 +37,21 @@ def predict(
|
|
|
36
37
|
|
|
37
38
|
try:
|
|
38
39
|
logger.info(f"Loading model from {checkpoint}...")
|
|
39
|
-
|
|
40
|
-
|
|
40
|
+
ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False)
|
|
41
|
+
|
|
42
|
+
if model_name is None:
|
|
43
|
+
logger.error("--model is required for predict command")
|
|
44
|
+
return 1
|
|
45
|
+
|
|
46
|
+
model_registry = _get_model_registry()
|
|
47
|
+
if model_name not in model_registry:
|
|
48
|
+
logger.error(f"Unknown model: {model_name}")
|
|
49
|
+
return 1
|
|
50
|
+
|
|
51
|
+
model_class = model_registry[model_name]
|
|
52
|
+
model = model_class(input_channels=1, output_size=4)
|
|
53
|
+
model.load_state_dict(ckpt["model_state_dict"])
|
|
54
|
+
model.eval()
|
|
41
55
|
|
|
42
56
|
logger.info(f"Loading input data from {input_path}...")
|
|
43
57
|
if input_file.suffix == ".npy":
|
|
@@ -61,11 +75,11 @@ def predict(
|
|
|
61
75
|
|
|
62
76
|
logger.info("Running inference...")
|
|
63
77
|
device = "cuda" if accelerator == "gpu" and torch.cuda.is_available() else "cpu"
|
|
64
|
-
|
|
78
|
+
model = model.to(device)
|
|
65
79
|
tensor = tensor.to(device)
|
|
66
80
|
|
|
67
81
|
with torch.no_grad():
|
|
68
|
-
outputs =
|
|
82
|
+
outputs = model(tensor)
|
|
69
83
|
if hasattr(outputs, "logits"):
|
|
70
84
|
outputs = outputs.logits
|
|
71
85
|
probabilities = torch.softmax(outputs, dim=-1)
|
|
@@ -5,10 +5,7 @@ from pathlib import Path
|
|
|
5
5
|
from typing import Any, Dict, Optional
|
|
6
6
|
|
|
7
7
|
import matplotlib
|
|
8
|
-
import pytorch_lightning as pl
|
|
9
8
|
import torch
|
|
10
|
-
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
|
|
11
|
-
from pytorch_lightning.loggers import CSVLogger
|
|
12
9
|
|
|
13
10
|
from deepecgkit.cli.evaluation import _save_multi_label_evaluation, _save_single_label_evaluation
|
|
14
11
|
from deepecgkit.cli.logger import CLILogger
|
|
@@ -22,7 +19,7 @@ from deepecgkit.cli.registry import (
|
|
|
22
19
|
)
|
|
23
20
|
from deepecgkit.datasets import ECGDataModule
|
|
24
21
|
from deepecgkit.evaluation.visualization import plot_training_curves
|
|
25
|
-
from deepecgkit.training import
|
|
22
|
+
from deepecgkit.training import ECGTrainer
|
|
26
23
|
from deepecgkit.utils.weights import load_pretrained_weights
|
|
27
24
|
|
|
28
25
|
|
|
@@ -65,7 +62,8 @@ def train(
|
|
|
65
62
|
logger.info(f"Available datasets: {', '.join(DATASET_NAMES)}")
|
|
66
63
|
return 1
|
|
67
64
|
|
|
68
|
-
|
|
65
|
+
ECGTrainer.seed_everything(seed)
|
|
66
|
+
torch.set_float32_matmul_precision("medium")
|
|
69
67
|
|
|
70
68
|
if output_dir is None:
|
|
71
69
|
timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
|
|
@@ -124,7 +122,6 @@ def train(
|
|
|
124
122
|
if binary_classification and dataset_name in af_datasets:
|
|
125
123
|
num_classes = 2
|
|
126
124
|
|
|
127
|
-
# BCE loss expects a single logit output for binary classification
|
|
128
125
|
model_output_size = (
|
|
129
126
|
1 if (binary_classification and dataset_name in af_datasets) else num_classes
|
|
130
127
|
)
|
|
@@ -177,41 +174,25 @@ def train(
|
|
|
177
174
|
"pos_weight": pos_weight,
|
|
178
175
|
}
|
|
179
176
|
|
|
180
|
-
|
|
177
|
+
trainer = ECGTrainer(model=model, train_config=train_config, device=accelerator)
|
|
181
178
|
|
|
182
|
-
|
|
183
|
-
|
|
184
|
-
filename=f"{model_name}-{dataset_name}-{{epoch:02d}}-{{val_loss:.4f}}",
|
|
185
|
-
monitor="val_loss",
|
|
186
|
-
mode="min",
|
|
187
|
-
save_top_k=3,
|
|
188
|
-
)
|
|
189
|
-
|
|
190
|
-
early_stopping = EarlyStopping(
|
|
191
|
-
monitor="val_loss",
|
|
192
|
-
patience=early_stopping_patience,
|
|
193
|
-
mode="min",
|
|
194
|
-
)
|
|
195
|
-
|
|
196
|
-
csv_logger = CSVLogger(save_dir=output_path, name="training_logs")
|
|
197
|
-
|
|
198
|
-
trainer = pl.Trainer(
|
|
199
|
-
max_epochs=epochs,
|
|
200
|
-
accelerator=accelerator,
|
|
201
|
-
devices=devices,
|
|
202
|
-
callbacks=[checkpoint_callback, early_stopping],
|
|
203
|
-
logger=csv_logger,
|
|
204
|
-
default_root_dir=output_path,
|
|
205
|
-
enable_progress_bar=not logger.quiet,
|
|
206
|
-
)
|
|
179
|
+
checkpoint_dir = output_path / "checkpoints"
|
|
180
|
+
log_dir = output_path / "training_logs"
|
|
207
181
|
|
|
208
182
|
try:
|
|
209
183
|
logger.info(f"Starting training for {epochs} epochs...")
|
|
210
|
-
trainer.fit(
|
|
184
|
+
trainer.fit(
|
|
185
|
+
data_module,
|
|
186
|
+
epochs=epochs,
|
|
187
|
+
early_stopping_patience=early_stopping_patience,
|
|
188
|
+
checkpoint_dir=str(checkpoint_dir),
|
|
189
|
+
log_dir=str(log_dir),
|
|
190
|
+
progress_bar=not logger.quiet,
|
|
191
|
+
)
|
|
211
192
|
|
|
212
|
-
if data_module.
|
|
193
|
+
if data_module.test_dataset is not None and len(data_module.test_dataset) > 0:
|
|
213
194
|
logger.info("Running test evaluation...")
|
|
214
|
-
trainer.test(
|
|
195
|
+
trainer.test(data_module)
|
|
215
196
|
|
|
216
197
|
logger.info(f"Training complete. Checkpoints saved to {output_path}")
|
|
217
198
|
|
|
@@ -222,7 +203,7 @@ def train(
|
|
|
222
203
|
if hasattr(data_module, "dataset") and hasattr(data_module.dataset, "class_names"):
|
|
223
204
|
class_names = data_module.dataset.class_names
|
|
224
205
|
|
|
225
|
-
metrics_file =
|
|
206
|
+
metrics_file = log_dir / "metrics.csv"
|
|
226
207
|
if metrics_file.exists():
|
|
227
208
|
plot_training_curves(
|
|
228
209
|
str(metrics_file),
|
|
@@ -231,7 +212,7 @@ def train(
|
|
|
231
212
|
logger.info(f"Loss plot saved to {output_path / 'loss.png'}")
|
|
232
213
|
logger.info(f"Accuracy plot saved to {output_path / 'accuracy.png'}")
|
|
233
214
|
|
|
234
|
-
y_pred, y_true, y_prob =
|
|
215
|
+
y_pred, y_true, y_prob = trainer.get_test_results()
|
|
235
216
|
|
|
236
217
|
if y_pred is not None:
|
|
237
218
|
is_multi_label = y_true.ndim > 1
|
|
@@ -273,6 +254,7 @@ def train(
|
|
|
273
254
|
|
|
274
255
|
def resume(
|
|
275
256
|
checkpoint: str,
|
|
257
|
+
model_name: str,
|
|
276
258
|
epochs: Optional[int] = None,
|
|
277
259
|
output_dir: Optional[str] = None,
|
|
278
260
|
accelerator: str = "auto",
|
|
@@ -290,42 +272,8 @@ def resume(
|
|
|
290
272
|
|
|
291
273
|
try:
|
|
292
274
|
logger.info(f"Loading checkpoint from {checkpoint}...")
|
|
293
|
-
|
|
294
|
-
|
|
295
|
-
output_path = Path(output_dir) if output_dir else checkpoint_path.parent
|
|
296
|
-
output_path.mkdir(parents=True, exist_ok=True)
|
|
297
|
-
|
|
298
|
-
checkpoint_callback = ModelCheckpoint(
|
|
299
|
-
dirpath=output_path,
|
|
300
|
-
filename="resumed-{epoch:02d}-{val_loss:.4f}",
|
|
301
|
-
monitor="val_loss",
|
|
302
|
-
mode="min",
|
|
303
|
-
save_top_k=3,
|
|
304
|
-
)
|
|
305
|
-
|
|
306
|
-
early_stopping = EarlyStopping(
|
|
307
|
-
monitor="val_loss",
|
|
308
|
-
patience=early_stopping_patience,
|
|
309
|
-
mode="min",
|
|
310
|
-
)
|
|
311
|
-
|
|
312
|
-
trainer_kwargs: Dict[str, Any] = {
|
|
313
|
-
"accelerator": accelerator,
|
|
314
|
-
"devices": devices,
|
|
315
|
-
"callbacks": [checkpoint_callback, early_stopping],
|
|
316
|
-
"default_root_dir": output_path,
|
|
317
|
-
"enable_progress_bar": not logger.quiet,
|
|
318
|
-
}
|
|
319
|
-
if epochs:
|
|
320
|
-
trainer_kwargs["max_epochs"] = epochs
|
|
321
|
-
|
|
322
|
-
trainer = pl.Trainer(**trainer_kwargs)
|
|
323
|
-
|
|
324
|
-
logger.info("Resuming training...")
|
|
325
|
-
trainer.fit(lit_model, ckpt_path=checkpoint)
|
|
326
|
-
|
|
327
|
-
logger.info(f"Training complete. Checkpoints saved to {output_path}")
|
|
328
|
-
return 0
|
|
275
|
+
logger.warning("Resume requires a data_module. Please use the train command instead.")
|
|
276
|
+
return 1
|
|
329
277
|
except Exception as e:
|
|
330
278
|
logger.error(f"Failed to resume training: {e}")
|
|
331
279
|
return 1
|
|
@@ -20,9 +20,9 @@ def main(argv: Optional[list] = None) -> int:
|
|
|
20
20
|
Examples:
|
|
21
21
|
deepecg train -m kanres -d af-classification
|
|
22
22
|
deepecg train -m kanres -d af-classification --force-download
|
|
23
|
-
deepecg evaluate --checkpoint model.ckpt -d af-classification
|
|
24
|
-
deepecg predict --checkpoint model.ckpt --input ecg.npy
|
|
25
|
-
deepecg resume --checkpoint model.ckpt --epochs 100
|
|
23
|
+
deepecg evaluate --checkpoint model.ckpt -m kanres -d af-classification
|
|
24
|
+
deepecg predict --checkpoint model.ckpt -m kanres --input ecg.npy
|
|
25
|
+
deepecg resume --checkpoint model.ckpt -m kanres --epochs 100
|
|
26
26
|
deepecg list-models
|
|
27
27
|
deepecg info -m kanres
|
|
28
28
|
|
|
@@ -190,6 +190,14 @@ Note: Datasets are automatically downloaded when first used.
|
|
|
190
190
|
required=True,
|
|
191
191
|
help="Path to model checkpoint",
|
|
192
192
|
)
|
|
193
|
+
eval_parser.add_argument(
|
|
194
|
+
"--model",
|
|
195
|
+
"-m",
|
|
196
|
+
type=str,
|
|
197
|
+
required=True,
|
|
198
|
+
choices=MODEL_NAMES,
|
|
199
|
+
help="Model architecture to load",
|
|
200
|
+
)
|
|
193
201
|
eval_parser.add_argument(
|
|
194
202
|
"--dataset",
|
|
195
203
|
"-d",
|
|
@@ -250,6 +258,14 @@ Note: Datasets are automatically downloaded when first used.
|
|
|
250
258
|
required=True,
|
|
251
259
|
help="Path to model checkpoint",
|
|
252
260
|
)
|
|
261
|
+
predict_parser.add_argument(
|
|
262
|
+
"--model",
|
|
263
|
+
"-m",
|
|
264
|
+
type=str,
|
|
265
|
+
required=True,
|
|
266
|
+
choices=MODEL_NAMES,
|
|
267
|
+
help="Model architecture to load",
|
|
268
|
+
)
|
|
253
269
|
predict_parser.add_argument(
|
|
254
270
|
"--input",
|
|
255
271
|
"-i",
|
|
@@ -286,6 +302,14 @@ Note: Datasets are automatically downloaded when first used.
|
|
|
286
302
|
required=True,
|
|
287
303
|
help="Path to checkpoint to resume from",
|
|
288
304
|
)
|
|
305
|
+
resume_parser.add_argument(
|
|
306
|
+
"--model",
|
|
307
|
+
"-m",
|
|
308
|
+
type=str,
|
|
309
|
+
required=True,
|
|
310
|
+
choices=MODEL_NAMES,
|
|
311
|
+
help="Model architecture to load",
|
|
312
|
+
)
|
|
289
313
|
resume_parser.add_argument(
|
|
290
314
|
"--epochs",
|
|
291
315
|
"-e",
|
|
@@ -384,6 +408,7 @@ Note: Datasets are automatically downloaded when first used.
|
|
|
384
408
|
return _cli.evaluate(
|
|
385
409
|
checkpoint=args.checkpoint,
|
|
386
410
|
dataset_name=args.dataset,
|
|
411
|
+
model_name=args.model,
|
|
387
412
|
data_dir=args.data_dir,
|
|
388
413
|
batch_size=args.batch_size,
|
|
389
414
|
num_workers=args.num_workers,
|
|
@@ -401,12 +426,14 @@ Note: Datasets are automatically downloaded when first used.
|
|
|
401
426
|
output_path=args.output,
|
|
402
427
|
batch_size=args.batch_size,
|
|
403
428
|
accelerator=args.accelerator,
|
|
429
|
+
model_name=args.model,
|
|
404
430
|
logger=logger,
|
|
405
431
|
)
|
|
406
432
|
|
|
407
433
|
elif args.command == "resume":
|
|
408
434
|
return _cli.resume(
|
|
409
435
|
checkpoint=args.checkpoint,
|
|
436
|
+
model_name=args.model,
|
|
410
437
|
epochs=args.epochs,
|
|
411
438
|
output_dir=args.output_dir,
|
|
412
439
|
accelerator=args.accelerator,
|
|
@@ -1,11 +1,10 @@
|
|
|
1
1
|
"""
|
|
2
|
-
|
|
2
|
+
Data module for ECG datasets.
|
|
3
3
|
"""
|
|
4
4
|
|
|
5
5
|
from typing import Dict, Optional, Type, Union
|
|
6
6
|
|
|
7
7
|
import numpy as np
|
|
8
|
-
import pytorch_lightning as pl
|
|
9
8
|
import torch
|
|
10
9
|
from torch.utils.data import DataLoader, Dataset
|
|
11
10
|
|
|
@@ -13,8 +12,11 @@ from deepecgkit.datasets.base import BaseECGDataset
|
|
|
13
12
|
from deepecgkit.datasets.splitting import DataSplitter
|
|
14
13
|
|
|
15
14
|
|
|
16
|
-
class ECGDataModule
|
|
17
|
-
"""
|
|
15
|
+
class ECGDataModule:
|
|
16
|
+
"""Data module for ECG datasets.
|
|
17
|
+
|
|
18
|
+
Handles dataset creation, train/val/test splitting, and DataLoader construction.
|
|
19
|
+
"""
|
|
18
20
|
|
|
19
21
|
def __init__(
|
|
20
22
|
self,
|
|
@@ -57,7 +59,6 @@ class ECGDataModule(pl.LightningDataModule):
|
|
|
57
59
|
stratify: Whether to use stratified splitting based on labels
|
|
58
60
|
dataset_kwargs: Additional keyword arguments to pass to the dataset class
|
|
59
61
|
"""
|
|
60
|
-
super().__init__()
|
|
61
62
|
self.dataset = dataset
|
|
62
63
|
self.dataset_class = dataset_class
|
|
63
64
|
self.data_dir = data_dir
|
|
@@ -16,9 +16,6 @@ from deepecgkit.evaluation.metrics import (
|
|
|
16
16
|
)
|
|
17
17
|
|
|
18
18
|
MAX_CLASSIFICATION_CLASSES = 20
|
|
19
|
-
TUPLE_LENGTH_2 = 2
|
|
20
|
-
TUPLE_LENGTH_3 = 3
|
|
21
|
-
MAX_BRANCHES = 12
|
|
22
19
|
|
|
23
20
|
|
|
24
21
|
class ECGEvaluator:
|
|
@@ -72,10 +69,7 @@ class ECGEvaluator:
|
|
|
72
69
|
Dictionary of metric values
|
|
73
70
|
"""
|
|
74
71
|
|
|
75
|
-
|
|
76
|
-
self.y_scores = y_scores
|
|
77
|
-
else:
|
|
78
|
-
self.y_scores = None
|
|
72
|
+
self.y_scores = y_scores
|
|
79
73
|
|
|
80
74
|
predictions, targets = self._process_input_data(model, test_data)
|
|
81
75
|
metrics = self._calculate_metrics(predictions, targets)
|
|
@@ -104,10 +98,10 @@ class ECGEvaluator:
|
|
|
104
98
|
if not isinstance(test_data, tuple):
|
|
105
99
|
return predictions, test_data
|
|
106
100
|
|
|
107
|
-
if len(test_data) ==
|
|
101
|
+
if len(test_data) == 3:
|
|
108
102
|
targets = test_data[1]
|
|
109
103
|
self.y_scores = test_data[2]
|
|
110
|
-
elif len(test_data) ==
|
|
104
|
+
elif len(test_data) == 2:
|
|
111
105
|
targets = test_data[1]
|
|
112
106
|
else:
|
|
113
107
|
targets = test_data
|
|
@@ -146,7 +140,7 @@ class ECGEvaluator:
|
|
|
146
140
|
|
|
147
141
|
if isinstance(test_data, torch.utils.data.DataLoader):
|
|
148
142
|
return self._process_dataloader(model, test_data)
|
|
149
|
-
elif isinstance(test_data, tuple) and len(test_data) ==
|
|
143
|
+
elif isinstance(test_data, tuple) and len(test_data) == 2:
|
|
150
144
|
return self._process_tuple_data(model, test_data)
|
|
151
145
|
elif isinstance(test_data, np.ndarray):
|
|
152
146
|
return self._process_numpy_data(model, test_data)
|
|
@@ -155,11 +149,11 @@ class ECGEvaluator:
|
|
|
155
149
|
|
|
156
150
|
def _handle_tuple_test_data(self, test_data: tuple) -> tuple:
|
|
157
151
|
"""Handle test data when it's a tuple."""
|
|
158
|
-
if len(test_data) ==
|
|
152
|
+
if len(test_data) == 3:
|
|
159
153
|
y_true, y_pred, y_scores = test_data
|
|
160
154
|
self.y_scores = y_scores
|
|
161
155
|
return y_pred, y_true
|
|
162
|
-
elif len(test_data) ==
|
|
156
|
+
elif len(test_data) == 2:
|
|
163
157
|
return test_data[0], test_data[1]
|
|
164
158
|
return test_data, None
|
|
165
159
|
|
|
@@ -171,7 +165,7 @@ class ECGEvaluator:
|
|
|
171
165
|
targets = []
|
|
172
166
|
with torch.no_grad():
|
|
173
167
|
for batch in test_data:
|
|
174
|
-
if isinstance(batch, (list, tuple)) and len(batch) ==
|
|
168
|
+
if isinstance(batch, (list, tuple)) and len(batch) == 2:
|
|
175
169
|
x, y = batch
|
|
176
170
|
else:
|
|
177
171
|
x, y = batch, None
|
|
@@ -2,7 +2,6 @@
|
|
|
2
2
|
Visualization functions for ECG evaluation.
|
|
3
3
|
"""
|
|
4
4
|
|
|
5
|
-
import warnings
|
|
6
5
|
from pathlib import Path
|
|
7
6
|
from typing import List, Optional
|
|
8
7
|
|
|
@@ -58,34 +57,30 @@ def plot_roc_curve(
|
|
|
58
57
|
save_path: Optional[str] = None,
|
|
59
58
|
):
|
|
60
59
|
"""Plot ROC curve."""
|
|
61
|
-
|
|
62
|
-
|
|
63
|
-
|
|
64
|
-
|
|
65
|
-
|
|
66
|
-
|
|
67
|
-
|
|
68
|
-
|
|
69
|
-
|
|
70
|
-
|
|
71
|
-
|
|
72
|
-
|
|
73
|
-
|
|
74
|
-
|
|
75
|
-
|
|
76
|
-
|
|
77
|
-
|
|
78
|
-
|
|
79
|
-
plt.legend(loc="lower right")
|
|
80
|
-
|
|
81
|
-
if save_path:
|
|
82
|
-
plt.savefig(save_path, dpi=300, bbox_inches="tight")
|
|
83
|
-
plt.close()
|
|
84
|
-
else:
|
|
85
|
-
plt.show()
|
|
60
|
+
fpr, tpr, _ = roc_curve(y_true, y_scores)
|
|
61
|
+
roc_auc = auc(fpr, tpr)
|
|
62
|
+
|
|
63
|
+
plt.figure(figsize=(8, 6))
|
|
64
|
+
plt.plot(
|
|
65
|
+
fpr,
|
|
66
|
+
tpr,
|
|
67
|
+
color="darkorange",
|
|
68
|
+
lw=2,
|
|
69
|
+
label=f"ROC curve (area = {roc_auc:.2f})",
|
|
70
|
+
)
|
|
71
|
+
plt.plot([0, 1], [0, 1], color="navy", lw=2, linestyle="--")
|
|
72
|
+
plt.xlim([0.0, 1.0])
|
|
73
|
+
plt.ylim([0.0, 1.05])
|
|
74
|
+
plt.xlabel("False Positive Rate")
|
|
75
|
+
plt.ylabel("True Positive Rate")
|
|
76
|
+
plt.title(title)
|
|
77
|
+
plt.legend(loc="lower right")
|
|
86
78
|
|
|
87
|
-
|
|
88
|
-
|
|
79
|
+
if save_path:
|
|
80
|
+
plt.savefig(save_path, dpi=300, bbox_inches="tight")
|
|
81
|
+
plt.close()
|
|
82
|
+
else:
|
|
83
|
+
plt.show()
|
|
89
84
|
|
|
90
85
|
|
|
91
86
|
def plot_predictions(
|
|
@@ -109,7 +109,7 @@ class TransformerECG(nn.Module):
|
|
|
109
109
|
def feature_dim(self) -> int:
|
|
110
110
|
return self._feature_dim
|
|
111
111
|
|
|
112
|
-
def
|
|
112
|
+
def extract_features(self, x: torch.Tensor) -> torch.Tensor:
|
|
113
113
|
x = x.transpose(1, 2)
|
|
114
114
|
x = self.input_projection(x)
|
|
115
115
|
x = self.positional_encoding(x)
|
|
@@ -119,9 +119,6 @@ class TransformerECG(nn.Module):
|
|
|
119
119
|
x = x.squeeze(-1)
|
|
120
120
|
return x
|
|
121
121
|
|
|
122
|
-
def extract_features(self, x: torch.Tensor) -> torch.Tensor:
|
|
123
|
-
return self._extract_encoder_features(x)
|
|
124
|
-
|
|
125
122
|
def forward(self, x):
|
|
126
123
|
x = self.extract_features(x)
|
|
127
124
|
return self.classifier(x)
|