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.
Files changed (78) hide show
  1. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/PKG-INFO +9 -5
  2. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/README.md +1 -1
  3. deepecgkit-0.2.0/deepecgkit/__init__.py +17 -0
  4. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/cli/__init__.py +2 -2
  5. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/cli/commands/evaluate.py +20 -14
  6. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/cli/commands/predict.py +19 -5
  7. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/cli/commands/train.py +21 -73
  8. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/cli/main.py +30 -3
  9. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/datasets/modules.py +6 -5
  10. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/evaluation/evaluator.py +7 -13
  11. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/evaluation/visualization.py +23 -28
  12. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/transformer_ecg.py +1 -4
  13. deepecgkit-0.2.0/deepecgkit/training/__init__.py +3 -0
  14. deepecgkit-0.2.0/deepecgkit/training/train.py +435 -0
  15. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit.egg-info/PKG-INFO +9 -5
  16. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit.egg-info/requires.txt +8 -3
  17. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/pyproject.toml +15 -4
  18. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/tests/test_cli.py +71 -13
  19. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/tests/test_deepecgkit.py +28 -21
  20. deepecgkit-0.2.0/tests/test_training.py +329 -0
  21. deepecgkit-0.1.0/deepecgkit/__init__.py +0 -42
  22. deepecgkit-0.1.0/deepecgkit/training/__init__.py +0 -9
  23. deepecgkit-0.1.0/deepecgkit/training/train.py +0 -171
  24. deepecgkit-0.1.0/tests/test_training.py +0 -457
  25. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/LICENSE +0 -0
  26. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/__main__.py +0 -0
  27. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/cli/commands/__init__.py +0 -0
  28. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/cli/commands/info.py +0 -0
  29. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/cli/config.py +0 -0
  30. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/cli/evaluation.py +0 -0
  31. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/cli/logger.py +0 -0
  32. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/cli/registry.py +0 -0
  33. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/datasets/__init__.py +0 -0
  34. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/datasets/af_classification.py +0 -0
  35. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/datasets/base.py +0 -0
  36. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/datasets/ltafdb.py +0 -0
  37. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/datasets/mitbih_afdb.py +0 -0
  38. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/datasets/preprocessing.py +0 -0
  39. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/datasets/ptbxl.py +0 -0
  40. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/datasets/splitting.py +0 -0
  41. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/datasets/unified_af.py +0 -0
  42. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/evaluation/__init__.py +0 -0
  43. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/evaluation/metrics.py +0 -0
  44. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/__init__.py +0 -0
  45. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/af_classifier.py +0 -0
  46. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/convnext_v2_1d.py +0 -0
  47. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/crnn.py +0 -0
  48. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/deep_res_cnn.py +0 -0
  49. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/ecg_dualnet.py +0 -0
  50. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/fcn_wang.py +0 -0
  51. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/gru_model.py +0 -0
  52. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/inception_time.py +0 -0
  53. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/kanres_wide_x.py +0 -0
  54. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/kanres_x.py +0 -0
  55. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/lstm_model.py +0 -0
  56. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/mamba1d.py +0 -0
  57. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/medformer.py +0 -0
  58. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/resnet1d.py +0 -0
  59. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/resnet1d_wang.py +0 -0
  60. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/se_resnet1d.py +0 -0
  61. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/simple_cnn.py +0 -0
  62. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/tcn.py +0 -0
  63. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/xresnet1d.py +0 -0
  64. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/models/xresnet1d_benchmark.py +0 -0
  65. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/registry.py +0 -0
  66. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/utils/__init__.py +0 -0
  67. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/utils/download.py +0 -0
  68. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit/utils/weights.py +0 -0
  69. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit.egg-info/SOURCES.txt +0 -0
  70. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit.egg-info/dependency_links.txt +0 -0
  71. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit.egg-info/entry_points.txt +0 -0
  72. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/deepecgkit.egg-info/top_level.txt +0 -0
  73. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/setup.cfg +0 -0
  74. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/tests/test_datasets.py +0 -0
  75. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/tests/test_evaluation.py +0 -0
  76. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/tests/test_models.py +0 -0
  77. {deepecgkit-0.1.0 → deepecgkit-0.2.0}/tests/test_registry.py +0 -0
  78. {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.1.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: pytorch-lightning>=1.5.0
23
- Requires-Dist: litmodels>=0.1.8
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 Lightning-based deep learning library for ECG analysis and arrhythmia classification. Provides a CLI for training, evaluation, and inference across multiple model architectures and ECG datasets.
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 Lightning-based deep learning library for ECG analysis and arrhythmia classification. Provides a CLI for training, evaluation, and inference across multiple model architectures and ECG datasets.
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 ECGLitModel
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
- "ECGLitModel",
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 DATASET_NAMES, _get_dataset_registry
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 ECGLitModel
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
- lit_model = ECGLitModel.load_from_checkpoint(checkpoint)
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(lit_model, data_module)
73
+ results = trainer.test(data_module)
68
74
  elif split == "val":
69
- results = trainer.validate(lit_model, data_module)
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[0].items():
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
- lit_model = _cli.ECGLitModel.load_from_checkpoint(checkpoint)
40
- lit_model.eval()
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
- lit_model = lit_model.to(device)
78
+ model = model.to(device)
65
79
  tensor = tensor.to(device)
66
80
 
67
81
  with torch.no_grad():
68
- outputs = lit_model(tensor)
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 ECGLitModel
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
- pl.seed_everything(seed)
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
- lit_model = ECGLitModel(model=model, train_config=train_config)
177
+ trainer = ECGTrainer(model=model, train_config=train_config, device=accelerator)
181
178
 
182
- checkpoint_callback = ModelCheckpoint(
183
- dirpath=output_path / "checkpoints",
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(lit_model, data_module)
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.test_dataloader() is not None:
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(lit_model, data_module)
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 = Path(csv_logger.log_dir) / "metrics.csv"
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 = lit_model.get_test_results()
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
- lit_model = ECGLitModel.load_from_checkpoint(checkpoint)
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
- PyTorch Lightning data module for ECG datasets.
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(pl.LightningDataModule):
17
- """PyTorch Lightning data module for ECG datasets."""
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
- if y_scores is not None:
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) == TUPLE_LENGTH_3:
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) == TUPLE_LENGTH_2:
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) == TUPLE_LENGTH_2:
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) == TUPLE_LENGTH_3:
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) == TUPLE_LENGTH_2:
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) == TUPLE_LENGTH_2:
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
- try:
62
- fpr, tpr, _ = roc_curve(y_true, y_scores)
63
- roc_auc = auc(fpr, tpr)
64
-
65
- plt.figure(figsize=(8, 6))
66
- plt.plot(
67
- fpr,
68
- tpr,
69
- color="darkorange",
70
- lw=2,
71
- label=f"ROC curve (area = {roc_auc:.2f})",
72
- )
73
- plt.plot([0, 1], [0, 1], color="navy", lw=2, linestyle="--")
74
- plt.xlim([0.0, 1.0])
75
- plt.ylim([0.0, 1.05])
76
- plt.xlabel("False Positive Rate")
77
- plt.ylabel("True Positive Rate")
78
- plt.title(title)
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
- except ImportError:
88
- warnings.warn("Could not plot ROC curve - missing dependencies", stacklevel=2)
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 _extract_encoder_features(self, x: torch.Tensor) -> torch.Tensor:
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)
@@ -0,0 +1,3 @@
1
+ from deepecgkit.training.train import ECGTrainer
2
+
3
+ __all__ = ["ECGTrainer"]