Upload folder using huggingface_hub
Browse filesThis view is limited to 50 files because it contains too many changes. See raw diff
- .gitattributes +62 -0
- .gitignore +15 -0
- .streamlit/config.toml +15 -0
- LICENSE +18 -0
- README.md +375 -0
- app.py +114 -0
- archive/classification/aggregate_subject_metrics.py +96 -0
- archive/classification/dataset_curate.py +66 -0
- archive/classification/feature_analysis.py +93 -0
- archive/classification/prepare_multiclass_data.py +71 -0
- archive/classification/re_plotter.py +494 -0
- archive/classification/run_benchmarking.py +64 -0
- archive/classification/run_classification.py +34 -0
- archive/classification/run_inference.py +297 -0
- archive/classification/run_lead_importance_test.py +128 -0
- archive/classification/train_and_save_models.py +427 -0
- archive/digitization/README.md +21 -0
- archive/digitization/run_ecg.py +302 -0
- assets/ETH_Z/303/274rich_Logo_black.svg.png +0 -0
- assets/Logo_Vanvitelli_university.svg.png +0 -0
- assets/eoc_logo.png +0 -0
- assets/logo.png +3 -0
- assets/scai_lab_logo.svg +90 -0
- assets/usi_logo.png +0 -0
- backend/__init__.py +1 -0
- backend/classification_runner.py +571 -0
- backend/digitization_runner.py +174 -0
- config.py +64 -0
- digitization.py +1784 -0
- environment.yml +194 -0
- example_inputs/test.jpg +3 -0
- models/classifier_models/ecg_surgery/inceptiontime_model.pkl +3 -0
- models/classifier_models/ecg_surgery/model_metadata.json +62 -0
- models/classifier_models/mi_vs_normal_segmented/arsenal_segmented_model.pkl +3 -0
- models/classifier_models/mi_vs_normal_segmented/model_metadata.json +61 -0
- models/classifier_models/omi_vs_nonomi/model_metadata.json +61 -0
- models/classifier_models/omi_vs_nonomi/rocket_model.pkl +3 -0
- models/digitization_models/yolo11_full/weights/best.pt +3 -0
- models/digitization_models/yolo11_lead/.ipynb_checkpoints/P_curve-checkpoint.png +3 -0
- models/digitization_models/yolo11_lead/.ipynb_checkpoints/args-checkpoint.yaml +105 -0
- models/digitization_models/yolo11_lead/.ipynb_checkpoints/confusion_matrix-checkpoint.png +3 -0
- models/digitization_models/yolo11_lead/.ipynb_checkpoints/confusion_matrix_normalized-checkpoint.png +3 -0
- models/digitization_models/yolo11_lead/.ipynb_checkpoints/labels-checkpoint.jpg +3 -0
- models/digitization_models/yolo11_lead/.ipynb_checkpoints/labels_correlogram-checkpoint.jpg +3 -0
- models/digitization_models/yolo11_lead/.ipynb_checkpoints/results-checkpoint.csv +101 -0
- models/digitization_models/yolo11_lead/.ipynb_checkpoints/results-checkpoint.png +3 -0
- models/digitization_models/yolo11_lead/.ipynb_checkpoints/val_batch0_labels-checkpoint.jpg +3 -0
- models/digitization_models/yolo11_lead/.ipynb_checkpoints/val_batch0_pred-checkpoint.jpg +3 -0
- models/digitization_models/yolo11_lead/.ipynb_checkpoints/val_batch1_labels-checkpoint.jpg +3 -0
- models/digitization_models/yolo11_lead/F1_curve.png +3 -0
.gitattributes
CHANGED
|
@@ -33,3 +33,65 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
| 36 |
+
assets/logo.png filter=lfs diff=lfs merge=lfs -text
|
| 37 |
+
example_inputs/test.jpg filter=lfs diff=lfs merge=lfs -text
|
| 38 |
+
models/digitization_models/yolo11_lead/.ipynb_checkpoints/confusion_matrix-checkpoint.png filter=lfs diff=lfs merge=lfs -text
|
| 39 |
+
models/digitization_models/yolo11_lead/.ipynb_checkpoints/confusion_matrix_normalized-checkpoint.png filter=lfs diff=lfs merge=lfs -text
|
| 40 |
+
models/digitization_models/yolo11_lead/.ipynb_checkpoints/labels-checkpoint.jpg filter=lfs diff=lfs merge=lfs -text
|
| 41 |
+
models/digitization_models/yolo11_lead/.ipynb_checkpoints/labels_correlogram-checkpoint.jpg filter=lfs diff=lfs merge=lfs -text
|
| 42 |
+
models/digitization_models/yolo11_lead/.ipynb_checkpoints/P_curve-checkpoint.png filter=lfs diff=lfs merge=lfs -text
|
| 43 |
+
models/digitization_models/yolo11_lead/.ipynb_checkpoints/results-checkpoint.png filter=lfs diff=lfs merge=lfs -text
|
| 44 |
+
models/digitization_models/yolo11_lead/.ipynb_checkpoints/val_batch0_labels-checkpoint.jpg filter=lfs diff=lfs merge=lfs -text
|
| 45 |
+
models/digitization_models/yolo11_lead/.ipynb_checkpoints/val_batch0_pred-checkpoint.jpg filter=lfs diff=lfs merge=lfs -text
|
| 46 |
+
models/digitization_models/yolo11_lead/.ipynb_checkpoints/val_batch1_labels-checkpoint.jpg filter=lfs diff=lfs merge=lfs -text
|
| 47 |
+
models/digitization_models/yolo11_lead/confusion_matrix.png filter=lfs diff=lfs merge=lfs -text
|
| 48 |
+
models/digitization_models/yolo11_lead/confusion_matrix_normalized.png filter=lfs diff=lfs merge=lfs -text
|
| 49 |
+
models/digitization_models/yolo11_lead/F1_curve.png filter=lfs diff=lfs merge=lfs -text
|
| 50 |
+
models/digitization_models/yolo11_lead/labels.jpg filter=lfs diff=lfs merge=lfs -text
|
| 51 |
+
models/digitization_models/yolo11_lead/labels_correlogram.jpg filter=lfs diff=lfs merge=lfs -text
|
| 52 |
+
models/digitization_models/yolo11_lead/P_curve.png filter=lfs diff=lfs merge=lfs -text
|
| 53 |
+
models/digitization_models/yolo11_lead/PR_curve.png filter=lfs diff=lfs merge=lfs -text
|
| 54 |
+
models/digitization_models/yolo11_lead/R_curve.png filter=lfs diff=lfs merge=lfs -text
|
| 55 |
+
models/digitization_models/yolo11_lead/results.png filter=lfs diff=lfs merge=lfs -text
|
| 56 |
+
models/digitization_models/yolo11_lead/train_batch0.jpg filter=lfs diff=lfs merge=lfs -text
|
| 57 |
+
models/digitization_models/yolo11_lead/train_batch1.jpg filter=lfs diff=lfs merge=lfs -text
|
| 58 |
+
models/digitization_models/yolo11_lead/train_batch2.jpg filter=lfs diff=lfs merge=lfs -text
|
| 59 |
+
models/digitization_models/yolo11_lead/train_batch237330.jpg filter=lfs diff=lfs merge=lfs -text
|
| 60 |
+
models/digitization_models/yolo11_lead/train_batch237331.jpg filter=lfs diff=lfs merge=lfs -text
|
| 61 |
+
models/digitization_models/yolo11_lead/train_batch237332.jpg filter=lfs diff=lfs merge=lfs -text
|
| 62 |
+
models/digitization_models/yolo11_lead/val_batch0_labels.jpg filter=lfs diff=lfs merge=lfs -text
|
| 63 |
+
models/digitization_models/yolo11_lead/val_batch0_pred.jpg filter=lfs diff=lfs merge=lfs -text
|
| 64 |
+
models/digitization_models/yolo11_lead/val_batch1_labels.jpg filter=lfs diff=lfs merge=lfs -text
|
| 65 |
+
models/digitization_models/yolo11_lead/val_batch1_pred.jpg filter=lfs diff=lfs merge=lfs -text
|
| 66 |
+
models/digitization_models/yolo11_lead/val_batch2_labels.jpg filter=lfs diff=lfs merge=lfs -text
|
| 67 |
+
models/digitization_models/yolo11_lead/val_batch2_pred.jpg filter=lfs diff=lfs merge=lfs -text
|
| 68 |
+
models/digitization_models/yolo11_patch/.ipynb_checkpoints/labels-checkpoint.jpg filter=lfs diff=lfs merge=lfs -text
|
| 69 |
+
models/digitization_models/yolo11_patch/.ipynb_checkpoints/train_batch0-checkpoint.jpg filter=lfs diff=lfs merge=lfs -text
|
| 70 |
+
models/digitization_models/yolo11_patch/.ipynb_checkpoints/train_batch1-checkpoint.jpg filter=lfs diff=lfs merge=lfs -text
|
| 71 |
+
models/digitization_models/yolo11_patch/.ipynb_checkpoints/train_batch2-checkpoint.jpg filter=lfs diff=lfs merge=lfs -text
|
| 72 |
+
models/digitization_models/yolo11_patch/labels.jpg filter=lfs diff=lfs merge=lfs -text
|
| 73 |
+
models/digitization_models/yolo11_patch/labels_correlogram.jpg filter=lfs diff=lfs merge=lfs -text
|
| 74 |
+
models/digitization_models/yolo11_patch/train_batch0.jpg filter=lfs diff=lfs merge=lfs -text
|
| 75 |
+
models/digitization_models/yolo11_patch/train_batch1.jpg filter=lfs diff=lfs merge=lfs -text
|
| 76 |
+
models/digitization_models/yolo11_patch/train_batch2.jpg filter=lfs diff=lfs merge=lfs -text
|
| 77 |
+
models/digitization_models/yolo11_pulse/.ipynb_checkpoints/labels-checkpoint.jpg filter=lfs diff=lfs merge=lfs -text
|
| 78 |
+
models/digitization_models/yolo11_pulse/.ipynb_checkpoints/results-checkpoint.png filter=lfs diff=lfs merge=lfs -text
|
| 79 |
+
models/digitization_models/yolo11_pulse/.ipynb_checkpoints/train_batch50402-checkpoint.jpg filter=lfs diff=lfs merge=lfs -text
|
| 80 |
+
models/digitization_models/yolo11_pulse/.ipynb_checkpoints/val_batch0_pred-checkpoint.jpg filter=lfs diff=lfs merge=lfs -text
|
| 81 |
+
models/digitization_models/yolo11_pulse/.ipynb_checkpoints/val_batch1_labels-checkpoint.jpg filter=lfs diff=lfs merge=lfs -text
|
| 82 |
+
models/digitization_models/yolo11_pulse/.ipynb_checkpoints/val_batch2_labels-checkpoint.jpg filter=lfs diff=lfs merge=lfs -text
|
| 83 |
+
models/digitization_models/yolo11_pulse/labels.jpg filter=lfs diff=lfs merge=lfs -text
|
| 84 |
+
models/digitization_models/yolo11_pulse/labels_correlogram.jpg filter=lfs diff=lfs merge=lfs -text
|
| 85 |
+
models/digitization_models/yolo11_pulse/results.png filter=lfs diff=lfs merge=lfs -text
|
| 86 |
+
models/digitization_models/yolo11_pulse/train_batch0.jpg filter=lfs diff=lfs merge=lfs -text
|
| 87 |
+
models/digitization_models/yolo11_pulse/train_batch1.jpg filter=lfs diff=lfs merge=lfs -text
|
| 88 |
+
models/digitization_models/yolo11_pulse/train_batch2.jpg filter=lfs diff=lfs merge=lfs -text
|
| 89 |
+
models/digitization_models/yolo11_pulse/train_batch50400.jpg filter=lfs diff=lfs merge=lfs -text
|
| 90 |
+
models/digitization_models/yolo11_pulse/train_batch50401.jpg filter=lfs diff=lfs merge=lfs -text
|
| 91 |
+
models/digitization_models/yolo11_pulse/train_batch50402.jpg filter=lfs diff=lfs merge=lfs -text
|
| 92 |
+
models/digitization_models/yolo11_pulse/val_batch0_labels.jpg filter=lfs diff=lfs merge=lfs -text
|
| 93 |
+
models/digitization_models/yolo11_pulse/val_batch0_pred.jpg filter=lfs diff=lfs merge=lfs -text
|
| 94 |
+
models/digitization_models/yolo11_pulse/val_batch1_labels.jpg filter=lfs diff=lfs merge=lfs -text
|
| 95 |
+
models/digitization_models/yolo11_pulse/val_batch1_pred.jpg filter=lfs diff=lfs merge=lfs -text
|
| 96 |
+
models/digitization_models/yolo11_pulse/val_batch2_labels.jpg filter=lfs diff=lfs merge=lfs -text
|
| 97 |
+
models/digitization_models/yolo11_pulse/val_batch2_pred.jpg filter=lfs diff=lfs merge=lfs -text
|
.gitignore
ADDED
|
@@ -0,0 +1,15 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Ignore model checkpoints (heavy weights)
|
| 2 |
+
models/
|
| 3 |
+
|
| 4 |
+
# Streamlit run output directories
|
| 5 |
+
output/
|
| 6 |
+
|
| 7 |
+
# Python cache files
|
| 8 |
+
__pycache__/
|
| 9 |
+
*.pyc
|
| 10 |
+
|
| 11 |
+
# Local Streamlit config files
|
| 12 |
+
.streamlit/
|
| 13 |
+
|
| 14 |
+
# Legacy/archive code directories (retained locally but excluded from version control)
|
| 15 |
+
archive/
|
.streamlit/config.toml
ADDED
|
@@ -0,0 +1,15 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
[theme]
|
| 2 |
+
primaryColor = "#E63946"
|
| 3 |
+
backgroundColor = "#FFFFFF"
|
| 4 |
+
secondaryBackgroundColor = "#F1F5F9"
|
| 5 |
+
textColor = "#1E293B"
|
| 6 |
+
font = "sans serif"
|
| 7 |
+
|
| 8 |
+
[server]
|
| 9 |
+
headless = true
|
| 10 |
+
enableCORS = false
|
| 11 |
+
enableXsrfProtection = false
|
| 12 |
+
maxUploadSize = 200
|
| 13 |
+
|
| 14 |
+
[runner]
|
| 15 |
+
fastReruns = true
|
LICENSE
ADDED
|
@@ -0,0 +1,18 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
NON-COMMERCIAL ACADEMIC AND RESEARCH LICENSE AGREEMENT
|
| 2 |
+
|
| 3 |
+
Copyright (c) 2026 SCAI Lab (Spinal Cord Injury and Artificial Intelligence Lab). All rights reserved.
|
| 4 |
+
|
| 5 |
+
Redistribution and use of this software (in source or binary forms, with or without modification) and its associated trained model weights/checkpoints are permitted solely for academic, educational, and non-profit research purposes, subject to the following terms and conditions:
|
| 6 |
+
|
| 7 |
+
1. Redistributions of source code must retain the above copyright notice, this list of conditions, and the following disclaimer.
|
| 8 |
+
|
| 9 |
+
2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions, and the following disclaimer in the documentation and/or other materials provided with the distribution.
|
| 10 |
+
|
| 11 |
+
3. The software, source code, binary forms, and all associated model checkpoints/weights (located under the `models/` directory or downloaded separately) may NOT be used for any commercial purpose, including but not limited to:
|
| 12 |
+
- Integration into a commercial or profit-generating product, platform, or service.
|
| 13 |
+
- Use in paid consulting, corporate research, clinical trial services, or commercial patient diagnostics.
|
| 14 |
+
- Sublicensing, renting, leasing, selling, or distributing the software or derivatives for financial gain.
|
| 15 |
+
|
| 16 |
+
4. Any publications, patents, research papers, presentations, or software derivatives that utilize this software or the associated digitization/classification pipeline must credit the original authors and cite this repository.
|
| 17 |
+
|
| 18 |
+
THIS SOFTWARE AND MODEL WEIGHTS ARE PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES (INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, AND NON-INFRINGEMENT) ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
README.md
ADDED
|
@@ -0,0 +1,375 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# ⚡ ECGLight: ECG Digitization & Classification Dashboard
|
| 2 |
+
|
| 3 |
+
<!--
|
| 4 |
+
<table align="center" border="0" style="border-collapse: collapse; border: none;">
|
| 5 |
+
<tr style="border: none;">
|
| 6 |
+
<td align="center" valign="middle" style="border: none; padding-right: 40px;">
|
| 7 |
+
<img src="assets/logo.png" alt="ECG Digitization & Classification Logo" width="180px" style="border-radius: 14px; box-shadow: 0 4px 16px rgba(15, 23, 42, 0.08);" />
|
| 8 |
+
</td>
|
| 9 |
+
<td align="center" valign="middle" style="border: none;">
|
| 10 |
+
<img src="assets/scai_lab_logo.svg" alt="SCAI Lab Logo" height="100px" />
|
| 11 |
+
</td>
|
| 12 |
+
</tr>
|
| 13 |
+
</table>
|
| 14 |
+
-->
|
| 15 |
+
|
| 16 |
+
An advanced, interactive Streamlit web workstation designed to convert printed/photographed 12-lead paper ECG reports into high-resolution digitized signals and carry out different types of classification over them. The suite is engineered to run with low computational resource requirements, operating seamlessly on a standard consumer laptop GPU (via CUDA) or running completely on CPU.
|
| 17 |
+
|
| 18 |
+
## <a id="table-of-contents"></a>📌 Table of Contents
|
| 19 |
+
|
| 20 |
+
- [🖥️ Web Dashboard Workstation Overview](#web-dashboard-workstation-overview)
|
| 21 |
+
- [🚀 Installation & Setup](#installation-setup)
|
| 22 |
+
- [🧠 Pre-Trained Classifiers & Tasks](#pre-trained-classifiers-tasks)
|
| 23 |
+
- [🚀 Command Line Usage](#command-line-usage)
|
| 24 |
+
- [📁 Repository Structure & Directory Organization](#repository-structure-directory-organization)
|
| 25 |
+
- [📷 How It Works: Signal Digitization](#how-it-works-signal-digitization)
|
| 26 |
+
- [📈 How It Works: Signal Analysis & Visualization](#how-it-works-signal-analysis-visualization)
|
| 27 |
+
- [⚡ How It Works: Heartbeat Segmentation](#how-it-works-heartbeat-segmentation)
|
| 28 |
+
- [🧠 How It Works: Cardiac Classification](#how-it-works-cardiac-classification)
|
| 29 |
+
- [🤝 Collaborating Institutions](#collaborating-institutions)
|
| 30 |
+
- [📄 Citation](#citation)
|
| 31 |
+
- [👥 Authors & Contact](#authors-contact)
|
| 32 |
+
- [📄 License](#license)
|
| 33 |
+
|
| 34 |
+
## <a id="web-dashboard-workstation-overview"></a>🖥️ Web Dashboard Workstation Overview
|
| 35 |
+
|
| 36 |
+
The dashboard provides a premium, responsive user interface designed for research, education, and clinical workflow exploration. It coordinates the digitization and classification pipelines into a unified, lightweight web application.
|
| 37 |
+
|
| 38 |
+
### Workstation Modules
|
| 39 |
+
|
| 40 |
+
1. **📷 ECG Image Digitizer**:
|
| 41 |
+
- Upload any scanned or photographed ECG image (`.png`, `.jpg`, `.jpeg`).
|
| 42 |
+
- Run the sequential YOLOv11 pipeline step-by-step with real-time progress indicators.
|
| 43 |
+
- Outputs a summary of detected leads and total samples.
|
| 44 |
+
- Automatically saves the digitized CSV to disk under `output/digitization/latest_digitized.csv` for downstream consumption.
|
| 45 |
+
|
| 46 |
+
2. **📈 ECG Signal Viewer**:
|
| 47 |
+
- Visualizes multi-channel ECG signals interactively using client-side native Streamlit line charts (supporting zoom, pan, and hover tooltips).
|
| 48 |
+
- Supports stacked subplots (with distinct clinical colors for each lead: clinical red, teal, deep blue, yellow, purple, etc.) or overlaid graphs.
|
| 49 |
+
- Displays statistical summaries (mean, standard deviation, min/max, range) and allows row-by-row signal previewing.
|
| 50 |
+
|
| 51 |
+
3. **❤️ ECG Classification**:
|
| 52 |
+
- Predicts cardiac conditions using pre-trained ensemble and deep learning classifiers.
|
| 53 |
+
- Automatically segments raw signals into heartbeats around R-peaks using the Pan-Tompkins algorithm before running inference.
|
| 54 |
+
- **Inference Mode (No Ground-Truth)**: If the uploaded CSV lacks diagnostic labels, the dashboard displays a downloadable **Predictions Table** detailing predicted class diagnoses and model confidence probabilities.
|
| 55 |
+
- **Evaluation Mode (With Ground-Truth)**: If labels are present, the page calculates and plots performance metrics (Accuracy, F1-Score, Sensitivity, Specificity, Confusion Matrix).
|
| 56 |
+
|
| 57 |
+
### 🔄 Workflow
|
| 58 |
+
|
| 59 |
+
The workstation coordinates the pipeline through four distinct steps: **Digitization**, **Analysis**, **Segmentation**, and **Classification**. The technical workflows for each process are documented below in their respective **How It Works** sections.
|
| 60 |
+
|
| 61 |
+
---
|
| 62 |
+
|
| 63 |
+
## <a id="installation-setup"></a>🚀 Installation & Setup
|
| 64 |
+
|
| 65 |
+
### Prerequisites
|
| 66 |
+
- Python 3.9
|
| 67 |
+
- CUDA-capable GPU recommended (automatically falls back to CPU if unavailable).
|
| 68 |
+
|
| 69 |
+
### Conda Environment & Model Setup
|
| 70 |
+
|
| 71 |
+
1. Clone the repository and navigate to the project directory:
|
| 72 |
+
```bash
|
| 73 |
+
git clone https://github.com/scai-lab/ECG-Digitization-Classification.git
|
| 74 |
+
cd ECG-Digitization-Classification
|
| 75 |
+
```
|
| 76 |
+
|
| 77 |
+
2. Create the conda environment using the provided `environment.yml` configuration:
|
| 78 |
+
```bash
|
| 79 |
+
conda env create -f environment.yml
|
| 80 |
+
conda activate infer
|
| 81 |
+
```
|
| 82 |
+
|
| 83 |
+
> [!IMPORTANT]
|
| 84 |
+
> **Windows Compatibility & TensorFlow Setup**:
|
| 85 |
+
> If you are on Windows and encounter native runtime loading failures (`ImportError: DLL load failed while importing _pywrap_tensorflow_internal: A dynamic link library (DLL) initialization routine failed`), you need to install a stable version pairing of TensorFlow and Protobuf:
|
| 86 |
+
> ```bash
|
| 87 |
+
> pip install tensorflow==2.15.0 protobuf==4.25.3
|
| 88 |
+
> ```
|
| 89 |
+
> *(Make sure no background Streamlit or Python tasks are running when executing this command, to prevent file locking issues on `.pyd` libraries).*
|
| 90 |
+
|
| 91 |
+
3. **Download Pre-Trained Model Weights**:
|
| 92 |
+
Due to their file sizes, the YOLO detection checkpoints and pre-trained classifiers are hosted externally. Download the `models/` directory from the link below and place it directly in the root of the project:
|
| 93 |
+
|
| 94 |
+
👉 **[Download Pre-Trained Models Directory (ETH Zürich Polybox)](https://polybox.ethz.ch/index.php/s/GDACstPtsoTrrWH)**
|
| 95 |
+
|
| 96 |
+
Once extracted, verify that the weights are located inside the directory tree structure:
|
| 97 |
+
```text
|
| 98 |
+
models/
|
| 99 |
+
├── digitization_models/
|
| 100 |
+
│ ├── yolo11_full/weights/best.pt
|
| 101 |
+
│ ├── yolo11_lead/weights/best.pt
|
| 102 |
+
│ ├── yolo11_pulse/weights/best.pt
|
| 103 |
+
│ └── yolo11_patch/weights/best.pt
|
| 104 |
+
└── classifier_models/
|
| 105 |
+
├── mi_vs_normal_segmented/
|
| 106 |
+
├── omi_vs_nonomi/
|
| 107 |
+
└── ecg_surgery/
|
| 108 |
+
```
|
| 109 |
+
|
| 110 |
+
Key packages installed by the environment: `torch 2.7`, `ultralytics 8.3`, `opencv-python 4.11`, `scikit-image 0.24`, `wfdb 4.3`, `patched-yolo-infer 1.3.8`, `sktime`, `streamlit`.
|
| 111 |
+
|
| 112 |
+
---
|
| 113 |
+
|
| 114 |
+
## <a id="pre-trained-classifiers-tasks"></a>🧠 Pre-Trained Classifiers & Tasks
|
| 115 |
+
|
| 116 |
+
The classification engine supports three diagnostic tasks using the pre-trained weights in `classifier_models/`:
|
| 117 |
+
|
| 118 |
+
| Classification Task | Model Type | Expected Input Shape | Test Accuracy | Positive Class |
|
| 119 |
+
| :--- | :--- | :--- | :---: | :--- |
|
| 120 |
+
| **Normal vs Myocardial Infarction (MI) - Segmented** | Arsenal | 12 leads × 140 timesteps | **92.3%** | `MYOCARDIAL_INFARCTION` |
|
| 121 |
+
| **Occlusive MI (OMI) vs non-OMI** | Rocket | 12 leads × 141 timesteps | **88.9%** | `OMI` |
|
| 122 |
+
| **Pre-Procedural vs Post-Procedural MI** | InceptionTime | 12 leads × 140 timesteps | **91.4%** | `pre-procedural MI` |
|
| 123 |
+
|
| 124 |
+
- **Arsenal**: An ensemble of ROCKET classifiers utilizing random convolutional kernels to extract feature representations combined with ridge regression.
|
| 125 |
+
- **Rocket**: Random Omni-directional Kernel Extraction (ROCKET) classifier, computing kernel convolutions quickly for high-dimensional time-series data.
|
| 126 |
+
- **InceptionTime**: A deep convolutional network ensemble modeled on the Inception architecture, extracting multi-scale temporal features.
|
| 127 |
+
|
| 128 |
+
---
|
| 129 |
+
|
| 130 |
+
## <a id="command-line-usage"></a>🚀 Command Line Usage
|
| 131 |
+
|
| 132 |
+
### Run Batch Digitization (`run_org.py`)
|
| 133 |
+
|
| 134 |
+
The batch processing script processes nested hospital directories, exporting structured folders of digitized CSVs:
|
| 135 |
+
|
| 136 |
+
1. Configure path variables at the top of `run_org.py`:
|
| 137 |
+
```python
|
| 138 |
+
ORGANIZED_DIR = "../ecg_files/ECG_organized_all" # Input dataset root
|
| 139 |
+
OUTPUT_DIR = "../ecg_files/ECG_digitized" # Mirrored CSV directory
|
| 140 |
+
CATEGORIES = ["pre", "index", "post"] # Categories to process
|
| 141 |
+
```
|
| 142 |
+
|
| 143 |
+
2. Run the script:
|
| 144 |
+
```bash
|
| 145 |
+
python run_org.py
|
| 146 |
+
```
|
| 147 |
+
|
| 148 |
+
### Run Model Inference (`run_inference.py`)
|
| 149 |
+
|
| 150 |
+
Execute predictions directly on digitized data from the command line using `run_inference.py` located in `archive/classification/`:
|
| 151 |
+
|
| 152 |
+
```bash
|
| 153 |
+
# MI vs Normal Segmented heartbeat classification
|
| 154 |
+
python archive/classification/run_inference.py --model mi_vs_normal_segmented --input data/ptb_xl/segmented_heartbeats.csv
|
| 155 |
+
|
| 156 |
+
# OMI vs non-OMI classification
|
| 157 |
+
python archive/classification/run_inference.py --model omi_vs_nonomi --input data/ecg_matrix_omi_segmented_50_150_90.csv
|
| 158 |
+
|
| 159 |
+
# Custom output file path
|
| 160 |
+
python archive/classification/run_inference.py --model ecg_surgery --input data/ecg_surgery_segmented_50_150_70.csv --output results/surgery_preds.csv
|
| 161 |
+
```
|
| 162 |
+
|
| 163 |
+
---
|
| 164 |
+
|
| 165 |
+
## <a id="repository-structure-directory-organization"></a>📁 Repository Structure & Directory Organization
|
| 166 |
+
|
| 167 |
+
The repository is structured to maintain a clean root directory, moving utility runners, UI views, model checkpoints, and legacy/training scripts into distinct modules:
|
| 168 |
+
|
| 169 |
+
```
|
| 170 |
+
.
|
| 171 |
+
├── app.py # Streamlit application main router
|
| 172 |
+
├── config.py # Centralized configuration and model registry
|
| 173 |
+
├── digitization.py # Core ECGImage extraction pipeline class
|
| 174 |
+
├── environment.yml # Conda environment dependency file
|
| 175 |
+
├── README.md # Comprehensive repository documentation
|
| 176 |
+
│
|
| 177 |
+
├── backend/ # Dashboard background execution adapters
|
| 178 |
+
│ ├── __init__.py # Backend package declaration
|
| 179 |
+
│ ├── digitization_runner.py # YOLO loader and single-image processor
|
| 180 |
+
│ └── classification_runner.py # Pre-trained model loader and preprocessor
|
| 181 |
+
│
|
| 182 |
+
├── utils/ # Streamlit front-end page components
|
| 183 |
+
│ ├── __init__.py # Utils package declaration
|
| 184 |
+
│ ├── branding.py # Sidebar titles, headers, and footer logos
|
| 185 |
+
│ ├── css.py # Custom clinical theme and grid background CSS
|
| 186 |
+
│ ├── hardware.py # Displays CPU/GPU hardware properties (cached)
|
| 187 |
+
│ ├── page_digitizer.py # Front-end for the ECG Digitizer page
|
| 188 |
+
│ ├── page_csv_viewer.py # Front-end for the interactive Signal Viewer
|
| 189 |
+
│ └── page_classifier.py # Front-end for the Classification workstation
|
| 190 |
+
│
|
| 191 |
+
├── models/ # Relocated YOLO checkpoints and classifiers
|
| 192 |
+
│ ├── digitization_models/ # YOLO v11 checkpoints for digitization
|
| 193 |
+
│ │ ├── yolo11_full/ # YOLO Bounding boxes
|
| 194 |
+
│ │ ├── yolo11_lead/ # YOLO Lead names
|
| 195 |
+
│ │ ├── yolo11_pulse/ # YOLO Reference pulses
|
| 196 |
+
│ │ └── yolo11_patch/ # YOLO Waveform segmentations
|
| 197 |
+
│ │
|
| 198 |
+
│ └── classifier_models/ # Bundled pre-trained diagnostic classifiers
|
| 199 |
+
│ ├── mi_vs_normal_segmented/ # Pre-trained Arsenal model (segmented beats)
|
| 200 |
+
│ ├── omi_vs_nonomi/ # Pre-trained Rocket model (segmented beats)
|
| 201 |
+
│ └── ecg_surgery/ # Pre-trained InceptionTime model (segmented beats)
|
| 202 |
+
│
|
| 203 |
+
└── archive/ # Archived developer, training, and legacy scripts
|
| 204 |
+
└── classification/
|
| 205 |
+
├── train_and_save_models.py# Script used to compile pre-trained models
|
| 206 |
+
├── run_inference.py # Independent CLI inference execution script
|
| 207 |
+
├── run_classification.py # Baseline MLP classifier pipeline
|
| 208 |
+
├── run_benchmarking.py # Comparison benchmarking suite
|
| 209 |
+
├── run_lead_importance_test.py# Individual lead performance evaluator
|
| 210 |
+
├── feature_analysis.py # Original all-in-one analysis script
|
| 211 |
+
├── aggregate_subject_metrics.py# Multi-subject performance aggregator
|
| 212 |
+
├── dataset_curate.py # Local dataset curation utility
|
| 213 |
+
└── re_plotter.py # Advanced Gaussian signal generator & visualizer
|
| 214 |
+
```
|
| 215 |
+
|
| 216 |
+
---
|
| 217 |
+
|
| 218 |
+
## <a id="how-it-works-signal-digitization"></a>📷 How It Works: Signal Digitization
|
| 219 |
+
|
| 220 |
+
The core class [digitization.py](file:///d:/Projects/ECGLight/digitization.py) operates a multi-stage sequential computer vision pipeline to translate raster images into digitized signals:
|
| 221 |
+
|
| 222 |
+
```mermaid
|
| 223 |
+
graph TD
|
| 224 |
+
A[ECG Image Upload] --> B[Preprocessing: Otsu & Blurring]
|
| 225 |
+
B --> C[YOLOv11 Detection & Segmentation]
|
| 226 |
+
subgraph YOLOv11 Models
|
| 227 |
+
C1[yolo11_full: Lead Boundaries]
|
| 228 |
+
C2[yolo11_lead: Text Name Labels]
|
| 229 |
+
C3[yolo11_pulse: Calibration Pulses]
|
| 230 |
+
C4[yolo11_patch: Waveform Segments]
|
| 231 |
+
end
|
| 232 |
+
C --> C1 & C2 & C3 & C4
|
| 233 |
+
C1 & C2 & C3 & C4 --> D[Hough Lines Calibration]
|
| 234 |
+
D --> E[K-Means Row & Column Grid Construction]
|
| 235 |
+
E --> F[Anti-Leakage Connected Components Filter]
|
| 236 |
+
F --> G[Centroid Trace & Resampling to 500Hz]
|
| 237 |
+
G --> H[Export latest_digitized.csv]
|
| 238 |
+
```
|
| 239 |
+
|
| 240 |
+
1. **Preprocessing**: Cleans the scanned image using shadow-removal masks, Otsu binarization, and Gaussian blurring to isolate ink lines from paper textures.
|
| 241 |
+
2. **YOLO Segmentation**: Applies a patched YOLO segmentation model at three crop scales (`4×`, `4.5×`, and `5×` height) to isolate individual lead waveform contours.
|
| 242 |
+
3. **Sequential Detections**: Runs three YOLO models in parallel:
|
| 243 |
+
- `yolo11_full`: Bounding boxes for the 12 lead channels.
|
| 244 |
+
- `yolo11_lead`: Text labels representing lead names (I, II, aVR...).
|
| 245 |
+
- `yolo11_pulse`: Bounding boxes for the calibration reference pulses (typically 1mV high, representing vertical scale).
|
| 246 |
+
4. **Scale Calibration**: Fits Hough lines to the calibration pulse boundaries. The pixel height determines the voltage scale (`volt/pixel`), while the width determines the time scale (`time/pixel`).
|
| 247 |
+
5. **Grid Construction**: Employs K-Means clustering on lead coordinates to map rows and columns, automatically parsing standard Cabrera orders and grid formats (3×4, 4×3, 6×2, 12×1).
|
| 248 |
+
6. **Signal Extraction & Post-Processing**: Traces contours to extract raw pixel centroids, performs baseline correction, applies linear interpolation to bridge gaps, and resamples to a standard **500 Hz** frequency calibrated in **millivolts (mV)**. In addition, an **anti-leakage component filter** is executed per cell crop using connected components analysis to automatically identify the primary waveform trace and strip out smaller, boundary-adjacent components (leaked signals from neighboring leads) that sit far from the row baseline.
|
| 249 |
+
|
| 250 |
+
---
|
| 251 |
+
|
| 252 |
+
## <a id="how-it-works-signal-analysis-visualization"></a>📈 How It Works: Signal Analysis & Visualization
|
| 253 |
+
|
| 254 |
+
Once continuous signals are extracted, the dashboard runs analytical tasks and displays interactive previews:
|
| 255 |
+
|
| 256 |
+
```mermaid
|
| 257 |
+
graph TD
|
| 258 |
+
A[Upload Digitized CSV] --> B[Parse Lead Voltages & Timestamps]
|
| 259 |
+
B --> C[Vega-Lite Interactive Visualizer]
|
| 260 |
+
C --> C1[Render Stacked Leads]
|
| 261 |
+
C --> C2[Render Overlaid Signals]
|
| 262 |
+
B --> D[Compute Signal Statistics: Mean, SD, Min/Max]
|
| 263 |
+
D --> E[Display Summary Dataframes & Row Previews]
|
| 264 |
+
```
|
| 265 |
+
|
| 266 |
+
1. **Parse Signals**: Reads digitized CSV format, validating lead names and timestamps.
|
| 267 |
+
2. **Vega-Lite Visualization**: Renders interactive charts supporting native client-side zoom, pan, and hover tooltips for all channels.
|
| 268 |
+
3. **Signal Statistics**: Automatically computes statistical characteristics (mean, standard deviation, min/max values) for each lead.
|
| 269 |
+
|
| 270 |
+
---
|
| 271 |
+
|
| 272 |
+
## <a id="how-it-works-heartbeat-segmentation"></a>⚡ How It Works: Heartbeat Segmentation
|
| 273 |
+
|
| 274 |
+
To prepare continuous digitized signals for the classification models, the pipeline runs the Pan-Tompkins R-peak detection algorithm:
|
| 275 |
+
|
| 276 |
+
```mermaid
|
| 277 |
+
graph TD
|
| 278 |
+
A[Digitized 500Hz Signal] --> B[Bandpass Filter 5-15Hz]
|
| 279 |
+
B --> C[Derivative Filter]
|
| 280 |
+
C --> D[Squaring Operation]
|
| 281 |
+
D --> E[Moving Window Integration]
|
| 282 |
+
E --> F[Adaptive Thresholding & R-Peak Search]
|
| 283 |
+
F --> G[Extract 140-sample Beats: 50ms pre-R, 150ms post-R]
|
| 284 |
+
G --> H[Max-Absolute Voltage Normalization]
|
| 285 |
+
```
|
| 286 |
+
|
| 287 |
+
1. **Filtering**: The Lead II signal is filtered via a bandpass filter (5-15 Hz) to suppress muscle noise, baseline wander, and T-wave interference.
|
| 288 |
+
2. **Differentiation**: Computes the slope of the signal to highlight the rapid change in the QRS complex.
|
| 289 |
+
3. **Squaring**: Performs point-by-point squaring to amplify QRS slopes while attenuating smaller waves.
|
| 290 |
+
4. **Integration**: A moving window integrator (typically 150ms wide) compiles the slope information into a peak window.
|
| 291 |
+
5. **Adaptive Thresholding & Peak Search**: Dynamically computes threshold constants based on average noise and signal levels, locating R-peaks.
|
| 292 |
+
6. **Beat Windowing**: Extracts a localized heartbeat around each R-peak (typically extending 50ms before and 150ms after the peak), normalizes the voltage per heartbeat using max-absolute scaling, and truncates/pads the resulting segments to the target model input width (e.g. 140 or 141 timesteps).
|
| 293 |
+
|
| 294 |
+
---
|
| 295 |
+
|
| 296 |
+
## <a id="how-it-works-cardiac-classification"></a>🧠 How It Works: Cardiac Classification
|
| 297 |
+
|
| 298 |
+
The heartbeat segment tensors are evaluated using pre-trained time-series classification models:
|
| 299 |
+
|
| 300 |
+
```mermaid
|
| 301 |
+
graph TD
|
| 302 |
+
A[Segmented Heartbeats] --> B[Numpy3D Reshaping: N_instances × 12_leads × N_timesteps]
|
| 303 |
+
B --> C[Select Classification Task]
|
| 304 |
+
subgraph Model Registry
|
| 305 |
+
C1[Normal vs MI: Arsenal]
|
| 306 |
+
C2[OMI vs non-OMI: Rocket]
|
| 307 |
+
C3[Pre vs Post-Procedural MI: InceptionTime]
|
| 308 |
+
end
|
| 309 |
+
C --> C1 & C2 & C3
|
| 310 |
+
C1 & C2 & C3 --> D[Load Pre-Trained Pickled Estimator]
|
| 311 |
+
D --> E[Predict Class Labels & Probabilities]
|
| 312 |
+
E --> F[Generate Downloadable Predictions CSV]
|
| 313 |
+
```
|
| 314 |
+
|
| 315 |
+
1. **Numpy3D Formatting**: Formats the heartbeat segments into a standard `sktime` `Numpy3D` tensor with shape `(N_instances, 12_leads, N_timesteps)`.
|
| 316 |
+
2. **Dynamic Task Selection**: Loads the pre-trained pickled model corresponding to the selected classification task.
|
| 317 |
+
3. **Model Inference**: Evaluates the model to compute class predictions and probability confidences.
|
| 318 |
+
4. **Result Generation**: Automatically builds downloadable prediction tables and calculates performance metrics if ground-truth labels are present in the dataset.
|
| 319 |
+
|
| 320 |
+
---
|
| 321 |
+
|
| 322 |
+
## <a id="collaborating-institutions"></a>🤝 Collaborating Institutions
|
| 323 |
+
|
| 324 |
+
This project was developed in collaboration with:
|
| 325 |
+
|
| 326 |
+
- [ETH Zürich](https://ethz.ch)
|
| 327 |
+
- [Istituto Cardiocentro Ticino (EOC)](https://www.cardiocentro.org)
|
| 328 |
+
- [Università della Svizzera italiana (USI)](https://www.usi.ch)
|
| 329 |
+
- [Università della Campania Luigi Vanvitelli](https://www.unicampania.it)
|
| 330 |
+
|
| 331 |
+
<p align="center">
|
| 332 |
+
<a href="https://ethz.ch" target="_blank">
|
| 333 |
+
<img src="assets/ETH_Zürich_Logo_black.svg.png" alt="ETH Zürich" height="30px" style="vertical-align: middle; margin: 0 15px;" />
|
| 334 |
+
</a>
|
| 335 |
+
|
| 336 |
+
<a href="https://www.cardiocentro.org" target="_blank">
|
| 337 |
+
<img src="assets/eoc_logo.png" alt="Istituto Cardiocentro Ticino (EOC)" height="35px" style="vertical-align: middle; margin: 0 15px;" />
|
| 338 |
+
</a>
|
| 339 |
+
|
| 340 |
+
<a href="https://www.usi.ch" target="_blank">
|
| 341 |
+
<img src="assets/usi_logo.png" alt="USI" height="35px" style="vertical-align: middle; margin: 0 15px;" />
|
| 342 |
+
</a>
|
| 343 |
+
|
| 344 |
+
<a href="https://www.unicampania.it" target="_blank">
|
| 345 |
+
<img src="assets/Logo_Vanvitelli_university.svg.png" alt="Università della Campania Luigi Vanvitelli" height="35px" style="vertical-align: middle; margin: 0 15px;" />
|
| 346 |
+
</a>
|
| 347 |
+
</p>
|
| 348 |
+
|
| 349 |
+
---
|
| 350 |
+
|
| 351 |
+
## <a id="citation"></a>📄 Citation
|
| 352 |
+
|
| 353 |
+
```bibtex
|
| 354 |
+
@article{natraj2026ecglight,
|
| 355 |
+
title={ECGLight: Compute-Light Framework For Paper ECG Digitization and Myocardial Infarction Screening},
|
| 356 |
+
author={Natraj, Shreyasvi and Achtari, Cyrus and Gragnano, Felice and Milzi, Andrea and Valgimigli, Marco and Paez-Granados, Diego},
|
| 357 |
+
journal={arXiv preprint arXiv:2607.07683},
|
| 358 |
+
year={2026},
|
| 359 |
+
url={https://arxiv.org/abs/2607.07683},
|
| 360 |
+
doi={10.48550/arXiv.2607.07683}
|
| 361 |
+
}
|
| 362 |
+
```
|
| 363 |
+
|
| 364 |
+
---
|
| 365 |
+
|
| 366 |
+
## <a id="authors-contact"></a>👥 Authors & Contact
|
| 367 |
+
|
| 368 |
+
- **Shreyasvi Natraj** — [snatraj@ethz.ch](mailto:snatraj@ethz.ch)
|
| 369 |
+
- **Cyrus Achtari**
|
| 370 |
+
|
| 371 |
+
---
|
| 372 |
+
|
| 373 |
+
## <a id="license"></a>📄 License
|
| 374 |
+
|
| 375 |
+
This project is released under the **Non-Commercial Academic and Research License Agreement**. Please refer to the [LICENSE](file:///d:/Projects/ECGLight/LICENSE) file in the repository root for the full licensing terms. The codebase and trained model weights are provided free of charge for personal, academic, and non-profit research use only. Commercial use is strictly prohibited.
|
app.py
ADDED
|
@@ -0,0 +1,114 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# app.py
|
| 2 |
+
"""
|
| 3 |
+
Main entry point for the ECGLight Dashboard.
|
| 4 |
+
Slim router that delegates to page modules in utils/.
|
| 5 |
+
"""
|
| 6 |
+
|
| 7 |
+
import os
|
| 8 |
+
import sys
|
| 9 |
+
import streamlit as st
|
| 10 |
+
|
| 11 |
+
# Add project root to path
|
| 12 |
+
sys.path.append(os.path.abspath(os.path.dirname(__file__)))
|
| 13 |
+
|
| 14 |
+
# Import configuration
|
| 15 |
+
import config
|
| 16 |
+
|
| 17 |
+
# Import utility modules (CSS, branding, hardware, pages)
|
| 18 |
+
from utils.css import inject_custom_css
|
| 19 |
+
from utils.branding import render_sidebar_logo, render_logo_footer
|
| 20 |
+
from utils.hardware import render_sidebar_hardware
|
| 21 |
+
from utils import page_digitizer, page_csv_viewer, page_classifier
|
| 22 |
+
|
| 23 |
+
# Backend runners are imported lazily inside page routing blocks below
|
| 24 |
+
# to avoid loading torch/sklearn/scipy when the user only needs the CSV viewer.
|
| 25 |
+
|
| 26 |
+
# ==============================================================================
|
| 27 |
+
# PAGE CONFIG (must be the first Streamlit command)
|
| 28 |
+
# ==============================================================================
|
| 29 |
+
st.set_page_config(
|
| 30 |
+
page_title="ECGLight",
|
| 31 |
+
page_icon="⚡",
|
| 32 |
+
layout="wide",
|
| 33 |
+
initial_sidebar_state="expanded"
|
| 34 |
+
)
|
| 35 |
+
|
| 36 |
+
# ==============================================================================
|
| 37 |
+
# INJECT CSS THEME
|
| 38 |
+
# ==============================================================================
|
| 39 |
+
inject_custom_css()
|
| 40 |
+
|
| 41 |
+
# ==============================================================================
|
| 42 |
+
# SIDEBAR
|
| 43 |
+
# ==============================================================================
|
| 44 |
+
# SCAI Lab logo at the top of the sidebar
|
| 45 |
+
render_sidebar_logo()
|
| 46 |
+
|
| 47 |
+
st.sidebar.markdown(
|
| 48 |
+
'<h1 style="font-size: 1.8rem; margin-bottom: 5px; font-weight: 800; letter-spacing: -0.5px; text-align: center; padding-right: 1.5rem;">'
|
| 49 |
+
'⚡ <span class="glow-text">ECGLight</span></h1>',
|
| 50 |
+
unsafe_allow_html=True
|
| 51 |
+
)
|
| 52 |
+
st.sidebar.markdown(
|
| 53 |
+
'<p class="section-subtitle" style="text-align: center;">Lightweight Digitization & Classification Suite</p>',
|
| 54 |
+
unsafe_allow_html=True
|
| 55 |
+
)
|
| 56 |
+
|
| 57 |
+
st.sidebar.markdown("---")
|
| 58 |
+
|
| 59 |
+
# Navigation (CSV Signal Viewer above Classification)
|
| 60 |
+
page = st.sidebar.radio(
|
| 61 |
+
"🧭 Navigation Pages",
|
| 62 |
+
["📷 ECG Image Digitizer", "📈 ECG Signal Viewer", "❤️ ECG Classification"],
|
| 63 |
+
index=0
|
| 64 |
+
)
|
| 65 |
+
|
| 66 |
+
# Keep track of active page to reset session state on page navigation
|
| 67 |
+
if "current_page" not in st.session_state:
|
| 68 |
+
st.session_state["current_page"] = page
|
| 69 |
+
|
| 70 |
+
if st.session_state["current_page"] != page:
|
| 71 |
+
# Clear session state keys to reset state on page change to prevent slowing things down,
|
| 72 |
+
# but we DO NOT clear "latest_digitized_path" to share the default digitized CSV.
|
| 73 |
+
keys_to_clear = [
|
| 74 |
+
"digitized_df", "digitized_name", "last_uploaded_name",
|
| 75 |
+
"clf_results", "clf_logs", "selected_task",
|
| 76 |
+
"csv_viewer_upload", "classifier_csv_upload"
|
| 77 |
+
]
|
| 78 |
+
for key in keys_to_clear:
|
| 79 |
+
st.session_state.pop(key, None)
|
| 80 |
+
|
| 81 |
+
st.session_state["current_page"] = page
|
| 82 |
+
st.rerun()
|
| 83 |
+
|
| 84 |
+
st.sidebar.markdown("---")
|
| 85 |
+
|
| 86 |
+
# Hardware info (cached — only runs once per session)
|
| 87 |
+
render_sidebar_hardware()
|
| 88 |
+
|
| 89 |
+
# ==============================================================================
|
| 90 |
+
# YOLO MODEL CHECK (used by Page 1)
|
| 91 |
+
# ==============================================================================
|
| 92 |
+
yolo_configured = all([
|
| 93 |
+
config.YOLO_BOX_MODEL_PATH,
|
| 94 |
+
config.YOLO_LEAD_NAME_MODEL_PATH,
|
| 95 |
+
config.YOLO_PULSE_MODEL_PATH,
|
| 96 |
+
config.YOLO_SEGMENTATION_MODEL_PATH
|
| 97 |
+
])
|
| 98 |
+
|
| 99 |
+
# ==============================================================================
|
| 100 |
+
# PAGE ROUTING (lazy backend imports — only load heavy modules when needed)
|
| 101 |
+
# ==============================================================================
|
| 102 |
+
if page == "📷 ECG Image Digitizer":
|
| 103 |
+
from backend import digitization_runner
|
| 104 |
+
page_digitizer.render(config, digitization_runner, yolo_configured)
|
| 105 |
+
render_logo_footer()
|
| 106 |
+
|
| 107 |
+
elif page == "📈 ECG Signal Viewer":
|
| 108 |
+
page_csv_viewer.render()
|
| 109 |
+
render_logo_footer()
|
| 110 |
+
|
| 111 |
+
elif page == "❤️ ECG Classification":
|
| 112 |
+
from backend import classification_runner
|
| 113 |
+
page_classifier.render(config, classification_runner)
|
| 114 |
+
render_logo_footer()
|
archive/classification/aggregate_subject_metrics.py
ADDED
|
@@ -0,0 +1,96 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import pandas as pd
|
| 2 |
+
import numpy as np
|
| 3 |
+
import os
|
| 4 |
+
import re
|
| 5 |
+
from sklearn.metrics import accuracy_score, precision_score, f1_score
|
| 6 |
+
|
| 7 |
+
def get_base_subject_id(subject_id_with_beat):
|
| 8 |
+
"""
|
| 9 |
+
Extracts the base subject ID by removing suffixes like _beatN, beatN, -index_beatN, etc.
|
| 10 |
+
"""
|
| 11 |
+
# Regex to catch variations like _beat1, beat1, -index_beat1, INDEX_beat1, etc.
|
| 12 |
+
# It looks for 'beat' preceded by common delimiters and captures everything before it.
|
| 13 |
+
match = re.search(r'^(.*?)(?:[_\s-](?:index|INDEX|post|POST|checkpoint|POST PCI|POST PCI \d|POST \d)*)?(?:[_\s-]beat\d+)$', str(subject_id_with_beat))
|
| 14 |
+
if match:
|
| 15 |
+
return match.group(1).strip()
|
| 16 |
+
|
| 17 |
+
# Fallback: just split by 'beat' if regex is too specific
|
| 18 |
+
if 'beat' in str(subject_id_with_beat):
|
| 19 |
+
return str(subject_id_with_beat).split('beat')[0].rstrip('_ -').strip()
|
| 20 |
+
|
| 21 |
+
return str(subject_id_with_beat).strip()
|
| 22 |
+
|
| 23 |
+
def calculate_metrics(df_subject):
|
| 24 |
+
"""
|
| 25 |
+
Calculates metrics for subject-level predictions.
|
| 26 |
+
"""
|
| 27 |
+
# Ground truth: since all beats for a subject have the same label, we just take the first
|
| 28 |
+
y_true = df_subject['true_label_numeric'].values
|
| 29 |
+
|
| 30 |
+
# Prediction based on averaged probability for Pre-Procedural MI
|
| 31 |
+
# Assuming probability_Pre-Procedural MI >= 0.5 is Pre-Procedural MI (1)
|
| 32 |
+
y_pred = (df_subject['avg_prob_pre_procedural'] >= 0.5).astype(int)
|
| 33 |
+
|
| 34 |
+
acc = accuracy_score(y_true, y_pred)
|
| 35 |
+
prec = precision_score(y_true, y_pred, zero_division=0)
|
| 36 |
+
f1 = f1_score(y_true, y_pred, zero_division=0)
|
| 37 |
+
|
| 38 |
+
return acc, prec, f1
|
| 39 |
+
|
| 40 |
+
def main():
|
| 41 |
+
prediction_dir = r'output/benchmark_results/ecg_matrix_dataset_segmented/model_predictions'
|
| 42 |
+
|
| 43 |
+
if not os.path.exists(prediction_dir):
|
| 44 |
+
print(f"Directory not found: {prediction_dir}")
|
| 45 |
+
return
|
| 46 |
+
|
| 47 |
+
# Filter out the combined file if it exists
|
| 48 |
+
files = [f for f in os.listdir(prediction_dir) if f.endswith('.csv') and f != 'all_models_predictions_combined.csv']
|
| 49 |
+
|
| 50 |
+
results = []
|
| 51 |
+
|
| 52 |
+
print(f"{'Model':<20} | {'Accuracy':<10} | {'Precision':<10} | {'F1 Score':<10}")
|
| 53 |
+
print("-" * 60)
|
| 54 |
+
|
| 55 |
+
for file in files:
|
| 56 |
+
model_name = file.replace('_predictions.csv', '')
|
| 57 |
+
file_path = os.path.join(prediction_dir, file)
|
| 58 |
+
|
| 59 |
+
try:
|
| 60 |
+
df = pd.read_csv(file_path)
|
| 61 |
+
|
| 62 |
+
# 1. Extract base subject ID
|
| 63 |
+
df['base_subject_id'] = df['subject_id'].apply(get_base_subject_id)
|
| 64 |
+
|
| 65 |
+
# 2. Group by base subject ID
|
| 66 |
+
# We average the probability and take the first true label (since it should be invariant for a subject)
|
| 67 |
+
subject_group = df.groupby('base_subject_id').agg({
|
| 68 |
+
'probability_Pre-Procedural MI': 'mean',
|
| 69 |
+
'true_label_numeric': 'first'
|
| 70 |
+
}).reset_index()
|
| 71 |
+
|
| 72 |
+
subject_group.rename(columns={'probability_Pre-Procedural MI': 'avg_prob_pre_procedural'}, inplace=True)
|
| 73 |
+
|
| 74 |
+
# 3. Calculate Metrics
|
| 75 |
+
acc, prec, f1 = calculate_metrics(subject_group)
|
| 76 |
+
|
| 77 |
+
results.append({
|
| 78 |
+
'Model': model_name,
|
| 79 |
+
'Accuracy': acc,
|
| 80 |
+
'Precision': prec,
|
| 81 |
+
'F1 Score': f1
|
| 82 |
+
})
|
| 83 |
+
|
| 84 |
+
print(f"{model_name:<20} | {acc:<10.4f} | {prec:<10.4f} | {f1:<10.4f}")
|
| 85 |
+
|
| 86 |
+
except Exception as e:
|
| 87 |
+
print(f"Error processing {file}: {e}")
|
| 88 |
+
|
| 89 |
+
# Save summary to CSV
|
| 90 |
+
summary_df = pd.DataFrame(results)
|
| 91 |
+
summary_path = os.path.join(prediction_dir, 'subject_level_metrics_summary.csv')
|
| 92 |
+
summary_df.to_csv(summary_path, index=False)
|
| 93 |
+
print(f"\nSummary saved to: {summary_path}")
|
| 94 |
+
|
| 95 |
+
if __name__ == "__main__":
|
| 96 |
+
main()
|
archive/classification/dataset_curate.py
ADDED
|
@@ -0,0 +1,66 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import os
|
| 2 |
+
import pandas as pd
|
| 3 |
+
from pathlib import Path
|
| 4 |
+
|
| 5 |
+
def process_ecg_data_wide_format(root_folder):
|
| 6 |
+
"""
|
| 7 |
+
Process ECG data in wide format (one row per timestamp with all leads as columns)
|
| 8 |
+
"""
|
| 9 |
+
lead_names = ['I', 'aVR', 'V1', 'V4', 'II', 'aVL', 'V2', 'V5', 'III', 'aVF', 'V3', 'V6']
|
| 10 |
+
|
| 11 |
+
all_data = []
|
| 12 |
+
base_path = Path(root_folder)
|
| 13 |
+
|
| 14 |
+
# Process both classes
|
| 15 |
+
for class_name, class_folder in [("Pre-Procedural MI", "Pre"),
|
| 16 |
+
("Post-Procedural MI", "Post")]:
|
| 17 |
+
class_path = base_path / class_folder
|
| 18 |
+
|
| 19 |
+
if class_path.exists():
|
| 20 |
+
for csv_file in class_path.glob("*.csv"):
|
| 21 |
+
try:
|
| 22 |
+
df = pd.read_csv(csv_file)
|
| 23 |
+
subject_id = csv_file.stem.replace('_hr', '')
|
| 24 |
+
|
| 25 |
+
# Add metadata columns
|
| 26 |
+
df['subject_id'] = subject_id
|
| 27 |
+
df['class'] = class_name
|
| 28 |
+
df['filename'] = csv_file.name
|
| 29 |
+
df['timestamp'] = range(len(df))
|
| 30 |
+
|
| 31 |
+
all_data.append(df)
|
| 32 |
+
print(f"Processed {csv_file.name}")
|
| 33 |
+
print(class_name)
|
| 34 |
+
|
| 35 |
+
except Exception as e:
|
| 36 |
+
print(f"Error processing {csv_file.name}: {str(e)}")
|
| 37 |
+
|
| 38 |
+
if all_data:
|
| 39 |
+
# Combine all DataFrames
|
| 40 |
+
combined_df = pd.concat(all_data, ignore_index=True)
|
| 41 |
+
|
| 42 |
+
# Reorder columns to have metadata first
|
| 43 |
+
metadata_cols = ['subject_id', 'timestamp', 'class', 'filename']
|
| 44 |
+
signal_cols = [col for col in combined_df.columns if col not in metadata_cols]
|
| 45 |
+
combined_df = combined_df[metadata_cols + signal_cols]
|
| 46 |
+
|
| 47 |
+
return combined_df
|
| 48 |
+
return None
|
| 49 |
+
|
| 50 |
+
# Usage for wide format:
|
| 51 |
+
if __name__ == "__main__":
|
| 52 |
+
folder_path = "."
|
| 53 |
+
|
| 54 |
+
if os.path.exists(folder_path):
|
| 55 |
+
print("Processing ECG data in wide format...")
|
| 56 |
+
wide_df = process_ecg_data_wide_format(folder_path)
|
| 57 |
+
|
| 58 |
+
if wide_df is not None:
|
| 59 |
+
wide_df.to_csv("combined_ecg_data_wide.csv", index=False)
|
| 60 |
+
print("Wide format data saved to combined_ecg_data_wide.csv")
|
| 61 |
+
print(f"Shape: {wide_df.shape}")
|
| 62 |
+
print(f"Columns: {list(wide_df.columns)}")
|
| 63 |
+
print("\nSample data:")
|
| 64 |
+
print(wide_df.head())
|
| 65 |
+
else:
|
| 66 |
+
print(f"Folder '{folder_path}' not found.")
|
archive/classification/feature_analysis.py
ADDED
|
@@ -0,0 +1,93 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import os
|
| 2 |
+
import pandas as pd
|
| 3 |
+
import numpy as np
|
| 4 |
+
from utils.mlp_feature_downsample import (
|
| 5 |
+
perform_comprehensive_heartbeat_analysis,
|
| 6 |
+
perform_full_signal_analysis,
|
| 7 |
+
DEFAULT_HEARTBEAT_OUTPUT_DIR,
|
| 8 |
+
DEFAULT_TARGET_FS,
|
| 9 |
+
)
|
| 10 |
+
|
| 11 |
+
if __name__ == "__main__":
|
| 12 |
+
|
| 13 |
+
# Create output directory
|
| 14 |
+
os.makedirs('output', exist_ok=True)
|
| 15 |
+
|
| 16 |
+
# Use the processed CSV from the previous step
|
| 17 |
+
combined_input_csv = 'data/ecg_dataset_processed.csv'
|
| 18 |
+
|
| 19 |
+
# Load and verify the data
|
| 20 |
+
df_combined = pd.read_csv(combined_input_csv)
|
| 21 |
+
print(f"Loaded combined dataset shape: {df_combined.shape}")
|
| 22 |
+
print(f"Unique subjects: {df_combined['subject_id'].nunique()}")
|
| 23 |
+
|
| 24 |
+
print("\n" + "="*50)
|
| 25 |
+
print("RUNNING SEGMENTED HEARTBEAT ANALYSIS (500 Hz)")
|
| 26 |
+
print("="*50)
|
| 27 |
+
|
| 28 |
+
heartbeat_results = perform_comprehensive_heartbeat_analysis(
|
| 29 |
+
input_csv=combined_input_csv,
|
| 30 |
+
output_dir=DEFAULT_HEARTBEAT_OUTPUT_DIR,
|
| 31 |
+
pca_components=50,
|
| 32 |
+
shap_samples=100,
|
| 33 |
+
shap_background_size=50,
|
| 34 |
+
hidden_layer_sizes=(100, 50),
|
| 35 |
+
max_iter=1000,
|
| 36 |
+
target_fs=500,
|
| 37 |
+
)
|
| 38 |
+
|
| 39 |
+
print("\n" + "="*50)
|
| 40 |
+
print("RUNNING FULL SIGNAL ANALYSIS (500 Hz)")
|
| 41 |
+
print("="*50)
|
| 42 |
+
|
| 43 |
+
full_signal_results = perform_full_signal_analysis(
|
| 44 |
+
input_csv=combined_input_csv,
|
| 45 |
+
output_dir='output/full_signal_analysis',
|
| 46 |
+
pca_components=50,
|
| 47 |
+
shap_samples=100,
|
| 48 |
+
shap_background_size=50,
|
| 49 |
+
hidden_layer_sizes=(100, 50),
|
| 50 |
+
max_iter=1000,
|
| 51 |
+
target_fs=500,
|
| 52 |
+
)
|
| 53 |
+
|
| 54 |
+
# Print Summary Results
|
| 55 |
+
print("\n" + "#" * 60)
|
| 56 |
+
print("FINAL SUMMARY RESULTS")
|
| 57 |
+
print("#" * 60)
|
| 58 |
+
|
| 59 |
+
if heartbeat_results is not None and 'metrics' in heartbeat_results:
|
| 60 |
+
m = heartbeat_results['metrics']
|
| 61 |
+
print("\n--- Segmented Heartbeat Analysis ---")
|
| 62 |
+
print(f" Accuracy: {m['accuracy']:.4f}")
|
| 63 |
+
print(f" Precision: {m['precision']:.4f}")
|
| 64 |
+
print(f" Recall: {m['recall']:.4f}")
|
| 65 |
+
print(f" F1-score: {m['f1']:.4f}")
|
| 66 |
+
print(f" Sensitivity: {m['sensitivity']:.4f}")
|
| 67 |
+
print(f" Specificity: {m['specificity']:.4f}")
|
| 68 |
+
print("=" * 50)
|
| 69 |
+
'''
|
| 70 |
+
############################################## Supervised Classifier Benchmarking ##################################
|
| 71 |
+
#print("\nPerforming classifier benchmarking...")
|
| 72 |
+
benchmark_results, models = perform_benchmark(combined_input_csv)
|
| 73 |
+
|
| 74 |
+
print("\nBenchmarking Results:")
|
| 75 |
+
print(benchmark_results)
|
| 76 |
+
|
| 77 |
+
# Save all results to a summary file
|
| 78 |
+
results_summary = {
|
| 79 |
+
#'kmeans_results': kmeans_results,
|
| 80 |
+
'mlp_results': mlp_results,
|
| 81 |
+
'benchmark_results': benchmark_results,
|
| 82 |
+
}
|
| 83 |
+
|
| 84 |
+
# Save summary to file
|
| 85 |
+
summary_df = pd.DataFrame([
|
| 86 |
+
#{'model': 'K-Means', 'details': str(kmeans_results)},
|
| 87 |
+
#{'model': 'MLP', 'details': str(mlp_results)},
|
| 88 |
+
{'model': 'Benchmark', 'details': str(benchmark_results)},
|
| 89 |
+
])
|
| 90 |
+
|
| 91 |
+
summary_df.to_csv('output/analysis_summary.csv', index=False)
|
| 92 |
+
print("\nAnalysis summary saved to 'output/analysis_summary.csv'")
|
| 93 |
+
'''
|
archive/classification/prepare_multiclass_data.py
ADDED
|
@@ -0,0 +1,71 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import pandas as pd
|
| 2 |
+
import os
|
| 3 |
+
|
| 4 |
+
if __name__ == "__main__":
|
| 5 |
+
# ── CONFIGURATION ──
|
| 6 |
+
# Input CSV file path
|
| 7 |
+
input_csv = 'data/combined_ecg_matrix_data_segmented.csv'
|
| 8 |
+
|
| 9 |
+
# Output CSV file path
|
| 10 |
+
output_csv = 'data/combined_ecg_matrix_data_filtered.csv'
|
| 11 |
+
|
| 12 |
+
# Option 1: Only include specific classes (set to None to use all)
|
| 13 |
+
# e.g. ['Pre', 'Post', 'Index']
|
| 14 |
+
include = None
|
| 15 |
+
|
| 16 |
+
# Option 2: Exclude specific classes (set to None to skip)
|
| 17 |
+
# e.g. ['Dimi']
|
| 18 |
+
exclude = None
|
| 19 |
+
|
| 20 |
+
# Option 3: Combine/merge classes under a new label (set to None to skip)
|
| 21 |
+
# e.g. {'Surgery': ['Pre', 'Post'], 'Other': ['Dimi', 'Discharge']}
|
| 22 |
+
combine = {'Post-Procedural MI': ['Discharge', 'Dimi', 'Post'], 'Pre-Procedural MI': ['Index', 'Pre']}
|
| 23 |
+
|
| 24 |
+
print(f"Loading data from '{input_csv}'...")
|
| 25 |
+
if not os.path.exists(input_csv):
|
| 26 |
+
print(f"Error: Input file '{input_csv}' not found.")
|
| 27 |
+
exit(1)
|
| 28 |
+
|
| 29 |
+
df = pd.read_csv(input_csv)
|
| 30 |
+
initial_rows = len(df)
|
| 31 |
+
print(f"Loaded {initial_rows} rows.")
|
| 32 |
+
|
| 33 |
+
if 'class' not in df.columns:
|
| 34 |
+
print("Error: 'class' column not found in input CSV.")
|
| 35 |
+
exit(1)
|
| 36 |
+
|
| 37 |
+
print(f"\nOriginal classes: {sorted(df['class'].dropna().unique())}")
|
| 38 |
+
for cls in sorted(df['class'].dropna().unique()):
|
| 39 |
+
print(f" {cls}: {len(df[df['class'] == cls])} rows")
|
| 40 |
+
|
| 41 |
+
# ── 1. Combine/Merge Classes ──
|
| 42 |
+
if combine:
|
| 43 |
+
print("\nApplying class combinations...")
|
| 44 |
+
for new_name, old_names in combine.items():
|
| 45 |
+
df.loc[df['class'].isin(old_names), 'class'] = new_name
|
| 46 |
+
print(f" Combined {old_names} -> '{new_name}'")
|
| 47 |
+
|
| 48 |
+
# ── 2. Include Classes ──
|
| 49 |
+
if include:
|
| 50 |
+
print(f"\nIncluding only: {include}")
|
| 51 |
+
df = df[df['class'].isin(include)]
|
| 52 |
+
|
| 53 |
+
# ── 3. Exclude Classes ──
|
| 54 |
+
if exclude:
|
| 55 |
+
print(f"\nExcluding: {exclude}")
|
| 56 |
+
df = df[~df['class'].isin(exclude)]
|
| 57 |
+
|
| 58 |
+
if len(df) == 0:
|
| 59 |
+
print("\nERROR: No data remaining after class filtering!")
|
| 60 |
+
exit(1)
|
| 61 |
+
|
| 62 |
+
print(f"\nData processing complete. Remaining rows: {len(df)} ({(len(df)/initial_rows)*100:.1f}%)")
|
| 63 |
+
|
| 64 |
+
final_classes = sorted(df['class'].dropna().unique())
|
| 65 |
+
print(f"\nFinal classes hierarchy: {final_classes}")
|
| 66 |
+
for cls in final_classes:
|
| 67 |
+
print(f" {cls}: {len(df[df['class'] == cls])} rows")
|
| 68 |
+
|
| 69 |
+
os.makedirs(os.path.dirname(output_csv), exist_ok=True)
|
| 70 |
+
df.to_csv(output_csv, index=False)
|
| 71 |
+
print(f"\nSaved filtered data to '{output_csv}'")
|
archive/classification/re_plotter.py
ADDED
|
@@ -0,0 +1,494 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import numpy as np
|
| 2 |
+
import pandas as pd
|
| 3 |
+
import matplotlib.pyplot as plt
|
| 4 |
+
import os
|
| 5 |
+
from matplotlib.gridspec import GridSpec
|
| 6 |
+
|
| 7 |
+
def replot_idealized_shap_all_leads(shap_values_csv, top_regions_csv, output_dir, n_timesteps=40, fs=500):
|
| 8 |
+
"""
|
| 9 |
+
Replot idealized heartbeat with SHAP values for all leads below it in order of importance
|
| 10 |
+
"""
|
| 11 |
+
|
| 12 |
+
# Create output directory
|
| 13 |
+
os.makedirs(output_dir, exist_ok=True)
|
| 14 |
+
|
| 15 |
+
# Load SHAP values and top regions
|
| 16 |
+
print("Loading SHAP values and top regions...")
|
| 17 |
+
shap_df = pd.read_csv(shap_values_csv)
|
| 18 |
+
top_regions_df = pd.read_csv(top_regions_csv)
|
| 19 |
+
|
| 20 |
+
print(f"SHAP DataFrame shape: {shap_df.shape}")
|
| 21 |
+
print(f"Number of timesteps inferred from columns: {len([col for col in shap_df.columns if '_t' in col]) // 12}") # 12 leads
|
| 22 |
+
|
| 23 |
+
# Calculate mean absolute SHAP values for each feature
|
| 24 |
+
mean_shap_abs = np.abs(shap_df).mean(axis=0)
|
| 25 |
+
|
| 26 |
+
# Extract unique leads and their overall importance
|
| 27 |
+
leads = []
|
| 28 |
+
lead_importance = {}
|
| 29 |
+
|
| 30 |
+
for feature_name in mean_shap_abs.index:
|
| 31 |
+
if '_t' in feature_name:
|
| 32 |
+
lead = feature_name.split('_t')[0]
|
| 33 |
+
if lead not in lead_importance:
|
| 34 |
+
lead_importance[lead] = 0
|
| 35 |
+
lead_importance[lead] += mean_shap_abs[feature_name]
|
| 36 |
+
|
| 37 |
+
# Sort leads by overall importance
|
| 38 |
+
sorted_leads = sorted(lead_importance.items(), key=lambda x: x[1], reverse=True)
|
| 39 |
+
lead_names = [lead for lead, importance in sorted_leads]
|
| 40 |
+
lead_importances = [importance for lead, importance in sorted_leads]
|
| 41 |
+
|
| 42 |
+
print("Lead importance ranking:")
|
| 43 |
+
for i, (lead, importance) in enumerate(sorted_leads):
|
| 44 |
+
print(f"{i+1}. {lead}: {importance:.4f}")
|
| 45 |
+
|
| 46 |
+
# Dynamically determine n_timesteps from the data
|
| 47 |
+
actual_timesteps = infer_timesteps_from_data(shap_df, lead_names[0])
|
| 48 |
+
print(f"Using {actual_timesteps} timesteps based on data")
|
| 49 |
+
|
| 50 |
+
# Create time points in percentage (0-100%)
|
| 51 |
+
time_points_pct = np.linspace(0, 100, actual_timesteps)
|
| 52 |
+
|
| 53 |
+
# Create main plot with idealized ECG at top and all leads below
|
| 54 |
+
create_main_comprehensive_plot(lead_names, shap_df, time_points_pct, output_dir, lead_importances, actual_timesteps)
|
| 55 |
+
|
| 56 |
+
return lead_names, lead_importances
|
| 57 |
+
|
| 58 |
+
def infer_timesteps_from_data(shap_df, sample_lead):
|
| 59 |
+
"""Infer number of timesteps from the data"""
|
| 60 |
+
lead_columns = [col for col in shap_df.columns if col.startswith(sample_lead + '_t')]
|
| 61 |
+
return len(lead_columns)
|
| 62 |
+
|
| 63 |
+
def create_main_comprehensive_plot(lead_names, shap_df, time_points_pct, output_dir, lead_importances, n_timesteps):
|
| 64 |
+
"""
|
| 65 |
+
Create main comprehensive plot with idealized ECG at top and all leads' SHAP values below
|
| 66 |
+
"""
|
| 67 |
+
n_leads = len(lead_names)
|
| 68 |
+
|
| 69 |
+
# Create larger figure with better spacing
|
| 70 |
+
fig = plt.figure(figsize=(18, 4 + n_leads * 1.5)) # Increased size
|
| 71 |
+
|
| 72 |
+
# Create grid with more space for idealized ECG and better spacing for leads
|
| 73 |
+
gs = GridSpec(n_leads + 1, 1, figure=fig, height_ratios=[3] + [1.2] * n_leads, hspace=0.4)
|
| 74 |
+
|
| 75 |
+
# Plot 1: Idealized ECG at the top
|
| 76 |
+
ax_ideal = fig.add_subplot(gs[0])
|
| 77 |
+
plot_idealized_heartbeat(ax_ideal, time_points_pct)
|
| 78 |
+
|
| 79 |
+
# Plot 2: SHAP values for each lead below (in order of importance)
|
| 80 |
+
for i, lead in enumerate(lead_names):
|
| 81 |
+
ax_shap = fig.add_subplot(gs[i + 1])
|
| 82 |
+
|
| 83 |
+
# Get SHAP values for this lead
|
| 84 |
+
lead_shap_values = extract_lead_shap_values(shap_df, lead, n_timesteps)
|
| 85 |
+
|
| 86 |
+
# Plot SHAP bars for this lead with cardiac cycle regions
|
| 87 |
+
plot_lead_shap_bars(ax_shap, time_points_pct, lead_shap_values, lead, lead_importances[i], i + 1, n_leads)
|
| 88 |
+
|
| 89 |
+
plt.tight_layout()
|
| 90 |
+
plt.savefig(f'{output_dir}/idealized_with_all_leads_shap.png',
|
| 91 |
+
dpi=300, bbox_inches='tight')
|
| 92 |
+
plt.savefig(f'{output_dir}/idealized_with_all_leads_shap.pdf',
|
| 93 |
+
bbox_inches='tight')
|
| 94 |
+
plt.close()
|
| 95 |
+
|
| 96 |
+
print(f"Main comprehensive plot saved: {output_dir}/idealized_with_all_leads_shap.png")
|
| 97 |
+
|
| 98 |
+
def extract_lead_shap_values(shap_df, lead, n_timesteps):
|
| 99 |
+
"""
|
| 100 |
+
Extract SHAP values for a specific lead across all timesteps
|
| 101 |
+
"""
|
| 102 |
+
lead_shap_values = np.zeros(n_timesteps)
|
| 103 |
+
found_features = 0
|
| 104 |
+
|
| 105 |
+
for t in range(n_timesteps):
|
| 106 |
+
feature_name = f"{lead}_t{t}"
|
| 107 |
+
if feature_name in shap_df.columns:
|
| 108 |
+
# Take mean absolute SHAP value for this time point
|
| 109 |
+
lead_shap_values[t] = np.abs(shap_df[feature_name]).mean()
|
| 110 |
+
found_features += 1
|
| 111 |
+
else:
|
| 112 |
+
print(f"Warning: Feature {feature_name} not found in SHAP DataFrame")
|
| 113 |
+
lead_shap_values[t] = 0
|
| 114 |
+
|
| 115 |
+
print(f"Lead {lead} - Found {found_features}/{n_timesteps} features, SHAP range: [{lead_shap_values.min():.6f}, {lead_shap_values.max():.6f}]")
|
| 116 |
+
return lead_shap_values
|
| 117 |
+
|
| 118 |
+
def plot_idealized_heartbeat(ax, time_points_pct):
|
| 119 |
+
"""
|
| 120 |
+
Plot accurate idealized ECG heartbeat on given axis with proper waveform characteristics
|
| 121 |
+
"""
|
| 122 |
+
# Create time array
|
| 123 |
+
t = time_points_pct
|
| 124 |
+
|
| 125 |
+
# Initialize ECG signal
|
| 126 |
+
ideal_ecg = np.zeros_like(t, dtype=float)
|
| 127 |
+
|
| 128 |
+
# Define standard ECG regions with PR segment removed and P-wave extended
|
| 129 |
+
ecg_regions = [
|
| 130 |
+
{'name': 'P Wave', 'description': 'Atrial Depolarization', 'start': 0, 'end': 25, 'center': 12.5, 'color': 'lightblue', 'alpha': 0.5},
|
| 131 |
+
{'name': 'QRS Complex', 'description': 'Ventricular Depolarization', 'start': 25, 'end': 42, 'center': 33.5, 'color': 'lightgreen', 'alpha': 0.5},
|
| 132 |
+
{'name': 'ST Segment', 'description': 'Ventricular Plateau', 'start': 42, 'end': 58, 'center': 50, 'color': 'orange', 'alpha': 0.5},
|
| 133 |
+
{'name': 'T Wave', 'description': 'Ventricular Repolarization', 'start': 58, 'end': 83, 'center': 70.5, 'color': 'violet', 'alpha': 0.5},
|
| 134 |
+
{'name': 'TP Segment', 'description': 'Ventricular Diastole', 'start': 83, 'end': 100, 'center': 91.5, 'color': 'saddlebrown', 'alpha': 0.5}
|
| 135 |
+
]
|
| 136 |
+
|
| 137 |
+
# Define precise landmark positions and values FIRST to ensure consistency
|
| 138 |
+
# We align the signal generation to these landmarks
|
| 139 |
+
landmark_data = [
|
| 140 |
+
# P wave peak - Center of P region
|
| 141 |
+
{'time': 12.5, 'label': 'P', 'value': 0.25, 'description': 'Atrial Depolarization'},
|
| 142 |
+
# Q wave nadir - Early QRS
|
| 143 |
+
{'time': 27.0, 'label': 'Q', 'value': -0.15, 'description': 'Start of Ventricular Depolarization'},
|
| 144 |
+
# R wave peak - Center of QRS
|
| 145 |
+
{'time': 33.5, 'label': 'R', 'value': 1.0, 'description': 'Ventricular Depolarization Peak'},
|
| 146 |
+
# S wave nadir - Late QRS
|
| 147 |
+
{'time': 39.0, 'label': 'S', 'value': -0.25, 'description': 'End of Ventricular Depolarization'},
|
| 148 |
+
# T wave peak - Center of T region
|
| 149 |
+
{'time': 70.5, 'label': 'T', 'value': 0.35, 'description': 'Ventricular Repolarization'}
|
| 150 |
+
]
|
| 151 |
+
|
| 152 |
+
# Helper to create Gaussian waves
|
| 153 |
+
def gaussian(x, mu, sig, amp):
|
| 154 |
+
return amp * np.exp(-np.power(x - mu, 2.) / (2 * np.power(sig, 2.)))
|
| 155 |
+
|
| 156 |
+
# Create accurate ECG waveform components using Gaussians for natural shape
|
| 157 |
+
for i, time_val in enumerate(t):
|
| 158 |
+
val = 0.0
|
| 159 |
+
|
| 160 |
+
# P Wave: centered at 12.5, wide
|
| 161 |
+
val += gaussian(time_val, 12.5, 4.0, 0.25)
|
| 162 |
+
|
| 163 |
+
# Q Wave: centered at 27.0, narrow
|
| 164 |
+
val += gaussian(time_val, 27.0, 1.2, -0.15)
|
| 165 |
+
|
| 166 |
+
# R Wave: centered at 33.5, narrow and tall
|
| 167 |
+
val += gaussian(time_val, 33.5, 1.5, 1.0)
|
| 168 |
+
|
| 169 |
+
# S Wave: centered at 39.0, narrow
|
| 170 |
+
val += gaussian(time_val, 39.0, 1.2, -0.25)
|
| 171 |
+
|
| 172 |
+
# T Wave: centered at 70.5, wide
|
| 173 |
+
val += gaussian(time_val, 70.5, 6.0, 0.35)
|
| 174 |
+
|
| 175 |
+
# Assign to array
|
| 176 |
+
ideal_ecg[i] = val
|
| 177 |
+
|
| 178 |
+
# Plot cardiac cycle regions with increased alpha
|
| 179 |
+
for region in ecg_regions:
|
| 180 |
+
ax.axvspan(region['start'], region['end'], alpha=region['alpha'], color=region['color'], zorder=0)
|
| 181 |
+
|
| 182 |
+
# Add region labels - name above and description below
|
| 183 |
+
ax.text(region['center'], -0.45, region['name'],
|
| 184 |
+
ha='center', va='center', fontsize=9, fontweight='bold',
|
| 185 |
+
bbox=dict(boxstyle="round,pad=0.2", facecolor='white', alpha=0.9, edgecolor=region['color']),
|
| 186 |
+
zorder=10)
|
| 187 |
+
ax.text(region['center'], -0.65, region['description'],
|
| 188 |
+
ha='center', va='center', fontsize=8,
|
| 189 |
+
bbox=dict(boxstyle="round,pad=0.2", facecolor='white', alpha=0.8, edgecolor=region['color']),
|
| 190 |
+
zorder=10)
|
| 191 |
+
|
| 192 |
+
# Plot idealized ECG
|
| 193 |
+
ax.plot(t, ideal_ecg, 'k-', linewidth=2, label='Idealized ECG', zorder=5)
|
| 194 |
+
ax.set_title('Idealized ECG with Standard Regions', fontsize=16, fontweight='bold', pad=20)
|
| 195 |
+
ax.set_ylabel('Normalized Amplitude', fontsize=12)
|
| 196 |
+
ax.grid(True, alpha=0.3, zorder=1)
|
| 197 |
+
|
| 198 |
+
# Plot ECG landmarks with annotations
|
| 199 |
+
for landmark in landmark_data:
|
| 200 |
+
# Calculate actual value on the curve at this time
|
| 201 |
+
# This ensures the point is exactly ON the line, even if Gaussian summation shifts it slightly
|
| 202 |
+
curve_val = 0.0
|
| 203 |
+
t_lm = landmark['time']
|
| 204 |
+
curve_val += gaussian(t_lm, 12.5, 4.0, 0.25)
|
| 205 |
+
curve_val += gaussian(t_lm, 27.0, 1.2, -0.15)
|
| 206 |
+
curve_val += gaussian(t_lm, 33.5, 1.5, 1.0)
|
| 207 |
+
curve_val += gaussian(t_lm, 39.0, 1.2, -0.25)
|
| 208 |
+
curve_val += gaussian(t_lm, 70.5, 6.0, 0.35)
|
| 209 |
+
|
| 210 |
+
ax.plot(t_lm, curve_val, 'ro', markersize=8, zorder=10)
|
| 211 |
+
|
| 212 |
+
# Position text based on whether it's a peak or trough
|
| 213 |
+
if landmark['value'] > 0:
|
| 214 |
+
text_y_offset = 0.15
|
| 215 |
+
va = 'bottom'
|
| 216 |
+
else:
|
| 217 |
+
text_y_offset = -0.15
|
| 218 |
+
va = 'top'
|
| 219 |
+
|
| 220 |
+
ax.annotate(landmark['label'],
|
| 221 |
+
(t_lm, curve_val),
|
| 222 |
+
xytext=(0, text_y_offset * 30),
|
| 223 |
+
textcoords='offset points',
|
| 224 |
+
ha='center', va=va,
|
| 225 |
+
fontweight='bold',
|
| 226 |
+
fontsize=12,
|
| 227 |
+
bbox=dict(boxstyle="round,pad=0.3", facecolor='white', alpha=0.9, edgecolor='red'),
|
| 228 |
+
zorder=15)
|
| 229 |
+
|
| 230 |
+
ax.set_ylim(-0.8, 1.5) # Increased range to accommodate annotations
|
| 231 |
+
ax.set_xlim(0, 100)
|
| 232 |
+
|
| 233 |
+
# Remove x-axis labels and ticks for the top plot
|
| 234 |
+
ax.set_xticklabels([])
|
| 235 |
+
ax.set_xticks([])
|
| 236 |
+
ax.set_xlabel('')
|
| 237 |
+
|
| 238 |
+
ax.tick_params(axis='y', which='major', labelsize=10)
|
| 239 |
+
|
| 240 |
+
# Add a legend for the waveform
|
| 241 |
+
ax.legend(loc='upper right', framealpha=0.9)
|
| 242 |
+
|
| 243 |
+
def plot_lead_shap_bars(ax, time_points_pct, shap_importance, lead, lead_importance, rank, total_leads):
|
| 244 |
+
"""
|
| 245 |
+
Plot SHAP importance bars for a single lead with cardiac cycle regions
|
| 246 |
+
"""
|
| 247 |
+
if len(shap_importance) == len(time_points_pct):
|
| 248 |
+
# Define cardiac cycle regions with PR segment removed and P-wave extended
|
| 249 |
+
ecg_regions = [
|
| 250 |
+
{'name': 'P Wave', 'start': 0, 'end': 25, 'color': 'lightblue', 'alpha': 0.5},
|
| 251 |
+
{'name': 'QRS Complex', 'start': 25, 'end': 42, 'color': 'lightgreen', 'alpha': 0.5},
|
| 252 |
+
{'name': 'ST Segment', 'start': 42, 'end': 58, 'color': 'orange', 'alpha': 0.5},
|
| 253 |
+
{'name': 'T Wave', 'start': 58, 'end': 83, 'color': 'violet', 'alpha': 0.5},
|
| 254 |
+
{'name': 'TP Segment', 'start': 83, 'end': 100, 'color': 'saddlebrown', 'alpha': 0.5}
|
| 255 |
+
]
|
| 256 |
+
|
| 257 |
+
# Plot cardiac cycle regions as background
|
| 258 |
+
for region in ecg_regions:
|
| 259 |
+
ax.axvspan(region['start'], region['end'], alpha=region['alpha'], color=region['color'], zorder=0)
|
| 260 |
+
|
| 261 |
+
bar_width = 100 / len(time_points_pct) * 0.8
|
| 262 |
+
|
| 263 |
+
# Debug info
|
| 264 |
+
max_shap = np.max(shap_importance) if len(shap_importance) > 0 else 0
|
| 265 |
+
print(f" {rank}. {lead} - Max SHAP: {max_shap:.6f}")
|
| 266 |
+
|
| 267 |
+
# Create bars with dark grey color
|
| 268 |
+
if max_shap > 0:
|
| 269 |
+
bars = ax.bar(time_points_pct, shap_importance, alpha=0.8, color='darkslategray',
|
| 270 |
+
width=bar_width, edgecolor='black', linewidth=0.5, zorder=5)
|
| 271 |
+
|
| 272 |
+
# Highlight top region
|
| 273 |
+
top_idx = np.argmax(shap_importance)
|
| 274 |
+
top_value = shap_importance[top_idx]
|
| 275 |
+
top_time = time_points_pct[top_idx]
|
| 276 |
+
|
| 277 |
+
# Highlight the top bar with red
|
| 278 |
+
bars[top_idx].set_color('red')
|
| 279 |
+
bars[top_idx].set_alpha(1.0)
|
| 280 |
+
bars[top_idx].set_edgecolor('darkred')
|
| 281 |
+
bars[top_idx].set_linewidth(1)
|
| 282 |
+
|
| 283 |
+
# Add annotation for top region
|
| 284 |
+
if top_value > max_shap * 0.1: # Only annotate if significant
|
| 285 |
+
ax.text(top_time, top_value * 0.9,
|
| 286 |
+
f'{top_time:.1f}%',
|
| 287 |
+
ha='center', va='top', fontsize=8, fontweight='bold', zorder=10,
|
| 288 |
+
bbox=dict(boxstyle="round,pad=0.2", facecolor='white', alpha=0.9))
|
| 289 |
+
else:
|
| 290 |
+
# If no SHAP values, plot zeros with different color
|
| 291 |
+
bars = ax.bar(time_points_pct, shap_importance, alpha=0.3, color='gray',
|
| 292 |
+
width=bar_width, edgecolor='darkgray', linewidth=0.5, zorder=5)
|
| 293 |
+
|
| 294 |
+
# Set labels with better formatting
|
| 295 |
+
ax.set_ylabel(f'{lead}\n({lead_importance:.3f})\nSHAP Value', fontsize=10, rotation=0,
|
| 296 |
+
ha='right', va='center', labelpad=10)
|
| 297 |
+
ax.grid(True, alpha=0.3, zorder=1)
|
| 298 |
+
ax.set_xlim(0, 100)
|
| 299 |
+
|
| 300 |
+
# Remove x-axis labels and ticks for all but bottom plot
|
| 301 |
+
if rank < total_leads:
|
| 302 |
+
ax.set_xticklabels([])
|
| 303 |
+
ax.set_xticks([])
|
| 304 |
+
else:
|
| 305 |
+
ax.set_xlabel('Cardiac Cycle (%)', fontsize=11)
|
| 306 |
+
ax.tick_params(axis='x', labelsize=9)
|
| 307 |
+
|
| 308 |
+
# Set y-axis limit based on max SHAP value with some margin
|
| 309 |
+
max_shap_val = max_shap if max_shap > 0 else 0.001
|
| 310 |
+
ax.set_ylim(0, max_shap_val * 1.3)
|
| 311 |
+
ax.tick_params(axis='y', labelsize=8)
|
| 312 |
+
|
| 313 |
+
# Remove y-axis ticks and labels (the numbers 1-12)
|
| 314 |
+
ax.set_yticklabels([])
|
| 315 |
+
ax.set_yticks([])
|
| 316 |
+
|
| 317 |
+
def create_lead_importance_table(lead_names, lead_importances, output_dir):
|
| 318 |
+
"""
|
| 319 |
+
Create a simple table showing lead importance ranking
|
| 320 |
+
"""
|
| 321 |
+
# Create importance DataFrame
|
| 322 |
+
importance_df = pd.DataFrame({
|
| 323 |
+
'Rank': range(1, len(lead_names) + 1),
|
| 324 |
+
'Lead': lead_names,
|
| 325 |
+
'Total_SHAP_Importance': lead_importances
|
| 326 |
+
})
|
| 327 |
+
|
| 328 |
+
# Save to CSV
|
| 329 |
+
importance_df.to_csv(f'{output_dir}/lead_importance_ranking.csv', index=False)
|
| 330 |
+
|
| 331 |
+
# Create simple text summary
|
| 332 |
+
with open(f'{output_dir}/lead_importance_summary.txt', 'w') as f:
|
| 333 |
+
f.write("Lead Importance Ranking (by Total SHAP Value)\n")
|
| 334 |
+
f.write("=" * 50 + "\n")
|
| 335 |
+
for i, (lead, importance) in enumerate(zip(lead_names, lead_importances)):
|
| 336 |
+
f.write(f"{i+1:2d}. {lead:4s}: {importance:.4f}\n")
|
| 337 |
+
|
| 338 |
+
print(f"Lead importance table saved: {output_dir}/lead_importance_ranking.csv")
|
| 339 |
+
return importance_df
|
| 340 |
+
|
| 341 |
+
# Function to fix other plots (avg_with_shap plots)
|
| 342 |
+
def fix_other_plots(analysis_output_dir, n_timesteps=40):
|
| 343 |
+
"""
|
| 344 |
+
Fix the other plots that are not showing SHAP values
|
| 345 |
+
"""
|
| 346 |
+
print("\nFixing other plots...")
|
| 347 |
+
|
| 348 |
+
# Load the necessary data
|
| 349 |
+
shap_values_csv = f'{analysis_output_dir}/shap/shap_values_detailed.csv'
|
| 350 |
+
top_regions_csv = f'{analysis_output_dir}/shap/top_shap_regions.csv'
|
| 351 |
+
|
| 352 |
+
if not os.path.exists(shap_values_csv) or not os.path.exists(top_regions_csv):
|
| 353 |
+
print("Required files not found. Skipping other plots.")
|
| 354 |
+
return
|
| 355 |
+
|
| 356 |
+
shap_df = pd.read_csv(shap_values_csv)
|
| 357 |
+
top_regions_df = pd.read_csv(top_regions_csv)
|
| 358 |
+
|
| 359 |
+
# Get lead names from top regions
|
| 360 |
+
lead_names = top_regions_df['feature'].unique()
|
| 361 |
+
|
| 362 |
+
# Create time points
|
| 363 |
+
time_points_pct = np.linspace(0, 100, n_timesteps)
|
| 364 |
+
|
| 365 |
+
# Recreate avg_with_shap plots for each lead
|
| 366 |
+
for lead in lead_names:
|
| 367 |
+
create_fixed_avg_shap_plot(lead, shap_df, time_points_pct, analysis_output_dir, n_timesteps)
|
| 368 |
+
|
| 369 |
+
def create_fixed_avg_shap_plot(lead, shap_df, time_points_pct, output_dir, n_timesteps):
|
| 370 |
+
"""
|
| 371 |
+
Create fixed avg_with_shap plot for a single lead
|
| 372 |
+
"""
|
| 373 |
+
# Extract SHAP values
|
| 374 |
+
lead_shap_values = extract_lead_shap_values(shap_df, lead, n_timesteps)
|
| 375 |
+
|
| 376 |
+
# Create figure
|
| 377 |
+
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(15, 8))
|
| 378 |
+
|
| 379 |
+
# Plot 1: Placeholder for average heartbeat (you would need the original data for this)
|
| 380 |
+
ax1.text(0.5, 0.5, f'Average Heartbeats - {lead}\n(Original data required)',
|
| 381 |
+
ha='center', va='center', transform=ax1.transAxes, fontsize=14)
|
| 382 |
+
ax1.set_title(f'Average Heartbeats - {lead}', fontsize=16)
|
| 383 |
+
ax1.set_ylabel('Normalized Amplitude')
|
| 384 |
+
ax1.grid(True, alpha=0.3)
|
| 385 |
+
|
| 386 |
+
# Plot 2: SHAP importance with updated regions
|
| 387 |
+
# Define cardiac cycle regions with PR segment removed and P-wave extended
|
| 388 |
+
ecg_regions = [
|
| 389 |
+
{'name': 'P Wave', 'start': 0, 'end': 25, 'color': 'lightblue', 'alpha': 0.3},
|
| 390 |
+
{'name': 'QRS Complex', 'start': 25, 'end': 42, 'color': 'lightgreen', 'alpha': 0.3},
|
| 391 |
+
{'name': 'ST Segment', 'start': 42, 'end': 58, 'color': 'orange', 'alpha': 0.3},
|
| 392 |
+
{'name': 'T Wave', 'start': 58, 'end': 83, 'color': 'violet', 'alpha': 0.3},
|
| 393 |
+
{'name': 'TP Segment', 'start': 83, 'end': 100, 'color': 'saddlebrown', 'alpha': 0.3}
|
| 394 |
+
]
|
| 395 |
+
|
| 396 |
+
# Plot cardiac cycle regions as background
|
| 397 |
+
for region in ecg_regions:
|
| 398 |
+
ax2.axvspan(region['start'], region['end'], alpha=region['alpha'], color=region['color'], zorder=0)
|
| 399 |
+
|
| 400 |
+
bar_width = 100 / len(time_points_pct) * 0.8
|
| 401 |
+
ax2.bar(time_points_pct, lead_shap_values, alpha=0.7, color='darkslategray',
|
| 402 |
+
width=bar_width, label='SHAP Importance', zorder=5)
|
| 403 |
+
ax2.set_title(f'SHAP Importance - {lead}', fontsize=16)
|
| 404 |
+
ax2.set_xlabel('Cardiac Cycle (%)')
|
| 405 |
+
ax2.set_ylabel('SHAP Value')
|
| 406 |
+
ax2.grid(True, alpha=0.3)
|
| 407 |
+
ax2.legend()
|
| 408 |
+
|
| 409 |
+
plt.tight_layout()
|
| 410 |
+
plt.savefig(f'{output_dir}/heartbeat_plots/avg_with_shap_{lead}_FIXED.png',
|
| 411 |
+
dpi=300, bbox_inches='tight')
|
| 412 |
+
plt.close()
|
| 413 |
+
|
| 414 |
+
print(f"Fixed plot saved: {output_dir}/heartbeat_plots/avg_with_shap_{lead}_FIXED.png")
|
| 415 |
+
|
| 416 |
+
# Additional function to create standalone idealized ECG plot
|
| 417 |
+
def create_standalone_idealized_ecg_plot(output_dir, n_timesteps=40):
|
| 418 |
+
"""
|
| 419 |
+
Create a standalone idealized ECG plot with the updated regions
|
| 420 |
+
"""
|
| 421 |
+
time_points_pct = np.linspace(0, 100, n_timesteps)
|
| 422 |
+
|
| 423 |
+
fig, ax = plt.subplots(figsize=(15, 6))
|
| 424 |
+
plot_idealized_heartbeat(ax, time_points_pct)
|
| 425 |
+
|
| 426 |
+
plt.tight_layout()
|
| 427 |
+
plt.savefig(f'{output_dir}/standalone_idealized_ecg.png',
|
| 428 |
+
dpi=300, bbox_inches='tight')
|
| 429 |
+
plt.savefig(f'{output_dir}/standalone_idealized_ecg.pdf',
|
| 430 |
+
bbox_inches='tight')
|
| 431 |
+
plt.close()
|
| 432 |
+
|
| 433 |
+
print(f"Standalone idealized ECG plot saved: {output_dir}/standalone_idealized_ecg.png")
|
| 434 |
+
|
| 435 |
+
# Main execution function
|
| 436 |
+
def main():
|
| 437 |
+
"""
|
| 438 |
+
Main function to run the SHAP replotting
|
| 439 |
+
"""
|
| 440 |
+
# Define paths
|
| 441 |
+
analysis_output_dir = 'results/heartbeat_analysis_500Hz_good'
|
| 442 |
+
shap_values_csv = f'{analysis_output_dir}/shap/shap_values_detailed.csv'
|
| 443 |
+
top_regions_csv = f'{analysis_output_dir}/shap/top_shap_regions.csv'
|
| 444 |
+
output_dir = f'{analysis_output_dir}/shap_replots'
|
| 445 |
+
|
| 446 |
+
# Check if files exist
|
| 447 |
+
if not os.path.exists(shap_values_csv):
|
| 448 |
+
print(f"Error: SHAP values file not found at {shap_values_csv}")
|
| 449 |
+
return
|
| 450 |
+
if not os.path.exists(top_regions_csv):
|
| 451 |
+
print(f"Error: Top regions file not found at {top_regions_csv}")
|
| 452 |
+
return
|
| 453 |
+
|
| 454 |
+
print("Starting SHAP replotting...")
|
| 455 |
+
print(f"SHAP values file: {shap_values_csv}")
|
| 456 |
+
print(f"Top regions file: {top_regions_csv}")
|
| 457 |
+
print(f"Output directory: {output_dir}")
|
| 458 |
+
|
| 459 |
+
# Replot all leads
|
| 460 |
+
lead_names, lead_importances = replot_idealized_shap_all_leads(
|
| 461 |
+
shap_values_csv, top_regions_csv, output_dir
|
| 462 |
+
)
|
| 463 |
+
|
| 464 |
+
# Create lead importance table
|
| 465 |
+
importance_df = create_lead_importance_table(lead_names, lead_importances, output_dir)
|
| 466 |
+
|
| 467 |
+
# Create standalone idealized ECG plot
|
| 468 |
+
create_standalone_idealized_ecg_plot(output_dir)
|
| 469 |
+
|
| 470 |
+
# Fix other plots
|
| 471 |
+
fix_other_plots(analysis_output_dir)
|
| 472 |
+
|
| 473 |
+
print("\n" + "="*60)
|
| 474 |
+
print("SHAP REPLOTTING COMPLETE!")
|
| 475 |
+
print("="*60)
|
| 476 |
+
print(f"Results saved to: {output_dir}")
|
| 477 |
+
print(f"Lead ranking (by importance):")
|
| 478 |
+
for i, (lead, importance) in enumerate(zip(lead_names, lead_importances)):
|
| 479 |
+
print(f" {i+1:2d}. {lead}: {importance:.4f}")
|
| 480 |
+
|
| 481 |
+
# Alternative: Function to use in your existing code
|
| 482 |
+
def replot_from_existing_analysis(analysis_output_dir):
|
| 483 |
+
"""
|
| 484 |
+
Replot from existing analysis output directory
|
| 485 |
+
"""
|
| 486 |
+
shap_values_csv = f'{analysis_output_dir}/shap/shap_values_detailed.csv'
|
| 487 |
+
top_regions_csv = f'{analysis_output_dir}/shap/top_shap_regions.csv'
|
| 488 |
+
output_dir = f'{analysis_output_dir}/shap_replots'
|
| 489 |
+
|
| 490 |
+
return replot_idealized_shap_all_leads(shap_values_csv, top_regions_csv, output_dir)
|
| 491 |
+
|
| 492 |
+
# Example usage:
|
| 493 |
+
if __name__ == "__main__":
|
| 494 |
+
main()
|
archive/classification/run_benchmarking.py
ADDED
|
@@ -0,0 +1,64 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import argparse
|
| 2 |
+
import os
|
| 3 |
+
from utils.classifier_benchmarking import perform_benchmark
|
| 4 |
+
|
| 5 |
+
def run_benchmarking(input_file, output_folder=None):
|
| 6 |
+
"""
|
| 7 |
+
Run benchmarking on the specified input file.
|
| 8 |
+
|
| 9 |
+
Args:
|
| 10 |
+
input_file: Path to input CSV
|
| 11 |
+
output_folder: Directory to save results. If None, derives from input filename.
|
| 12 |
+
"""
|
| 13 |
+
if not os.path.exists(input_file):
|
| 14 |
+
print(f"Error: Input file '{input_file}' not found.")
|
| 15 |
+
return
|
| 16 |
+
|
| 17 |
+
# Determine output folder if not specified
|
| 18 |
+
if output_folder is None:
|
| 19 |
+
base_name = os.path.splitext(os.path.basename(input_file))[0]
|
| 20 |
+
output_folder = f'output/benchmark_results/{base_name}'
|
| 21 |
+
|
| 22 |
+
print(f"\n{'='*60}")
|
| 23 |
+
print(f"STARTING BENCHMARKING")
|
| 24 |
+
print(f"Input: {input_file}")
|
| 25 |
+
print(f"Output: {output_folder}")
|
| 26 |
+
print(f"{'='*60}")
|
| 27 |
+
|
| 28 |
+
try:
|
| 29 |
+
perform_benchmark(input_file, output_dir=output_folder)
|
| 30 |
+
print(f"\nBenchmarking complete for {input_file}")
|
| 31 |
+
print(f"Results saved to: {output_folder}")
|
| 32 |
+
except Exception as e:
|
| 33 |
+
print(f"\nError running benchmark: {e}")
|
| 34 |
+
import traceback
|
| 35 |
+
traceback.print_exc()
|
| 36 |
+
|
| 37 |
+
if __name__ == "__main__":
|
| 38 |
+
parser = argparse.ArgumentParser(description="Run ECG Classifier Benchmarking")
|
| 39 |
+
parser.add_argument('--input', type=str, help='Path to input CSV file (optional)')
|
| 40 |
+
parser.add_argument('--all', action='store_true', help='Run on both standard datasets (segmented and full)')
|
| 41 |
+
|
| 42 |
+
args = parser.parse_args()
|
| 43 |
+
|
| 44 |
+
# Default files
|
| 45 |
+
segmented_file = 'data/segmented_heartbeats.csv'
|
| 46 |
+
processed_file = 'data/ecg_dataset_processed.csv'
|
| 47 |
+
|
| 48 |
+
# Define explicit output paths for standard runs
|
| 49 |
+
output_segmented = 'output/benchmark_results/segmented_heartbeats'
|
| 50 |
+
output_processed = 'output/benchmark_results/ecg_dataset_processed'
|
| 51 |
+
|
| 52 |
+
if args.input:
|
| 53 |
+
# Run on user specified file (folder derived from filename)
|
| 54 |
+
run_benchmarking(args.input)
|
| 55 |
+
elif args.all:
|
| 56 |
+
# Run on both with specific folders
|
| 57 |
+
print("Running comprehensive benchmarking on ALL datasets...")
|
| 58 |
+
run_benchmarking(segmented_file, output_folder=output_segmented)
|
| 59 |
+
run_benchmarking(processed_file, output_folder=output_processed)
|
| 60 |
+
else:
|
| 61 |
+
# Default run
|
| 62 |
+
print("No input specified. Defaulting to processed ECG data.")
|
| 63 |
+
print(f"Usage: python run_benchmarking.py --input <path> OR --all")
|
| 64 |
+
run_benchmarking(processed_file, output_folder=output_processed)
|
archive/classification/run_classification.py
ADDED
|
@@ -0,0 +1,34 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from utils.ecg_classification import perform_classification_analysis
|
| 2 |
+
import os
|
| 3 |
+
# Import the universally defined target frequency
|
| 4 |
+
from run_segmentation import TARGET_FS
|
| 5 |
+
|
| 6 |
+
if __name__ == "__main__":
|
| 7 |
+
# Define paths
|
| 8 |
+
input_csv = 'data/ecg_matrix_dataset_segmented.csv'
|
| 9 |
+
output_dir = 'output/heartbeat_analysis_matrix'
|
| 10 |
+
|
| 11 |
+
# Check if input exists
|
| 12 |
+
if not os.path.exists(input_csv):
|
| 13 |
+
print(f"Error: Input file {input_csv} not found.")
|
| 14 |
+
print("Please run run_segmentation.py first.")
|
| 15 |
+
exit(1)
|
| 16 |
+
|
| 17 |
+
print(f"Starting classification analysis...")
|
| 18 |
+
print(f"Input: {input_csv}")
|
| 19 |
+
print(f"Output Directory: {output_dir}")
|
| 20 |
+
print(f"Target Frequency: {TARGET_FS} Hz (imported from run_segmentation.py)")
|
| 21 |
+
|
| 22 |
+
# Run analysis
|
| 23 |
+
# Normalization happens inside this function
|
| 24 |
+
results = perform_classification_analysis(
|
| 25 |
+
segmented_csv=input_csv,
|
| 26 |
+
output_dir=output_dir,
|
| 27 |
+
target_fs=TARGET_FS
|
| 28 |
+
)
|
| 29 |
+
|
| 30 |
+
if results:
|
| 31 |
+
print("\nAnalysis successful!")
|
| 32 |
+
print(f"Results available in {output_dir}")
|
| 33 |
+
else:
|
| 34 |
+
print("\nAnalysis failed.")
|
archive/classification/run_inference.py
ADDED
|
@@ -0,0 +1,297 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Run inference on new ECG data using a previously saved model.
|
| 3 |
+
|
| 4 |
+
Loads the serialised model and its metadata, applies identical preprocessing
|
| 5 |
+
(per-subject normalisation, reshape to numpy3D), and outputs per-subject
|
| 6 |
+
predictions with confidence scores.
|
| 7 |
+
|
| 8 |
+
Usage:
|
| 9 |
+
# MI vs Normal
|
| 10 |
+
python run_inference.py --model mi_vs_normal --input data/new_ecg.csv
|
| 11 |
+
|
| 12 |
+
# OMI vs non-OMI
|
| 13 |
+
python run_inference.py --model omi_vs_nonomi --input data/new_ecg.csv
|
| 14 |
+
|
| 15 |
+
# Custom output path
|
| 16 |
+
python run_inference.py --model mi_vs_normal --input data/new_ecg.csv --output results/preds.csv
|
| 17 |
+
"""
|
| 18 |
+
|
| 19 |
+
import os
|
| 20 |
+
import argparse
|
| 21 |
+
import pickle
|
| 22 |
+
import json
|
| 23 |
+
import numpy as np
|
| 24 |
+
import pandas as pd
|
| 25 |
+
from tqdm.auto import tqdm
|
| 26 |
+
|
| 27 |
+
# ─── Configuration ───────────────────────────────────────────────────────────
|
| 28 |
+
|
| 29 |
+
MODELS_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../models/classifier_models'))
|
| 30 |
+
|
| 31 |
+
# ─── Model Loading ───────────────────────────────────────────────────────────
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
def load_model(use_case):
|
| 35 |
+
"""
|
| 36 |
+
Load a saved model and its metadata from the models directory.
|
| 37 |
+
|
| 38 |
+
Args:
|
| 39 |
+
use_case: One of 'mi_vs_normal' or 'omi_vs_nonomi'.
|
| 40 |
+
|
| 41 |
+
Returns:
|
| 42 |
+
(model, metadata) tuple.
|
| 43 |
+
"""
|
| 44 |
+
model_dir = os.path.join(MODELS_DIR, use_case)
|
| 45 |
+
|
| 46 |
+
# Metadata
|
| 47 |
+
metadata_path = os.path.join(model_dir, 'model_metadata.json')
|
| 48 |
+
if not os.path.exists(metadata_path):
|
| 49 |
+
raise FileNotFoundError(
|
| 50 |
+
f"Model metadata not found: {metadata_path}\n"
|
| 51 |
+
f"Have you run train_and_save_models.py first?")
|
| 52 |
+
|
| 53 |
+
with open(metadata_path, 'r') as f:
|
| 54 |
+
metadata = json.load(f)
|
| 55 |
+
|
| 56 |
+
# Model pickle
|
| 57 |
+
model_path = os.path.join(model_dir, metadata['model_file'])
|
| 58 |
+
if not os.path.exists(model_path):
|
| 59 |
+
raise FileNotFoundError(
|
| 60 |
+
f"Model file not found: {model_path}\n"
|
| 61 |
+
f"Have you run train_and_save_models.py first?")
|
| 62 |
+
|
| 63 |
+
print(f"Loading {metadata['model_name']} model for {metadata['use_case']} ...")
|
| 64 |
+
with open(model_path, 'rb') as f:
|
| 65 |
+
model = pickle.load(f)
|
| 66 |
+
|
| 67 |
+
print(f" Model loaded successfully.")
|
| 68 |
+
print(f" Trained on: {metadata['training_data']}")
|
| 69 |
+
print(f" Training date: {metadata['training_date']}")
|
| 70 |
+
print(f" Expected input: {metadata['n_features']} leads x "
|
| 71 |
+
f"{metadata['n_timesteps']} timesteps")
|
| 72 |
+
print(f" Performance metrics on held-out test set:")
|
| 73 |
+
metrics = metadata.get('test_metrics', {})
|
| 74 |
+
print(f" - Accuracy: {metrics.get('accuracy', 0.0):.4f}")
|
| 75 |
+
print(f" - Precision: {metrics.get('precision', 0.0):.4f}")
|
| 76 |
+
print(f" - Recall: {metrics.get('recall', 0.0):.4f}")
|
| 77 |
+
print(f" - F1 Score: {metrics.get('f1', 0.0):.4f}")
|
| 78 |
+
print(f" - Sensitivity: {metrics.get('sensitivity', 0.0):.4f}")
|
| 79 |
+
print(f" - Specificity: {metrics.get('specificity', 0.0):.4f}")
|
| 80 |
+
|
| 81 |
+
return model, metadata
|
| 82 |
+
|
| 83 |
+
|
| 84 |
+
# ─── Preprocessing ───────────────────────────────────────────────────────────
|
| 85 |
+
|
| 86 |
+
|
| 87 |
+
def preprocess_data(input_csv, metadata):
|
| 88 |
+
"""
|
| 89 |
+
Preprocess new ECG data for inference using the same pipeline as training.
|
| 90 |
+
|
| 91 |
+
Steps (mirrors prepare_data() in classifier_benchmarking.py):
|
| 92 |
+
1. Per-subject normalisation (divide each lead by its max absolute value)
|
| 93 |
+
2. Reshape to numpy3D (n_subjects, n_features, n_timesteps)
|
| 94 |
+
3. Truncate or zero-pad each subject to the training n_timesteps
|
| 95 |
+
|
| 96 |
+
Args:
|
| 97 |
+
input_csv: Path to the input CSV. Must contain columns
|
| 98 |
+
'subject_id', 'timestamp', and the 12 ECG lead columns.
|
| 99 |
+
metadata: The model_metadata dict loaded from JSON.
|
| 100 |
+
|
| 101 |
+
Returns:
|
| 102 |
+
(X, valid_subject_ids) where X has shape
|
| 103 |
+
(n_subjects, n_features, n_timesteps).
|
| 104 |
+
"""
|
| 105 |
+
features = metadata['features']
|
| 106 |
+
n_timesteps = metadata['n_timesteps']
|
| 107 |
+
n_features = metadata['n_features']
|
| 108 |
+
|
| 109 |
+
# ── Load ──────────────────────────────────────────────────────────────
|
| 110 |
+
print(f"\nLoading input data from {input_csv} ...")
|
| 111 |
+
df = pd.read_csv(input_csv, index_col=['subject_id', 'timestamp'])
|
| 112 |
+
|
| 113 |
+
missing = [f for f in features if f not in df.columns]
|
| 114 |
+
if missing:
|
| 115 |
+
raise ValueError(f"Missing required lead columns in input: {missing}")
|
| 116 |
+
|
| 117 |
+
# ── Normalise ─────────────────────────────────────────────────────────
|
| 118 |
+
def safe_normalize(x):
|
| 119 |
+
max_val = x.abs().max()
|
| 120 |
+
return x / max_val if max_val > 0 else x
|
| 121 |
+
|
| 122 |
+
print("Normalising features per subject ...")
|
| 123 |
+
df_X = df[features].copy()
|
| 124 |
+
for col in tqdm(df_X.columns, desc="Normalising"):
|
| 125 |
+
df_X[col] = df_X.groupby(level='subject_id')[col].transform(safe_normalize)
|
| 126 |
+
|
| 127 |
+
# ── Reshape ───────────────────────────────────────────────────────────
|
| 128 |
+
subject_ids = df_X.index.get_level_values('subject_id').unique()
|
| 129 |
+
timesteps_per_subject = df_X.groupby(level='subject_id').size()
|
| 130 |
+
|
| 131 |
+
print(f" Subjects found: {len(subject_ids)}")
|
| 132 |
+
print(f" Timesteps per subject — "
|
| 133 |
+
f"min: {timesteps_per_subject.min()}, "
|
| 134 |
+
f"max: {timesteps_per_subject.max()}")
|
| 135 |
+
print(f" Target timesteps (train): {n_timesteps}")
|
| 136 |
+
|
| 137 |
+
X_list = []
|
| 138 |
+
valid_subject_ids = []
|
| 139 |
+
skipped = 0
|
| 140 |
+
|
| 141 |
+
for subject_id in tqdm(subject_ids, desc="Reshaping subjects"):
|
| 142 |
+
try:
|
| 143 |
+
subject_data = df_X.xs(subject_id, level='subject_id')
|
| 144 |
+
actual_len = len(subject_data)
|
| 145 |
+
|
| 146 |
+
if actual_len >= n_timesteps:
|
| 147 |
+
# Truncate to training length
|
| 148 |
+
arr = subject_data.iloc[:n_timesteps].values.T
|
| 149 |
+
else:
|
| 150 |
+
# Zero-pad short subjects
|
| 151 |
+
arr = np.zeros((n_features, n_timesteps))
|
| 152 |
+
arr[:, :actual_len] = subject_data.values.T
|
| 153 |
+
print(f" [warn] Subject {subject_id}: {actual_len} timesteps "
|
| 154 |
+
f"→ padded to {n_timesteps}")
|
| 155 |
+
|
| 156 |
+
if arr.shape == (n_features, n_timesteps):
|
| 157 |
+
X_list.append(arr)
|
| 158 |
+
valid_subject_ids.append(subject_id)
|
| 159 |
+
else:
|
| 160 |
+
skipped += 1
|
| 161 |
+
except Exception as e:
|
| 162 |
+
print(f" [warn] Skipped subject {subject_id}: {e}")
|
| 163 |
+
skipped += 1
|
| 164 |
+
|
| 165 |
+
if not X_list:
|
| 166 |
+
raise ValueError("No valid subjects found in the input data!")
|
| 167 |
+
|
| 168 |
+
X = np.array(X_list)
|
| 169 |
+
print(f"\n Preprocessed shape: {X.shape} "
|
| 170 |
+
f"({len(valid_subject_ids)} subjects, {skipped} skipped)")
|
| 171 |
+
|
| 172 |
+
return X, valid_subject_ids
|
| 173 |
+
|
| 174 |
+
|
| 175 |
+
# ─── Inference ───────────────────────────────────────────────────────────────
|
| 176 |
+
|
| 177 |
+
|
| 178 |
+
def run_predictions(model, X, metadata, subject_ids):
|
| 179 |
+
"""
|
| 180 |
+
Run model predictions and package them into a DataFrame.
|
| 181 |
+
|
| 182 |
+
Returns:
|
| 183 |
+
pd.DataFrame with columns:
|
| 184 |
+
subject_id, predicted_class, confidence,
|
| 185 |
+
probability_<class0>, probability_<class1>
|
| 186 |
+
"""
|
| 187 |
+
print(f"\nRunning inference on {len(subject_ids)} subjects ...")
|
| 188 |
+
|
| 189 |
+
# Predict labels and probabilities in threading backend to bypass Numba read-only array issue
|
| 190 |
+
from joblib import parallel_backend
|
| 191 |
+
with parallel_backend('threading', n_jobs=-1):
|
| 192 |
+
y_pred = model.predict(X)
|
| 193 |
+
|
| 194 |
+
# Predict probabilities (Arsenal & Rocket both support this)
|
| 195 |
+
has_proba = False
|
| 196 |
+
y_proba = None
|
| 197 |
+
try:
|
| 198 |
+
y_proba = model.predict_proba(X)
|
| 199 |
+
has_proba = True
|
| 200 |
+
except (AttributeError, NotImplementedError):
|
| 201 |
+
pass
|
| 202 |
+
|
| 203 |
+
# ── Build output DataFrame ────────────────────────────────────────────
|
| 204 |
+
results = {
|
| 205 |
+
'subject_id': subject_ids,
|
| 206 |
+
'predicted_class': y_pred,
|
| 207 |
+
}
|
| 208 |
+
|
| 209 |
+
if has_proba:
|
| 210 |
+
class_labels = list(model.classes_)
|
| 211 |
+
for i, cls in enumerate(class_labels):
|
| 212 |
+
results[f'probability_{cls}'] = y_proba[:, i]
|
| 213 |
+
|
| 214 |
+
# Confidence = probability of the predicted class
|
| 215 |
+
confidence = []
|
| 216 |
+
for j in range(len(y_pred)):
|
| 217 |
+
pred_idx = class_labels.index(y_pred[j])
|
| 218 |
+
confidence.append(float(y_proba[j, pred_idx]))
|
| 219 |
+
results['confidence'] = confidence
|
| 220 |
+
|
| 221 |
+
return pd.DataFrame(results)
|
| 222 |
+
|
| 223 |
+
|
| 224 |
+
# ─── Entry Point ──────────────────────────────────────────────────────────────
|
| 225 |
+
|
| 226 |
+
|
| 227 |
+
def main():
|
| 228 |
+
parser = argparse.ArgumentParser(
|
| 229 |
+
description='Run ECG classification inference with a saved model',
|
| 230 |
+
formatter_class=argparse.RawDescriptionHelpFormatter,
|
| 231 |
+
epilog="""
|
| 232 |
+
Examples:
|
| 233 |
+
python run_inference.py --model mi_vs_normal_processed --input data/ptb_xl/ecg_dataset_processed.csv
|
| 234 |
+
python run_inference.py --model mi_vs_normal_segmented --input data/ptb_xl/segmented_heartbeats.csv
|
| 235 |
+
python run_inference.py --model omi_vs_nonomi --input data/ecg_matrix_omi_segmented_50_150_90.csv
|
| 236 |
+
python run_inference.py --model ecg_surgery --input data/ecg_surgery_segmented_50_150_70.csv
|
| 237 |
+
""")
|
| 238 |
+
|
| 239 |
+
parser.add_argument('--model', type=str, required=True,
|
| 240 |
+
choices=['mi_vs_normal_processed', 'mi_vs_normal_segmented', 'omi_vs_nonomi', 'ecg_surgery'],
|
| 241 |
+
help='Which saved model to use')
|
| 242 |
+
parser.add_argument('--input', type=str, required=True,
|
| 243 |
+
help='Path to input CSV (subject_id, timestamp, 12 leads)')
|
| 244 |
+
parser.add_argument('--output', type=str, default=None,
|
| 245 |
+
help='Path to save predictions CSV '
|
| 246 |
+
'(default: results/<model>_predictions.csv)')
|
| 247 |
+
|
| 248 |
+
args = parser.parse_args()
|
| 249 |
+
|
| 250 |
+
# Default output
|
| 251 |
+
if args.output is None:
|
| 252 |
+
os.makedirs('results', exist_ok=True)
|
| 253 |
+
args.output = f'results/{args.model}_predictions.csv'
|
| 254 |
+
|
| 255 |
+
print('=' * 60)
|
| 256 |
+
print(' ECG CLASSIFICATION - INFERENCE')
|
| 257 |
+
print('=' * 60)
|
| 258 |
+
|
| 259 |
+
# ── Load model ────────────────────────────────────────────────────────
|
| 260 |
+
model, metadata = load_model(args.model)
|
| 261 |
+
|
| 262 |
+
# ── Preprocess ────────────────────────────────────────────────────────
|
| 263 |
+
X, subject_ids = preprocess_data(args.input, metadata)
|
| 264 |
+
|
| 265 |
+
# ── Predict ───────────────────────────────────────────────────────────
|
| 266 |
+
results_df = run_predictions(model, X, metadata, subject_ids)
|
| 267 |
+
|
| 268 |
+
# ── Save ──────────────────────────────────────────────────────────────
|
| 269 |
+
out_dir = os.path.dirname(args.output)
|
| 270 |
+
if out_dir:
|
| 271 |
+
os.makedirs(out_dir, exist_ok=True)
|
| 272 |
+
results_df.to_csv(args.output, index=False)
|
| 273 |
+
|
| 274 |
+
# ── Summary ───────────────────────────────────────────────────────────
|
| 275 |
+
print(f"\n{'=' * 60}")
|
| 276 |
+
print(f" INFERENCE RESULTS")
|
| 277 |
+
print(f"{'=' * 60}")
|
| 278 |
+
print(f" Model: {metadata['model_name']} ({metadata['use_case']})")
|
| 279 |
+
print(f" Total subjects: {len(results_df)}")
|
| 280 |
+
print(f"\n Predicted class distribution:")
|
| 281 |
+
for cls, count in results_df['predicted_class'].value_counts().items():
|
| 282 |
+
print(f" {cls}: {count}")
|
| 283 |
+
|
| 284 |
+
if 'confidence' in results_df.columns:
|
| 285 |
+
conf = results_df['confidence']
|
| 286 |
+
print(f"\n Confidence statistics:")
|
| 287 |
+
print(f" Mean: {conf.mean():.4f}")
|
| 288 |
+
print(f" Median: {conf.median():.4f}")
|
| 289 |
+
print(f" Min: {conf.min():.4f}")
|
| 290 |
+
print(f" Max: {conf.max():.4f}")
|
| 291 |
+
|
| 292 |
+
print(f"\n Results saved to: {args.output}")
|
| 293 |
+
print(f"{'=' * 60}")
|
| 294 |
+
|
| 295 |
+
|
| 296 |
+
if __name__ == '__main__':
|
| 297 |
+
main()
|
archive/classification/run_lead_importance_test.py
ADDED
|
@@ -0,0 +1,128 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import numpy as np
|
| 2 |
+
import pandas as pd
|
| 3 |
+
import matplotlib.pyplot as plt
|
| 4 |
+
import os
|
| 5 |
+
from sklearn.model_selection import train_test_split
|
| 6 |
+
from sklearn.neural_network import MLPClassifier
|
| 7 |
+
from sklearn.preprocessing import StandardScaler
|
| 8 |
+
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
|
| 9 |
+
from tqdm.auto import tqdm
|
| 10 |
+
|
| 11 |
+
# Import common configuration
|
| 12 |
+
try:
|
| 13 |
+
from run_segmentation import TARGET_FS
|
| 14 |
+
except ImportError:
|
| 15 |
+
TARGET_FS = 500 # Default fallback
|
| 16 |
+
|
| 17 |
+
def run_lead_importance_test(input_csv, output_dir):
|
| 18 |
+
"""
|
| 19 |
+
Test the accuracy of identifying MI using each ECG lead individually.
|
| 20 |
+
"""
|
| 21 |
+
os.makedirs(output_dir, exist_ok=True)
|
| 22 |
+
|
| 23 |
+
# Load segmented data
|
| 24 |
+
print(f"Loading data from {input_csv}...")
|
| 25 |
+
df_segmented = pd.read_csv(input_csv, index_col=['subject_id', 'timestamp'])
|
| 26 |
+
|
| 27 |
+
if df_segmented.empty:
|
| 28 |
+
print("Error: Input data is empty.")
|
| 29 |
+
return
|
| 30 |
+
|
| 31 |
+
# Define features/leads
|
| 32 |
+
leads = ['I', 'aVR', 'V1', 'V4', 'II', 'aVL', 'V2', 'V5', 'III', 'aVF', 'V3', 'V6']
|
| 33 |
+
label_map = {'Post-Procedural MI': 0, 'Pre-Procedural MI': 1}
|
| 34 |
+
|
| 35 |
+
results = []
|
| 36 |
+
|
| 37 |
+
print(f"\nEvaluating {len(leads)} leads individually at {TARGET_FS} Hz...")
|
| 38 |
+
|
| 39 |
+
for lead in leads:
|
| 40 |
+
print(f"\n>>> Testing Lead {lead}...")
|
| 41 |
+
|
| 42 |
+
# 1. Prepare data for this lead
|
| 43 |
+
df_lead = df_segmented[[lead, 'class']].copy()
|
| 44 |
+
|
| 45 |
+
# Normalize per heartbeat
|
| 46 |
+
df_lead[lead] = df_lead.groupby(level='subject_id')[lead].transform(
|
| 47 |
+
lambda x: x / x.abs().max() if x.abs().max() > 0 else x
|
| 48 |
+
)
|
| 49 |
+
|
| 50 |
+
# Reshape to (heartbeats, timesteps)
|
| 51 |
+
beat_ids = df_lead.index.get_level_values('subject_id').unique()
|
| 52 |
+
timesteps_per_beat = df_lead.groupby(level='subject_id').size().min()
|
| 53 |
+
|
| 54 |
+
X_list = []
|
| 55 |
+
y_list = []
|
| 56 |
+
|
| 57 |
+
for beat_id in beat_ids:
|
| 58 |
+
beat_data = df_lead.xs(beat_id, level='subject_id')
|
| 59 |
+
if len(beat_data) >= timesteps_per_beat:
|
| 60 |
+
X_list.append(beat_data[lead].values[:timesteps_per_beat])
|
| 61 |
+
y_list.append(label_map[beat_data['class'].iloc[0]])
|
| 62 |
+
|
| 63 |
+
X = np.array(X_list)
|
| 64 |
+
y = np.array(y_list)
|
| 65 |
+
|
| 66 |
+
# 2. Train and Evaluate
|
| 67 |
+
X_train, X_test, y_train, y_test = train_test_split(
|
| 68 |
+
X, y, test_size=0.25, random_state=42, stratify=y
|
| 69 |
+
)
|
| 70 |
+
|
| 71 |
+
scaler = StandardScaler()
|
| 72 |
+
X_train_scaled = scaler.fit_transform(X_train)
|
| 73 |
+
X_test_scaled = scaler.transform(X_test)
|
| 74 |
+
|
| 75 |
+
clf = MLPClassifier(
|
| 76 |
+
hidden_layer_sizes=(100,),
|
| 77 |
+
max_iter=1000,
|
| 78 |
+
random_state=42
|
| 79 |
+
)
|
| 80 |
+
clf.fit(X_train_scaled, y_train)
|
| 81 |
+
|
| 82 |
+
y_pred = clf.predict(X_test_scaled)
|
| 83 |
+
|
| 84 |
+
# Collect metrics
|
| 85 |
+
metrics = {
|
| 86 |
+
'Lead': lead,
|
| 87 |
+
'Accuracy': accuracy_score(y_test, y_pred),
|
| 88 |
+
'Precision': precision_score(y_test, y_pred, zero_division=0),
|
| 89 |
+
'Recall': recall_score(y_test, y_pred, zero_division=0),
|
| 90 |
+
'F1_Score': f1_score(y_test, y_pred, zero_division=0)
|
| 91 |
+
}
|
| 92 |
+
results.append(metrics)
|
| 93 |
+
print(f" Accuracy: {metrics['Accuracy']:.4f}")
|
| 94 |
+
|
| 95 |
+
# 3. Save and Visualize Results
|
| 96 |
+
results_df = pd.DataFrame(results).sort_values(by='Accuracy', ascending=False)
|
| 97 |
+
results_df.to_csv(f'{output_dir}/lead_accuracies.csv', index=False)
|
| 98 |
+
|
| 99 |
+
print("\n" + "="*40)
|
| 100 |
+
print("LEAD IMPORTANCE TEST RESULTS")
|
| 101 |
+
print("="*40)
|
| 102 |
+
print(results_df.to_string(index=False))
|
| 103 |
+
|
| 104 |
+
# Plot results
|
| 105 |
+
plt.figure(figsize=(12, 6))
|
| 106 |
+
plt.bar(results_df['Lead'], results_df['Accuracy'], color='skyblue', edgecolor='navy')
|
| 107 |
+
plt.axhline(y=0.5, color='red', linestyle='--', label='Baseline (Random)')
|
| 108 |
+
plt.title(f'MLP Accuracy per Single ECG Lead ({TARGET_FS} Hz)', fontsize=15, fontweight='bold')
|
| 109 |
+
plt.xlabel('ECG Lead', fontsize=12)
|
| 110 |
+
plt.ylabel('Accuracy', fontsize=12)
|
| 111 |
+
plt.ylim(0, 1.0)
|
| 112 |
+
plt.grid(axis='y', alpha=0.3)
|
| 113 |
+
plt.legend()
|
| 114 |
+
plt.tight_layout()
|
| 115 |
+
plt.savefig(f'{output_dir}/lead_accuracies_plot.png', dpi=300)
|
| 116 |
+
plt.close()
|
| 117 |
+
|
| 118 |
+
print(f"\nResults saved to {output_dir}")
|
| 119 |
+
|
| 120 |
+
if __name__ == "__main__":
|
| 121 |
+
#input_csv = 'data/segmented_heartbeats.csv'
|
| 122 |
+
input_csv = 'data/ecg_matrix_dataset_segmented.csv'
|
| 123 |
+
output_dir = 'output/lead_importance_test'
|
| 124 |
+
|
| 125 |
+
if not os.path.exists(input_csv):
|
| 126 |
+
print(f"Error: {input_csv} not found. Run segmentation first.")
|
| 127 |
+
else:
|
| 128 |
+
run_lead_importance_test(input_csv, output_dir)
|
archive/classification/train_and_save_models.py
ADDED
|
@@ -0,0 +1,427 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import os
|
| 2 |
+
import argparse
|
| 3 |
+
import pickle
|
| 4 |
+
import json
|
| 5 |
+
import time
|
| 6 |
+
import numpy as np
|
| 7 |
+
import pandas as pd
|
| 8 |
+
from sklearn.model_selection import train_test_split
|
| 9 |
+
from imblearn.combine import SMOTETomek
|
| 10 |
+
from sklearn.metrics import (accuracy_score, precision_score,
|
| 11 |
+
recall_score, f1_score, confusion_matrix)
|
| 12 |
+
from sktime.classification.kernel_based import Arsenal, RocketClassifier
|
| 13 |
+
from sktime.classification.deep_learning.inceptiontime import InceptionTimeClassifier
|
| 14 |
+
from tqdm.auto import tqdm
|
| 15 |
+
import warnings
|
| 16 |
+
warnings.filterwarnings('ignore')
|
| 17 |
+
|
| 18 |
+
|
| 19 |
+
# ─── Configuration ───────────────────────────────────────────────────────────
|
| 20 |
+
|
| 21 |
+
MODELS_DIR = 'models'
|
| 22 |
+
|
| 23 |
+
FEATURES = ['I', 'aVR', 'V1', 'V4', 'II', 'aVL', 'V2', 'V5',
|
| 24 |
+
'III', 'aVF', 'V3', 'V6']
|
| 25 |
+
|
| 26 |
+
USE_CASES = {
|
| 27 |
+
'mi_vs_normal_processed': {
|
| 28 |
+
'name': 'MI vs Normal (Processed)',
|
| 29 |
+
'input_csv': 'data/ptb_xl/ecg_dataset_processed.csv',
|
| 30 |
+
'model_class': Arsenal,
|
| 31 |
+
'model_name': 'Arsenal',
|
| 32 |
+
'model_file': 'arsenal_processed_model.pkl',
|
| 33 |
+
'positive_class': 'MYOCARDIAL_INFARCTION',
|
| 34 |
+
'negative_class': 'NORMAL',
|
| 35 |
+
'hyperparameters': {
|
| 36 |
+
'num_kernels': 10000,
|
| 37 |
+
'random_state': 42,
|
| 38 |
+
'n_jobs': 1,
|
| 39 |
+
},
|
| 40 |
+
},
|
| 41 |
+
'mi_vs_normal_segmented': {
|
| 42 |
+
'name': 'MI vs Normal (Segmented)',
|
| 43 |
+
'input_csv': 'data/ptb_xl/segmented_heartbeats.csv',
|
| 44 |
+
'model_class': Arsenal,
|
| 45 |
+
'model_name': 'Arsenal',
|
| 46 |
+
'model_file': 'arsenal_segmented_model.pkl',
|
| 47 |
+
'positive_class': 'MYOCARDIAL_INFARCTION',
|
| 48 |
+
'negative_class': 'NORMAL',
|
| 49 |
+
'hyperparameters': {
|
| 50 |
+
'num_kernels': 20000,
|
| 51 |
+
'random_state': 42,
|
| 52 |
+
'n_jobs': 1,
|
| 53 |
+
},
|
| 54 |
+
},
|
| 55 |
+
'omi_vs_nonomi': {
|
| 56 |
+
'name': 'OMI vs non-OMI',
|
| 57 |
+
'input_csv': 'data/ecg_matrix_omi_segmented_50_150_90.csv',
|
| 58 |
+
'model_class': RocketClassifier,
|
| 59 |
+
'model_name': 'Rocket',
|
| 60 |
+
'model_file': 'rocket_model.pkl',
|
| 61 |
+
'positive_class': 'OMI',
|
| 62 |
+
'negative_class': 'non-OMI',
|
| 63 |
+
'hyperparameters': {
|
| 64 |
+
'num_kernels': 50000,
|
| 65 |
+
'random_state': 42,
|
| 66 |
+
'n_jobs': -1,
|
| 67 |
+
},
|
| 68 |
+
},
|
| 69 |
+
'ecg_surgery': {
|
| 70 |
+
'name': 'Pre vs Post Procedural MI (Surgery)',
|
| 71 |
+
'input_csv': 'data/ecg_surgery_segmented_50_150_70.csv',
|
| 72 |
+
'model_class': InceptionTimeClassifier,
|
| 73 |
+
'model_name': 'InceptionTime',
|
| 74 |
+
'model_file': 'inceptiontime_model.pkl',
|
| 75 |
+
'positive_class': 'pre-procedural MI',
|
| 76 |
+
'negative_class': 'post-procedural MI',
|
| 77 |
+
'hyperparameters': {
|
| 78 |
+
'n_epochs': 200,
|
| 79 |
+
'batch_size': 64,
|
| 80 |
+
'random_state': 42,
|
| 81 |
+
'verbose': True,
|
| 82 |
+
},
|
| 83 |
+
},
|
| 84 |
+
}
|
| 85 |
+
|
| 86 |
+
# ─── Data Preprocessing ──────────────────────────────────────────────────────
|
| 87 |
+
|
| 88 |
+
def prepare_data(df):
|
| 89 |
+
"""Prepare data for benchmarking in sktime compatible format (numpy3D)"""
|
| 90 |
+
features = ['I', 'aVR', 'V1', 'V4', 'II', 'aVL', 'V2', 'V5', 'III', 'aVF', 'V3', 'V6']
|
| 91 |
+
|
| 92 |
+
def safe_normalize(x):
|
| 93 |
+
max_val = x.abs().max()
|
| 94 |
+
# Prevent division by zero which yields NaNs
|
| 95 |
+
return x / max_val if max_val > 0 else x
|
| 96 |
+
|
| 97 |
+
print(" Normalizing features...")
|
| 98 |
+
# Normalize features within each subject
|
| 99 |
+
df_X = df[features].copy()
|
| 100 |
+
for col in tqdm(df_X.columns, desc="Normalizing features"):
|
| 101 |
+
df_X[col] = df_X.groupby(level='subject_id')[col].transform(safe_normalize)
|
| 102 |
+
|
| 103 |
+
# Get unique subjects
|
| 104 |
+
subject_ids = df_X.index.get_level_values('subject_id').unique()
|
| 105 |
+
|
| 106 |
+
# Check timesteps per subject
|
| 107 |
+
timesteps_per_subject = df_X.groupby(level='subject_id').size()
|
| 108 |
+
print(f" Timesteps per subject stats:")
|
| 109 |
+
print(f" Min: {timesteps_per_subject.min()}, Max: {timesteps_per_subject.max()}, Mean: {timesteps_per_subject.mean():.1f}")
|
| 110 |
+
|
| 111 |
+
# Use minimum timesteps to ensure consistent shape
|
| 112 |
+
n_timesteps = timesteps_per_subject.min()
|
| 113 |
+
n_features = len(features)
|
| 114 |
+
|
| 115 |
+
print(f" Using {n_timesteps} timesteps per subject (minimum)")
|
| 116 |
+
|
| 117 |
+
# Manually reshape to numpy3D format: (n_instances, n_features, n_timesteps)
|
| 118 |
+
print(" Reshaping data to numpy3D format...")
|
| 119 |
+
X_list = []
|
| 120 |
+
y_list = []
|
| 121 |
+
valid_subject_ids = []
|
| 122 |
+
|
| 123 |
+
for subject_id in tqdm(subject_ids, desc="Processing subjects"):
|
| 124 |
+
try:
|
| 125 |
+
subject_data = df_X.xs(subject_id, level='subject_id')
|
| 126 |
+
|
| 127 |
+
# Ensure we have enough timesteps
|
| 128 |
+
if len(subject_data) >= n_timesteps:
|
| 129 |
+
# Take first n_timesteps and transpose to (n_features, n_timesteps)
|
| 130 |
+
subject_array = subject_data.iloc[:n_timesteps].values.T
|
| 131 |
+
|
| 132 |
+
# Check shape
|
| 133 |
+
if subject_array.shape == (n_features, n_timesteps):
|
| 134 |
+
X_list.append(subject_array)
|
| 135 |
+
|
| 136 |
+
# Get label for this subject
|
| 137 |
+
subject_label = df.loc[subject_id].iloc[0]['class']
|
| 138 |
+
y_list.append(subject_label)
|
| 139 |
+
valid_subject_ids.append(subject_id)
|
| 140 |
+
else:
|
| 141 |
+
continue
|
| 142 |
+
|
| 143 |
+
except Exception as e:
|
| 144 |
+
continue
|
| 145 |
+
|
| 146 |
+
# Convert to numpy arrays
|
| 147 |
+
X = np.array(X_list) # Shape: (n_instances, n_features, n_timesteps)
|
| 148 |
+
y = pd.Series(y_list, index=valid_subject_ids)
|
| 149 |
+
|
| 150 |
+
print(f" Final data shape: {X.shape}")
|
| 151 |
+
print(f" Class distribution:")
|
| 152 |
+
print(y.value_counts())
|
| 153 |
+
|
| 154 |
+
# Split into train/test
|
| 155 |
+
print(" Splitting data into train/test sets...")
|
| 156 |
+
X_train, X_test, y_train, y_test = train_test_split(
|
| 157 |
+
X, y, test_size=0.25, random_state=42, stratify=y
|
| 158 |
+
)
|
| 159 |
+
|
| 160 |
+
# --- SMOTE-TOMEK INJECTION ---
|
| 161 |
+
print(" Executing SMOTE-Tomek to isolate and synthetically balance training geometries...")
|
| 162 |
+
try:
|
| 163 |
+
# Flatten 3D arrays to 2D for SMOTETomek: (N, features, timesteps) -> (N, features * timesteps)
|
| 164 |
+
n_train, n_feat, n_time = X_train.shape
|
| 165 |
+
X_train_flat = X_train.reshape(n_train, n_feat * n_time)
|
| 166 |
+
|
| 167 |
+
# Balance classes
|
| 168 |
+
smote_tomek = SMOTETomek(random_state=42)
|
| 169 |
+
X_resampled_flat, y_resampled = smote_tomek.fit_resample(X_train_flat, y_train)
|
| 170 |
+
|
| 171 |
+
# Deflate back out to precise 3D format for Sktime models
|
| 172 |
+
n_new = X_resampled_flat.shape[0]
|
| 173 |
+
X_train = X_resampled_flat.reshape(n_new, n_feat, n_time)
|
| 174 |
+
y_train = y_resampled
|
| 175 |
+
|
| 176 |
+
print(f" -> SMOTE Success! Training geometries augmented from {n_train} to {n_new} samples.")
|
| 177 |
+
except Exception as e:
|
| 178 |
+
print(f" [!] SMOTETomek failed (are you missing imbalanced-learn?): {e}")
|
| 179 |
+
# -----------------------------
|
| 180 |
+
|
| 181 |
+
print(f" Final Training set: {X_train.shape}")
|
| 182 |
+
print(f" Final Test set: {X_test.shape}")
|
| 183 |
+
print(f" y_train balanced distribution: {y_train.value_counts().to_dict()}")
|
| 184 |
+
print(f" y_test strictly isolated distribution: {y_test.value_counts().to_dict()}")
|
| 185 |
+
|
| 186 |
+
return X_train, X_test, y_train, y_test
|
| 187 |
+
|
| 188 |
+
|
| 189 |
+
# ─── Training Logic ──────────────────────────────────────────────────────────
|
| 190 |
+
|
| 191 |
+
|
| 192 |
+
def train_and_save(use_case_key, config):
|
| 193 |
+
"""
|
| 194 |
+
Train a single model for the given use case, evaluate it, and persist
|
| 195 |
+
both the serialised model and a metadata JSON to disk.
|
| 196 |
+
|
| 197 |
+
Returns True on success, False on failure.
|
| 198 |
+
"""
|
| 199 |
+
|
| 200 |
+
print(f"\n{'=' * 60}")
|
| 201 |
+
print(f" TRAINING: {config['name']}")
|
| 202 |
+
print(f" Model: {config['model_name']}")
|
| 203 |
+
print(f"{'=' * 60}")
|
| 204 |
+
|
| 205 |
+
input_csv = config['input_csv']
|
| 206 |
+
if not os.path.exists(input_csv):
|
| 207 |
+
print(f"[ERROR] Input file not found: {input_csv}")
|
| 208 |
+
return False
|
| 209 |
+
|
| 210 |
+
# ── Load data ──────────────────────────────────────────────────────────
|
| 211 |
+
print(f"\nLoading data from {input_csv} ...")
|
| 212 |
+
df = pd.read_csv(input_csv, index_col=['subject_id', 'timestamp'])
|
| 213 |
+
|
| 214 |
+
# Map labels to Pre/Post Procedural MI if it's the surgery dataset
|
| 215 |
+
if use_case_key == 'ecg_surgery':
|
| 216 |
+
print(" Mapping class labels to 'pre-procedural MI' and 'post-procedural MI' ...")
|
| 217 |
+
df['class'] = df['class'].replace({
|
| 218 |
+
'Pre-Surgery MI': 'pre-procedural MI',
|
| 219 |
+
'Post-Surgery MI': 'post-procedural MI'
|
| 220 |
+
})
|
| 221 |
+
|
| 222 |
+
print(f" DataFrame shape: {df.shape}")
|
| 223 |
+
|
| 224 |
+
subject_class = df.groupby(level='subject_id').first()['class']
|
| 225 |
+
print(f" Class distribution (subjects):\n{subject_class.value_counts().to_string()}")
|
| 226 |
+
|
| 227 |
+
# ── Prepare data (normalize → reshape → split → SMOTE-Tomek) ──────────
|
| 228 |
+
print("\nPreparing data (normalize, reshape, split, SMOTE-Tomek) ...")
|
| 229 |
+
X_train, X_test, y_train, y_test = prepare_data(df)
|
| 230 |
+
|
| 231 |
+
n_timesteps = X_train.shape[2]
|
| 232 |
+
n_features = X_train.shape[1]
|
| 233 |
+
print(f" n_timesteps = {n_timesteps}")
|
| 234 |
+
print(f" n_features = {n_features}")
|
| 235 |
+
print(f" X_train shape: {X_train.shape} | X_test shape: {X_test.shape}")
|
| 236 |
+
|
| 237 |
+
# ── Initialise & train model ──────────────────────────────────────────
|
| 238 |
+
model = config['model_class'](**config['hyperparameters'])
|
| 239 |
+
|
| 240 |
+
params_str = ""
|
| 241 |
+
if 'num_kernels' in config['hyperparameters']:
|
| 242 |
+
params_str = f" (num_kernels={config['hyperparameters']['num_kernels']})"
|
| 243 |
+
elif 'n_epochs' in config['hyperparameters']:
|
| 244 |
+
params_str = f" (n_epochs={config['hyperparameters']['n_epochs']})"
|
| 245 |
+
if 'batch_size' in config['hyperparameters']:
|
| 246 |
+
params_str += f" (batch_size={config['hyperparameters']['batch_size']})"
|
| 247 |
+
|
| 248 |
+
print(f"\nTraining {config['model_name']}{params_str} ...")
|
| 249 |
+
train_start = time.time()
|
| 250 |
+
|
| 251 |
+
from joblib import parallel_backend
|
| 252 |
+
with parallel_backend('threading'):
|
| 253 |
+
model.fit(X_train, y_train)
|
| 254 |
+
train_time = time.time() - train_start
|
| 255 |
+
print(f" Training completed in {train_time:.1f}s")
|
| 256 |
+
|
| 257 |
+
# ── Evaluate on test set ──────────────────────────────────────────────
|
| 258 |
+
print("\nEvaluating on held-out test set ...")
|
| 259 |
+
y_pred = model.predict(X_test)
|
| 260 |
+
|
| 261 |
+
positive = config['positive_class']
|
| 262 |
+
y_test_num = (y_test == positive).astype(int)
|
| 263 |
+
y_pred_num = (y_pred == positive).astype(int)
|
| 264 |
+
|
| 265 |
+
accuracy = accuracy_score(y_test_num, y_pred_num)
|
| 266 |
+
precision = precision_score(y_test_num, y_pred_num, zero_division=0)
|
| 267 |
+
recall = recall_score(y_test_num, y_pred_num, zero_division=0)
|
| 268 |
+
f1 = f1_score(y_test_num, y_pred_num, zero_division=0)
|
| 269 |
+
|
| 270 |
+
cm = confusion_matrix(y_test_num, y_pred_num)
|
| 271 |
+
if cm.shape == (2, 2):
|
| 272 |
+
tn, fp, fn, tp = cm.flatten()
|
| 273 |
+
sensitivity = tp / (tp + fn) if (tp + fn) > 0 else 0
|
| 274 |
+
specificity = tn / (tn + fp) if (tn + fp) > 0 else 0
|
| 275 |
+
else:
|
| 276 |
+
sensitivity = specificity = 0.0
|
| 277 |
+
|
| 278 |
+
test_metrics = {
|
| 279 |
+
'accuracy': float(accuracy),
|
| 280 |
+
'precision': float(precision),
|
| 281 |
+
'recall': float(recall),
|
| 282 |
+
'f1': float(f1),
|
| 283 |
+
'sensitivity': float(sensitivity),
|
| 284 |
+
'specificity': float(specificity),
|
| 285 |
+
'confusion_matrix': cm.tolist(),
|
| 286 |
+
}
|
| 287 |
+
|
| 288 |
+
print(f"\n{'-' * 40}")
|
| 289 |
+
print(f" TEST-SET RESULTS")
|
| 290 |
+
print(f"{'-' * 40}")
|
| 291 |
+
print(f" Accuracy: {accuracy:.4f}")
|
| 292 |
+
print(f" Precision: {precision:.4f}")
|
| 293 |
+
print(f" Recall: {recall:.4f}")
|
| 294 |
+
print(f" F1 Score: {f1:.4f}")
|
| 295 |
+
print(f" Sensitivity: {sensitivity:.4f}")
|
| 296 |
+
print(f" Specificity: {specificity:.4f}")
|
| 297 |
+
print(f" Confusion Matrix:\n {cm}")
|
| 298 |
+
|
| 299 |
+
# ── Save model (pickle) ──────────────────────────────────────────────
|
| 300 |
+
output_dir = os.path.join(MODELS_DIR, use_case_key)
|
| 301 |
+
os.makedirs(output_dir, exist_ok=True)
|
| 302 |
+
|
| 303 |
+
model_path = os.path.join(output_dir, config['model_file'])
|
| 304 |
+
print(f"\nSaving model to {model_path} ...")
|
| 305 |
+
with open(model_path, 'wb') as f:
|
| 306 |
+
pickle.dump(model, f)
|
| 307 |
+
|
| 308 |
+
model_size_mb = os.path.getsize(model_path) / (1024 * 1024)
|
| 309 |
+
print(f" Model saved ({model_size_mb:.1f} MB)")
|
| 310 |
+
|
| 311 |
+
# ── Save metadata (JSON) ─────────────────────────────────────────────
|
| 312 |
+
metadata = {
|
| 313 |
+
'model_name': config['model_name'],
|
| 314 |
+
'model_file': config['model_file'],
|
| 315 |
+
'use_case': config['name'],
|
| 316 |
+
'positive_class': config['positive_class'],
|
| 317 |
+
'negative_class': config['negative_class'],
|
| 318 |
+
'class_labels': [config['negative_class'], config['positive_class']],
|
| 319 |
+
'features': FEATURES,
|
| 320 |
+
'n_features': int(n_features),
|
| 321 |
+
'n_timesteps': int(n_timesteps),
|
| 322 |
+
'hyperparameters': config['hyperparameters'],
|
| 323 |
+
'training_data': input_csv,
|
| 324 |
+
'training_date': time.strftime('%Y-%m-%d %H:%M:%S'),
|
| 325 |
+
'training_time_seconds': round(train_time, 2),
|
| 326 |
+
'preprocessing': {
|
| 327 |
+
'normalization': 'per_subject_max_abs',
|
| 328 |
+
'smote_tomek': True,
|
| 329 |
+
'test_size': 0.25,
|
| 330 |
+
'random_state': 42,
|
| 331 |
+
},
|
| 332 |
+
'training_samples': int(X_train.shape[0]),
|
| 333 |
+
'test_samples': int(X_test.shape[0]),
|
| 334 |
+
'test_metrics': test_metrics,
|
| 335 |
+
}
|
| 336 |
+
|
| 337 |
+
metadata_path = os.path.join(output_dir, 'model_metadata.json')
|
| 338 |
+
with open(metadata_path, 'w') as f:
|
| 339 |
+
json.dump(metadata, f, indent=2)
|
| 340 |
+
print(f" Metadata saved to {metadata_path}")
|
| 341 |
+
|
| 342 |
+
print(f"\n[OK] {config['name']} - done.")
|
| 343 |
+
return True
|
| 344 |
+
|
| 345 |
+
|
| 346 |
+
# ─── Entry Point ──────────────────────────────────────────────────────────────
|
| 347 |
+
|
| 348 |
+
def show_saved_metrics():
|
| 349 |
+
"""
|
| 350 |
+
Search the MODELS_DIR directory for any saved model metadata JSON files
|
| 351 |
+
and print a beautiful summary of their performance metrics.
|
| 352 |
+
"""
|
| 353 |
+
print('=' * 60)
|
| 354 |
+
print(' TRAINED ECG MODEL METRICS SUMMARY')
|
| 355 |
+
print('=' * 60)
|
| 356 |
+
|
| 357 |
+
found = False
|
| 358 |
+
if os.path.exists(MODELS_DIR):
|
| 359 |
+
for item in sorted(os.listdir(MODELS_DIR)):
|
| 360 |
+
meta_path = os.path.join(MODELS_DIR, item, 'model_metadata.json')
|
| 361 |
+
if os.path.exists(meta_path):
|
| 362 |
+
try:
|
| 363 |
+
with open(meta_path, 'r') as f:
|
| 364 |
+
meta = json.load(f)
|
| 365 |
+
|
| 366 |
+
found = True
|
| 367 |
+
print(f"\nUse Case: {meta.get('use_case', 'Unknown')}")
|
| 368 |
+
print(f"Model Class: {meta.get('model_name', 'Unknown')}")
|
| 369 |
+
print(f"Model File: {os.path.join(MODELS_DIR, item, meta.get('model_file', ''))}")
|
| 370 |
+
print(f"Trained On: {meta.get('training_data', 'Unknown')}")
|
| 371 |
+
print(f"Training Date: {meta.get('training_date', 'Unknown')}")
|
| 372 |
+
if 'training_time_seconds' in meta:
|
| 373 |
+
print(f"Train Time: {meta['training_time_seconds']:.1f}s")
|
| 374 |
+
|
| 375 |
+
metrics = meta.get('test_metrics', {})
|
| 376 |
+
print("Test Metrics:")
|
| 377 |
+
print(f" - Accuracy: {metrics.get('accuracy', 0.0):.4f}")
|
| 378 |
+
print(f" - Precision: {metrics.get('precision', 0.0):.4f}")
|
| 379 |
+
print(f" - Recall: {metrics.get('recall', 0.0):.4f}")
|
| 380 |
+
print(f" - F1 Score: {metrics.get('f1', 0.0):.4f}")
|
| 381 |
+
print(f" - Sensitivity: {metrics.get('sensitivity', 0.0):.4f}")
|
| 382 |
+
print(f" - Specificity: {metrics.get('specificity', 0.0):.4f}")
|
| 383 |
+
|
| 384 |
+
cm = metrics.get('confusion_matrix', None)
|
| 385 |
+
if cm:
|
| 386 |
+
print(f" - Confusion Matrix:\n {np.array(cm)}")
|
| 387 |
+
print("-" * 60)
|
| 388 |
+
except Exception as e:
|
| 389 |
+
print(f"[Error loading metadata from {meta_path}]: {e}")
|
| 390 |
+
|
| 391 |
+
if not found:
|
| 392 |
+
print(f"\nNo saved model metadata files found under '{MODELS_DIR}/'.")
|
| 393 |
+
print("Please train models first to generate performance metrics.")
|
| 394 |
+
print('=' * 60)
|
| 395 |
+
|
| 396 |
+
|
| 397 |
+
if __name__ == '__main__':
|
| 398 |
+
parser = argparse.ArgumentParser(
|
| 399 |
+
description='Train and save the best ECG classification models')
|
| 400 |
+
parser.add_argument(
|
| 401 |
+
'--use-case', type=str,
|
| 402 |
+
choices=['mi_vs_normal_processed', 'mi_vs_normal_segmented', 'omi_vs_nonomi', 'ecg_surgery', 'all'],
|
| 403 |
+
default='all',
|
| 404 |
+
help='Which use case to train (default: all)')
|
| 405 |
+
parser.add_argument(
|
| 406 |
+
'--show-metrics', action='store_true',
|
| 407 |
+
help='Display performance metrics of already trained/saved models and exit')
|
| 408 |
+
args = parser.parse_args()
|
| 409 |
+
|
| 410 |
+
if args.show_metrics:
|
| 411 |
+
show_saved_metrics()
|
| 412 |
+
else:
|
| 413 |
+
cases = list(USE_CASES.keys()) if args.use_case == 'all' else [args.use_case]
|
| 414 |
+
|
| 415 |
+
print('=' * 60)
|
| 416 |
+
print(' ECG MODEL TRAINING & SERIALISATION')
|
| 417 |
+
print('=' * 60)
|
| 418 |
+
|
| 419 |
+
for key in cases:
|
| 420 |
+
success = train_and_save(key, USE_CASES[key])
|
| 421 |
+
if not success:
|
| 422 |
+
print(f"\n[FAIL] Could not train {USE_CASES[key]['name']}")
|
| 423 |
+
|
| 424 |
+
print(f"\n{'=' * 60}")
|
| 425 |
+
print(f" ALL TRAINING COMPLETE")
|
| 426 |
+
print(f" Models saved to: {MODELS_DIR}/")
|
| 427 |
+
print(f"{'=' * 60}")
|
archive/digitization/README.md
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# archive/digitization
|
| 2 |
+
|
| 3 |
+
Legacy ECG image digitization scripts that have been **superseded** by the
|
| 4 |
+
Streamlit dashboard in the workspace root.
|
| 5 |
+
|
| 6 |
+
## Contents
|
| 7 |
+
|
| 8 |
+
| File | Status | Reason for archival |
|
| 9 |
+
| --- | --- | --- |
|
| 10 |
+
| `run_ecg.py` | Legacy batch digitizer | Processes a directory tree of ECG paper scans (one CSV per image). Uses Linux-only `SIGALRM` timeouts and hard-coded model paths under `Submission code/...`. Replaced by the interactive pipeline in `backend/digitization_runner.py` invoked from `app.py`. |
|
| 11 |
+
|
| 12 |
+
## Current digitization entry point
|
| 13 |
+
|
| 14 |
+
Use the Streamlit dashboard:
|
| 15 |
+
|
| 16 |
+
```bash
|
| 17 |
+
streamlit run app.py
|
| 18 |
+
```
|
| 19 |
+
|
| 20 |
+
Then open the **"📷 Image Digitizer"** (now **"ECG Image Digitizer"**) page and
|
| 21 |
+
either upload an ECG scan or check the sample image option.
|
archive/digitization/run_ecg.py
ADDED
|
@@ -0,0 +1,302 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Force all BLAS / OpenMP thread pools to a single thread before any other
|
| 2 |
+
# import. This must happen before numpy, torch, or cv2 are imported, otherwise
|
| 3 |
+
# the libraries have already spawned their thread pools and the env vars have
|
| 4 |
+
# no effect. Without this, BLAS and PyTorch worker threads compete for CPU
|
| 5 |
+
# cores and can deadlock on multi-core servers.
|
| 6 |
+
import os
|
| 7 |
+
os.environ["OPENBLAS_NUM_THREADS"] = "1"
|
| 8 |
+
os.environ["OMP_NUM_THREADS"] = "1"
|
| 9 |
+
os.environ["MKL_NUM_THREADS"] = "1"
|
| 10 |
+
|
| 11 |
+
import glob
|
| 12 |
+
import signal
|
| 13 |
+
import torch
|
| 14 |
+
import pandas as pd
|
| 15 |
+
from ultralytics import YOLO
|
| 16 |
+
from digitization import ECGImage
|
| 17 |
+
|
| 18 |
+
# -------------------------
|
| 19 |
+
# Configuration
|
| 20 |
+
# -------------------------
|
| 21 |
+
|
| 22 |
+
# Root of the organised input dataset. Each immediate subdirectory is treated
|
| 23 |
+
# as one hospital and must itself contain subfolders named after CATEGORIES.
|
| 24 |
+
ORGANIZED_DIR = "../ecg_files/ECG_organized_all"
|
| 25 |
+
|
| 26 |
+
# Root where output CSVs are written, mirroring the hospital/category structure
|
| 27 |
+
# of ORGANIZED_DIR. Created automatically if it does not exist.
|
| 28 |
+
OUTPUT_DIR = "../ecg_files/ECG_digitized"
|
| 29 |
+
|
| 30 |
+
# Subfolders to process within each hospital directory.
|
| 31 |
+
# Any subfolder not listed here is silently skipped.
|
| 32 |
+
# Add or remove entries to match your dataset's naming convention.
|
| 33 |
+
CATEGORIES = ["pre", "index", "post", "discharge", "dimi"]
|
| 34 |
+
|
| 35 |
+
# Per-image wall-clock time limit in seconds, enforced via SIGALRM.
|
| 36 |
+
# Images that exceed this budget are logged as timeouts and skipped.
|
| 37 |
+
# Increase for very large images or slow CPU-only machines.
|
| 38 |
+
# Note: SIGALRM is only available on Linux.
|
| 39 |
+
TIMEOUT = 600
|
| 40 |
+
|
| 41 |
+
# -------------------------
|
| 42 |
+
# Torch / GPU settings
|
| 43 |
+
# -------------------------
|
| 44 |
+
|
| 45 |
+
if torch.cuda.is_available():
|
| 46 |
+
# benchmark=True lets cuDNN auto-tune convolution algorithms at the first
|
| 47 |
+
# forward pass. Speeds up subsequent passes when input sizes are fixed,
|
| 48 |
+
# which is the case here (images are resized to a fixed height).
|
| 49 |
+
torch.backends.cudnn.benchmark = True
|
| 50 |
+
|
| 51 |
+
# deterministic=False allows cuDNN to use non-deterministic (but faster)
|
| 52 |
+
# algorithms. Set to True if exact reproducibility is required.
|
| 53 |
+
torch.backends.cudnn.deterministic = False
|
| 54 |
+
|
| 55 |
+
print(f"GPU: {torch.cuda.get_device_name(0)}", flush=True)
|
| 56 |
+
print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB\n",
|
| 57 |
+
flush=True)
|
| 58 |
+
else:
|
| 59 |
+
print("No GPU found — running on CPU.\n", flush=True)
|
| 60 |
+
|
| 61 |
+
# -------------------------
|
| 62 |
+
# Timeout (Linux only)
|
| 63 |
+
# -------------------------
|
| 64 |
+
|
| 65 |
+
class TimeoutException(Exception):
|
| 66 |
+
"""Raised by timeout_handler when SIGALRM fires."""
|
| 67 |
+
pass
|
| 68 |
+
|
| 69 |
+
def timeout_handler(signum, frame):
|
| 70 |
+
"""Signal handler that converts SIGALRM into a TimeoutException.
|
| 71 |
+
|
| 72 |
+
Registered below with signal.signal() so that any blocking call
|
| 73 |
+
inside process_folder() can be interrupted after TIMEOUT seconds.
|
| 74 |
+
"""
|
| 75 |
+
raise TimeoutException("Processing timed out")
|
| 76 |
+
|
| 77 |
+
# Register the handler for SIGALRM. signal.alarm(n) starts a countdown;
|
| 78 |
+
# if the process has not called signal.alarm(0) within n seconds, this
|
| 79 |
+
# handler fires and raises TimeoutException.
|
| 80 |
+
signal.signal(signal.SIGALRM, timeout_handler)
|
| 81 |
+
|
| 82 |
+
# -------------------------
|
| 83 |
+
# Load models once (on GPU if available)
|
| 84 |
+
# -------------------------
|
| 85 |
+
|
| 86 |
+
# All four models are loaded once before the processing loop so that the
|
| 87 |
+
# weights are only read from disk once, regardless of how many images are
|
| 88 |
+
# processed. Loading inside the loop would add several seconds of overhead
|
| 89 |
+
# per image.
|
| 90 |
+
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
| 91 |
+
|
| 92 |
+
print("Loading models...", flush=True)
|
| 93 |
+
|
| 94 |
+
# Detection model for full lead bounding boxes (one box per lead region).
|
| 95 |
+
box_model = YOLO("Submission code/3. Model training/runs/yolo11_full/weights/best.pt")
|
| 96 |
+
|
| 97 |
+
# Classification model for lead name labels (I, II, III, aVR, aVL, aVF, V1–V6).
|
| 98 |
+
lead_name_model = YOLO("Submission code/3. Model training/runs/yolo11_lead/weights/best.pt")
|
| 99 |
+
|
| 100 |
+
# Detection model for calibration (reference) pulse boxes.
|
| 101 |
+
pulse_model = YOLO("Submission code/3. Model training/runs/yolo11_pulse/weights/best.pt")
|
| 102 |
+
|
| 103 |
+
box_model.to(DEVICE)
|
| 104 |
+
lead_name_model.to(DEVICE)
|
| 105 |
+
pulse_model.to(DEVICE)
|
| 106 |
+
|
| 107 |
+
# The segmentation model is kept as a file path rather than a loaded model
|
| 108 |
+
# object because it is consumed by patched_yolo_infer.MakeCropsDetectThem,
|
| 109 |
+
# which loads it internally on each call. This avoids a second copy of the
|
| 110 |
+
# weights living in VRAM alongside the three detection models above.
|
| 111 |
+
segmentation_model = "Submission code/3. Model training/runs/yolo11_patch/weights/best.pt"
|
| 112 |
+
|
| 113 |
+
print(f"Models loaded on {DEVICE.upper()}.\n", flush=True)
|
| 114 |
+
|
| 115 |
+
# -------------------------
|
| 116 |
+
# Process a single folder
|
| 117 |
+
# -------------------------
|
| 118 |
+
|
| 119 |
+
def process_folder(input_dir, output_dir):
|
| 120 |
+
"""Digitize all ECG images in input_dir and save CSVs to output_dir.
|
| 121 |
+
|
| 122 |
+
Iterates over all .png / .jpg / .jpeg files in input_dir (sorted
|
| 123 |
+
alphabetically). Each image is processed by ECGImage.run_full_pipeline()
|
| 124 |
+
and the resulting signals are exported as a CSV file with the same base
|
| 125 |
+
name as the image.
|
| 126 |
+
|
| 127 |
+
Parameters
|
| 128 |
+
----------
|
| 129 |
+
input_dir : str
|
| 130 |
+
Directory containing the input ECG images.
|
| 131 |
+
output_dir : str
|
| 132 |
+
Directory where output CSV files are written (created if absent).
|
| 133 |
+
|
| 134 |
+
Returns
|
| 135 |
+
-------
|
| 136 |
+
success : int
|
| 137 |
+
Number of images successfully digitized.
|
| 138 |
+
total : int
|
| 139 |
+
Total number of images found in input_dir.
|
| 140 |
+
"""
|
| 141 |
+
# Collect all supported image formats; sort for reproducible ordering.
|
| 142 |
+
files = []
|
| 143 |
+
for ext in ("*.png", "*.jpg", "*.jpeg"):
|
| 144 |
+
files.extend(glob.glob(os.path.join(input_dir, ext)))
|
| 145 |
+
files.sort()
|
| 146 |
+
|
| 147 |
+
# Return early without creating the output directory if the folder is empty.
|
| 148 |
+
if not files:
|
| 149 |
+
return 0, 0
|
| 150 |
+
|
| 151 |
+
os.makedirs(output_dir, exist_ok=True)
|
| 152 |
+
success = errors = timeouts = 0
|
| 153 |
+
|
| 154 |
+
for i, file_path in enumerate(files):
|
| 155 |
+
base_name = os.path.splitext(os.path.basename(file_path))[0]
|
| 156 |
+
print(f" [{i+1}/{len(files)}] {base_name}", flush=True)
|
| 157 |
+
try:
|
| 158 |
+
# Start the per-image countdown. If the block below takes longer
|
| 159 |
+
# than TIMEOUT seconds, SIGALRM fires and TimeoutException is raised.
|
| 160 |
+
signal.alarm(TIMEOUT)
|
| 161 |
+
|
| 162 |
+
# Instantiate the digitizer for this image. Models are passed in
|
| 163 |
+
# so they are not reloaded on every iteration.
|
| 164 |
+
ecg = ECGImage(
|
| 165 |
+
box_model=box_model,
|
| 166 |
+
segmentation_model=segmentation_model,
|
| 167 |
+
lead_name_model=lead_name_model,
|
| 168 |
+
pulse_model=pulse_model,
|
| 169 |
+
image_path=file_path,
|
| 170 |
+
)
|
| 171 |
+
|
| 172 |
+
# Run all pipeline stages: preprocessing → segmentation → YOLO
|
| 173 |
+
# detection (×3) → scale calibration → grid construction →
|
| 174 |
+
# signal extraction. Internally retries at different image sizes
|
| 175 |
+
# if any detection returns empty.
|
| 176 |
+
ecg.run_full_pipeline()
|
| 177 |
+
|
| 178 |
+
# Write one CSV per image: columns = lead names, rows = time samples.
|
| 179 |
+
ecg.save_signals_as_csv(base_name, directory=output_dir)
|
| 180 |
+
success += 1
|
| 181 |
+
|
| 182 |
+
except TimeoutException:
|
| 183 |
+
# The image took longer than TIMEOUT seconds. Counted separately
|
| 184 |
+
# from errors so they can be investigated independently.
|
| 185 |
+
timeouts += 1
|
| 186 |
+
print(f" ⏰ TIMEOUT: {base_name}", flush=True)
|
| 187 |
+
|
| 188 |
+
except Exception as e:
|
| 189 |
+
# Any other failure (bad image, detection failure, calibration
|
| 190 |
+
# error, etc.) is caught here so the loop continues with the
|
| 191 |
+
# remaining images.
|
| 192 |
+
errors += 1
|
| 193 |
+
print(f" ❌ ERROR: {base_name} → {e}", flush=True)
|
| 194 |
+
|
| 195 |
+
finally:
|
| 196 |
+
# Always cancel the alarm, even if an exception was raised, to
|
| 197 |
+
# prevent a stale countdown from interrupting the next image.
|
| 198 |
+
signal.alarm(0)
|
| 199 |
+
|
| 200 |
+
# Release VRAM after each image to prevent fragmentation across
|
| 201 |
+
# a long batch run. synchronize() ensures all GPU ops are complete
|
| 202 |
+
# before empty_cache() frees the memory.
|
| 203 |
+
if torch.cuda.is_available():
|
| 204 |
+
try:
|
| 205 |
+
torch.cuda.synchronize()
|
| 206 |
+
torch.cuda.empty_cache()
|
| 207 |
+
except Exception:
|
| 208 |
+
pass
|
| 209 |
+
|
| 210 |
+
total = len(files)
|
| 211 |
+
print(f" → {success}/{total} ok | {errors} errors | {timeouts} timeouts\n",
|
| 212 |
+
flush=True)
|
| 213 |
+
return success, total
|
| 214 |
+
|
| 215 |
+
# -------------------------
|
| 216 |
+
# Main loop
|
| 217 |
+
# -------------------------
|
| 218 |
+
|
| 219 |
+
# Discover hospital directories. Each immediate child directory of
|
| 220 |
+
# ORGANIZED_DIR is treated as one hospital; files at the root level are ignored.
|
| 221 |
+
hospitals = sorted([
|
| 222 |
+
name for name in os.listdir(ORGANIZED_DIR)
|
| 223 |
+
if os.path.isdir(os.path.join(ORGANIZED_DIR, name))
|
| 224 |
+
])
|
| 225 |
+
|
| 226 |
+
# Accumulate per-hospital/category counts for the summary report.
|
| 227 |
+
log_rows = []
|
| 228 |
+
|
| 229 |
+
print(f"Found {len(hospitals)} hospital(s). Starting digitization...\n", flush=True)
|
| 230 |
+
|
| 231 |
+
for hospital in hospitals:
|
| 232 |
+
print(f"{'='*60}", flush=True)
|
| 233 |
+
print(f"Hospital: {hospital}", flush=True)
|
| 234 |
+
print(f"{'='*60}", flush=True)
|
| 235 |
+
|
| 236 |
+
for cat in CATEGORIES:
|
| 237 |
+
input_dir = os.path.join(ORGANIZED_DIR, hospital, cat)
|
| 238 |
+
output_dir = os.path.join(OUTPUT_DIR, hospital, cat)
|
| 239 |
+
|
| 240 |
+
# Skip silently if this category does not exist for this hospital.
|
| 241 |
+
if not os.path.isdir(input_dir):
|
| 242 |
+
continue
|
| 243 |
+
|
| 244 |
+
print(f" [{cat.upper()}]", flush=True)
|
| 245 |
+
success, total = process_folder(input_dir, output_dir)
|
| 246 |
+
|
| 247 |
+
# Only log categories that had at least one image, to keep the
|
| 248 |
+
# summary table free of zero-row entries.
|
| 249 |
+
if total > 0:
|
| 250 |
+
log_rows.append({
|
| 251 |
+
"hospital": hospital,
|
| 252 |
+
"category": cat,
|
| 253 |
+
"total" : total,
|
| 254 |
+
"success" : success,
|
| 255 |
+
"failed" : total - success,
|
| 256 |
+
"rate" : f"{100 * success / total:.1f}%",
|
| 257 |
+
})
|
| 258 |
+
|
| 259 |
+
# -------------------------
|
| 260 |
+
# Summary report
|
| 261 |
+
# -------------------------
|
| 262 |
+
|
| 263 |
+
print("\n" + "="*60, flush=True)
|
| 264 |
+
print("DIGITIZATION SUMMARY", flush=True)
|
| 265 |
+
print("="*60, flush=True)
|
| 266 |
+
|
| 267 |
+
df = pd.DataFrame(log_rows)
|
| 268 |
+
|
| 269 |
+
if df.empty:
|
| 270 |
+
print("No images were processed.", flush=True)
|
| 271 |
+
else:
|
| 272 |
+
# Flat per-hospital/category breakdown.
|
| 273 |
+
print("\nPer hospital / category:\n")
|
| 274 |
+
print(df.to_string(index=False))
|
| 275 |
+
|
| 276 |
+
# Pivot to a hospital × category matrix of successful digitizations,
|
| 277 |
+
# with a TOTAL row and column for quick inspection.
|
| 278 |
+
print("\n\nSuccessful digitizations — distribution:\n")
|
| 279 |
+
success_pivot = (
|
| 280 |
+
df.pivot_table(index="hospital", columns="category",
|
| 281 |
+
values="success", aggfunc="sum", fill_value=0)
|
| 282 |
+
.reindex(columns=CATEGORIES, fill_value=0) # preserve CATEGORIES column order
|
| 283 |
+
)
|
| 284 |
+
success_pivot["TOTAL"] = success_pivot.sum(axis=1) # row totals
|
| 285 |
+
success_pivot.loc["TOTAL"] = success_pivot.sum(axis=0) # grand total row
|
| 286 |
+
print(success_pivot.to_string())
|
| 287 |
+
|
| 288 |
+
# Overall counts across the entire dataset.
|
| 289 |
+
total_imgs = df["total"].sum()
|
| 290 |
+
total_ok = df["success"].sum()
|
| 291 |
+
total_failed = df["failed"].sum()
|
| 292 |
+
print(f"\n✅ Overall: {total_ok}/{total_imgs} successful "
|
| 293 |
+
f"({100*total_ok/total_imgs:.1f}%) | "
|
| 294 |
+
f"{total_failed} failed", flush=True)
|
| 295 |
+
|
| 296 |
+
# Persist the pivot table as a CSV for later analysis.
|
| 297 |
+
summary_path = os.path.join(OUTPUT_DIR, "digitization_summary.csv")
|
| 298 |
+
os.makedirs(OUTPUT_DIR, exist_ok=True)
|
| 299 |
+
success_pivot.to_csv(summary_path)
|
| 300 |
+
print(f"📄 Summary saved to {summary_path}", flush=True)
|
| 301 |
+
|
| 302 |
+
print("\nAll processing complete.", flush=True)
|
assets/ETH_Z/303/274rich_Logo_black.svg.png
ADDED
|
assets/Logo_Vanvitelli_university.svg.png
ADDED
|
assets/eoc_logo.png
ADDED
|
assets/logo.png
ADDED
|
Git LFS Details
|
assets/scai_lab_logo.svg
ADDED
|
|
assets/usi_logo.png
ADDED
|
backend/__init__.py
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
# backend package
|
backend/classification_runner.py
ADDED
|
@@ -0,0 +1,571 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# backend/classification_runner.py
|
| 2 |
+
"""
|
| 3 |
+
Backend adapter for ECG Classification.
|
| 4 |
+
Handles heartbeat segmentation, data preparation for sktime models, and model evaluation.
|
| 5 |
+
"""
|
| 6 |
+
|
| 7 |
+
import os
|
| 8 |
+
import sys
|
| 9 |
+
import numpy as np
|
| 10 |
+
import pandas as pd
|
| 11 |
+
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, confusion_matrix
|
| 12 |
+
|
| 13 |
+
# Add parent directory to path to ensure digitization and other local modules are importable
|
| 14 |
+
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
| 15 |
+
|
| 16 |
+
# Custom Exceptions for dependencies
|
| 17 |
+
class DependencyMissingError(Exception):
|
| 18 |
+
"""Raised when a required third-party library is not installed."""
|
| 19 |
+
pass
|
| 20 |
+
|
| 21 |
+
class EmptyDatasetError(Exception):
|
| 22 |
+
"""Raised when the uploaded dataset is empty or invalid."""
|
| 23 |
+
pass
|
| 24 |
+
|
| 25 |
+
# Helper to check dependencies
|
| 26 |
+
def check_classification_dependencies():
|
| 27 |
+
"""
|
| 28 |
+
Checks if sktime is installed in the current environment.
|
| 29 |
+
Raises DependencyMissingError if it is not.
|
| 30 |
+
"""
|
| 31 |
+
missing = []
|
| 32 |
+
try:
|
| 33 |
+
import sktime
|
| 34 |
+
except ImportError:
|
| 35 |
+
missing.append("sktime")
|
| 36 |
+
|
| 37 |
+
if missing:
|
| 38 |
+
raise DependencyMissingError(
|
| 39 |
+
f"The following libraries are required for classification but are missing in this environment:\n"
|
| 40 |
+
f"{', '.join(missing)}\n\n"
|
| 41 |
+
f"Please run: pip install sktime"
|
| 42 |
+
)
|
| 43 |
+
|
| 44 |
+
# ==============================================================================
|
| 45 |
+
# HEARTBEAT SEGMENTATION
|
| 46 |
+
# ==============================================================================
|
| 47 |
+
# ==============================================================================
|
| 48 |
+
# LOCAL HEARTBEAT SEGMENTATION IMPLEMENTATION (copied from ecg_segmentation.py)
|
| 49 |
+
# ==============================================================================
|
| 50 |
+
|
| 51 |
+
def bandpass_filter(signal_data, lowcut=5, highcut=15, fs=500, order=2):
|
| 52 |
+
"""Bandpass filter for ECG signal"""
|
| 53 |
+
from scipy.signal import butter, filtfilt
|
| 54 |
+
try:
|
| 55 |
+
nyquist = 0.5 * fs
|
| 56 |
+
low = lowcut / nyquist
|
| 57 |
+
high = highcut / nyquist
|
| 58 |
+
b, a = butter(order, [low, high], btype='band')
|
| 59 |
+
return filtfilt(b, a, signal_data)
|
| 60 |
+
except Exception:
|
| 61 |
+
return signal_data
|
| 62 |
+
|
| 63 |
+
|
| 64 |
+
def find_r_peaks(ecg_signal, min_heart_rate=40, max_heart_rate=180, fs=500):
|
| 65 |
+
"""Find R-peaks in ECG signal"""
|
| 66 |
+
from scipy.signal import find_peaks
|
| 67 |
+
# Calculate expected distance between peaks based on heart rate
|
| 68 |
+
min_distance = int((60 / max_heart_rate) * fs) # samples
|
| 69 |
+
max_distance = int((60 / min_heart_rate) * fs) # samples
|
| 70 |
+
|
| 71 |
+
# Find peaks
|
| 72 |
+
peaks, properties = find_peaks(ecg_signal,
|
| 73 |
+
height=np.percentile(ecg_signal, 80),
|
| 74 |
+
distance=min_distance,
|
| 75 |
+
prominence=np.std(ecg_signal) * 0.3)
|
| 76 |
+
|
| 77 |
+
# Filter peaks based on reasonable heart rate
|
| 78 |
+
if len(peaks) == 0:
|
| 79 |
+
return peaks
|
| 80 |
+
|
| 81 |
+
valid_peaks = [peaks[0]]
|
| 82 |
+
for i in range(1, len(peaks)):
|
| 83 |
+
distance = peaks[i] - valid_peaks[-1]
|
| 84 |
+
if min_distance <= distance <= max_distance:
|
| 85 |
+
valid_peaks.append(peaks[i])
|
| 86 |
+
|
| 87 |
+
return np.array(valid_peaks)
|
| 88 |
+
|
| 89 |
+
|
| 90 |
+
def downsample_beat(beat_data, features, downsampling_factor):
|
| 91 |
+
"""
|
| 92 |
+
Downsample a single heartbeat by taking every nth sample.
|
| 93 |
+
Uses vectorized iloc slicing instead of row-by-row iteration.
|
| 94 |
+
"""
|
| 95 |
+
if downsampling_factor <= 1:
|
| 96 |
+
return beat_data
|
| 97 |
+
try:
|
| 98 |
+
downsampled = beat_data.iloc[::downsampling_factor]
|
| 99 |
+
return downsampled if len(downsampled) > 0 else None
|
| 100 |
+
except Exception as e:
|
| 101 |
+
print(f"Error in downsampling: {e}")
|
| 102 |
+
return None
|
| 103 |
+
|
| 104 |
+
|
| 105 |
+
def segment_beats_around_r_peaks(subject_data, r_peaks, features, fs=500, target_fs=500, pre_r_ms=150, post_r_ms=250):
|
| 106 |
+
"""
|
| 107 |
+
Segment heartbeats around R-peaks with downsampling
|
| 108 |
+
"""
|
| 109 |
+
heartbeats = []
|
| 110 |
+
total_samples = len(subject_data)
|
| 111 |
+
|
| 112 |
+
# Calculate downsampling factor
|
| 113 |
+
downsampling_factor = fs // target_fs
|
| 114 |
+
if downsampling_factor < 1:
|
| 115 |
+
downsampling_factor = 1
|
| 116 |
+
|
| 117 |
+
# Convert milliseconds to samples at original and target frequencies
|
| 118 |
+
pre_r_original = int((pre_r_ms / 1000) * fs)
|
| 119 |
+
post_r_original = int((post_r_ms / 1000) * fs)
|
| 120 |
+
|
| 121 |
+
min_beat_len = int((pre_r_original + post_r_original) * 0.7)
|
| 122 |
+
|
| 123 |
+
for r_peak in r_peaks:
|
| 124 |
+
start_idx = max(0, r_peak - pre_r_original)
|
| 125 |
+
end_idx = min(total_samples, r_peak + post_r_original)
|
| 126 |
+
|
| 127 |
+
# Ensure we have enough samples for a complete beat
|
| 128 |
+
if (end_idx - start_idx) >= min_beat_len:
|
| 129 |
+
beat_data_original = subject_data.iloc[start_idx:end_idx]
|
| 130 |
+
|
| 131 |
+
# Downsample the beat
|
| 132 |
+
beat_data_downsampled = downsample_beat(beat_data_original, features, downsampling_factor)
|
| 133 |
+
|
| 134 |
+
if beat_data_downsampled is not None:
|
| 135 |
+
heartbeats.append((beat_data_downsampled, r_peak))
|
| 136 |
+
|
| 137 |
+
return heartbeats
|
| 138 |
+
|
| 139 |
+
|
| 140 |
+
def local_segment_heartbeats(input_csv, output_csv=None, min_heart_rate=40, max_heart_rate=180,
|
| 141 |
+
fs=500, target_fs=500):
|
| 142 |
+
"""
|
| 143 |
+
Segment ECG data into individual heartbeats for each subject with downsampling.
|
| 144 |
+
Self-contained implementation.
|
| 145 |
+
Optimized: uses vectorized pandas concat instead of row-by-row dict assembly.
|
| 146 |
+
"""
|
| 147 |
+
print("Loading ECG data for heartbeat segmentation...")
|
| 148 |
+
df = pd.read_csv(input_csv, index_col=['subject_id', 'timestamp'])
|
| 149 |
+
features = ['I', 'aVR', 'V1', 'V4', 'II', 'aVL', 'V2', 'V5', 'III', 'aVF', 'V3', 'V6']
|
| 150 |
+
|
| 151 |
+
# Fill NaN values in features using linear interpolation and fallback to 0.0
|
| 152 |
+
existing_features = [f for f in features if f in df.columns]
|
| 153 |
+
if existing_features:
|
| 154 |
+
df[existing_features] = df[existing_features].interpolate(method='linear', limit_direction='both').fillna(0.0)
|
| 155 |
+
|
| 156 |
+
# Calculate downsampling factor
|
| 157 |
+
downsampling_factor = fs // target_fs
|
| 158 |
+
print(f"Downsampling from {fs} Hz to {target_fs} Hz (factor: {downsampling_factor})")
|
| 159 |
+
|
| 160 |
+
# Collect beat DataFrames in a list for a single pd.concat at the end
|
| 161 |
+
beat_frames = []
|
| 162 |
+
|
| 163 |
+
# Get unique subjects
|
| 164 |
+
subject_ids = df.index.get_level_values('subject_id').unique()
|
| 165 |
+
|
| 166 |
+
for subject_id in subject_ids:
|
| 167 |
+
# Get subject data
|
| 168 |
+
subject_data = df.xs(subject_id, level='subject_id')
|
| 169 |
+
subject_class = subject_data['class'].iloc[0]
|
| 170 |
+
|
| 171 |
+
# Use Lead II for R-peak detection (commonly used)
|
| 172 |
+
ecg_signal = subject_data['II'].values
|
| 173 |
+
|
| 174 |
+
# Skip if signal is too short
|
| 175 |
+
if len(ecg_signal) < fs * 2: # Less than 2 seconds
|
| 176 |
+
continue
|
| 177 |
+
|
| 178 |
+
# Preprocess signal (apply bandpass filter for robust R-peak detection)
|
| 179 |
+
try:
|
| 180 |
+
ecg_filtered = bandpass_filter(ecg_signal, lowcut=5, highcut=15, fs=fs, order=2)
|
| 181 |
+
except Exception:
|
| 182 |
+
ecg_filtered = ecg_signal
|
| 183 |
+
|
| 184 |
+
# Find R-peaks
|
| 185 |
+
r_peaks = find_r_peaks(ecg_filtered, min_heart_rate=min_heart_rate, max_heart_rate=max_heart_rate, fs=fs)
|
| 186 |
+
|
| 187 |
+
if len(r_peaks) < 2:
|
| 188 |
+
# Fallback: if filtered failed or returned no peaks, try on unfiltered signal
|
| 189 |
+
try:
|
| 190 |
+
r_peaks = find_r_peaks(ecg_signal, min_heart_rate=min_heart_rate, max_heart_rate=max_heart_rate, fs=fs)
|
| 191 |
+
except Exception:
|
| 192 |
+
pass
|
| 193 |
+
|
| 194 |
+
if len(r_peaks) < 2:
|
| 195 |
+
print(f" Warning: Only {len(r_peaks)} R-peaks found for subject {subject_id}")
|
| 196 |
+
continue
|
| 197 |
+
|
| 198 |
+
# Segment heartbeats around R-peaks
|
| 199 |
+
heartbeats = segment_beats_around_r_peaks(subject_data, r_peaks, features, fs=fs, target_fs=target_fs)
|
| 200 |
+
|
| 201 |
+
# Vectorized beat assembly: build each beat as a small DataFrame, then concat once
|
| 202 |
+
for i, (beat_data, r_peak_pos) in enumerate(heartbeats):
|
| 203 |
+
beat_id = f"{subject_id}_beat{i+1}"
|
| 204 |
+
n_rows = len(beat_data)
|
| 205 |
+
|
| 206 |
+
# Build the beat frame from the existing beat_data slice
|
| 207 |
+
beat_frame = beat_data[existing_features].copy()
|
| 208 |
+
beat_frame = beat_frame.reset_index(drop=True)
|
| 209 |
+
beat_frame['subject_id'] = beat_id
|
| 210 |
+
beat_frame['original_subject'] = subject_id
|
| 211 |
+
beat_frame['beat_number'] = i + 1
|
| 212 |
+
beat_frame['r_peak_position'] = int(r_peak_pos)
|
| 213 |
+
beat_frame['class'] = subject_class
|
| 214 |
+
beat_frame['timestamp'] = np.arange(n_rows)
|
| 215 |
+
|
| 216 |
+
beat_frames.append(beat_frame)
|
| 217 |
+
|
| 218 |
+
# Single concat at the end instead of thousands of dict appends
|
| 219 |
+
if not beat_frames:
|
| 220 |
+
return pd.DataFrame()
|
| 221 |
+
|
| 222 |
+
segmented_df = pd.concat(beat_frames, ignore_index=True)
|
| 223 |
+
|
| 224 |
+
# Set multi-index
|
| 225 |
+
if not segmented_df.empty:
|
| 226 |
+
segmented_df.set_index(['subject_id', 'timestamp'], inplace=True)
|
| 227 |
+
|
| 228 |
+
# Save if output path provided
|
| 229 |
+
if output_csv and not segmented_df.empty:
|
| 230 |
+
os.makedirs(os.path.dirname(output_csv), exist_ok=True)
|
| 231 |
+
segmented_df.to_csv(output_csv)
|
| 232 |
+
print(f"Segmented data saved to: {output_csv}")
|
| 233 |
+
|
| 234 |
+
print(f"\nSegmentation complete:")
|
| 235 |
+
print(f"Total heartbeats: {segmented_df.index.get_level_values('subject_id').nunique() if not segmented_df.empty else 0}")
|
| 236 |
+
|
| 237 |
+
return segmented_df
|
| 238 |
+
|
| 239 |
+
|
| 240 |
+
def segment_uploaded_csv(input_csv_path, output_csv_path, target_fs=500):
|
| 241 |
+
"""
|
| 242 |
+
Segments a raw digitized dataset CSV into individual heartbeats.
|
| 243 |
+
Uses the local self-contained heartbeat segmentation implementation.
|
| 244 |
+
"""
|
| 245 |
+
# Read the first few lines to inspect columns
|
| 246 |
+
df_raw = pd.read_csv(input_csv_path)
|
| 247 |
+
if df_raw.empty:
|
| 248 |
+
raise EmptyDatasetError("The uploaded CSV file is empty.")
|
| 249 |
+
|
| 250 |
+
# Standardize input formatting for the segmentation module.
|
| 251 |
+
# The segmentation module expects 'subject_id', 'timestamp', 'class' columns.
|
| 252 |
+
required_cols = ['subject_id', 'timestamp', 'class']
|
| 253 |
+
needs_formatting = not all(col in df_raw.columns for col in required_cols)
|
| 254 |
+
|
| 255 |
+
formatted_csv_path = input_csv_path
|
| 256 |
+
if needs_formatting:
|
| 257 |
+
print("Formatting raw CSV columns to match expected segmentation schema...")
|
| 258 |
+
# Check if we have lead columns
|
| 259 |
+
lead_cols = ['I', 'aVR', 'V1', 'V4', 'II', 'aVL', 'V2', 'V5', 'III', 'aVF', 'V3', 'V6']
|
| 260 |
+
existing_leads = [col for col in lead_cols if col in df_raw.columns]
|
| 261 |
+
if not existing_leads:
|
| 262 |
+
raise ValueError(
|
| 263 |
+
f"The uploaded CSV does not contain standard ECG lead columns.\n"
|
| 264 |
+
f"Expected some of: {lead_cols}"
|
| 265 |
+
)
|
| 266 |
+
|
| 267 |
+
# Format df_raw
|
| 268 |
+
df_formatted = df_raw.copy()
|
| 269 |
+
if 'subject_id' not in df_formatted.columns:
|
| 270 |
+
df_formatted['subject_id'] = 'subject_1'
|
| 271 |
+
if 'class' not in df_formatted.columns:
|
| 272 |
+
# Default placeholder label
|
| 273 |
+
df_formatted['class'] = 'Pre-Procedural MI'
|
| 274 |
+
if 'timestamp' not in df_formatted.columns:
|
| 275 |
+
df_formatted['timestamp'] = range(len(df_formatted))
|
| 276 |
+
|
| 277 |
+
# Ensure subject_id and class are strings
|
| 278 |
+
df_formatted['subject_id'] = df_formatted['subject_id'].astype(str)
|
| 279 |
+
df_formatted['class'] = df_formatted['class'].astype(str)
|
| 280 |
+
|
| 281 |
+
# Save to temporary path for segmentation input
|
| 282 |
+
formatted_csv_path = input_csv_path + "_formatted.csv"
|
| 283 |
+
df_formatted.to_csv(formatted_csv_path, index=False)
|
| 284 |
+
|
| 285 |
+
# Run the local segmentation module
|
| 286 |
+
try:
|
| 287 |
+
segmented_df = local_segment_heartbeats(
|
| 288 |
+
formatted_csv_path,
|
| 289 |
+
output_csv=output_csv_path,
|
| 290 |
+
fs=target_fs,
|
| 291 |
+
target_fs=target_fs
|
| 292 |
+
)
|
| 293 |
+
finally:
|
| 294 |
+
# Clean up temporary formatted file if created
|
| 295 |
+
if needs_formatting and os.path.exists(formatted_csv_path):
|
| 296 |
+
try:
|
| 297 |
+
os.remove(formatted_csv_path)
|
| 298 |
+
except Exception:
|
| 299 |
+
pass
|
| 300 |
+
|
| 301 |
+
if segmented_df is None or segmented_df.empty:
|
| 302 |
+
raise ValueError(
|
| 303 |
+
"Heartbeat segmentation returned no beats. This can happen if the signal is too short or R-peaks were not detected."
|
| 304 |
+
)
|
| 305 |
+
|
| 306 |
+
return segmented_df
|
| 307 |
+
|
| 308 |
+
|
| 309 |
+
#
|
| 310 |
+
|
| 311 |
+
|
| 312 |
+
# ==============================================================================
|
| 313 |
+
# PRE-TRAINED MODEL LOADING & INFERENCE SUPPORT
|
| 314 |
+
# ==============================================================================
|
| 315 |
+
|
| 316 |
+
def _load_model_and_metadata_impl(model_key):
|
| 317 |
+
"""
|
| 318 |
+
Internal implementation: loads a pre-trained model and its metadata from
|
| 319 |
+
models/classifier_models/{model_key}.
|
| 320 |
+
"""
|
| 321 |
+
import json
|
| 322 |
+
import pickle
|
| 323 |
+
|
| 324 |
+
model_dir = os.path.join("models", "classifier_models", model_key)
|
| 325 |
+
metadata_path = os.path.join(model_dir, "model_metadata.json")
|
| 326 |
+
|
| 327 |
+
if not os.path.exists(metadata_path):
|
| 328 |
+
raise FileNotFoundError(
|
| 329 |
+
f"Pre-trained model metadata not found: {metadata_path}\n"
|
| 330 |
+
f"Please verify models/classifier_models/ directory contents."
|
| 331 |
+
)
|
| 332 |
+
|
| 333 |
+
with open(metadata_path, 'r') as f:
|
| 334 |
+
metadata = json.load(f)
|
| 335 |
+
|
| 336 |
+
model_path = os.path.join(model_dir, metadata['model_file'])
|
| 337 |
+
if not os.path.exists(model_path):
|
| 338 |
+
raise FileNotFoundError(
|
| 339 |
+
f"Pre-trained model pickle not found: {model_path}\n"
|
| 340 |
+
f"Please verify models/classifier_models/ directory contents."
|
| 341 |
+
)
|
| 342 |
+
|
| 343 |
+
print(f"Loading {metadata['model_name']} pre-trained model for {metadata['use_case']}...")
|
| 344 |
+
with open(model_path, 'rb') as f:
|
| 345 |
+
model = pickle.load(f)
|
| 346 |
+
|
| 347 |
+
# Restore missing classes_ on unpickled scikit-learn models due to version inconsistencies
|
| 348 |
+
def _fix_unpickled_classes(obj, visited=None):
|
| 349 |
+
if obj is None:
|
| 350 |
+
return
|
| 351 |
+
if visited is None:
|
| 352 |
+
visited = set()
|
| 353 |
+
obj_id = id(obj)
|
| 354 |
+
if obj_id in visited:
|
| 355 |
+
return
|
| 356 |
+
visited.add(obj_id)
|
| 357 |
+
|
| 358 |
+
if hasattr(obj, '_label_binarizer') and hasattr(obj._label_binarizer, 'classes_'):
|
| 359 |
+
if not hasattr(obj, 'classes_') or obj.classes_ is None:
|
| 360 |
+
try:
|
| 361 |
+
obj.classes_ = obj._label_binarizer.classes_
|
| 362 |
+
except Exception:
|
| 363 |
+
pass
|
| 364 |
+
|
| 365 |
+
if isinstance(obj, (list, tuple)):
|
| 366 |
+
for item in obj:
|
| 367 |
+
_fix_unpickled_classes(item, visited)
|
| 368 |
+
elif isinstance(obj, dict):
|
| 369 |
+
for val in obj.values():
|
| 370 |
+
_fix_unpickled_classes(val, visited)
|
| 371 |
+
else:
|
| 372 |
+
for attr in ['estimator', 'estimator_', 'classifier', 'classifier_',
|
| 373 |
+
'estimators', 'estimators_', 'transformers', 'transformers_',
|
| 374 |
+
'steps', 'steps_']:
|
| 375 |
+
if hasattr(obj, attr):
|
| 376 |
+
try:
|
| 377 |
+
val = getattr(obj, attr)
|
| 378 |
+
_fix_unpickled_classes(val, visited)
|
| 379 |
+
except Exception:
|
| 380 |
+
pass
|
| 381 |
+
|
| 382 |
+
_fix_unpickled_classes(model)
|
| 383 |
+
|
| 384 |
+
return model, metadata
|
| 385 |
+
|
| 386 |
+
|
| 387 |
+
# Try to use Streamlit cache; fall back to plain function if running outside Streamlit
|
| 388 |
+
try:
|
| 389 |
+
import streamlit as st
|
| 390 |
+
|
| 391 |
+
@st.cache_resource(show_spinner="Loading pre-trained model...")
|
| 392 |
+
def load_pretrained_model_and_metadata(model_key):
|
| 393 |
+
"""Cached wrapper — model pickle is only deserialized once per model_key."""
|
| 394 |
+
return _load_model_and_metadata_impl(model_key)
|
| 395 |
+
except Exception:
|
| 396 |
+
# Fallback when Streamlit is not available (e.g. tests, scripts)
|
| 397 |
+
load_pretrained_model_and_metadata = _load_model_and_metadata_impl
|
| 398 |
+
|
| 399 |
+
|
| 400 |
+
def preprocess_dataframe_for_inference(df, metadata):
|
| 401 |
+
"""
|
| 402 |
+
Preprocesses raw or segmented DataFrame into numpy3D array for the pre-trained model.
|
| 403 |
+
Reshapes each subject's signal to match the exact target length (n_timesteps),
|
| 404 |
+
applies max absolute voltage normalization per subject, and returns the 3D numpy array.
|
| 405 |
+
Optimized: single grouped normalization instead of per-column loop.
|
| 406 |
+
"""
|
| 407 |
+
features = metadata['features']
|
| 408 |
+
n_timesteps = metadata['n_timesteps']
|
| 409 |
+
n_features = metadata['n_features']
|
| 410 |
+
|
| 411 |
+
# Standardize input formatting
|
| 412 |
+
# If index is multi-index with subject_id, or if subject_id/timestamp are columns
|
| 413 |
+
df_reset = df.copy()
|
| 414 |
+
if 'subject_id' in df_reset.index.names:
|
| 415 |
+
df_reset = df_reset.reset_index()
|
| 416 |
+
elif 'subject_id' not in df_reset.columns:
|
| 417 |
+
df_reset['subject_id'] = 'subject_1'
|
| 418 |
+
|
| 419 |
+
if 'timestamp' in df_reset.index.names:
|
| 420 |
+
df_reset = df_reset.reset_index()
|
| 421 |
+
elif 'timestamp' not in df_reset.columns:
|
| 422 |
+
df_reset['timestamp'] = df_reset.groupby('subject_id').cumcount()
|
| 423 |
+
|
| 424 |
+
df_reset = df_reset.set_index(['subject_id', 'timestamp'])
|
| 425 |
+
|
| 426 |
+
# Verify required features exist
|
| 427 |
+
missing = [f for f in features if f not in df_reset.columns]
|
| 428 |
+
if missing:
|
| 429 |
+
# Fallback to standard channels if names differ slightly (e.g. casing)
|
| 430 |
+
df_reset.columns = [c.strip() for c in df_reset.columns]
|
| 431 |
+
missing = [f for f in features if f not in df_reset.columns]
|
| 432 |
+
if missing:
|
| 433 |
+
raise ValueError(f"Missing required lead columns in input: {missing}")
|
| 434 |
+
|
| 435 |
+
# Vectorized normalization: single grouped abs-max across all feature columns at once
|
| 436 |
+
df_X = df_reset[features].copy()
|
| 437 |
+
grouped_max = df_X.abs().groupby(level='subject_id').transform('max')
|
| 438 |
+
# Avoid division by zero
|
| 439 |
+
grouped_max = grouped_max.replace(0, 1.0)
|
| 440 |
+
df_X = df_X / grouped_max
|
| 441 |
+
|
| 442 |
+
# Reshape to 3D array: (n_subjects, n_features, n_timesteps)
|
| 443 |
+
subject_ids = df_X.index.get_level_values('subject_id').unique()
|
| 444 |
+
X_list = []
|
| 445 |
+
valid_subject_ids = []
|
| 446 |
+
y_true_str = []
|
| 447 |
+
|
| 448 |
+
for subject_id in subject_ids:
|
| 449 |
+
try:
|
| 450 |
+
subject_data = df_X.xs(subject_id, level='subject_id')
|
| 451 |
+
actual_len = len(subject_data)
|
| 452 |
+
|
| 453 |
+
if actual_len >= n_timesteps:
|
| 454 |
+
# Truncate
|
| 455 |
+
arr = subject_data.iloc[:n_timesteps].values.T
|
| 456 |
+
else:
|
| 457 |
+
# Zero-pad
|
| 458 |
+
arr = np.zeros((n_features, n_timesteps))
|
| 459 |
+
arr[:, :actual_len] = subject_data.values.T
|
| 460 |
+
|
| 461 |
+
if arr.shape == (n_features, n_timesteps):
|
| 462 |
+
X_list.append(arr)
|
| 463 |
+
valid_subject_ids.append(subject_id)
|
| 464 |
+
|
| 465 |
+
# Capture ground truth class label if it exists in the source DataFrame
|
| 466 |
+
if 'class' in df_reset.columns:
|
| 467 |
+
lbl = df_reset.loc[subject_id].iloc[0]['class']
|
| 468 |
+
y_true_str.append(str(lbl))
|
| 469 |
+
except Exception:
|
| 470 |
+
continue
|
| 471 |
+
|
| 472 |
+
if not X_list:
|
| 473 |
+
raise ValueError(
|
| 474 |
+
f"No valid subjects/heartbeats found after preprocessing. "
|
| 475 |
+
f"Ensure data contains at least one subject with some timesteps."
|
| 476 |
+
)
|
| 477 |
+
|
| 478 |
+
return np.array(X_list), valid_subject_ids, y_true_str
|
| 479 |
+
|
| 480 |
+
|
| 481 |
+
def run_pretrained_inference(model, X, metadata, subject_ids):
|
| 482 |
+
"""
|
| 483 |
+
Runs pre-trained model predictions and returns a results dictionary.
|
| 484 |
+
"""
|
| 485 |
+
from joblib import parallel_backend
|
| 486 |
+
|
| 487 |
+
with parallel_backend('threading', n_jobs=-1):
|
| 488 |
+
y_pred = model.predict(X)
|
| 489 |
+
|
| 490 |
+
has_proba = False
|
| 491 |
+
y_proba = None
|
| 492 |
+
try:
|
| 493 |
+
y_proba = model.predict_proba(X)
|
| 494 |
+
has_proba = True
|
| 495 |
+
except (AttributeError, NotImplementedError):
|
| 496 |
+
pass
|
| 497 |
+
|
| 498 |
+
results = {
|
| 499 |
+
'subject_id': subject_ids,
|
| 500 |
+
'predicted_class': y_pred.tolist(),
|
| 501 |
+
}
|
| 502 |
+
|
| 503 |
+
if has_proba and y_proba is not None:
|
| 504 |
+
class_labels = list(model.classes_)
|
| 505 |
+
results['class_labels'] = class_labels
|
| 506 |
+
results['probabilities'] = y_proba.tolist()
|
| 507 |
+
|
| 508 |
+
# Calculate confidence as the probability of the predicted class
|
| 509 |
+
confidence = []
|
| 510 |
+
for j in range(len(y_pred)):
|
| 511 |
+
try:
|
| 512 |
+
pred_idx = class_labels.index(y_pred[j])
|
| 513 |
+
confidence.append(float(y_proba[j, pred_idx]))
|
| 514 |
+
except ValueError:
|
| 515 |
+
confidence.append(1.0)
|
| 516 |
+
results['confidence'] = confidence
|
| 517 |
+
|
| 518 |
+
return results
|
| 519 |
+
|
| 520 |
+
|
| 521 |
+
def calculate_evaluation_metrics(y_true_str, y_pred_str, class_labels, positive_class):
|
| 522 |
+
"""
|
| 523 |
+
Calculates accuracy, precision, recall, f1, sensitivity, specificity, and confusion matrix.
|
| 524 |
+
Maps ground-truth and prediction string labels to binary (0/1).
|
| 525 |
+
"""
|
| 526 |
+
# Map strings to binary values (0/1) based on positive class
|
| 527 |
+
pos_str = str(positive_class).strip().lower()
|
| 528 |
+
|
| 529 |
+
# We do a lenient comparison (check substring or exact match)
|
| 530 |
+
y_true_num = [1 if pos_str in str(y).strip().lower() else 0 for y in y_true_str]
|
| 531 |
+
y_pred_num = [1 if pos_str in str(y).strip().lower() else 0 for y in y_pred_str]
|
| 532 |
+
|
| 533 |
+
# Handle cases where all targets are of one class
|
| 534 |
+
if len(np.unique(y_true_num)) <= 1 and len(np.unique(y_pred_num)) <= 1:
|
| 535 |
+
# Avoid standard metrics crashing
|
| 536 |
+
accuracy = float(accuracy_score(y_true_num, y_pred_num))
|
| 537 |
+
return {
|
| 538 |
+
"Accuracy": accuracy,
|
| 539 |
+
"Precision": 1.0 if accuracy == 1.0 else 0.0,
|
| 540 |
+
"Recall": 1.0 if accuracy == 1.0 else 0.0,
|
| 541 |
+
"F1": 1.0 if accuracy == 1.0 else 0.0,
|
| 542 |
+
"Sensitivity": 1.0 if accuracy == 1.0 else 0.0,
|
| 543 |
+
"Specificity": 1.0 if accuracy == 1.0 else 0.0,
|
| 544 |
+
"Confusion Matrix": confusion_matrix(y_true_num, y_pred_num, labels=[0, 1]).tolist(),
|
| 545 |
+
"TN": int(accuracy * len(y_true_num)) if y_true_num[0] == 0 else 0,
|
| 546 |
+
"FP": 0,
|
| 547 |
+
"FN": 0,
|
| 548 |
+
"TP": int(accuracy * len(y_true_num)) if y_true_num[0] == 1 else 0
|
| 549 |
+
}
|
| 550 |
+
|
| 551 |
+
accuracy = float(accuracy_score(y_true_num, y_pred_num))
|
| 552 |
+
precision = float(precision_score(y_true_num, y_pred_num, zero_division=0))
|
| 553 |
+
recall = float(recall_score(y_true_num, y_pred_num, zero_division=0))
|
| 554 |
+
f1 = float(f1_score(y_true_num, y_pred_num, zero_division=0))
|
| 555 |
+
|
| 556 |
+
cm = confusion_matrix(y_true_num, y_pred_num, labels=[0, 1])
|
| 557 |
+
tn, fp, fn, tp = cm.flatten()
|
| 558 |
+
|
| 559 |
+
sensitivity = float(tp / (tp + fn) if (tp + fn) > 0 else 0)
|
| 560 |
+
specificity = float(tn / (tn + fp) if (tn + fp) > 0 else 0)
|
| 561 |
+
|
| 562 |
+
return {
|
| 563 |
+
"Accuracy": accuracy,
|
| 564 |
+
"Precision": precision,
|
| 565 |
+
"Recall": recall,
|
| 566 |
+
"F1": f1,
|
| 567 |
+
"Sensitivity": sensitivity,
|
| 568 |
+
"Specificity": specificity,
|
| 569 |
+
"Confusion Matrix": cm.tolist(),
|
| 570 |
+
"TN": int(tn), "FP": int(fp), "FN": int(fn), "TP": int(tp)
|
| 571 |
+
}
|
backend/digitization_runner.py
ADDED
|
@@ -0,0 +1,174 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# backend/digitization_runner.py
|
| 2 |
+
"""
|
| 3 |
+
Backend adapter for ECG Digitization.
|
| 4 |
+
Wraps the ECGImage pipeline from digitization.py and manages model caching.
|
| 5 |
+
"""
|
| 6 |
+
|
| 7 |
+
import os
|
| 8 |
+
import sys
|
| 9 |
+
import pandas as pd
|
| 10 |
+
import numpy as np
|
| 11 |
+
import streamlit as st
|
| 12 |
+
import torch
|
| 13 |
+
|
| 14 |
+
# Add parent directory to path to ensure digitization and other local modules are importable
|
| 15 |
+
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
| 16 |
+
|
| 17 |
+
try:
|
| 18 |
+
from digitization import ECGImage
|
| 19 |
+
from ultralytics import YOLO
|
| 20 |
+
except ImportError as e:
|
| 21 |
+
# We will log the error but keep running so the dashboard can load and display elegant warnings
|
| 22 |
+
pass
|
| 23 |
+
|
| 24 |
+
class ModelPathNotSetError(Exception):
|
| 25 |
+
"""Raised when one or more YOLO model paths are not set in the configuration."""
|
| 26 |
+
pass
|
| 27 |
+
|
| 28 |
+
class ModelFileNotFoundError(Exception):
|
| 29 |
+
"""Raised when a specified YOLO model file does not exist on disk."""
|
| 30 |
+
pass
|
| 31 |
+
|
| 32 |
+
def check_models_configured(config):
|
| 33 |
+
"""
|
| 34 |
+
Verify that all YOLO model paths are set in the configuration.
|
| 35 |
+
"""
|
| 36 |
+
missing_configs = []
|
| 37 |
+
if not config.YOLO_BOX_MODEL_PATH:
|
| 38 |
+
missing_configs.append("YOLO Box Model (YOLO_BOX_MODEL_PATH)")
|
| 39 |
+
if not config.YOLO_LEAD_NAME_MODEL_PATH:
|
| 40 |
+
missing_configs.append("YOLO Lead Name Model (YOLO_LEAD_NAME_MODEL_PATH)")
|
| 41 |
+
if not config.YOLO_PULSE_MODEL_PATH:
|
| 42 |
+
missing_configs.append("YOLO Pulse Model (YOLO_PULSE_MODEL_PATH)")
|
| 43 |
+
if not config.YOLO_SEGMENTATION_MODEL_PATH:
|
| 44 |
+
missing_configs.append("YOLO Segmentation Model (YOLO_SEGMENTATION_MODEL_PATH)")
|
| 45 |
+
|
| 46 |
+
if missing_configs:
|
| 47 |
+
raise ModelPathNotSetError(
|
| 48 |
+
f"The following YOLO model paths are not set in config.py:\n" +
|
| 49 |
+
"\n".join([f"- {m}" for m in missing_configs]) +
|
| 50 |
+
"\n\nPlease configure these paths in config.py before running digitization."
|
| 51 |
+
)
|
| 52 |
+
|
| 53 |
+
def check_model_files_exist(config):
|
| 54 |
+
"""
|
| 55 |
+
Verify that the configured model files actually exist on disk.
|
| 56 |
+
"""
|
| 57 |
+
missing_files = []
|
| 58 |
+
for name, path in [
|
| 59 |
+
("Box Model", config.YOLO_BOX_MODEL_PATH),
|
| 60 |
+
("Lead Name Model", config.YOLO_LEAD_NAME_MODEL_PATH),
|
| 61 |
+
("Pulse Model", config.YOLO_PULSE_MODEL_PATH),
|
| 62 |
+
("Segmentation Model", config.YOLO_SEGMENTATION_MODEL_PATH)
|
| 63 |
+
]:
|
| 64 |
+
if path and not os.path.exists(path):
|
| 65 |
+
missing_files.append(f"{name} at path: '{path}'")
|
| 66 |
+
|
| 67 |
+
if missing_files:
|
| 68 |
+
raise ModelFileNotFoundError(
|
| 69 |
+
f"The following model files were not found on disk:\n" +
|
| 70 |
+
"\n".join([f"- {f}" for f in missing_files]) +
|
| 71 |
+
"\n\nPlease verify that the files are downloaded and located at these paths."
|
| 72 |
+
)
|
| 73 |
+
|
| 74 |
+
@st.cache_resource
|
| 75 |
+
def load_yolo_models(box_path, lead_name_path, pulse_path):
|
| 76 |
+
"""
|
| 77 |
+
Load YOLO models once and cache them in Streamlit memory.
|
| 78 |
+
Returns a dict containing the loaded model objects.
|
| 79 |
+
"""
|
| 80 |
+
device = "cuda" if torch.cuda.is_available() else "cpu"
|
| 81 |
+
|
| 82 |
+
# Load each model
|
| 83 |
+
box_model = YOLO(box_path)
|
| 84 |
+
lead_name_model = YOLO(lead_name_path)
|
| 85 |
+
pulse_model = YOLO(pulse_path)
|
| 86 |
+
|
| 87 |
+
# Move to GPU if available
|
| 88 |
+
box_model.to(device)
|
| 89 |
+
lead_name_model.to(device)
|
| 90 |
+
pulse_model.to(device)
|
| 91 |
+
|
| 92 |
+
return {
|
| 93 |
+
"box_model": box_model,
|
| 94 |
+
"lead_name_model": lead_name_model,
|
| 95 |
+
"pulse_model": pulse_model,
|
| 96 |
+
"device": device
|
| 97 |
+
}
|
| 98 |
+
|
| 99 |
+
def run_digitization_pipeline(image_path, config):
|
| 100 |
+
"""
|
| 101 |
+
Runs the full digitization pipeline on a single ECG image.
|
| 102 |
+
|
| 103 |
+
Parameters
|
| 104 |
+
----------
|
| 105 |
+
image_path : str
|
| 106 |
+
Path to the uploaded image file.
|
| 107 |
+
config : module
|
| 108 |
+
The config module containing paths and settings.
|
| 109 |
+
|
| 110 |
+
Returns
|
| 111 |
+
-------
|
| 112 |
+
df : pandas.DataFrame
|
| 113 |
+
Digitized ECG signals with lead names as columns and timesteps as rows.
|
| 114 |
+
ecg_instance : ECGImage
|
| 115 |
+
The raw ECGImage instance for accessing signal_grid, masks, or metadata.
|
| 116 |
+
"""
|
| 117 |
+
# 1. Validation
|
| 118 |
+
check_models_configured(config)
|
| 119 |
+
check_model_files_exist(config)
|
| 120 |
+
|
| 121 |
+
# 2. Load models (cached via st.cache_resource)
|
| 122 |
+
models = load_yolo_models(
|
| 123 |
+
config.YOLO_BOX_MODEL_PATH,
|
| 124 |
+
config.YOLO_LEAD_NAME_MODEL_PATH,
|
| 125 |
+
config.YOLO_PULSE_MODEL_PATH
|
| 126 |
+
)
|
| 127 |
+
|
| 128 |
+
# 3. Instantiate ECGImage
|
| 129 |
+
ecg = ECGImage(
|
| 130 |
+
box_model=models["box_model"],
|
| 131 |
+
segmentation_model=config.YOLO_SEGMENTATION_MODEL_PATH,
|
| 132 |
+
lead_name_model=models["lead_name_model"],
|
| 133 |
+
pulse_model=models["pulse_model"],
|
| 134 |
+
image_path=image_path
|
| 135 |
+
)
|
| 136 |
+
|
| 137 |
+
# 4. Run the pipeline (without Linux-specific SIGALRM timeouts)
|
| 138 |
+
ecg.run_full_pipeline()
|
| 139 |
+
|
| 140 |
+
# 5. Extract DataFrame from signal_grid
|
| 141 |
+
df = extract_dataframe_from_ecg(ecg)
|
| 142 |
+
|
| 143 |
+
# 6. Clear GPU cache if needed
|
| 144 |
+
if torch.cuda.is_available():
|
| 145 |
+
try:
|
| 146 |
+
torch.cuda.synchronize()
|
| 147 |
+
torch.cuda.empty_cache()
|
| 148 |
+
except Exception:
|
| 149 |
+
pass
|
| 150 |
+
|
| 151 |
+
return df, ecg
|
| 152 |
+
|
| 153 |
+
def extract_dataframe_from_ecg(ecg_instance):
|
| 154 |
+
"""
|
| 155 |
+
Extracts signal data from an ECGImage instance and converts it to a pandas DataFrame.
|
| 156 |
+
"""
|
| 157 |
+
signals, lead_names, max_length = [], [], 0
|
| 158 |
+
for row in ecg_instance.signal_grid:
|
| 159 |
+
for cell in row:
|
| 160 |
+
if 'signal' in cell and 'lead' in cell:
|
| 161 |
+
sig = np.asarray(cell['signal'], dtype=float)
|
| 162 |
+
signals.append(sig)
|
| 163 |
+
lead_names.append(cell['lead'])
|
| 164 |
+
max_length = max(max_length, len(sig))
|
| 165 |
+
|
| 166 |
+
if not signals:
|
| 167 |
+
raise ValueError("No valid lead signals were extracted from the ECG image.")
|
| 168 |
+
|
| 169 |
+
# Pad shorter signals with NaN to make them uniform length
|
| 170 |
+
padded = [np.pad(s, (0, max_length - len(s)), mode='constant', constant_values=np.nan) for s in signals]
|
| 171 |
+
|
| 172 |
+
# Build DataFrame
|
| 173 |
+
df = pd.DataFrame(np.vstack(padded).T, columns=lead_names)
|
| 174 |
+
return df
|
config.py
ADDED
|
@@ -0,0 +1,64 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# config.py
|
| 2 |
+
"""
|
| 3 |
+
Centralized configuration settings for the ECG Digitization and Classification Dashboard.
|
| 4 |
+
This holds paths, hyperparameter defaults, and metadata for classification tasks.
|
| 5 |
+
"""
|
| 6 |
+
|
| 7 |
+
import os
|
| 8 |
+
|
| 9 |
+
# ==============================================================================
|
| 10 |
+
# YOLO Models Settings (For Digitization)
|
| 11 |
+
# ==============================================================================
|
| 12 |
+
# These paths are initially left empty as requested. The user will set them in their environment
|
| 13 |
+
# or directly in this file. The Streamlit app will warn if they are empty.
|
| 14 |
+
YOLO_BOX_MODEL_PATH = "models/digitization_models/yolo11_full/weights/best.pt" # Bounding boxes for full lead regions
|
| 15 |
+
YOLO_LEAD_NAME_MODEL_PATH = "models/digitization_models/yolo11_lead/weights/best.pt" # Classification for lead labels (I, II, etc.)
|
| 16 |
+
YOLO_PULSE_MODEL_PATH = "models/digitization_models/yolo11_pulse/weights/best.pt" # Detection for calibration reference pulses
|
| 17 |
+
YOLO_SEGMENTATION_MODEL_PATH = "models/digitization_models/yolo11_patch/weights/best.pt" # Path to the YOLO patch segmentation model checkpoint
|
| 18 |
+
|
| 19 |
+
# ==============================================================================
|
| 20 |
+
# ECG Signal Metadata
|
| 21 |
+
# ==============================================================================
|
| 22 |
+
# Target sampling frequency for digitization and classification (in Hz)
|
| 23 |
+
TARGET_FS = 500
|
| 24 |
+
|
| 25 |
+
# Ordered standard list of 12 ECG Leads
|
| 26 |
+
LEAD_NAMES = ['I', 'aVR', 'V1', 'V4', 'II', 'aVL', 'V2', 'V5', 'III', 'aVF', 'V3', 'V6']
|
| 27 |
+
|
| 28 |
+
# ==============================================================================
|
| 29 |
+
# Classification Task Mapping
|
| 30 |
+
# ==============================================================================
|
| 31 |
+
# Maps the task selected in the Streamlit UI to its associated model and target class labels.
|
| 32 |
+
CLASSIFICATION_TASKS = {
|
| 33 |
+
"Normal vs Myocardial Infarction (MI) - Segmented": {
|
| 34 |
+
"model_dir": "mi_vs_normal_segmented",
|
| 35 |
+
"model": "Arsenal",
|
| 36 |
+
"labels": {0: "NORMAL", 1: "MYOCARDIAL_INFARCTION"},
|
| 37 |
+
"description": "Differentiates normal healthy heartbeats from Myocardial Infarction using the pre-trained Arsenal ensemble classifier (Segmented Heartbeats)."
|
| 38 |
+
},
|
| 39 |
+
"Occlusive Myocardial Infarction (OMI) vs non-OMI": {
|
| 40 |
+
"model_dir": "omi_vs_nonomi",
|
| 41 |
+
"model": "Rocket",
|
| 42 |
+
"labels": {0: "non-OMI", 1: "OMI"},
|
| 43 |
+
"description": "Identifies Occlusive Myocardial Infarction (OMI) vs Non-Occlusive Myocardial Infarction/controls using the pre-trained Rocket classifier."
|
| 44 |
+
},
|
| 45 |
+
"Pre-Procedural vs Post-Procedural MI": {
|
| 46 |
+
"model_dir": "ecg_surgery",
|
| 47 |
+
"model": "InceptionTime",
|
| 48 |
+
"labels": {0: "post-procedural MI", 1: "pre-procedural MI"},
|
| 49 |
+
"description": "Classifies pre-procedural MI vs post-procedural MI using the pre-trained InceptionTime deep learning time-series classifier."
|
| 50 |
+
}
|
| 51 |
+
}
|
| 52 |
+
|
| 53 |
+
|
| 54 |
+
# ==============================================================================
|
| 55 |
+
# File Management
|
| 56 |
+
# ==============================================================================
|
| 57 |
+
# Base output directories for storing run results
|
| 58 |
+
OUTPUT_BASE_DIR = "output"
|
| 59 |
+
DIGITIZATION_OUTPUT_DIR = os.path.join(OUTPUT_BASE_DIR, "digitization")
|
| 60 |
+
CLASSIFICATION_OUTPUT_DIR = os.path.join(OUTPUT_BASE_DIR, "classification")
|
| 61 |
+
|
| 62 |
+
# Ensure base directories exist
|
| 63 |
+
os.makedirs(DIGITIZATION_OUTPUT_DIR, exist_ok=True)
|
| 64 |
+
os.makedirs(CLASSIFICATION_OUTPUT_DIR, exist_ok=True)
|
digitization.py
ADDED
|
@@ -0,0 +1,1784 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
digitization.py
|
| 3 |
+
===============
|
| 4 |
+
Core ECG image-to-signal conversion module.
|
| 5 |
+
|
| 6 |
+
The public interface is the :class:`ECGImage` class. Instantiate it with
|
| 7 |
+
four pre-loaded YOLO models and the path to an ECG image, then call
|
| 8 |
+
:meth:`ECGImage.run_full_pipeline` followed by
|
| 9 |
+
:meth:`ECGImage.save_signals_as_csv` (or :meth:`ECGImage.save_signals_as_wfdb`).
|
| 10 |
+
|
| 11 |
+
Module-level helpers
|
| 12 |
+
--------------------
|
| 13 |
+
plot_image — Quick matplotlib preview of a grayscale array.
|
| 14 |
+
shadow_removal — Morphological background subtraction.
|
| 15 |
+
line_length — Euclidean distance between two 2-D points.
|
| 16 |
+
parse_layout_from_folder — Parse rows/cols/layout flags from a folder name.
|
| 17 |
+
"""
|
| 18 |
+
|
| 19 |
+
import os
|
| 20 |
+
os.environ["OPENBLAS_NUM_THREADS"] = "1"
|
| 21 |
+
os.environ["OMP_NUM_THREADS"] = "1"
|
| 22 |
+
os.environ["MKL_NUM_THREADS"] = "1"
|
| 23 |
+
|
| 24 |
+
import numpy as np
|
| 25 |
+
import cv2
|
| 26 |
+
cv2.setNumThreads(0) # disable OpenCV's internal thread pool
|
| 27 |
+
|
| 28 |
+
import matplotlib.pyplot as plt
|
| 29 |
+
from ultralytics import YOLO
|
| 30 |
+
from PIL import Image
|
| 31 |
+
from sklearn.cluster import KMeans
|
| 32 |
+
from sklearn.metrics import silhouette_score, mean_squared_error
|
| 33 |
+
from scipy.interpolate import interp1d
|
| 34 |
+
import wfdb
|
| 35 |
+
from scipy import signal
|
| 36 |
+
from skimage import morphology, segmentation
|
| 37 |
+
from scipy.signal import savgol_filter, find_peaks
|
| 38 |
+
from scipy.stats import pearsonr
|
| 39 |
+
from skimage.filters import threshold_multiotsu
|
| 40 |
+
from concurrent.futures import ThreadPoolExecutor
|
| 41 |
+
import re
|
| 42 |
+
import pandas as pd
|
| 43 |
+
import torch
|
| 44 |
+
import time
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
from patched_yolo_infer import (
|
| 48 |
+
MakeCropsDetectThem,
|
| 49 |
+
CombineDetections,
|
| 50 |
+
visualize_results,
|
| 51 |
+
)
|
| 52 |
+
|
| 53 |
+
|
| 54 |
+
def plot_image(img, title="Image Plot", size=(12, 12), show_axis=False):
|
| 55 |
+
"""Display a grayscale image with matplotlib.
|
| 56 |
+
|
| 57 |
+
Parameters
|
| 58 |
+
----------
|
| 59 |
+
img : np.ndarray
|
| 60 |
+
Grayscale image array (H × W).
|
| 61 |
+
title : str, optional
|
| 62 |
+
Figure title (default ``"Image Plot"``).
|
| 63 |
+
size : tuple[int, int], optional
|
| 64 |
+
Figure size in inches ``(width, height)`` (default ``(12, 12)``).
|
| 65 |
+
show_axis : bool, optional
|
| 66 |
+
Whether to draw axis ticks (default ``False``).
|
| 67 |
+
"""
|
| 68 |
+
plt.figure(figsize=size)
|
| 69 |
+
plt.imshow(img, cmap='gray')
|
| 70 |
+
plt.title(title)
|
| 71 |
+
if not show_axis:
|
| 72 |
+
plt.axis('off')
|
| 73 |
+
plt.show()
|
| 74 |
+
|
| 75 |
+
|
| 76 |
+
def shadow_removal(img):
|
| 77 |
+
"""Remove uneven illumination / shadow from a grayscale image.
|
| 78 |
+
|
| 79 |
+
Uses morphological dilation followed by median blur to estimate the
|
| 80 |
+
background, then subtracts it from the original to yield a
|
| 81 |
+
normalised, shadow-free image.
|
| 82 |
+
|
| 83 |
+
Parameters
|
| 84 |
+
----------
|
| 85 |
+
img : np.ndarray
|
| 86 |
+
Grayscale uint8 image.
|
| 87 |
+
|
| 88 |
+
Returns
|
| 89 |
+
-------
|
| 90 |
+
np.ndarray
|
| 91 |
+
Shadow-corrected uint8 image, intensity range [0, 255].
|
| 92 |
+
"""
|
| 93 |
+
dilated_img = cv2.dilate(img, np.ones((7, 7), np.uint8))
|
| 94 |
+
bg_img = cv2.medianBlur(dilated_img, 15)
|
| 95 |
+
diff_img = 255 - cv2.absdiff(img, bg_img)
|
| 96 |
+
norm_img = cv2.normalize(diff_img, None, alpha=0, beta=255,
|
| 97 |
+
norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8UC1)
|
| 98 |
+
return norm_img
|
| 99 |
+
|
| 100 |
+
|
| 101 |
+
def line_length(x1, y1, x2, y2):
|
| 102 |
+
"""Return the Euclidean length of a line segment.
|
| 103 |
+
|
| 104 |
+
Parameters
|
| 105 |
+
----------
|
| 106 |
+
x1, y1 : float
|
| 107 |
+
Coordinates of the first endpoint.
|
| 108 |
+
x2, y2 : float
|
| 109 |
+
Coordinates of the second endpoint.
|
| 110 |
+
|
| 111 |
+
Returns
|
| 112 |
+
-------
|
| 113 |
+
float
|
| 114 |
+
Length of the segment in the same units as the inputs.
|
| 115 |
+
"""
|
| 116 |
+
return np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
|
| 117 |
+
|
| 118 |
+
|
| 119 |
+
def parse_layout_from_folder(folder_path):
|
| 120 |
+
"""Parse ECG layout metadata encoded in a folder name.
|
| 121 |
+
|
| 122 |
+
Expected naming convention::
|
| 123 |
+
|
| 124 |
+
ecg_signals_<rows>x<cols>_<Rythm|None>_<Cabrera|Normal>
|
| 125 |
+
|
| 126 |
+
Parameters
|
| 127 |
+
----------
|
| 128 |
+
folder_path : str
|
| 129 |
+
Absolute or relative path whose final component encodes the layout.
|
| 130 |
+
|
| 131 |
+
Returns
|
| 132 |
+
-------
|
| 133 |
+
layout_key : tuple[int, int, bool] or None
|
| 134 |
+
``(rows, cols, is_cabrera)``. ``None`` if the name does not match.
|
| 135 |
+
calibration : bool or None
|
| 136 |
+
``True`` when a rhythm (calibration) strip is present, ``None`` on
|
| 137 |
+
parse failure.
|
| 138 |
+
"""
|
| 139 |
+
base_name = folder_path.split('/')[-1]
|
| 140 |
+
match = re.search(r'ecg_signals_(\d+)x(\d+)_(None|Rythm)_(Cabrera|Normal)', base_name)
|
| 141 |
+
if match:
|
| 142 |
+
rows = int(match.group(1))
|
| 143 |
+
cols = int(match.group(2))
|
| 144 |
+
calibration = (match.group(3) == 'Rythm')
|
| 145 |
+
cabrera_flag = (match.group(4) == 'Cabrera')
|
| 146 |
+
layout_key = (rows, cols, cabrera_flag)
|
| 147 |
+
return layout_key, calibration
|
| 148 |
+
else:
|
| 149 |
+
return None, None
|
| 150 |
+
|
| 151 |
+
|
| 152 |
+
class ECGImage:
|
| 153 |
+
"""End-to-end digitizer for a single 12-lead ECG image.
|
| 154 |
+
|
| 155 |
+
The class encapsulates every processing stage — from raw pixel
|
| 156 |
+
loading to per-lead voltage time-series export. All intermediate
|
| 157 |
+
results (masks, bounding boxes, calibration constants, signal grid)
|
| 158 |
+
are stored as instance attributes so that individual stages can be
|
| 159 |
+
inspected or visualised after the pipeline runs.
|
| 160 |
+
|
| 161 |
+
Typical usage
|
| 162 |
+
-------------
|
| 163 |
+
>>> ecg = ECGImage(box_model, seg_model, name_model, pulse_model, "ecg.png")
|
| 164 |
+
>>> ecg.run_full_pipeline()
|
| 165 |
+
>>> ecg.save_signals_as_csv("record_001", directory="./output")
|
| 166 |
+
|
| 167 |
+
Parameters
|
| 168 |
+
----------
|
| 169 |
+
box_model : ultralytics.YOLO
|
| 170 |
+
Pre-loaded YOLO model that detects lead bounding boxes.
|
| 171 |
+
segmentation_model : str
|
| 172 |
+
Path to the YOLO segmentation model checkpoint. Loaded
|
| 173 |
+
internally by ``patched_yolo_infer.MakeCropsDetectThem`` so
|
| 174 |
+
that it can operate on image patches.
|
| 175 |
+
lead_name_model : ultralytics.YOLO
|
| 176 |
+
Pre-loaded YOLO model that classifies lead name labels
|
| 177 |
+
(I, II, III, aVR, aVL, aVF, V1–V6).
|
| 178 |
+
pulse_model : ultralytics.YOLO
|
| 179 |
+
Pre-loaded YOLO model that detects calibration pulse boxes.
|
| 180 |
+
image_path : str
|
| 181 |
+
Path to the ECG image (.png, .jpg, or .jpeg).
|
| 182 |
+
wfdb_path : str, optional
|
| 183 |
+
Path to a WFDB record (without extension) used for ground-truth
|
| 184 |
+
comparison in :meth:`calculate_metrics_ptb`. Not required for
|
| 185 |
+
normal digitization.
|
| 186 |
+
|
| 187 |
+
Attributes set by the pipeline
|
| 188 |
+
--------------------------------
|
| 189 |
+
image : np.ndarray
|
| 190 |
+
Loaded and padded grayscale image.
|
| 191 |
+
processed_image : np.ndarray
|
| 192 |
+
Shadow-removed, blurred version used for all model inference.
|
| 193 |
+
lead_segmentation : list
|
| 194 |
+
``CombineDetections`` objects from the three segmentation passes.
|
| 195 |
+
mask_image : np.ndarray
|
| 196 |
+
Binary mask (H × W, uint8) combining all segmented lead polygons.
|
| 197 |
+
row_centers : np.ndarray
|
| 198 |
+
Y-pixel positions of detected lead row centres.
|
| 199 |
+
roi : tuple[float, float]
|
| 200 |
+
``(min_y, max_y)`` bounding the active ECG area.
|
| 201 |
+
lead_bboxes : list[list[float]]
|
| 202 |
+
``[x1, y1, x2, y2]`` boxes from ``box_model``.
|
| 203 |
+
lead_name_bboxes : list[dict]
|
| 204 |
+
``{'bbox': [...], 'class_name': str}`` from ``lead_name_model``.
|
| 205 |
+
reference_pulses : list[dict]
|
| 206 |
+
``{'bbox': [...], 'image': np.ndarray}`` from ``pulse_model``.
|
| 207 |
+
volt_per_pixel : float
|
| 208 |
+
Calibration constant: millivolts per pixel (vertical axis).
|
| 209 |
+
time_per_pixel : float
|
| 210 |
+
Calibration constant: seconds per pixel (horizontal axis).
|
| 211 |
+
is_cabrera : bool
|
| 212 |
+
Whether the ECG uses Cabrera lead ordering.
|
| 213 |
+
has_calibration_pulse : bool
|
| 214 |
+
Whether a rhythm/calibration strip row is present.
|
| 215 |
+
layout : list[list[str]]
|
| 216 |
+
2-D grid of lead names matching the detected row/column layout.
|
| 217 |
+
grid : list[list[dict]]
|
| 218 |
+
Mask grid — each cell contains ``'lead'`` (str) and ``'signal'``
|
| 219 |
+
(np.ndarray of the binary mask slice).
|
| 220 |
+
signal_grid : list[list[dict]]
|
| 221 |
+
After :meth:`extract_signals`: cells additionally contain
|
| 222 |
+
``'signal'`` as a list of float mV values.
|
| 223 |
+
"""
|
| 224 |
+
|
| 225 |
+
def __init__(self, box_model, segmentation_model, lead_name_model, pulse_model,
|
| 226 |
+
image_path, wfdb_path=""):
|
| 227 |
+
self.image_path = image_path # stored for retry in run_full_pipeline
|
| 228 |
+
self.load_image(image_path)
|
| 229 |
+
self.wfdb_path = wfdb_path
|
| 230 |
+
|
| 231 |
+
self.box_model = box_model
|
| 232 |
+
self.segmentation_model = segmentation_model
|
| 233 |
+
self.lead_name_model = lead_name_model
|
| 234 |
+
self.pulse_model = pulse_model
|
| 235 |
+
|
| 236 |
+
self.is_cabrera = None
|
| 237 |
+
self.has_calibration_pulse = None
|
| 238 |
+
|
| 239 |
+
self.standard_layouts = {
|
| 240 |
+
# Standard layouts
|
| 241 |
+
(12, 1, False): [['I'], ['II'], ['III'], ['aVR'], ['aVL'], ['aVF'],
|
| 242 |
+
['V1'], ['V2'], ['V3'], ['V4'], ['V5'], ['V6']],
|
| 243 |
+
(6, 2, False): [['I', 'V1'], ['II', 'V2'], ['III', 'V3'],
|
| 244 |
+
['aVR', 'V4'], ['aVL', 'V5'], ['aVF', 'V6']],
|
| 245 |
+
(4, 3, False): [['I', 'II', 'III'], ['aVR', 'aVL', 'aVF'],
|
| 246 |
+
['V1', 'V2', 'V3'], ['V4', 'V5', 'V6']],
|
| 247 |
+
(3, 4, False): [['I', 'aVR', 'V1', 'V4'], ['II', 'aVL', 'V2', 'V5'],
|
| 248 |
+
['III', 'aVF', 'V3', 'V6']],
|
| 249 |
+
# Cabrera layouts
|
| 250 |
+
(12, 1, True): [['aVL'], ['I'], ['aVR'], ['II'], ['aVF'], ['III'],
|
| 251 |
+
['V1'], ['V2'], ['V3'], ['V4'], ['V5'], ['V6']],
|
| 252 |
+
(6, 2, True): [['aVL', 'V1'], ['I', 'V2'], ['aVR', 'V3'],
|
| 253 |
+
['II', 'V4'], ['aVF', 'V5'], ['III', 'V6']],
|
| 254 |
+
(4, 3, True): [['aVL', 'I', 'aVR'], ['II', 'aVF', 'III'],
|
| 255 |
+
['V1', 'V2', 'V3'], ['V4', 'V5', 'V6']],
|
| 256 |
+
(3, 4, True): [['aVL', 'II', 'V1', 'V4'], ['I', 'aVF', 'V2', 'V5'],
|
| 257 |
+
['aVR', 'III', 'V3', 'V6']],
|
| 258 |
+
}
|
| 259 |
+
|
| 260 |
+
# ------------------------------------------------------------------
|
| 261 |
+
# Image loading & preprocessing
|
| 262 |
+
# ------------------------------------------------------------------
|
| 263 |
+
|
| 264 |
+
def load_image(self, path, target_size=2100):
|
| 265 |
+
"""Load, resize, and pad an ECG image.
|
| 266 |
+
|
| 267 |
+
The image is scaled so that its height equals *target_size* (cubic
|
| 268 |
+
interpolation) and then padded with a 20-pixel white border on all
|
| 269 |
+
sides. The result is stored in ``self.image``.
|
| 270 |
+
|
| 271 |
+
Parameters
|
| 272 |
+
----------
|
| 273 |
+
path : str
|
| 274 |
+
Path to the image file.
|
| 275 |
+
target_size : int, optional
|
| 276 |
+
Target image height in pixels (default ``2100``).
|
| 277 |
+
|
| 278 |
+
Raises
|
| 279 |
+
------
|
| 280 |
+
FileNotFoundError
|
| 281 |
+
If OpenCV cannot open the file at *path*.
|
| 282 |
+
"""
|
| 283 |
+
img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
|
| 284 |
+
if img is None:
|
| 285 |
+
raise FileNotFoundError(f"Cannot read image: {path}")
|
| 286 |
+
resample_factor = target_size / img.shape[0]
|
| 287 |
+
img = cv2.resize(img,
|
| 288 |
+
(int(img.shape[1] * resample_factor),
|
| 289 |
+
int(img.shape[0] * resample_factor)),
|
| 290 |
+
interpolation=cv2.INTER_CUBIC)
|
| 291 |
+
self.image = cv2.copyMakeBorder(img, 20, 20, 20, 20,
|
| 292 |
+
cv2.BORDER_CONSTANT, value=[255, 255, 255])
|
| 293 |
+
|
| 294 |
+
def preprocess_image(self):
|
| 295 |
+
"""Apply shadow removal and Gaussian smoothing to ``self.image``.
|
| 296 |
+
|
| 297 |
+
The result is stored in ``self.processed_image`` and is used as
|
| 298 |
+
input for all subsequent YOLO inference calls.
|
| 299 |
+
"""
|
| 300 |
+
rem = shadow_removal(self.image)
|
| 301 |
+
self.processed_image = cv2.GaussianBlur(rem, (3, 3), 0)
|
| 302 |
+
|
| 303 |
+
# ------------------------------------------------------------------
|
| 304 |
+
# Segmentation
|
| 305 |
+
# ------------------------------------------------------------------
|
| 306 |
+
|
| 307 |
+
def segment_leads(self):
|
| 308 |
+
"""Run patched YOLO instance segmentation at three crop scales.
|
| 309 |
+
|
| 310 |
+
The segmentation model is applied to overlapping image patches at
|
| 311 |
+
crop factors of 4, 4.5, and 5 (relative to image height).
|
| 312 |
+
Detections from each pass are combined with NMS
|
| 313 |
+
(``nms_threshold=0.5``) and stored in ``self.lead_segmentation``
|
| 314 |
+
as a list of three ``CombineDetections`` objects.
|
| 315 |
+
"""
|
| 316 |
+
segmentations = []
|
| 317 |
+
for shape in [4, 4.5, 5]:
|
| 318 |
+
element_crops = MakeCropsDetectThem(
|
| 319 |
+
image=cv2.cvtColor(self.processed_image, cv2.COLOR_GRAY2BGR),
|
| 320 |
+
model_path=self.segmentation_model,
|
| 321 |
+
segment=True,
|
| 322 |
+
show_crops=False,
|
| 323 |
+
shape_x=int(self.processed_image.shape[0] // shape),
|
| 324 |
+
shape_y=int(self.processed_image.shape[0] // shape),
|
| 325 |
+
overlap_x=50,
|
| 326 |
+
overlap_y=50,
|
| 327 |
+
conf=0.8,
|
| 328 |
+
iou=0.7,
|
| 329 |
+
classes_list=[0],
|
| 330 |
+
)
|
| 331 |
+
segmentations.append(CombineDetections(element_crops, nms_threshold=0.5))
|
| 332 |
+
self.lead_segmentation = segmentations
|
| 333 |
+
|
| 334 |
+
def make_segmentation_mask(self):
|
| 335 |
+
"""Rasterise all segmented polygons into a single binary mask.
|
| 336 |
+
|
| 337 |
+
Polygons from all three segmentation passes in
|
| 338 |
+
``self.lead_segmentation`` are filled and merged into a combined
|
| 339 |
+
mask. A morphological opening (5 × 5 kernel) removes small
|
| 340 |
+
spurious regions. The result is stored in ``self.mask_image``
|
| 341 |
+
(uint8, values 0 or 255).
|
| 342 |
+
"""
|
| 343 |
+
height, width = self.image.shape[:2]
|
| 344 |
+
combined_mask = np.zeros((height, width), dtype=np.uint8)
|
| 345 |
+
|
| 346 |
+
for segmentation in self.lead_segmentation:
|
| 347 |
+
for poly in segmentation.filtered_polygons:
|
| 348 |
+
pts = np.array(poly, dtype=np.int32)
|
| 349 |
+
if pts.ndim == 2:
|
| 350 |
+
pts = [pts]
|
| 351 |
+
cv2.fillPoly(combined_mask, pts, color=255)
|
| 352 |
+
|
| 353 |
+
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
|
| 354 |
+
self.mask_image = cv2.morphologyEx(combined_mask, cv2.MORPH_OPEN, kernel)
|
| 355 |
+
|
| 356 |
+
# ------------------------------------------------------------------
|
| 357 |
+
# Row detection
|
| 358 |
+
# ------------------------------------------------------------------
|
| 359 |
+
|
| 360 |
+
def find_row_centers(self):
|
| 361 |
+
"""Detect the vertical centre of each ECG lead row.
|
| 362 |
+
|
| 363 |
+
Projects the binary mask onto the vertical axis (row sums) and
|
| 364 |
+
applies ``scipy.signal.find_peaks`` twice: once to identify all
|
| 365 |
+
peaks, and a second time with an adaptive minimum distance
|
| 366 |
+
(⅔ × mean inter-peak spacing) to consolidate merged rows.
|
| 367 |
+
|
| 368 |
+
Sets
|
| 369 |
+
----
|
| 370 |
+
self.row_centers : np.ndarray
|
| 371 |
+
Y-pixel indices of each row centre.
|
| 372 |
+
self.first_peak_start : int
|
| 373 |
+
Y-pixel where the topmost row's signal begins.
|
| 374 |
+
self.last_peak_end : int
|
| 375 |
+
Y-pixel where the bottommost row's signal ends.
|
| 376 |
+
"""
|
| 377 |
+
image_height, image_width = self.mask_image.shape[:2]
|
| 378 |
+
proj = np.sum(self.mask_image, 1)
|
| 379 |
+
|
| 380 |
+
height = (image_width // 10) * 255
|
| 381 |
+
distance = image_height // 30
|
| 382 |
+
|
| 383 |
+
peaks, _ = find_peaks(proj, height=height, distance=distance)
|
| 384 |
+
row_centers, _ = find_peaks(
|
| 385 |
+
proj, height=height,
|
| 386 |
+
distance=int(np.mean(np.diff(peaks)) * (2 / 3))
|
| 387 |
+
)
|
| 388 |
+
self.row_centers = row_centers
|
| 389 |
+
|
| 390 |
+
if len(peaks) > 0:
|
| 391 |
+
first_peak = peaks[0]
|
| 392 |
+
start_index = first_peak
|
| 393 |
+
zero_gap = 0
|
| 394 |
+
for i in range(first_peak - 1, -1, -1):
|
| 395 |
+
if proj[i] == 0:
|
| 396 |
+
zero_gap += 1
|
| 397 |
+
if zero_gap > 2:
|
| 398 |
+
break
|
| 399 |
+
else:
|
| 400 |
+
zero_gap = 0
|
| 401 |
+
start_index = i
|
| 402 |
+
self.first_peak_start = start_index
|
| 403 |
+
|
| 404 |
+
last_peak = peaks[-1]
|
| 405 |
+
end_index = last_peak
|
| 406 |
+
zero_gap = 0
|
| 407 |
+
for i in range(last_peak + 1, image_height):
|
| 408 |
+
if proj[i] == 0:
|
| 409 |
+
zero_gap += 1
|
| 410 |
+
if zero_gap > 2:
|
| 411 |
+
break
|
| 412 |
+
else:
|
| 413 |
+
zero_gap = 0
|
| 414 |
+
end_index = i
|
| 415 |
+
self.last_peak_end = end_index
|
| 416 |
+
|
| 417 |
+
def get_roi(self):
|
| 418 |
+
"""Compute the vertical region-of-interest (ROI) spanning all lead rows.
|
| 419 |
+
|
| 420 |
+
Extends ⅔ of the mean row spacing above the first row centre and
|
| 421 |
+
below the last. Stores the result as ``self.roi = (min_y, max_y)``.
|
| 422 |
+
"""
|
| 423 |
+
spacing = 2 / 3 * np.mean(np.diff(self.row_centers))
|
| 424 |
+
min_y = max(0, self.row_centers[0] - spacing)
|
| 425 |
+
max_y = min(self.image.shape[0], self.row_centers[-1] + spacing)
|
| 426 |
+
self.roi = (min_y, max_y)
|
| 427 |
+
|
| 428 |
+
# ------------------------------------------------------------------
|
| 429 |
+
# YOLO inference helpers — GPU-safe wrappers
|
| 430 |
+
# ------------------------------------------------------------------
|
| 431 |
+
|
| 432 |
+
def _predict_safe(self, model, image_bgr, **kwargs):
|
| 433 |
+
"""Run a YOLO model prediction in a thread-safe manner.
|
| 434 |
+
|
| 435 |
+
On CUDA: wraps the call in a dedicated ``torch.cuda.Stream`` and
|
| 436 |
+
synchronises before returning, preventing contention on the
|
| 437 |
+
default stream when multiple threads call this simultaneously.
|
| 438 |
+
On CPU: wraps with ``torch.no_grad()`` only.
|
| 439 |
+
|
| 440 |
+
Parameters
|
| 441 |
+
----------
|
| 442 |
+
model : ultralytics.YOLO
|
| 443 |
+
The detection/classification model to run.
|
| 444 |
+
image_bgr : np.ndarray
|
| 445 |
+
BGR image array as expected by Ultralytics.
|
| 446 |
+
**kwargs
|
| 447 |
+
Forwarded to ``model.predict()`` (e.g. ``conf``, ``iou``).
|
| 448 |
+
|
| 449 |
+
Returns
|
| 450 |
+
-------
|
| 451 |
+
list
|
| 452 |
+
Ultralytics ``Results`` objects.
|
| 453 |
+
"""
|
| 454 |
+
if torch.cuda.is_available():
|
| 455 |
+
stream = torch.cuda.Stream()
|
| 456 |
+
with torch.no_grad():
|
| 457 |
+
with torch.cuda.stream(stream):
|
| 458 |
+
results = model.predict(image_bgr, verbose=False, **kwargs)
|
| 459 |
+
stream.synchronize()
|
| 460 |
+
else:
|
| 461 |
+
with torch.no_grad():
|
| 462 |
+
results = model.predict(image_bgr, verbose=False, **kwargs)
|
| 463 |
+
return results
|
| 464 |
+
|
| 465 |
+
# ------------------------------------------------------------------
|
| 466 |
+
# Timing helper
|
| 467 |
+
# ------------------------------------------------------------------
|
| 468 |
+
|
| 469 |
+
@staticmethod
|
| 470 |
+
def _timed(label, fn, *args, **kwargs):
|
| 471 |
+
"""Call *fn* with *args*/*kwargs*, print its wall-clock duration, and return its result.
|
| 472 |
+
|
| 473 |
+
Parameters
|
| 474 |
+
----------
|
| 475 |
+
label : str
|
| 476 |
+
Human-readable stage name printed alongside the elapsed time.
|
| 477 |
+
fn : callable
|
| 478 |
+
The stage function to time.
|
| 479 |
+
*args, **kwargs
|
| 480 |
+
Forwarded to *fn*.
|
| 481 |
+
|
| 482 |
+
Returns
|
| 483 |
+
-------
|
| 484 |
+
Any
|
| 485 |
+
Whatever *fn* returns.
|
| 486 |
+
"""
|
| 487 |
+
t = time.time()
|
| 488 |
+
result = fn(*args, **kwargs)
|
| 489 |
+
print(f" ⏱ {label}: {time.time() - t:.1f}s", flush=True)
|
| 490 |
+
return result
|
| 491 |
+
|
| 492 |
+
# ------------------------------------------------------------------
|
| 493 |
+
# Box & name extraction
|
| 494 |
+
# ------------------------------------------------------------------
|
| 495 |
+
|
| 496 |
+
def extract_lead_boxes(self):
|
| 497 |
+
"""Detect lead bounding boxes with ``box_model`` and filter to the ROI.
|
| 498 |
+
|
| 499 |
+
Runs ``_predict_safe`` on the preprocessed BGR image at confidence
|
| 500 |
+
0.8, then discards boxes whose ``y1``/``y2`` coordinates fall
|
| 501 |
+
outside ``self.roi``. Sets ``self.lead_bboxes`` as a list of
|
| 502 |
+
``[x1, y1, x2, y2]`` float lists.
|
| 503 |
+
"""
|
| 504 |
+
image_bgr = cv2.cvtColor(self.processed_image, cv2.COLOR_GRAY2BGR)
|
| 505 |
+
results = self._predict_safe(self.box_model, image_bgr, conf=0.8)
|
| 506 |
+
|
| 507 |
+
min_y, max_y = self.roi
|
| 508 |
+
lead_boxes = []
|
| 509 |
+
for r in results:
|
| 510 |
+
for box in r.boxes:
|
| 511 |
+
x1, y1, x2, y2 = box.xyxy.cpu().numpy()[0].tolist()
|
| 512 |
+
if y1 >= min_y and y2 <= max_y:
|
| 513 |
+
lead_boxes.append([x1, y1, x2, y2])
|
| 514 |
+
self.lead_bboxes = lead_boxes
|
| 515 |
+
|
| 516 |
+
def extract_lead_name_boxes(self):
|
| 517 |
+
"""Detect and classify lead name labels with ``lead_name_model``.
|
| 518 |
+
|
| 519 |
+
Runs at confidence 0.8, filters to the ROI, and stores results in
|
| 520 |
+
``self.lead_name_bboxes`` as a list of
|
| 521 |
+
``{'bbox': [x1, y1, x2, y2], 'class_name': str}`` dicts.
|
| 522 |
+
"""
|
| 523 |
+
image_bgr = cv2.cvtColor(self.processed_image, cv2.COLOR_GRAY2BGR)
|
| 524 |
+
results = self._predict_safe(self.lead_name_model, image_bgr, conf=0.8)
|
| 525 |
+
|
| 526 |
+
min_y, max_y = self.roi
|
| 527 |
+
name_boxes = []
|
| 528 |
+
for r in results:
|
| 529 |
+
for box in r.boxes:
|
| 530 |
+
x1, y1, x2, y2 = box.xyxy.cpu().numpy()[0].tolist()
|
| 531 |
+
if y1 >= min_y and y2 <= max_y:
|
| 532 |
+
cls_id = int(box.cls.cpu().numpy()[0])
|
| 533 |
+
cls_name = self.lead_name_model.names[cls_id]
|
| 534 |
+
name_boxes.append({'bbox': [x1, y1, x2, y2], 'class_name': cls_name})
|
| 535 |
+
self.lead_name_bboxes = name_boxes
|
| 536 |
+
|
| 537 |
+
def extract_reference_pulses(self):
|
| 538 |
+
"""Detect calibration pulse boxes with ``pulse_model``.
|
| 539 |
+
|
| 540 |
+
Runs at confidence 0.7 (lower than lead detection to improve
|
| 541 |
+
recall on small pulses). For each detection, stores the
|
| 542 |
+
corresponding image crop alongside the bounding-box coordinates.
|
| 543 |
+
Sets ``self.reference_pulses`` as a list of
|
| 544 |
+
``{'bbox': [x1, y1, x2, y2], 'image': np.ndarray}`` dicts.
|
| 545 |
+
"""
|
| 546 |
+
image_bgr = cv2.cvtColor(self.processed_image, cv2.COLOR_GRAY2BGR)
|
| 547 |
+
results = self._predict_safe(self.pulse_model, image_bgr, conf=0.7)
|
| 548 |
+
|
| 549 |
+
pulse_boxes = []
|
| 550 |
+
for r in results:
|
| 551 |
+
for box in r.boxes:
|
| 552 |
+
coord = box.xyxy.cpu().numpy()[0].tolist()
|
| 553 |
+
pulse_boxes.append({
|
| 554 |
+
'image': self.image[int(coord[1]) - 5:int(coord[3]) + 5,
|
| 555 |
+
int(coord[0]) - 5:int(coord[2]) + 5],
|
| 556 |
+
'bbox': coord,
|
| 557 |
+
})
|
| 558 |
+
self.reference_pulses = pulse_boxes
|
| 559 |
+
|
| 560 |
+
# ------------------------------------------------------------------
|
| 561 |
+
# Visualisation helpers
|
| 562 |
+
# ------------------------------------------------------------------
|
| 563 |
+
|
| 564 |
+
def visualize_boxes(self, task='Lead name', show_axis=False):
|
| 565 |
+
"""Overlay detection bounding boxes on the original image.
|
| 566 |
+
|
| 567 |
+
Parameters
|
| 568 |
+
----------
|
| 569 |
+
task : {'Lead name', 'Lead box', 'Reference pulse'}
|
| 570 |
+
Which set of boxes to draw.
|
| 571 |
+
show_axis : bool, optional
|
| 572 |
+
Whether to display axis ticks (default ``False``).
|
| 573 |
+
"""
|
| 574 |
+
img_copy = cv2.cvtColor(self.image.copy(), cv2.COLOR_GRAY2BGR)
|
| 575 |
+
|
| 576 |
+
if task == 'Lead name':
|
| 577 |
+
if not self.lead_name_bboxes:
|
| 578 |
+
print("Lead name boxes not extracted.")
|
| 579 |
+
return
|
| 580 |
+
for box in self.lead_name_bboxes:
|
| 581 |
+
x1, y1, x2, y2 = map(int, box['bbox'])
|
| 582 |
+
cv2.rectangle(img_copy, (x1, y1), (x2, y2), (255, 0, 0), 2)
|
| 583 |
+
cv2.putText(img_copy, box['class_name'], (x1, y1 - 10),
|
| 584 |
+
cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 3)
|
| 585 |
+
|
| 586 |
+
elif task == 'Lead box':
|
| 587 |
+
if not self.lead_bboxes:
|
| 588 |
+
print("Lead boxes not extracted.")
|
| 589 |
+
return
|
| 590 |
+
for bbox in self.lead_bboxes:
|
| 591 |
+
x1, y1, x2, y2 = map(int, bbox)
|
| 592 |
+
cv2.rectangle(img_copy, (x1, y1), (x2, y2), (255, 0, 0), 2)
|
| 593 |
+
|
| 594 |
+
elif task == 'Reference pulse':
|
| 595 |
+
if not self.reference_pulses:
|
| 596 |
+
print("Reference pulses not extracted.")
|
| 597 |
+
return
|
| 598 |
+
for bbox in self.reference_pulses:
|
| 599 |
+
x1, y1, x2, y2 = map(int, bbox['bbox'])
|
| 600 |
+
cv2.rectangle(img_copy, (x1, y1), (x2, y2), (255, 0, 0), 2)
|
| 601 |
+
|
| 602 |
+
else:
|
| 603 |
+
print(f"Unknown task: {task}")
|
| 604 |
+
return
|
| 605 |
+
|
| 606 |
+
plt.figure(figsize=(12, 10))
|
| 607 |
+
plt.imshow(cv2.cvtColor(img_copy, cv2.COLOR_BGR2RGB))
|
| 608 |
+
if not show_axis:
|
| 609 |
+
plt.axis('off')
|
| 610 |
+
plt.show()
|
| 611 |
+
|
| 612 |
+
def visualize_segmentation(self, show_boxes=False, show_axis=False,
|
| 613 |
+
fill_mask=True, thickness=1):
|
| 614 |
+
"""Render all segmented lead polygons on the original image.
|
| 615 |
+
|
| 616 |
+
Aggregates results from all three segmentation passes in
|
| 617 |
+
``self.lead_segmentation`` and passes them to
|
| 618 |
+
``patched_yolo_infer.visualize_results``.
|
| 619 |
+
|
| 620 |
+
Parameters
|
| 621 |
+
----------
|
| 622 |
+
show_boxes : bool, optional
|
| 623 |
+
Also draw bounding boxes around each polygon (default ``False``).
|
| 624 |
+
show_axis : bool, optional
|
| 625 |
+
Display axis ticks (default ``False``).
|
| 626 |
+
fill_mask : bool, optional
|
| 627 |
+
Fill polygon interiors with a translucent colour (default ``True``).
|
| 628 |
+
thickness : int, optional
|
| 629 |
+
Polygon outline thickness in pixels (default ``1``).
|
| 630 |
+
"""
|
| 631 |
+
all_confidences, all_boxes, all_polygons = [], [], []
|
| 632 |
+
all_classes_ids, all_classes_names = [], []
|
| 633 |
+
|
| 634 |
+
for seg in self.lead_segmentation:
|
| 635 |
+
all_confidences.extend(seg.filtered_confidences)
|
| 636 |
+
all_boxes.extend(seg.filtered_boxes)
|
| 637 |
+
all_polygons.extend(seg.filtered_polygons)
|
| 638 |
+
all_classes_ids.extend(seg.filtered_classes_id)
|
| 639 |
+
all_classes_names.extend(seg.filtered_classes_names)
|
| 640 |
+
|
| 641 |
+
visualize_results(
|
| 642 |
+
img=cv2.cvtColor(self.image, cv2.COLOR_GRAY2RGB),
|
| 643 |
+
confidences=all_confidences,
|
| 644 |
+
boxes=all_boxes,
|
| 645 |
+
polygons=all_polygons,
|
| 646 |
+
classes_ids=all_classes_ids,
|
| 647 |
+
classes_names=all_classes_names,
|
| 648 |
+
segment=True,
|
| 649 |
+
thickness=thickness,
|
| 650 |
+
fill_mask=fill_mask,
|
| 651 |
+
show_boxes=show_boxes,
|
| 652 |
+
show_class=False,
|
| 653 |
+
axis_off=(not show_axis),
|
| 654 |
+
)
|
| 655 |
+
|
| 656 |
+
# ------------------------------------------------------------------
|
| 657 |
+
# Calibration (volt/pixel & time/pixel)
|
| 658 |
+
# ------------------------------------------------------------------
|
| 659 |
+
|
| 660 |
+
def get_reference_scale(self):
|
| 661 |
+
"""Derive pixel-to-physical calibration constants from the calibration pulses.
|
| 662 |
+
|
| 663 |
+
For each pulse in ``self.reference_pulses``, the method:
|
| 664 |
+
|
| 665 |
+
1. Enhances contrast with CLAHE and binarises with multi-Otsu thresholding.
|
| 666 |
+
2. Detects horizontal and vertical lines via morphological filtering and
|
| 667 |
+
Hough transform to measure the pulse amplitude (voltage) and width (time).
|
| 668 |
+
3. Refines the time calibration using the Line Segment Detector (LSD).
|
| 669 |
+
|
| 670 |
+
Sets
|
| 671 |
+
----
|
| 672 |
+
self.volt_per_pixel : float
|
| 673 |
+
Millivolts per pixel (vertical axis). Calibrated to the standard
|
| 674 |
+
1 mV / 10 mm calibration pulse.
|
| 675 |
+
self.time_per_pixel : float
|
| 676 |
+
Seconds per pixel (horizontal axis). Calibrated to 0.2 s between
|
| 677 |
+
the two inner vertical edges of the calibration pulse.
|
| 678 |
+
|
| 679 |
+
Raises
|
| 680 |
+
------
|
| 681 |
+
RuntimeError
|
| 682 |
+
If no valid pulse yields voltage or time measurements.
|
| 683 |
+
"""
|
| 684 |
+
voltages, times, dist = [], [], []
|
| 685 |
+
|
| 686 |
+
def _line_length(x1, y1, x2, y2):
|
| 687 |
+
return np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
|
| 688 |
+
|
| 689 |
+
for pulse_idx, pulse in enumerate(self.reference_pulses):
|
| 690 |
+
try:
|
| 691 |
+
img = pulse['image']
|
| 692 |
+
if img is None or img.size == 0:
|
| 693 |
+
continue
|
| 694 |
+
h, w = img.shape
|
| 695 |
+
if h < 10 or w < 10:
|
| 696 |
+
continue
|
| 697 |
+
|
| 698 |
+
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
|
| 699 |
+
enhanced = clahe.apply(img)
|
| 700 |
+
blurred = cv2.GaussianBlur(enhanced, (3, 3), 0)
|
| 701 |
+
|
| 702 |
+
thresholds = threshold_multiotsu(blurred, classes=2)
|
| 703 |
+
regions = np.digitize(blurred, bins=thresholds)
|
| 704 |
+
binary = (regions == 0).astype(np.uint8) * 255
|
| 705 |
+
|
| 706 |
+
h_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (25, 1))
|
| 707 |
+
v_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 25))
|
| 708 |
+
h_lines_img = cv2.morphologyEx(binary, cv2.MORPH_OPEN, h_kernel)
|
| 709 |
+
v_lines_img = cv2.morphologyEx(binary, cv2.MORPH_OPEN, v_kernel)
|
| 710 |
+
|
| 711 |
+
h_lines_raw = cv2.HoughLinesP(
|
| 712 |
+
h_lines_img, rho=1, theta=np.pi / 180,
|
| 713 |
+
threshold=w // 4, minLineLength=w // 2, maxLineGap=1
|
| 714 |
+
)
|
| 715 |
+
v_lines_raw = cv2.HoughLinesP(
|
| 716 |
+
v_lines_img, rho=1, theta=np.pi / 180,
|
| 717 |
+
threshold=h // 4, minLineLength=h // 2, maxLineGap=1
|
| 718 |
+
)
|
| 719 |
+
|
| 720 |
+
combined_lines = []
|
| 721 |
+
if h_lines_raw is not None:
|
| 722 |
+
combined_lines.extend(h_lines_raw)
|
| 723 |
+
if v_lines_raw is not None:
|
| 724 |
+
combined_lines.extend(v_lines_raw)
|
| 725 |
+
|
| 726 |
+
horizontal_lines, vertical_lines = [], []
|
| 727 |
+
for line in combined_lines:
|
| 728 |
+
x1, y1, x2, y2 = line[0]
|
| 729 |
+
angle = np.arctan2(abs(y2 - y1), abs(x2 - x1)) * 180 / np.pi
|
| 730 |
+
if angle < 10 or angle > 170:
|
| 731 |
+
horizontal_lines.append((x1, y1, x2, y2))
|
| 732 |
+
elif 80 < angle < 100:
|
| 733 |
+
vertical_lines.append((x1, y1, x2, y2))
|
| 734 |
+
|
| 735 |
+
if len(vertical_lines) >= 2:
|
| 736 |
+
longest_v = sorted(vertical_lines,
|
| 737 |
+
key=lambda l: _line_length(*l), reverse=True)[:2]
|
| 738 |
+
xc1 = (longest_v[0][0] + longest_v[0][2]) / 2
|
| 739 |
+
xc2 = (longest_v[1][0] + longest_v[1][2]) / 2
|
| 740 |
+
vertical_spacing = abs(xc1 - xc2)
|
| 741 |
+
|
| 742 |
+
volt_vals = sorted([_line_length(*l) for l in vertical_lines],
|
| 743 |
+
reverse=True)[:2]
|
| 744 |
+
if volt_vals:
|
| 745 |
+
voltages.append(volt_vals)
|
| 746 |
+
if vertical_spacing > 5:
|
| 747 |
+
times.append(vertical_spacing)
|
| 748 |
+
|
| 749 |
+
# LSD for more precise time calibration
|
| 750 |
+
lsd = cv2.createLineSegmentDetector(refine=2)
|
| 751 |
+
lsd_lines, _, _, _ = lsd.detect(v_lines_img)
|
| 752 |
+
|
| 753 |
+
min_length = h / 3
|
| 754 |
+
angle_tol_rad = np.deg2rad(5)
|
| 755 |
+
filtered = []
|
| 756 |
+
|
| 757 |
+
if lsd_lines is not None:
|
| 758 |
+
for line in lsd_lines:
|
| 759 |
+
x1, y1, x2, y2 = line[0]
|
| 760 |
+
length = np.hypot(x2 - x1, y2 - y1)
|
| 761 |
+
if length >= min_length:
|
| 762 |
+
angle = np.arctan2(y2 - y1, x2 - x1)
|
| 763 |
+
if np.abs(np.abs(angle) - np.pi / 2) <= angle_tol_rad:
|
| 764 |
+
filtered.append([[x1, y1, x2, y2]])
|
| 765 |
+
|
| 766 |
+
if len(filtered) >= 4:
|
| 767 |
+
filtered.sort(
|
| 768 |
+
key=lambda l: np.hypot(l[0][2] - l[0][0], l[0][3] - l[0][1]),
|
| 769 |
+
reverse=True
|
| 770 |
+
)
|
| 771 |
+
top4 = filtered[:4]
|
| 772 |
+
top4.sort(key=lambda l: (l[0][0] + l[0][2]) / 2)
|
| 773 |
+
|
| 774 |
+
midpoints = []
|
| 775 |
+
for l1, l2 in [(top4[0], top4[1]), (top4[2], top4[3])]:
|
| 776 |
+
xc1 = (l1[0][0] + l1[0][2]) / 2
|
| 777 |
+
xc2 = (l2[0][0] + l2[0][2]) / 2
|
| 778 |
+
midpoints.append((xc1 + xc2) / 2)
|
| 779 |
+
|
| 780 |
+
dist.append(midpoints[1] - midpoints[0])
|
| 781 |
+
|
| 782 |
+
except Exception as e:
|
| 783 |
+
print(f" ⚠ Skipping pulse {pulse_idx}: {e}", flush=True)
|
| 784 |
+
continue
|
| 785 |
+
|
| 786 |
+
if not voltages:
|
| 787 |
+
raise RuntimeError("No valid reference pulses found for voltage calibration.")
|
| 788 |
+
if not dist:
|
| 789 |
+
raise RuntimeError("No valid reference pulses found for time calibration.")
|
| 790 |
+
|
| 791 |
+
self.volt_per_pixel = 1 / np.mean(voltages)
|
| 792 |
+
self.time_per_pixel = 0.2 / np.mean(dist)
|
| 793 |
+
|
| 794 |
+
# ------------------------------------------------------------------
|
| 795 |
+
# Layout helpers
|
| 796 |
+
# ------------------------------------------------------------------
|
| 797 |
+
|
| 798 |
+
def make_bounding_box_features(self, box, axis):
|
| 799 |
+
"""Extract three positional features from a bounding box along one axis.
|
| 800 |
+
|
| 801 |
+
Parameters
|
| 802 |
+
----------
|
| 803 |
+
box : list[float]
|
| 804 |
+
Bounding box as ``[x1, y1, x2, y2]``.
|
| 805 |
+
axis : {'x', 'y'}
|
| 806 |
+
Axis along which to compute features.
|
| 807 |
+
|
| 808 |
+
Returns
|
| 809 |
+
-------
|
| 810 |
+
list[float]
|
| 811 |
+
``[axis_min, axis_center, axis_max]``.
|
| 812 |
+
"""
|
| 813 |
+
if axis == 'y':
|
| 814 |
+
axis_min, axis_max = box[1], box[3]
|
| 815 |
+
else:
|
| 816 |
+
axis_min, axis_max = box[0], box[2]
|
| 817 |
+
axis_center = (axis_min + axis_max) / 2
|
| 818 |
+
return [axis_min, axis_center, axis_max]
|
| 819 |
+
|
| 820 |
+
def bounding_boxes_kmeans(self, bounding_boxes, axis='y',
|
| 821 |
+
k_min=1, k_max=13, return_model=True):
|
| 822 |
+
"""Cluster bounding boxes along one axis using K-Means.
|
| 823 |
+
|
| 824 |
+
Tries all values of *k* in ``[k_min, k_max]`` and selects the one
|
| 825 |
+
that maximises the silhouette score. Cluster labels are sorted by
|
| 826 |
+
mean position along *axis* (ascending).
|
| 827 |
+
|
| 828 |
+
Parameters
|
| 829 |
+
----------
|
| 830 |
+
bounding_boxes : list[list[float]]
|
| 831 |
+
List of ``[x1, y1, x2, y2]`` boxes.
|
| 832 |
+
axis : {'x', 'y'}, optional
|
| 833 |
+
Axis along which to cluster (default ``'y'``).
|
| 834 |
+
k_min : int, optional
|
| 835 |
+
Minimum number of clusters (default ``1``).
|
| 836 |
+
k_max : int, optional
|
| 837 |
+
Maximum number of clusters (default ``13``).
|
| 838 |
+
return_model : bool, optional
|
| 839 |
+
If ``True``, also return the fitted ``KMeans`` model and the
|
| 840 |
+
label remapping dict (default ``True``).
|
| 841 |
+
|
| 842 |
+
Returns
|
| 843 |
+
-------
|
| 844 |
+
sorted_labels : np.ndarray of int
|
| 845 |
+
Cluster index (sorted) for each input box.
|
| 846 |
+
best_k : int
|
| 847 |
+
Optimal number of clusters.
|
| 848 |
+
sorted_centers : np.ndarray
|
| 849 |
+
Mean positional value of each cluster, sorted ascending.
|
| 850 |
+
label_map : dict
|
| 851 |
+
Mapping from original K-Means label → sorted label.
|
| 852 |
+
best_model : sklearn.cluster.KMeans
|
| 853 |
+
Fitted model (only when ``return_model=True``).
|
| 854 |
+
|
| 855 |
+
Raises
|
| 856 |
+
------
|
| 857 |
+
ValueError
|
| 858 |
+
If *axis* is not ``'x'`` or ``'y'``, or if there are fewer
|
| 859 |
+
boxes than *k_min*.
|
| 860 |
+
"""
|
| 861 |
+
if axis not in ('x', 'y'):
|
| 862 |
+
raise ValueError("Axis must be 'x' or 'y'")
|
| 863 |
+
if len(bounding_boxes) < k_min:
|
| 864 |
+
raise ValueError("Not enough bounding boxes to cluster")
|
| 865 |
+
|
| 866 |
+
features = np.array([self.make_bounding_box_features(b, axis)
|
| 867 |
+
for b in bounding_boxes])
|
| 868 |
+
|
| 869 |
+
best_score, best_k = -1, k_min
|
| 870 |
+
best_labels = best_centers = best_model = None
|
| 871 |
+
|
| 872 |
+
for k in range(k_min, min(k_max + 1, len(bounding_boxes))):
|
| 873 |
+
kmeans = KMeans(n_clusters=k, random_state=42, n_init="auto")
|
| 874 |
+
labels = kmeans.fit_predict(features)
|
| 875 |
+
score = silhouette_score(features, labels)
|
| 876 |
+
if score > best_score:
|
| 877 |
+
best_score = score
|
| 878 |
+
best_k = k
|
| 879 |
+
best_labels = labels
|
| 880 |
+
best_centers = kmeans.cluster_centers_
|
| 881 |
+
best_model = kmeans
|
| 882 |
+
|
| 883 |
+
cluster_avgs = best_centers.mean(axis=1)
|
| 884 |
+
sorted_indices = np.argsort(cluster_avgs)
|
| 885 |
+
label_map = {old: new for new, old in enumerate(sorted_indices)}
|
| 886 |
+
sorted_labels = np.array([label_map[l] for l in best_labels])
|
| 887 |
+
sorted_centers = cluster_avgs[sorted_indices]
|
| 888 |
+
|
| 889 |
+
if return_model:
|
| 890 |
+
return sorted_labels, best_k, sorted_centers, label_map, best_model
|
| 891 |
+
return sorted_labels, best_k, sorted_centers, label_map
|
| 892 |
+
|
| 893 |
+
def check_cabrera(self, num_rows, num_cols):
|
| 894 |
+
"""Determine whether the ECG uses Cabrera lead ordering.
|
| 895 |
+
|
| 896 |
+
The heuristic varies by layout:
|
| 897 |
+
|
| 898 |
+
* **12- or 6-row layouts:** compares the vertical spacing of
|
| 899 |
+
augmented-limb leads (aVR/aVL/aVF) against precordial leads.
|
| 900 |
+
Cabrera re-orders the limb leads so their spacing differs.
|
| 901 |
+
* **4-row, 3-col / 5-row layouts:** checks the vertical spread
|
| 902 |
+
of augmented leads (high std → interleaved → Cabrera).
|
| 903 |
+
* **3-row / 4-col layouts:** checks the horizontal spread of
|
| 904 |
+
augmented leads.
|
| 905 |
+
|
| 906 |
+
Sets ``self.is_cabrera`` as a side-effect.
|
| 907 |
+
|
| 908 |
+
Parameters
|
| 909 |
+
----------
|
| 910 |
+
num_rows : int
|
| 911 |
+
Number of detected lead rows.
|
| 912 |
+
num_cols : int
|
| 913 |
+
Number of detected lead columns.
|
| 914 |
+
|
| 915 |
+
Returns
|
| 916 |
+
-------
|
| 917 |
+
bool
|
| 918 |
+
``True`` if Cabrera ordering is detected.
|
| 919 |
+
"""
|
| 920 |
+
if num_rows in [13, 12, 7, 6]:
|
| 921 |
+
av_leads = [b for b in self.lead_name_bboxes
|
| 922 |
+
if b['class_name'] in {'aVR', 'aVL', 'aVF'}]
|
| 923 |
+
v_leads = [b for b in self.lead_name_bboxes
|
| 924 |
+
if b['class_name'] in {'V1', 'V2', 'V3', 'V4', 'V5', 'V6'}]
|
| 925 |
+
|
| 926 |
+
if not av_leads or not v_leads:
|
| 927 |
+
return False
|
| 928 |
+
|
| 929 |
+
y_v = sorted([(b['bbox'][1] + b['bbox'][3]) / 2 for b in v_leads])
|
| 930 |
+
y_av = sorted([(b['bbox'][1] + b['bbox'][3]) / 2 for b in av_leads])
|
| 931 |
+
|
| 932 |
+
threshold = 30
|
| 933 |
+
diff_v = np.diff(y_v)
|
| 934 |
+
diff_av = np.diff(y_av)
|
| 935 |
+
diff_v = diff_v[np.abs(diff_v) > threshold]
|
| 936 |
+
diff_av = diff_av[np.abs(diff_av) > threshold]
|
| 937 |
+
|
| 938 |
+
if len(diff_v) == 0 or len(diff_av) == 0:
|
| 939 |
+
return False
|
| 940 |
+
|
| 941 |
+
if abs(np.min(diff_v) - np.min(diff_av)) > (0.25 * np.min(diff_v)):
|
| 942 |
+
self.is_cabrera = True
|
| 943 |
+
return True
|
| 944 |
+
self.is_cabrera = False
|
| 945 |
+
return False
|
| 946 |
+
|
| 947 |
+
elif (num_rows == 4 and num_cols == 3) or num_rows == 5:
|
| 948 |
+
av_leads = [b for b in self.lead_name_bboxes
|
| 949 |
+
if b['class_name'] in {'aVR', 'aVL', 'aVF'}]
|
| 950 |
+
if not av_leads:
|
| 951 |
+
return False
|
| 952 |
+
y_coords = [(b['bbox'][1] + b['bbox'][3]) / 2 for b in av_leads]
|
| 953 |
+
self.is_cabrera = np.std(y_coords) > 25
|
| 954 |
+
return self.is_cabrera
|
| 955 |
+
|
| 956 |
+
elif (num_rows == 4 and num_cols == 4) or num_rows == 3:
|
| 957 |
+
av_leads = [b for b in self.lead_name_bboxes
|
| 958 |
+
if b['class_name'] in {'aVR', 'aVL', 'aVF'}]
|
| 959 |
+
if not av_leads:
|
| 960 |
+
return False
|
| 961 |
+
x_coords = [(b['bbox'][0] + b['bbox'][2]) / 2 for b in av_leads]
|
| 962 |
+
self.is_cabrera = np.std(x_coords) > 25
|
| 963 |
+
return self.is_cabrera
|
| 964 |
+
|
| 965 |
+
def get_layout(self, num_rows):
|
| 966 |
+
"""Map the detected row count to a standard ECG layout.
|
| 967 |
+
|
| 968 |
+
Calls :meth:`check_cabrera` to determine lead ordering, sets
|
| 969 |
+
``self.layout`` (2-D list of lead name strings) and
|
| 970 |
+
``self.has_calibration_pulse``, and returns the number of columns.
|
| 971 |
+
|
| 972 |
+
Supported row counts
|
| 973 |
+
--------------------
|
| 974 |
+
* 13 → 12×1 with calibration pulse
|
| 975 |
+
* 12 → 12×1 without calibration pulse
|
| 976 |
+
* 7 → 6×2 with calibration pulse
|
| 977 |
+
* 6 → 6×2 without calibration pulse
|
| 978 |
+
* 5 → 4×3 with calibration pulse
|
| 979 |
+
* 4 → 3×4 or 4×3 (disambiguated by V-lead spatial distribution)
|
| 980 |
+
* 3 → 3×4 without calibration pulse
|
| 981 |
+
|
| 982 |
+
Parameters
|
| 983 |
+
----------
|
| 984 |
+
num_rows : int
|
| 985 |
+
Number of detected lead rows.
|
| 986 |
+
|
| 987 |
+
Returns
|
| 988 |
+
-------
|
| 989 |
+
int
|
| 990 |
+
Number of lead columns in the detected layout.
|
| 991 |
+
"""
|
| 992 |
+
if num_rows == 13:
|
| 993 |
+
num_cols = 1
|
| 994 |
+
self.has_calibration_pulse = True
|
| 995 |
+
cabrera = self.check_cabrera(num_rows, num_cols)
|
| 996 |
+
self.layout = self.standard_layouts[(12, 1, cabrera)]
|
| 997 |
+
|
| 998 |
+
elif num_rows == 12:
|
| 999 |
+
num_cols = 1
|
| 1000 |
+
self.has_calibration_pulse = False
|
| 1001 |
+
cabrera = self.check_cabrera(num_rows, num_cols)
|
| 1002 |
+
self.layout = self.standard_layouts[(12, 1, cabrera)]
|
| 1003 |
+
|
| 1004 |
+
elif num_rows == 7:
|
| 1005 |
+
num_cols = 2
|
| 1006 |
+
self.has_calibration_pulse = True
|
| 1007 |
+
cabrera = self.check_cabrera(num_rows, num_cols)
|
| 1008 |
+
self.layout = self.standard_layouts[(6, 2, cabrera)]
|
| 1009 |
+
|
| 1010 |
+
elif num_rows == 6:
|
| 1011 |
+
num_cols = 2
|
| 1012 |
+
self.has_calibration_pulse = False
|
| 1013 |
+
cabrera = self.check_cabrera(num_rows, num_cols)
|
| 1014 |
+
self.layout = self.standard_layouts[(6, 2, cabrera)]
|
| 1015 |
+
|
| 1016 |
+
elif num_rows == 5:
|
| 1017 |
+
num_cols = 3
|
| 1018 |
+
self.has_calibration_pulse = True
|
| 1019 |
+
cabrera = self.check_cabrera(num_rows, num_cols)
|
| 1020 |
+
self.layout = self.standard_layouts[(4, 3, cabrera)]
|
| 1021 |
+
|
| 1022 |
+
elif num_rows == 4:
|
| 1023 |
+
v_leads1 = [b for b in self.lead_name_bboxes
|
| 1024 |
+
if b['class_name'] in {'V1', 'V2', 'V3'}]
|
| 1025 |
+
v_leads2 = [b for b in self.lead_name_bboxes
|
| 1026 |
+
if b['class_name'] in {'V4', 'V5', 'V6'}]
|
| 1027 |
+
|
| 1028 |
+
def centers(boxes, axis):
|
| 1029 |
+
if axis == 'x':
|
| 1030 |
+
return [(b['bbox'][0] + b['bbox'][2]) / 2 for b in boxes]
|
| 1031 |
+
return [(b['bbox'][1] + b['bbox'][3]) / 2 for b in boxes]
|
| 1032 |
+
|
| 1033 |
+
x_std1 = np.std(centers(v_leads1, 'x'))
|
| 1034 |
+
y_std1 = np.std(centers(v_leads1, 'y'))
|
| 1035 |
+
x_std2 = np.std(centers(v_leads2, 'x'))
|
| 1036 |
+
y_std2 = np.std(centers(v_leads2, 'y'))
|
| 1037 |
+
|
| 1038 |
+
if x_std1 < y_std1 and x_std2 < y_std2:
|
| 1039 |
+
num_cols = 4
|
| 1040 |
+
self.has_calibration_pulse = True
|
| 1041 |
+
elif y_std1 < x_std1 and y_std2 < x_std2:
|
| 1042 |
+
num_cols = 3
|
| 1043 |
+
self.has_calibration_pulse = False
|
| 1044 |
+
else:
|
| 1045 |
+
num_cols = 4
|
| 1046 |
+
self.has_calibration_pulse = True
|
| 1047 |
+
|
| 1048 |
+
cabrera = self.check_cabrera(num_rows, num_cols)
|
| 1049 |
+
key = (3, 4, cabrera) if num_cols == 4 else (4, 3, cabrera)
|
| 1050 |
+
self.layout = self.standard_layouts[key]
|
| 1051 |
+
|
| 1052 |
+
elif num_rows == 3:
|
| 1053 |
+
num_cols = 4
|
| 1054 |
+
self.has_calibration_pulse = False
|
| 1055 |
+
cabrera = self.check_cabrera(num_rows, num_cols)
|
| 1056 |
+
self.layout = self.standard_layouts[(3, 4, cabrera)]
|
| 1057 |
+
|
| 1058 |
+
return num_cols
|
| 1059 |
+
|
| 1060 |
+
def clean_lead_mask(self, img, baseline_y):
|
| 1061 |
+
"""Remove leaked segments of neighboring leads from a single lead cell binary mask.
|
| 1062 |
+
|
| 1063 |
+
Parameters
|
| 1064 |
+
----------
|
| 1065 |
+
img : np.ndarray
|
| 1066 |
+
Binary mask for a single lead cell (H × W).
|
| 1067 |
+
baseline_y : float
|
| 1068 |
+
Y-coordinate of the row baseline relative to the cropped row mask.
|
| 1069 |
+
|
| 1070 |
+
Returns
|
| 1071 |
+
-------
|
| 1072 |
+
np.ndarray
|
| 1073 |
+
Cleaned binary mask.
|
| 1074 |
+
"""
|
| 1075 |
+
height, width = img.shape
|
| 1076 |
+
if height < 10 or width < 10:
|
| 1077 |
+
return img
|
| 1078 |
+
|
| 1079 |
+
num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(img)
|
| 1080 |
+
if num_labels <= 2:
|
| 1081 |
+
return img
|
| 1082 |
+
|
| 1083 |
+
# Find the primary component (largest horizontal span)
|
| 1084 |
+
primary_c = -1
|
| 1085 |
+
max_w = -1
|
| 1086 |
+
for c in range(1, num_labels):
|
| 1087 |
+
w = stats[c, cv2.CC_STAT_WIDTH]
|
| 1088 |
+
if w > max_w:
|
| 1089 |
+
max_w = w
|
| 1090 |
+
primary_c = c
|
| 1091 |
+
|
| 1092 |
+
if primary_c == -1:
|
| 1093 |
+
return img
|
| 1094 |
+
|
| 1095 |
+
primary_area = stats[primary_c, cv2.CC_STAT_AREA]
|
| 1096 |
+
cleaned_img = img.copy()
|
| 1097 |
+
|
| 1098 |
+
for c in range(1, num_labels):
|
| 1099 |
+
if c == primary_c:
|
| 1100 |
+
continue
|
| 1101 |
+
|
| 1102 |
+
c_top = stats[c, cv2.CC_STAT_TOP]
|
| 1103 |
+
c_height = stats[c, cv2.CC_STAT_HEIGHT]
|
| 1104 |
+
c_bottom = c_top + c_height
|
| 1105 |
+
c_width = stats[c, cv2.CC_STAT_WIDTH]
|
| 1106 |
+
c_area = stats[c, cv2.CC_STAT_AREA]
|
| 1107 |
+
c_centroid_y = centroids[c][1]
|
| 1108 |
+
|
| 1109 |
+
# Heuristics to identify a leak:
|
| 1110 |
+
# 1. It is relatively small (either narrow or small area)
|
| 1111 |
+
is_small = (c_width < 0.35 * width) or (c_area < 0.25 * primary_area)
|
| 1112 |
+
|
| 1113 |
+
# 2. It is located near the vertical boundaries of the cell
|
| 1114 |
+
near_boundary = (c_top < 0.15 * height) or (c_bottom > 0.85 * height)
|
| 1115 |
+
|
| 1116 |
+
# 3. It doesn't cross the baseline (with safety margin)
|
| 1117 |
+
baseline_margin = max(5, int(0.08 * height))
|
| 1118 |
+
crosses_baseline = (c_top - baseline_margin <= baseline_y <= c_bottom + baseline_margin)
|
| 1119 |
+
|
| 1120 |
+
# 4. It is far from the baseline
|
| 1121 |
+
far_from_baseline = abs(c_centroid_y - baseline_y) > 0.2 * height
|
| 1122 |
+
|
| 1123 |
+
if is_small and near_boundary and not crosses_baseline and far_from_baseline:
|
| 1124 |
+
cleaned_img[labels == c] = 0
|
| 1125 |
+
|
| 1126 |
+
return cleaned_img
|
| 1127 |
+
|
| 1128 |
+
# ------------------------------------------------------------------
|
| 1129 |
+
# Grid construction
|
| 1130 |
+
# ------------------------------------------------------------------
|
| 1131 |
+
|
| 1132 |
+
def make_grid(self, padding=0):
|
| 1133 |
+
"""Build the lead mask grid from segmentation polygons and lead boxes.
|
| 1134 |
+
|
| 1135 |
+
Steps:
|
| 1136 |
+
|
| 1137 |
+
1. Assigns each segmented polygon to the nearest row centre by
|
| 1138 |
+
polygon centroid Y-coordinate.
|
| 1139 |
+
2. Crops each row mask to its vertical extent.
|
| 1140 |
+
3. Clusters lead bounding boxes into *num_cols* columns with K-Means.
|
| 1141 |
+
4. Slices each (row, col) cell from the corresponding row mask.
|
| 1142 |
+
5. Applies a morphological opening to each cell slice.
|
| 1143 |
+
|
| 1144 |
+
Sets
|
| 1145 |
+
----
|
| 1146 |
+
self.grid : list[list[dict]]
|
| 1147 |
+
``grid[row][col]`` contains ``{'lead': str, 'signal': np.ndarray}``,
|
| 1148 |
+
where ``'signal'`` is currently a binary mask slice.
|
| 1149 |
+
self.baseline : list[float]
|
| 1150 |
+
Relative Y-position of the isoelectric baseline within each row
|
| 1151 |
+
(pixels from the top of the cropped row mask).
|
| 1152 |
+
self.row : dict[int, np.ndarray]
|
| 1153 |
+
Cropped binary mask for each row index.
|
| 1154 |
+
|
| 1155 |
+
Parameters
|
| 1156 |
+
----------
|
| 1157 |
+
padding : int, optional
|
| 1158 |
+
Extra pixel padding (currently unused, reserved for future use).
|
| 1159 |
+
"""
|
| 1160 |
+
image_height, image_width = self.image.shape[:2]
|
| 1161 |
+
num_rows = len(self.row_centers)
|
| 1162 |
+
num_cols = self.get_layout(num_rows)
|
| 1163 |
+
|
| 1164 |
+
row_masks = {i: np.zeros((image_height, image_width), dtype=np.uint8)
|
| 1165 |
+
for i in range(num_rows)}
|
| 1166 |
+
row_polygons = {i: [] for i in range(num_rows)}
|
| 1167 |
+
row_limits = {i: [] for i in range(num_rows)}
|
| 1168 |
+
|
| 1169 |
+
for seg in self.lead_segmentation:
|
| 1170 |
+
for box, polygon in zip(seg.filtered_boxes, seg.filtered_polygons):
|
| 1171 |
+
poly_np = np.array(polygon, dtype=np.int32).reshape((-1, 1, 2))
|
| 1172 |
+
temp_mask = np.zeros((image_height, image_width), dtype=np.uint8)
|
| 1173 |
+
cv2.fillPoly(temp_mask, [poly_np], color=1)
|
| 1174 |
+
|
| 1175 |
+
temp_proj = np.sum(temp_mask, axis=1)
|
| 1176 |
+
total_weight_y = np.sum(temp_proj)
|
| 1177 |
+
if total_weight_y == 0:
|
| 1178 |
+
continue
|
| 1179 |
+
|
| 1180 |
+
centroid_y = int(np.sum(np.arange(temp_proj.shape[0]) * temp_proj)
|
| 1181 |
+
/ total_weight_y)
|
| 1182 |
+
|
| 1183 |
+
y_vals = [pt[1] for pt in polygon]
|
| 1184 |
+
min_y, max_y = min(y_vals), max(y_vals)
|
| 1185 |
+
|
| 1186 |
+
if max_y < self.first_peak_start or min_y > self.last_peak_end:
|
| 1187 |
+
continue
|
| 1188 |
+
|
| 1189 |
+
closest_idx = int(np.argmin(np.abs(self.row_centers - centroid_y)))
|
| 1190 |
+
|
| 1191 |
+
if self.has_calibration_pulse and closest_idx == num_rows - 1:
|
| 1192 |
+
continue
|
| 1193 |
+
|
| 1194 |
+
row_polygons[closest_idx].append(polygon)
|
| 1195 |
+
cv2.fillPoly(row_masks[closest_idx], [poly_np], color=1)
|
| 1196 |
+
|
| 1197 |
+
if not row_limits[closest_idx]:
|
| 1198 |
+
row_limits[closest_idx] = [min_y, max_y]
|
| 1199 |
+
else:
|
| 1200 |
+
row_limits[closest_idx][0] = min(row_limits[closest_idx][0], min_y)
|
| 1201 |
+
row_limits[closest_idx][1] = max(row_limits[closest_idx][1], max_y)
|
| 1202 |
+
|
| 1203 |
+
cropped_row_masks = {}
|
| 1204 |
+
for i in range(num_rows):
|
| 1205 |
+
if self.has_calibration_pulse and i == num_rows - 1:
|
| 1206 |
+
continue
|
| 1207 |
+
if not row_limits[i]:
|
| 1208 |
+
continue
|
| 1209 |
+
min_y, max_y = row_limits[i]
|
| 1210 |
+
cropped_row_masks[i] = row_masks[i][min_y:max_y + 1, :]
|
| 1211 |
+
|
| 1212 |
+
lead_boxes = [
|
| 1213 |
+
box for box in self.lead_bboxes
|
| 1214 |
+
if not (
|
| 1215 |
+
self.has_calibration_pulse and
|
| 1216 |
+
int(np.argmin(np.abs(self.row_centers - (box[1] + box[3]) / 2)))
|
| 1217 |
+
== num_rows - 1
|
| 1218 |
+
)
|
| 1219 |
+
] if self.has_calibration_pulse else self.lead_bboxes
|
| 1220 |
+
|
| 1221 |
+
self.row = cropped_row_masks
|
| 1222 |
+
|
| 1223 |
+
if num_cols != 1:
|
| 1224 |
+
labels_cols, _, _, _ = self.bounding_boxes_kmeans(
|
| 1225 |
+
lead_boxes, axis='x', k_min=num_cols, k_max=num_cols, return_model=False
|
| 1226 |
+
)
|
| 1227 |
+
else:
|
| 1228 |
+
labels_cols = np.zeros(len(lead_boxes), dtype=int)
|
| 1229 |
+
|
| 1230 |
+
boxes_arr = np.array(lead_boxes)
|
| 1231 |
+
min_x_per_col, max_x_per_col = [], []
|
| 1232 |
+
for col_label in range(num_cols):
|
| 1233 |
+
col_boxes = boxes_arr[labels_cols == col_label]
|
| 1234 |
+
min_x_per_col.append(col_boxes[:, 0].min())
|
| 1235 |
+
max_x_per_col.append(col_boxes[:, 2].max())
|
| 1236 |
+
|
| 1237 |
+
mask_grid, relative_baselines = [], []
|
| 1238 |
+
for row_idx, row_slice in cropped_row_masks.items():
|
| 1239 |
+
row_cells = []
|
| 1240 |
+
for col_idx in range(num_cols):
|
| 1241 |
+
x_min = max(0, int(min_x_per_col[col_idx]))
|
| 1242 |
+
x_max = min(image_width, int(max_x_per_col[col_idx]))
|
| 1243 |
+
cell_slice = row_slice[:, x_min:x_max]
|
| 1244 |
+
|
| 1245 |
+
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
|
| 1246 |
+
cell_slice = cv2.morphologyEx(cell_slice, cv2.MORPH_OPEN, kernel)
|
| 1247 |
+
|
| 1248 |
+
# Post-processing: Clean up leaked segments from neighboring leads
|
| 1249 |
+
baseline_y = self.row_centers[row_idx] - row_limits[row_idx][0]
|
| 1250 |
+
cell_slice = self.clean_lead_mask(cell_slice, baseline_y)
|
| 1251 |
+
|
| 1252 |
+
row_cells.append({
|
| 1253 |
+
'lead': self.layout[row_idx][col_idx],
|
| 1254 |
+
'signal': cell_slice,
|
| 1255 |
+
})
|
| 1256 |
+
|
| 1257 |
+
relative_baselines.append(
|
| 1258 |
+
self.row_centers[row_idx] - row_limits[row_idx][0]
|
| 1259 |
+
)
|
| 1260 |
+
mask_grid.append(row_cells)
|
| 1261 |
+
|
| 1262 |
+
self.grid = mask_grid
|
| 1263 |
+
self.baseline = relative_baselines
|
| 1264 |
+
|
| 1265 |
+
def visualize_grid(self, figsize=(15, 10)):
|
| 1266 |
+
"""Display the lead mask grid — one subplot per (row, col) cell.
|
| 1267 |
+
|
| 1268 |
+
Parameters
|
| 1269 |
+
----------
|
| 1270 |
+
figsize : tuple[int, int], optional
|
| 1271 |
+
Overall figure size in inches (default ``(15, 10)``).
|
| 1272 |
+
|
| 1273 |
+
Raises
|
| 1274 |
+
------
|
| 1275 |
+
ValueError
|
| 1276 |
+
If :meth:`make_grid` has not been called yet.
|
| 1277 |
+
"""
|
| 1278 |
+
if not hasattr(self, 'grid') or not self.grid:
|
| 1279 |
+
raise ValueError("Grid not generated. Call make_grid() first.")
|
| 1280 |
+
|
| 1281 |
+
num_rows = len(self.grid)
|
| 1282 |
+
num_cols = len(self.grid[0]) if isinstance(self.grid[0], list) else 1
|
| 1283 |
+
|
| 1284 |
+
fig, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
|
| 1285 |
+
if num_cols == 1:
|
| 1286 |
+
axes = np.atleast_2d(axes).T
|
| 1287 |
+
elif num_rows == 1:
|
| 1288 |
+
axes = np.atleast_2d(axes)
|
| 1289 |
+
|
| 1290 |
+
for row_idx in range(num_rows):
|
| 1291 |
+
row = self.grid[row_idx] if isinstance(self.grid[row_idx], list) \
|
| 1292 |
+
else [self.grid[row_idx]]
|
| 1293 |
+
for col_idx, cell in enumerate(row):
|
| 1294 |
+
ax = axes[row_idx][col_idx]
|
| 1295 |
+
ax.imshow(cell['signal'], cmap='gray', aspect='auto')
|
| 1296 |
+
ax.set_title(cell['lead'], fontsize=10)
|
| 1297 |
+
ax.axis('off')
|
| 1298 |
+
|
| 1299 |
+
plt.tight_layout()
|
| 1300 |
+
plt.show()
|
| 1301 |
+
|
| 1302 |
+
# ------------------------------------------------------------------
|
| 1303 |
+
# Signal extraction
|
| 1304 |
+
# ------------------------------------------------------------------
|
| 1305 |
+
|
| 1306 |
+
def binarize_signal(self, img, window_length=11, polyorder=2):
|
| 1307 |
+
"""Extract a 1-D signal trace from a binary lead mask column by column.
|
| 1308 |
+
|
| 1309 |
+
For each image column, computes the intensity-weighted centroid of
|
| 1310 |
+
non-zero pixels. Uses the second derivative of the initial centroid
|
| 1311 |
+
trace (Savitzky-Golay filtered) to disambiguate ambiguous columns:
|
| 1312 |
+
upward concavity → take the topmost pixels; downward concavity →
|
| 1313 |
+
take the bottommost pixels; otherwise use all non-zero pixels.
|
| 1314 |
+
|
| 1315 |
+
Parameters
|
| 1316 |
+
----------
|
| 1317 |
+
img : np.ndarray
|
| 1318 |
+
Binary uint8 mask for a single lead cell (H × W).
|
| 1319 |
+
window_length : int, optional
|
| 1320 |
+
Savitzky-Golay filter window (must be odd, default ``11``).
|
| 1321 |
+
polyorder : int, optional
|
| 1322 |
+
Savitzky-Golay polynomial order (default ``2``).
|
| 1323 |
+
|
| 1324 |
+
Returns
|
| 1325 |
+
-------
|
| 1326 |
+
x_coords : np.ndarray of int
|
| 1327 |
+
Column indices where a signal was detected.
|
| 1328 |
+
final_signal : np.ndarray of float
|
| 1329 |
+
Weighted-centroid Y-coordinate for each column in *x_coords*.
|
| 1330 |
+
"""
|
| 1331 |
+
height, width = img.shape
|
| 1332 |
+
x_coords = []
|
| 1333 |
+
initial_signal = []
|
| 1334 |
+
|
| 1335 |
+
for col in range(width):
|
| 1336 |
+
column = img[:, col]
|
| 1337 |
+
if np.sum(column) == 0:
|
| 1338 |
+
continue
|
| 1339 |
+
|
| 1340 |
+
y_indices = np.arange(height)
|
| 1341 |
+
weights = column.astype(float)
|
| 1342 |
+
centroid = np.average(y_indices, weights=weights)
|
| 1343 |
+
initial_signal.append(centroid)
|
| 1344 |
+
x_coords.append(col)
|
| 1345 |
+
|
| 1346 |
+
x_coords = np.array(x_coords)
|
| 1347 |
+
initial_signal = np.array(initial_signal)
|
| 1348 |
+
|
| 1349 |
+
second_deriv = np.gradient(np.gradient(savgol_filter(initial_signal, window_length, polyorder)))
|
| 1350 |
+
second_deriv = savgol_filter(second_deriv, 11, 3)
|
| 1351 |
+
|
| 1352 |
+
x_out, final_signal = [], []
|
| 1353 |
+
for idx, col in enumerate(x_coords):
|
| 1354 |
+
column = img[:, col]
|
| 1355 |
+
nz_idx = np.where(column > 0)[0]
|
| 1356 |
+
if nz_idx.size == 0:
|
| 1357 |
+
continue
|
| 1358 |
+
|
| 1359 |
+
if second_deriv[idx] > 0.5:
|
| 1360 |
+
sel_idx = nz_idx[:5]
|
| 1361 |
+
elif second_deriv[idx] < -0.5:
|
| 1362 |
+
sel_idx = nz_idx[-5:]
|
| 1363 |
+
else:
|
| 1364 |
+
sel_idx = nz_idx
|
| 1365 |
+
|
| 1366 |
+
weights = column[sel_idx].astype(float)
|
| 1367 |
+
centroid = np.average(sel_idx, weights=weights)
|
| 1368 |
+
x_out.append(col)
|
| 1369 |
+
final_signal.append(centroid)
|
| 1370 |
+
|
| 1371 |
+
return np.array(x_out), np.array(final_signal)
|
| 1372 |
+
|
| 1373 |
+
def fill_gaps(self, x_coords, y_coords, method='linear'):
|
| 1374 |
+
"""Interpolate missing columns to produce a dense, evenly-spaced signal.
|
| 1375 |
+
|
| 1376 |
+
Parameters
|
| 1377 |
+
----------
|
| 1378 |
+
x_coords : np.ndarray of int
|
| 1379 |
+
Sparse column indices from :meth:`binarize_signal`.
|
| 1380 |
+
y_coords : np.ndarray of float
|
| 1381 |
+
Corresponding centroid Y-values.
|
| 1382 |
+
method : str, optional
|
| 1383 |
+
Interpolation kind passed to ``scipy.interpolate.interp1d``
|
| 1384 |
+
(default ``'linear'``).
|
| 1385 |
+
|
| 1386 |
+
Returns
|
| 1387 |
+
-------
|
| 1388 |
+
x_full : np.ndarray of int
|
| 1389 |
+
Dense array ``[x_coords[0], …, x_coords[-1]]``.
|
| 1390 |
+
y_interp : np.ndarray of float
|
| 1391 |
+
Interpolated Y-values at every integer column.
|
| 1392 |
+
"""
|
| 1393 |
+
x_full = np.arange(x_coords[0], x_coords[-1] + 1)
|
| 1394 |
+
interpolator = interp1d(x_coords, y_coords, kind=method, fill_value="extrapolate")
|
| 1395 |
+
y_interp = interpolator(x_full)
|
| 1396 |
+
return x_full, y_interp
|
| 1397 |
+
|
| 1398 |
+
def smooth_signal(self, sig, window_length=7, polyorder=4):
|
| 1399 |
+
"""Apply a Savitzky-Golay low-pass filter to a 1-D signal.
|
| 1400 |
+
|
| 1401 |
+
Parameters
|
| 1402 |
+
----------
|
| 1403 |
+
sig : np.ndarray
|
| 1404 |
+
Input signal array.
|
| 1405 |
+
window_length : int, optional
|
| 1406 |
+
Filter window size in samples (must be odd, default ``7``).
|
| 1407 |
+
polyorder : int, optional
|
| 1408 |
+
Polynomial order (default ``4``).
|
| 1409 |
+
|
| 1410 |
+
Returns
|
| 1411 |
+
-------
|
| 1412 |
+
np.ndarray
|
| 1413 |
+
Smoothed signal of the same length as *sig*.
|
| 1414 |
+
"""
|
| 1415 |
+
return savgol_filter(sig, window_length, polyorder)
|
| 1416 |
+
|
| 1417 |
+
def extract_signals(self):
|
| 1418 |
+
"""Convert mask grid cells into calibrated voltage time-series.
|
| 1419 |
+
|
| 1420 |
+
For each cell in ``self.grid``:
|
| 1421 |
+
|
| 1422 |
+
1. Calls :meth:`binarize_signal` to get column-centroid Y-values.
|
| 1423 |
+
2. Trims 5 edge samples, fills gaps, and smooths with Savitzky-Golay.
|
| 1424 |
+
3. Converts pixel Y-coordinates to millivolts using
|
| 1425 |
+
``self.baseline`` and ``self.volt_per_pixel``.
|
| 1426 |
+
4. Converts column indices to seconds using ``self.time_per_pixel``.
|
| 1427 |
+
5. Resamples uniformly at 500 Hz via linear interpolation.
|
| 1428 |
+
|
| 1429 |
+
Sets ``self.signal_grid`` with the same structure as ``self.grid``
|
| 1430 |
+
but each cell additionally contains:
|
| 1431 |
+
|
| 1432 |
+
* ``'time'`` – np.ndarray of time values in seconds.
|
| 1433 |
+
* ``'signal'`` – np.ndarray of voltage values in millivolts.
|
| 1434 |
+
"""
|
| 1435 |
+
signal_grid = []
|
| 1436 |
+
sample_rate = 500
|
| 1437 |
+
|
| 1438 |
+
for row_idx, row in enumerate(self.grid):
|
| 1439 |
+
baseline_y = self.baseline[row_idx]
|
| 1440 |
+
row_signals = []
|
| 1441 |
+
|
| 1442 |
+
for cell in row:
|
| 1443 |
+
x_coords, sig = self.binarize_signal(cell['signal'])
|
| 1444 |
+
|
| 1445 |
+
x_coords = x_coords[5:-5]
|
| 1446 |
+
sig = sig[5:-5]
|
| 1447 |
+
|
| 1448 |
+
x_coords, sig = self.fill_gaps(x_coords, sig, method='linear')
|
| 1449 |
+
sig = self.smooth_signal(sig)
|
| 1450 |
+
|
| 1451 |
+
signal_volts = (baseline_y - sig) * self.volt_per_pixel
|
| 1452 |
+
x_seconds = x_coords * self.time_per_pixel
|
| 1453 |
+
|
| 1454 |
+
duration = x_seconds[-1] - x_seconds[0]
|
| 1455 |
+
num_samples = round(duration * sample_rate) + 1
|
| 1456 |
+
resampled_time = np.linspace(x_seconds[0], x_seconds[-1], num_samples)
|
| 1457 |
+
|
| 1458 |
+
interpolator = interp1d(x_seconds, signal_volts, kind='linear',
|
| 1459 |
+
fill_value="extrapolate")
|
| 1460 |
+
resampled_signal = interpolator(resampled_time)
|
| 1461 |
+
|
| 1462 |
+
row_signals.append({
|
| 1463 |
+
'lead': cell['lead'],
|
| 1464 |
+
'time': resampled_time,
|
| 1465 |
+
'signal': resampled_signal
|
| 1466 |
+
})
|
| 1467 |
+
|
| 1468 |
+
signal_grid.append(row_signals)
|
| 1469 |
+
|
| 1470 |
+
self.signal_grid = signal_grid
|
| 1471 |
+
|
| 1472 |
+
# ------------------------------------------------------------------
|
| 1473 |
+
# Metrics & comparison
|
| 1474 |
+
# ------------------------------------------------------------------
|
| 1475 |
+
|
| 1476 |
+
def sliding_metrics(self, signal_a, signal_b, return_aligned_signals=False):
|
| 1477 |
+
"""Compute alignment-aware similarity metrics between two signals.
|
| 1478 |
+
|
| 1479 |
+
Slides the shorter signal across the longer one and finds the offset
|
| 1480 |
+
that maximises the Pearson correlation. Reports RMSE and SNR at
|
| 1481 |
+
the best-aligned position.
|
| 1482 |
+
|
| 1483 |
+
Parameters
|
| 1484 |
+
----------
|
| 1485 |
+
signal_a, signal_b : array-like
|
| 1486 |
+
The two signals to compare (need not be the same length).
|
| 1487 |
+
return_aligned_signals : bool, optional
|
| 1488 |
+
If ``True``, also return the aligned signal pair (default ``False``).
|
| 1489 |
+
|
| 1490 |
+
Returns
|
| 1491 |
+
-------
|
| 1492 |
+
max_corr : float
|
| 1493 |
+
Maximum Pearson correlation found across all offsets.
|
| 1494 |
+
rmse : float
|
| 1495 |
+
Root mean squared error at the best offset.
|
| 1496 |
+
snr : float
|
| 1497 |
+
Signal-to-noise ratio in dB at the best offset
|
| 1498 |
+
(``inf`` if noise power is zero).
|
| 1499 |
+
short_signal : np.ndarray
|
| 1500 |
+
The shorter signal (only when *return_aligned_signals* is ``True``).
|
| 1501 |
+
best_window : np.ndarray
|
| 1502 |
+
The aligned window from the longer signal
|
| 1503 |
+
(only when *return_aligned_signals* is ``True``).
|
| 1504 |
+
"""
|
| 1505 |
+
len_a, len_b = len(signal_a), len(signal_b)
|
| 1506 |
+
if len_a > len_b:
|
| 1507 |
+
long_signal, short_signal = signal_a, signal_b
|
| 1508 |
+
else:
|
| 1509 |
+
long_signal, short_signal = signal_b, signal_a
|
| 1510 |
+
|
| 1511 |
+
len_long, len_short = len(long_signal), len(short_signal)
|
| 1512 |
+
max_corr, best_corr_offset = -1, 0
|
| 1513 |
+
|
| 1514 |
+
for i in range(len_long - len_short + 1):
|
| 1515 |
+
window = long_signal[i:i + len_short]
|
| 1516 |
+
corr, _ = pearsonr(window, short_signal)
|
| 1517 |
+
if corr > max_corr:
|
| 1518 |
+
max_corr = corr
|
| 1519 |
+
best_corr_offset = i
|
| 1520 |
+
|
| 1521 |
+
best_window = long_signal[best_corr_offset:best_corr_offset + len_short]
|
| 1522 |
+
rmse = mean_squared_error(best_window, short_signal)
|
| 1523 |
+
signal_power = np.mean(np.square(short_signal))
|
| 1524 |
+
noise_power = np.mean(np.square(best_window - short_signal))
|
| 1525 |
+
snr = (10 * np.log10(signal_power / noise_power)
|
| 1526 |
+
if noise_power > 0 else np.inf)
|
| 1527 |
+
|
| 1528 |
+
if return_aligned_signals:
|
| 1529 |
+
return max_corr, rmse, snr, short_signal, best_window
|
| 1530 |
+
return max_corr, rmse, snr
|
| 1531 |
+
|
| 1532 |
+
def calculate_metrics_ptb(self, plot_signals=True, per_lead_scores=None):
|
| 1533 |
+
"""Evaluate digitization quality against a WFDB ground-truth record.
|
| 1534 |
+
|
| 1535 |
+
For each lead in ``self.signal_grid``, aligns the extracted signal
|
| 1536 |
+
to the corresponding WFDB signal using :meth:`sliding_metrics` and
|
| 1537 |
+
records Pearson r, RMSE, SNR, and p-value. Requires ``self.wfdb_path``
|
| 1538 |
+
to point to a valid WFDB record.
|
| 1539 |
+
|
| 1540 |
+
Sets the following instance attributes after completion:
|
| 1541 |
+
|
| 1542 |
+
* ``self.average_pearson`` – mean Pearson r across all leads.
|
| 1543 |
+
* ``self.average_rmse`` – mean RMSE (mV).
|
| 1544 |
+
* ``self.average_snr`` – mean SNR (dB).
|
| 1545 |
+
* ``self.average_pval`` – mean p-value.
|
| 1546 |
+
|
| 1547 |
+
Parameters
|
| 1548 |
+
----------
|
| 1549 |
+
plot_signals : bool, optional
|
| 1550 |
+
Unused placeholder (default ``True``).
|
| 1551 |
+
per_lead_scores : dict or None, optional
|
| 1552 |
+
If provided, per-lead metrics are appended here for leads where
|
| 1553 |
+
Pearson r > 0.60. Expected structure:
|
| 1554 |
+
``{lead_name: {'pearson': [], 'rmse': [], 'snr': [], 'pval': []}}``.
|
| 1555 |
+
"""
|
| 1556 |
+
record = wfdb.rdrecord(self.wfdb_path)
|
| 1557 |
+
avg_pearson, avg_rmse, avg_snr, avg_pval = [], [], [], []
|
| 1558 |
+
|
| 1559 |
+
for row in self.signal_grid:
|
| 1560 |
+
for cell in row:
|
| 1561 |
+
if 'lead' not in cell or 'signal' not in cell:
|
| 1562 |
+
continue
|
| 1563 |
+
try:
|
| 1564 |
+
lead_index = record.sig_name.index(cell['lead'])
|
| 1565 |
+
except ValueError:
|
| 1566 |
+
continue
|
| 1567 |
+
|
| 1568 |
+
wfdb_signal = record.p_signal[:, lead_index]
|
| 1569 |
+
wfdb_signal = wfdb_signal[~np.isnan(wfdb_signal)]
|
| 1570 |
+
voltage_signal = np.array(cell['signal'])
|
| 1571 |
+
|
| 1572 |
+
pearson_val, rmse, snr, sig1, sig2 = self.sliding_metrics(
|
| 1573 |
+
voltage_signal, wfdb_signal, return_aligned_signals=True
|
| 1574 |
+
)
|
| 1575 |
+
|
| 1576 |
+
try:
|
| 1577 |
+
pearson, pval = pearsonr(sig1, sig2)
|
| 1578 |
+
except Exception:
|
| 1579 |
+
pearson, pval = np.nan, np.nan
|
| 1580 |
+
|
| 1581 |
+
cell.update({'pearson': pearson, 'rmse': rmse,
|
| 1582 |
+
'snr': snr, 'pval': pval})
|
| 1583 |
+
|
| 1584 |
+
avg_pearson.append(pearson)
|
| 1585 |
+
avg_rmse.append(rmse)
|
| 1586 |
+
avg_snr.append(snr)
|
| 1587 |
+
avg_pval.append(pval)
|
| 1588 |
+
|
| 1589 |
+
if per_lead_scores is not None and pearson > 0.60:
|
| 1590 |
+
lead = cell['lead']
|
| 1591 |
+
if lead not in per_lead_scores:
|
| 1592 |
+
per_lead_scores[lead] = {
|
| 1593 |
+
'pearson': [], 'rmse': [], 'snr': [], 'pval': []
|
| 1594 |
+
}
|
| 1595 |
+
per_lead_scores[lead]['pearson'].append(pearson)
|
| 1596 |
+
per_lead_scores[lead]['rmse'].append(rmse)
|
| 1597 |
+
per_lead_scores[lead]['snr'].append(snr)
|
| 1598 |
+
per_lead_scores[lead]['pval'].append(pval)
|
| 1599 |
+
|
| 1600 |
+
self.average_pearson = np.mean(avg_pearson)
|
| 1601 |
+
self.average_rmse = np.mean(avg_rmse)
|
| 1602 |
+
self.average_snr = np.mean(avg_snr)
|
| 1603 |
+
self.average_pval = np.mean(avg_pval)
|
| 1604 |
+
|
| 1605 |
+
def plot_signals(self, title='', plot_wfdb=False):
|
| 1606 |
+
"""Plot extracted signals, optionally overlaid with the WFDB ground truth.
|
| 1607 |
+
|
| 1608 |
+
Generates one figure per lead. When *plot_wfdb* is ``True``, the
|
| 1609 |
+
extracted signal and the best-aligned WFDB window are plotted
|
| 1610 |
+
together for visual comparison.
|
| 1611 |
+
|
| 1612 |
+
Parameters
|
| 1613 |
+
----------
|
| 1614 |
+
title : str, optional
|
| 1615 |
+
Figure title applied to every subplot (default ``''``).
|
| 1616 |
+
plot_wfdb : bool, optional
|
| 1617 |
+
If ``True``, overlay the ground-truth WFDB signal. Requires
|
| 1618 |
+
``self.wfdb_path`` and a loaded ``record`` in scope
|
| 1619 |
+
(default ``False``).
|
| 1620 |
+
"""
|
| 1621 |
+
for row in self.signal_grid:
|
| 1622 |
+
for cell in row:
|
| 1623 |
+
voltage_signal = cell['signal']
|
| 1624 |
+
plt.figure(figsize=(10, 4))
|
| 1625 |
+
if not plot_wfdb:
|
| 1626 |
+
plt.plot(voltage_signal, linewidth=1.5)
|
| 1627 |
+
else:
|
| 1628 |
+
lead_index = record.sig_name.index(cell['lead'])
|
| 1629 |
+
wfdb_signal = record.p_signal[:, lead_index]
|
| 1630 |
+
wfdb_signal = [x for x in wfdb_signal if not np.isnan(x)]
|
| 1631 |
+
_, _, _, sig1, sig2 = self.sliding_metrics(
|
| 1632 |
+
voltage_signal, wfdb_signal, return_aligned_signals=True
|
| 1633 |
+
)
|
| 1634 |
+
plt.plot(sig1, label='Extracted Signal', linewidth=1.5)
|
| 1635 |
+
plt.plot(sig2, label='Ground Truth', linewidth=1.5)
|
| 1636 |
+
plt.title(title)
|
| 1637 |
+
plt.legend()
|
| 1638 |
+
plt.xlabel("Time (ms)")
|
| 1639 |
+
plt.ylabel("Voltage (mV)")
|
| 1640 |
+
plt.tight_layout()
|
| 1641 |
+
plt.show()
|
| 1642 |
+
|
| 1643 |
+
# ------------------------------------------------------------------
|
| 1644 |
+
# Export
|
| 1645 |
+
# ------------------------------------------------------------------
|
| 1646 |
+
|
| 1647 |
+
def save_signals_as_wfdb(self, record_name, directory='.'):
|
| 1648 |
+
"""Export extracted signals as a WFDB record (.hea + .dat files).
|
| 1649 |
+
|
| 1650 |
+
Pads all lead signals to the same length with NaN (converted to 0),
|
| 1651 |
+
then writes a 500 Hz, 16-bit WFDB record in millivolts.
|
| 1652 |
+
|
| 1653 |
+
Parameters
|
| 1654 |
+
----------
|
| 1655 |
+
record_name : str
|
| 1656 |
+
Base name for the output files (no extension).
|
| 1657 |
+
directory : str, optional
|
| 1658 |
+
Destination directory (default ``'.'``).
|
| 1659 |
+
"""
|
| 1660 |
+
signals, lead_names, max_length = [], [], 0
|
| 1661 |
+
for row in self.signal_grid:
|
| 1662 |
+
for cell in row:
|
| 1663 |
+
if 'signal' in cell and 'lead' in cell:
|
| 1664 |
+
sig = np.array(cell['signal'])
|
| 1665 |
+
signals.append(sig)
|
| 1666 |
+
lead_names.append(cell['lead'])
|
| 1667 |
+
max_length = max(max_length, len(sig))
|
| 1668 |
+
|
| 1669 |
+
padded = [np.pad(s, (0, max_length - len(s)),
|
| 1670 |
+
mode='constant', constant_values=np.nan)
|
| 1671 |
+
for s in signals]
|
| 1672 |
+
signal_array = np.nan_to_num(np.array(padded).T, nan=0.0)
|
| 1673 |
+
|
| 1674 |
+
wfdb.wrsamp(
|
| 1675 |
+
record_name=record_name,
|
| 1676 |
+
fs=500,
|
| 1677 |
+
units=['mV'] * len(lead_names),
|
| 1678 |
+
sig_name=lead_names,
|
| 1679 |
+
p_signal=signal_array,
|
| 1680 |
+
fmt=['16'] * len(lead_names),
|
| 1681 |
+
write_dir=directory,
|
| 1682 |
+
)
|
| 1683 |
+
|
| 1684 |
+
def save_signals_as_csv(self, record_name, directory='.'):
|
| 1685 |
+
"""
|
| 1686 |
+
Save ECG signals to CSV. Each column = one lead, each row = one time sample.
|
| 1687 |
+
"""
|
| 1688 |
+
signals, lead_names, max_length = [], [], 0
|
| 1689 |
+
for row in self.signal_grid:
|
| 1690 |
+
for cell in row:
|
| 1691 |
+
if 'signal' in cell and 'lead' in cell:
|
| 1692 |
+
sig = np.asarray(cell['signal'], dtype=float)
|
| 1693 |
+
signals.append(sig)
|
| 1694 |
+
lead_names.append(cell['lead'])
|
| 1695 |
+
max_length = max(max_length, len(sig))
|
| 1696 |
+
|
| 1697 |
+
if not signals:
|
| 1698 |
+
raise ValueError("No valid signals found in signal_grid.")
|
| 1699 |
+
|
| 1700 |
+
padded = [np.pad(s, (0, max_length - len(s)),
|
| 1701 |
+
mode='constant', constant_values=np.nan)
|
| 1702 |
+
for s in signals]
|
| 1703 |
+
df = pd.DataFrame(np.vstack(padded).T, columns=lead_names)
|
| 1704 |
+
|
| 1705 |
+
os.makedirs(directory, exist_ok=True)
|
| 1706 |
+
csv_path = os.path.join(directory, f"{record_name}.csv")
|
| 1707 |
+
df.to_csv(csv_path, index=False, float_format="%.6f")
|
| 1708 |
+
|
| 1709 |
+
# ------------------------------------------------------------------
|
| 1710 |
+
# Full pipeline
|
| 1711 |
+
# ------------------------------------------------------------------
|
| 1712 |
+
|
| 1713 |
+
def run_full_pipeline(self):
|
| 1714 |
+
"""Run all digitization stages end-to-end for a single ECG image.
|
| 1715 |
+
|
| 1716 |
+
Executes the following stages in order, with per-stage wall-clock
|
| 1717 |
+
timing printed to stdout via :meth:`_timed`:
|
| 1718 |
+
|
| 1719 |
+
1. **Image loading & preprocessing** — ``load_image`` + ``preprocess_image``.
|
| 1720 |
+
2. **Segmentation** — ``segment_leads`` → ``make_segmentation_mask``
|
| 1721 |
+
→ ``find_row_centers`` → ``get_roi``.
|
| 1722 |
+
3. **Sequential YOLO detections** — ``extract_lead_boxes``,
|
| 1723 |
+
``extract_lead_name_boxes``, ``extract_reference_pulses``.
|
| 1724 |
+
(Sequential, not parallel, to avoid a PyTorch thread-pool deadlock.)
|
| 1725 |
+
4. **Retry loop** — if any of the three detections is empty, the
|
| 1726 |
+
image is reloaded at the next fallback target size from
|
| 1727 |
+
``[2000, 2100, 1900, 2200, 1800, 1700]`` px and stages 1–3 are
|
| 1728 |
+
re-run. Raises ``RuntimeError`` if all sizes fail.
|
| 1729 |
+
5. **Calibration** — ``get_reference_scale``.
|
| 1730 |
+
6. **Grid & signal extraction** — ``make_grid`` → ``extract_signals``.
|
| 1731 |
+
|
| 1732 |
+
Raises
|
| 1733 |
+
------
|
| 1734 |
+
RuntimeError
|
| 1735 |
+
If all fallback target sizes yield at least one empty detection.
|
| 1736 |
+
"""
|
| 1737 |
+
FALLBACK_SIZES = [2000, 2100, 1900, 2200, 1800, 1700]
|
| 1738 |
+
|
| 1739 |
+
for target_size in FALLBACK_SIZES:
|
| 1740 |
+
|
| 1741 |
+
# -- Image loading & preprocessing --
|
| 1742 |
+
self._timed("load_image", self.load_image, self.image_path,
|
| 1743 |
+
target_size=target_size)
|
| 1744 |
+
self._timed("preprocess", self.preprocess_image)
|
| 1745 |
+
|
| 1746 |
+
# -- Segmentation (sequential: each step depends on the previous) --
|
| 1747 |
+
self._timed("segment_leads", self.segment_leads)
|
| 1748 |
+
self._timed("make_mask", self.make_segmentation_mask)
|
| 1749 |
+
self._timed("find_row_centers", self.find_row_centers)
|
| 1750 |
+
self._timed("get_roi", self.get_roi)
|
| 1751 |
+
|
| 1752 |
+
# -- Three YOLO calls sequentially --
|
| 1753 |
+
self._timed("extract_lead_boxes", self.extract_lead_boxes)
|
| 1754 |
+
self._timed("extract_lead_name_boxes", self.extract_lead_name_boxes)
|
| 1755 |
+
self._timed("extract_reference_pulses", self.extract_reference_pulses)
|
| 1756 |
+
|
| 1757 |
+
# -- Check all three detections; retry at next size if any is empty --
|
| 1758 |
+
missing = []
|
| 1759 |
+
if len(self.reference_pulses) == 0:
|
| 1760 |
+
missing.append("pulses")
|
| 1761 |
+
if len(self.lead_bboxes) == 0:
|
| 1762 |
+
missing.append("lead boxes")
|
| 1763 |
+
if len(self.lead_name_bboxes) == 0:
|
| 1764 |
+
missing.append("lead names")
|
| 1765 |
+
|
| 1766 |
+
if not missing:
|
| 1767 |
+
if target_size != FALLBACK_SIZES[0]:
|
| 1768 |
+
print(f" ↳ All detections found at target_size={target_size}",
|
| 1769 |
+
flush=True)
|
| 1770 |
+
break
|
| 1771 |
+
else:
|
| 1772 |
+
print(f" ↳ Missing {', '.join(missing)} at target_size={target_size},"
|
| 1773 |
+
f" retrying...", flush=True)
|
| 1774 |
+
|
| 1775 |
+
else:
|
| 1776 |
+
raise RuntimeError(
|
| 1777 |
+
f"Detections incomplete at all target sizes {FALLBACK_SIZES}. "
|
| 1778 |
+
f"Last missing: {', '.join(missing)}."
|
| 1779 |
+
)
|
| 1780 |
+
|
| 1781 |
+
# -- Remainder of pipeline (sequential) --
|
| 1782 |
+
self._timed("get_reference_scale", self.get_reference_scale)
|
| 1783 |
+
self._timed("make_grid", self.make_grid)
|
| 1784 |
+
self._timed("extract_signals", self.extract_signals)
|
environment.yml
ADDED
|
@@ -0,0 +1,194 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
name: infer
|
| 2 |
+
channels:
|
| 3 |
+
- main
|
| 4 |
+
- r
|
| 5 |
+
- conda-forge
|
| 6 |
+
- bioconda
|
| 7 |
+
- pytorch
|
| 8 |
+
dependencies:
|
| 9 |
+
- _libgcc_mutex=0.1=main
|
| 10 |
+
- _openmp_mutex=5.1=1_gnu
|
| 11 |
+
- aiohappyeyeballs=2.4.4=py39h06a4308_0
|
| 12 |
+
- aiohttp=3.11.10=py39h5eee18b_0
|
| 13 |
+
- aiosignal=1.2.0=pyhd3eb1b0_0
|
| 14 |
+
- aom=3.6.0=h6a678d5_0
|
| 15 |
+
- async-timeout=5.0.1=py39h06a4308_0
|
| 16 |
+
- attrs=24.3.0=py39h06a4308_0
|
| 17 |
+
- blas=1.0=openblas
|
| 18 |
+
- blosc=1.21.3=h6a678d5_0
|
| 19 |
+
- bottleneck=1.4.2=py39ha9d4c09_0
|
| 20 |
+
- brotli=1.0.9=h5eee18b_9
|
| 21 |
+
- brotli-bin=1.0.9=h5eee18b_9
|
| 22 |
+
- brotli-python=1.0.9=py39h6a678d5_9
|
| 23 |
+
- brunsli=0.1=h2531618_0
|
| 24 |
+
- bzip2=1.0.8=h5eee18b_6
|
| 25 |
+
- c-ares=1.19.1=h5eee18b_0
|
| 26 |
+
- ca-certificates=2025.2.25=h06a4308_0
|
| 27 |
+
- certifi=2025.4.26=py39h06a4308_0
|
| 28 |
+
- cffi=1.17.1=py39h1fdaa30_1
|
| 29 |
+
- cfitsio=3.470=h5893167_7
|
| 30 |
+
- charls=2.2.0=h2531618_0
|
| 31 |
+
- dav1d=1.2.1=h5eee18b_0
|
| 32 |
+
- freetype=2.13.3=h4a9f257_0
|
| 33 |
+
- frozenlist=1.5.0=py39h5eee18b_0
|
| 34 |
+
- fsspec=2025.3.2=py39he3bba80_0
|
| 35 |
+
- gettext=0.21.0=hedfda30_2
|
| 36 |
+
- giflib=5.2.2=h5eee18b_0
|
| 37 |
+
- icu=73.1=h6a678d5_0
|
| 38 |
+
- imagecodecs=2024.9.22=py39h652d938_1
|
| 39 |
+
- imageio=2.37.0=py39h06a4308_0
|
| 40 |
+
- importlib_resources=6.4.0=py39h06a4308_0
|
| 41 |
+
- joblib=1.4.2=py39h06a4308_0
|
| 42 |
+
- jpeg=9e=h5eee18b_3
|
| 43 |
+
- jxrlib=1.1=h7b6447c_2
|
| 44 |
+
- krb5=1.20.1=h143b758_1
|
| 45 |
+
- lazy_loader=0.4=py39h06a4308_0
|
| 46 |
+
- lcms2=2.16=h92b89f2_1
|
| 47 |
+
- ld_impl_linux-64=2.40=h12ee557_0
|
| 48 |
+
- lerc=4.0.0=h6a678d5_0
|
| 49 |
+
- libaec=1.1.3=h6a678d5_0
|
| 50 |
+
- libavif=1.1.1=h5eee18b_0
|
| 51 |
+
- libbrotlicommon=1.0.9=h5eee18b_9
|
| 52 |
+
- libbrotlidec=1.0.9=h5eee18b_9
|
| 53 |
+
- libbrotlienc=1.0.9=h5eee18b_9
|
| 54 |
+
- libcurl=8.12.1=hc9e6f67_0
|
| 55 |
+
- libdeflate=1.22=h5eee18b_0
|
| 56 |
+
- libedit=3.1.20230828=h5eee18b_0
|
| 57 |
+
- libev=4.33=h7f8727e_1
|
| 58 |
+
- libffi=3.4.4=h6a678d5_1
|
| 59 |
+
- libflac=1.3.4=h27087fc_0
|
| 60 |
+
- libgcc-ng=11.2.0=h1234567_1
|
| 61 |
+
- libgfortran-ng=11.2.0=h00389a5_1
|
| 62 |
+
- libgfortran5=11.2.0=h1234567_1
|
| 63 |
+
- libgomp=11.2.0=h1234567_1
|
| 64 |
+
- libnghttp2=1.57.0=h2d74bed_0
|
| 65 |
+
- libogg=1.3.5=h27cfd23_1
|
| 66 |
+
- libopenblas=0.3.29=ha39b09d_0
|
| 67 |
+
- libopus=1.3.1=h5eee18b_1
|
| 68 |
+
- libpng=1.6.39=h5eee18b_0
|
| 69 |
+
- libsndfile=1.0.31=h9c3ff4c_1
|
| 70 |
+
- libssh2=1.11.1=h251f7ec_0
|
| 71 |
+
- libstdcxx-ng=11.2.0=h1234567_1
|
| 72 |
+
- libtiff=4.7.0=hde9077f_0
|
| 73 |
+
- libvorbis=1.3.7=h7b6447c_0
|
| 74 |
+
- libwebp-base=1.3.2=h5eee18b_1
|
| 75 |
+
- libxml2=2.13.8=hfdd30dd_0
|
| 76 |
+
- libzopfli=1.0.3=he6710b0_0
|
| 77 |
+
- lz4-c=1.9.4=h6a678d5_1
|
| 78 |
+
- matplotlib-base=3.9.2=py39hbfdbfaf_1
|
| 79 |
+
- multidict=6.1.0=py39h5eee18b_0
|
| 80 |
+
- ncurses=6.4=h6a678d5_0
|
| 81 |
+
- networkx=3.2.1=py39h06a4308_0
|
| 82 |
+
- numexpr=2.10.1=py39hd28fd6d_0
|
| 83 |
+
- numpy=1.26.4=py39heeff2f4_0
|
| 84 |
+
- numpy-base=1.26.4=py39h8a23956_0
|
| 85 |
+
- openjpeg=2.5.2=h0d4d230_1
|
| 86 |
+
- openssl=3.0.16=h5eee18b_0
|
| 87 |
+
- pandas=2.2.3=py39h6a678d5_0
|
| 88 |
+
- pip=25.1=pyhc872135_2
|
| 89 |
+
- propcache=0.3.1=py39h5eee18b_0
|
| 90 |
+
- pybind11-abi=4=hd3eb1b0_1
|
| 91 |
+
- pycparser=2.21=pyhd3eb1b0_0
|
| 92 |
+
- pysocks=1.7.1=py39h06a4308_0
|
| 93 |
+
- pysoundfile=0.11.0=pyhd8ed1ab_0
|
| 94 |
+
- python=3.9.21=he870216_1
|
| 95 |
+
- python-dateutil=2.9.0post0=py39h06a4308_2
|
| 96 |
+
- python-tzdata=2025.2=pyhd3eb1b0_0
|
| 97 |
+
- readline=8.2=h5eee18b_0
|
| 98 |
+
- requests=2.32.3=py39h06a4308_1
|
| 99 |
+
- scikit-image=0.24.0=py39h1128e8f_0
|
| 100 |
+
- scikit-learn=1.6.1=py39h6a678d5_0
|
| 101 |
+
- scipy=1.13.1=py39heeff2f4_1
|
| 102 |
+
- setuptools=78.1.1=py39h06a4308_0
|
| 103 |
+
- six=1.17.0=py39h06a4308_0
|
| 104 |
+
- snappy=1.2.1=h6a678d5_0
|
| 105 |
+
- sqlite=3.45.3=h5eee18b_0
|
| 106 |
+
- threadpoolctl=3.5.0=py39h2f386ee_0
|
| 107 |
+
- tifffile=2023.4.12=py39h06a4308_0
|
| 108 |
+
- tk=8.6.14=h39e8969_0
|
| 109 |
+
- typing_extensions=4.12.2=py39h06a4308_0
|
| 110 |
+
- tzdata=2025b=h04d1e81_0
|
| 111 |
+
- unicodedata2=15.1.0=py39h5eee18b_1
|
| 112 |
+
- wfdb=4.3.0=pyhd8ed1ab_0
|
| 113 |
+
- wheel=0.45.1=py39h06a4308_0
|
| 114 |
+
- xz=5.6.4=h5eee18b_1
|
| 115 |
+
- yarl=1.18.0=py39h5eee18b_0
|
| 116 |
+
- zfp=1.0.0=h6a678d5_0
|
| 117 |
+
- zipp=3.21.0=py39h06a4308_0
|
| 118 |
+
- zlib=1.2.13=h5eee18b_1
|
| 119 |
+
- zlib-ng=2.0.7=h5eee18b_0
|
| 120 |
+
- zstd=1.5.6=hc292b87_0
|
| 121 |
+
- pip:
|
| 122 |
+
- asttokens==3.0.0
|
| 123 |
+
- charset-normalizer==3.4.2
|
| 124 |
+
- comm==0.2.2
|
| 125 |
+
- contourpy==1.3.0
|
| 126 |
+
- cycler==0.12.1
|
| 127 |
+
- debugpy==1.8.14
|
| 128 |
+
- decorator==5.2.1
|
| 129 |
+
- exceptiongroup==1.3.0
|
| 130 |
+
- executing==2.2.0
|
| 131 |
+
- filelock==3.18.0
|
| 132 |
+
- fonttools==4.58.0
|
| 133 |
+
- idna==3.10
|
| 134 |
+
- importlib-metadata==8.7.0
|
| 135 |
+
- importlib-resources==6.5.2
|
| 136 |
+
- ipykernel==6.29.5
|
| 137 |
+
- ipython==8.18.1
|
| 138 |
+
- jedi==0.19.2
|
| 139 |
+
- jinja2==3.1.6
|
| 140 |
+
- jupyter-client==8.6.3
|
| 141 |
+
- jupyter-core==5.7.2
|
| 142 |
+
- kiwisolver==1.4.7
|
| 143 |
+
- markupsafe==3.0.2
|
| 144 |
+
- matplotlib==3.9.4
|
| 145 |
+
- matplotlib-inline==0.1.7
|
| 146 |
+
- mpmath==1.3.0
|
| 147 |
+
- nest-asyncio==1.6.0
|
| 148 |
+
- nvidia-cublas-cu12==12.6.4.1
|
| 149 |
+
- nvidia-cuda-cupti-cu12==12.6.80
|
| 150 |
+
- nvidia-cuda-nvrtc-cu12==12.6.77
|
| 151 |
+
- nvidia-cuda-runtime-cu12==12.6.77
|
| 152 |
+
- nvidia-cudnn-cu12==9.5.1.17
|
| 153 |
+
- nvidia-cufft-cu12==11.3.0.4
|
| 154 |
+
- nvidia-cufile-cu12==1.11.1.6
|
| 155 |
+
- nvidia-curand-cu12==10.3.7.77
|
| 156 |
+
- nvidia-cusolver-cu12==11.7.1.2
|
| 157 |
+
- nvidia-cusparse-cu12==12.5.4.2
|
| 158 |
+
- nvidia-cusparselt-cu12==0.6.3
|
| 159 |
+
- nvidia-nccl-cu12==2.26.2
|
| 160 |
+
- nvidia-nvjitlink-cu12==12.6.85
|
| 161 |
+
- nvidia-nvtx-cu12==12.6.77
|
| 162 |
+
- opencv-python==4.11.0.86
|
| 163 |
+
- packaging==25.0
|
| 164 |
+
- parso==0.8.4
|
| 165 |
+
- patched-yolo-infer==1.3.8
|
| 166 |
+
- pexpect==4.9.0
|
| 167 |
+
- pillow==11.2.1
|
| 168 |
+
- platformdirs==4.3.8
|
| 169 |
+
- prompt-toolkit==3.0.51
|
| 170 |
+
- psutil==7.0.0
|
| 171 |
+
- ptyprocess==0.7.0
|
| 172 |
+
- pure-eval==0.2.3
|
| 173 |
+
- py-cpuinfo==9.0.0
|
| 174 |
+
- pygments==2.19.1
|
| 175 |
+
- pyparsing==3.2.3
|
| 176 |
+
- pytz==2025.2
|
| 177 |
+
- pyyaml==6.0.2
|
| 178 |
+
- pyzmq==26.4.0
|
| 179 |
+
- seaborn==0.13.2
|
| 180 |
+
- simpleitk==2.5.0
|
| 181 |
+
- stack-data==0.6.3
|
| 182 |
+
- sympy==1.14.0
|
| 183 |
+
- torch==2.7.0
|
| 184 |
+
- torchvision==0.22.0
|
| 185 |
+
- tornado==6.4.2
|
| 186 |
+
- tqdm==4.67.1
|
| 187 |
+
- traitlets==5.14.3
|
| 188 |
+
- triton==3.3.0
|
| 189 |
+
- typing-extensions==4.13.2
|
| 190 |
+
- ultralytics==8.3.133
|
| 191 |
+
- ultralytics-thop==2.0.14
|
| 192 |
+
- urllib3==2.4.0
|
| 193 |
+
- wcwidth==0.2.13
|
| 194 |
+
prefix: /cluster/apps/scaimed/cachtari/conda/envs/infer
|
example_inputs/test.jpg
ADDED
|
Git LFS Details
|
models/classifier_models/ecg_surgery/inceptiontime_model.pkl
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:f6d52babc1f5c750f89c9a442ccd5fdb7835e6aac2cba31c39c184d59305c282
|
| 3 |
+
size 2153
|
models/classifier_models/ecg_surgery/model_metadata.json
ADDED
|
@@ -0,0 +1,62 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"model_name": "InceptionTime",
|
| 3 |
+
"model_file": "inceptiontime_model.pkl",
|
| 4 |
+
"use_case": "Pre vs Post Procedural MI (Surgery)",
|
| 5 |
+
"positive_class": "pre-procedural MI",
|
| 6 |
+
"negative_class": "post-procedural MI",
|
| 7 |
+
"class_labels": [
|
| 8 |
+
"post-procedural MI",
|
| 9 |
+
"pre-procedural MI"
|
| 10 |
+
],
|
| 11 |
+
"features": [
|
| 12 |
+
"I",
|
| 13 |
+
"aVR",
|
| 14 |
+
"V1",
|
| 15 |
+
"V4",
|
| 16 |
+
"II",
|
| 17 |
+
"aVL",
|
| 18 |
+
"V2",
|
| 19 |
+
"V5",
|
| 20 |
+
"III",
|
| 21 |
+
"aVF",
|
| 22 |
+
"V3",
|
| 23 |
+
"V6"
|
| 24 |
+
],
|
| 25 |
+
"n_features": 12,
|
| 26 |
+
"n_timesteps": 140,
|
| 27 |
+
"hyperparameters": {
|
| 28 |
+
"n_epochs": 200,
|
| 29 |
+
"batch_size": 64,
|
| 30 |
+
"random_state": 42,
|
| 31 |
+
"verbose": true
|
| 32 |
+
},
|
| 33 |
+
"training_data": "data/ecg_surgery_segmented_50_150_70.csv",
|
| 34 |
+
"training_date": "2026-06-02 20:26:15",
|
| 35 |
+
"training_time_seconds": 251.83,
|
| 36 |
+
"preprocessing": {
|
| 37 |
+
"normalization": "per_subject_max_abs",
|
| 38 |
+
"smote_tomek": true,
|
| 39 |
+
"test_size": 0.25,
|
| 40 |
+
"random_state": 42
|
| 41 |
+
},
|
| 42 |
+
"training_samples": 1104,
|
| 43 |
+
"test_samples": 338,
|
| 44 |
+
"test_metrics": {
|
| 45 |
+
"accuracy": 0.9142011834319527,
|
| 46 |
+
"precision": 0.8974358974358975,
|
| 47 |
+
"recall": 0.9150326797385621,
|
| 48 |
+
"f1": 0.9061488673139159,
|
| 49 |
+
"sensitivity": 0.9150326797385621,
|
| 50 |
+
"specificity": 0.9135135135135135,
|
| 51 |
+
"confusion_matrix": [
|
| 52 |
+
[
|
| 53 |
+
169,
|
| 54 |
+
16
|
| 55 |
+
],
|
| 56 |
+
[
|
| 57 |
+
13,
|
| 58 |
+
140
|
| 59 |
+
]
|
| 60 |
+
]
|
| 61 |
+
}
|
| 62 |
+
}
|
models/classifier_models/mi_vs_normal_segmented/arsenal_segmented_model.pkl
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:fb22793560c0939a574e1cdab6e59b3e732cd8592be10cad7dd3fc46e46134c8
|
| 3 |
+
size 112024703
|
models/classifier_models/mi_vs_normal_segmented/model_metadata.json
ADDED
|
@@ -0,0 +1,61 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"model_name": "Arsenal",
|
| 3 |
+
"model_file": "arsenal_segmented_model.pkl",
|
| 4 |
+
"use_case": "MI vs Normal (Segmented)",
|
| 5 |
+
"positive_class": "MYOCARDIAL_INFARCTION",
|
| 6 |
+
"negative_class": "NORMAL",
|
| 7 |
+
"class_labels": [
|
| 8 |
+
"NORMAL",
|
| 9 |
+
"MYOCARDIAL_INFARCTION"
|
| 10 |
+
],
|
| 11 |
+
"features": [
|
| 12 |
+
"I",
|
| 13 |
+
"aVR",
|
| 14 |
+
"V1",
|
| 15 |
+
"V4",
|
| 16 |
+
"II",
|
| 17 |
+
"aVL",
|
| 18 |
+
"V2",
|
| 19 |
+
"V5",
|
| 20 |
+
"III",
|
| 21 |
+
"aVF",
|
| 22 |
+
"V3",
|
| 23 |
+
"V6"
|
| 24 |
+
],
|
| 25 |
+
"n_features": 12,
|
| 26 |
+
"n_timesteps": 140,
|
| 27 |
+
"hyperparameters": {
|
| 28 |
+
"num_kernels": 20000,
|
| 29 |
+
"random_state": 42,
|
| 30 |
+
"n_jobs": 1
|
| 31 |
+
},
|
| 32 |
+
"training_data": "data/ptb_xl/segmented_heartbeats.csv",
|
| 33 |
+
"training_date": "2026-06-03 03:58:21",
|
| 34 |
+
"training_time_seconds": 10302.47,
|
| 35 |
+
"preprocessing": {
|
| 36 |
+
"normalization": "per_subject_max_abs",
|
| 37 |
+
"smote_tomek": true,
|
| 38 |
+
"test_size": 0.25,
|
| 39 |
+
"random_state": 42
|
| 40 |
+
},
|
| 41 |
+
"training_samples": 6954,
|
| 42 |
+
"test_samples": 2349,
|
| 43 |
+
"test_metrics": {
|
| 44 |
+
"accuracy": 0.922520221370796,
|
| 45 |
+
"precision": 0.940864960282436,
|
| 46 |
+
"recall": 0.9026248941574937,
|
| 47 |
+
"f1": 0.9213483146067416,
|
| 48 |
+
"sensitivity": 0.9026248941574937,
|
| 49 |
+
"specificity": 0.9426369863013698,
|
| 50 |
+
"confusion_matrix": [
|
| 51 |
+
[
|
| 52 |
+
1101,
|
| 53 |
+
67
|
| 54 |
+
],
|
| 55 |
+
[
|
| 56 |
+
115,
|
| 57 |
+
1066
|
| 58 |
+
]
|
| 59 |
+
]
|
| 60 |
+
}
|
| 61 |
+
}
|
models/classifier_models/omi_vs_nonomi/model_metadata.json
ADDED
|
@@ -0,0 +1,61 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"model_name": "Rocket",
|
| 3 |
+
"model_file": "rocket_model.pkl",
|
| 4 |
+
"use_case": "OMI vs non-OMI",
|
| 5 |
+
"positive_class": "OMI",
|
| 6 |
+
"negative_class": "non-OMI",
|
| 7 |
+
"class_labels": [
|
| 8 |
+
"non-OMI",
|
| 9 |
+
"OMI"
|
| 10 |
+
],
|
| 11 |
+
"features": [
|
| 12 |
+
"I",
|
| 13 |
+
"aVR",
|
| 14 |
+
"V1",
|
| 15 |
+
"V4",
|
| 16 |
+
"II",
|
| 17 |
+
"aVL",
|
| 18 |
+
"V2",
|
| 19 |
+
"V5",
|
| 20 |
+
"III",
|
| 21 |
+
"aVF",
|
| 22 |
+
"V3",
|
| 23 |
+
"V6"
|
| 24 |
+
],
|
| 25 |
+
"n_features": 12,
|
| 26 |
+
"n_timesteps": 141,
|
| 27 |
+
"hyperparameters": {
|
| 28 |
+
"num_kernels": 50000,
|
| 29 |
+
"random_state": 42,
|
| 30 |
+
"n_jobs": -1
|
| 31 |
+
},
|
| 32 |
+
"training_data": "data/ecg_matrix_omi_segmented_50_150_90.csv",
|
| 33 |
+
"training_date": "2026-06-02 16:41:15",
|
| 34 |
+
"training_time_seconds": 132.78,
|
| 35 |
+
"preprocessing": {
|
| 36 |
+
"normalization": "per_subject_max_abs",
|
| 37 |
+
"smote_tomek": true,
|
| 38 |
+
"test_size": 0.25,
|
| 39 |
+
"random_state": 42
|
| 40 |
+
},
|
| 41 |
+
"training_samples": 2084,
|
| 42 |
+
"test_samples": 666,
|
| 43 |
+
"test_metrics": {
|
| 44 |
+
"accuracy": 0.8888888888888888,
|
| 45 |
+
"precision": 0.8571428571428571,
|
| 46 |
+
"recall": 0.9171974522292994,
|
| 47 |
+
"f1": 0.8861538461538462,
|
| 48 |
+
"sensitivity": 0.9171974522292994,
|
| 49 |
+
"specificity": 0.8636363636363636,
|
| 50 |
+
"confusion_matrix": [
|
| 51 |
+
[
|
| 52 |
+
304,
|
| 53 |
+
48
|
| 54 |
+
],
|
| 55 |
+
[
|
| 56 |
+
26,
|
| 57 |
+
288
|
| 58 |
+
]
|
| 59 |
+
]
|
| 60 |
+
}
|
| 61 |
+
}
|
models/classifier_models/omi_vs_nonomi/rocket_model.pkl
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:d35281f43a57aae562b0f929c7912e5bb8c17efad4dc3a50907fe93cf4f40318
|
| 3 |
+
size 11204282
|
models/digitization_models/yolo11_full/weights/best.pt
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:2a9adc74478e051e36652555cc81e35d787435a4799fa7a8abde6c6fa117a68a
|
| 3 |
+
size 124776986
|
models/digitization_models/yolo11_lead/.ipynb_checkpoints/P_curve-checkpoint.png
ADDED
|
Git LFS Details
|
models/digitization_models/yolo11_lead/.ipynb_checkpoints/args-checkpoint.yaml
ADDED
|
@@ -0,0 +1,105 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
task: detect
|
| 2 |
+
mode: train
|
| 3 |
+
model: yolo11x.pt
|
| 4 |
+
data: /cluster/home/cachtari/dataset_leadname.yaml
|
| 5 |
+
epochs: 100
|
| 6 |
+
time: null
|
| 7 |
+
patience: 100
|
| 8 |
+
batch: 6
|
| 9 |
+
imgsz: 1500
|
| 10 |
+
save: true
|
| 11 |
+
save_period: -1
|
| 12 |
+
cache: false
|
| 13 |
+
device: 0,1
|
| 14 |
+
workers: 8
|
| 15 |
+
project: /cluster/home/cachtari/runs
|
| 16 |
+
name: yolo11_1500_leadname_class
|
| 17 |
+
exist_ok: false
|
| 18 |
+
pretrained: true
|
| 19 |
+
optimizer: auto
|
| 20 |
+
verbose: true
|
| 21 |
+
seed: 0
|
| 22 |
+
deterministic: true
|
| 23 |
+
single_cls: false
|
| 24 |
+
rect: false
|
| 25 |
+
cos_lr: false
|
| 26 |
+
close_mosaic: 10
|
| 27 |
+
resume: false
|
| 28 |
+
amp: true
|
| 29 |
+
fraction: 1.0
|
| 30 |
+
profile: false
|
| 31 |
+
freeze: null
|
| 32 |
+
multi_scale: false
|
| 33 |
+
overlap_mask: true
|
| 34 |
+
mask_ratio: 4
|
| 35 |
+
dropout: 0.0
|
| 36 |
+
val: true
|
| 37 |
+
split: val
|
| 38 |
+
save_json: false
|
| 39 |
+
conf: null
|
| 40 |
+
iou: 0.7
|
| 41 |
+
max_det: 300
|
| 42 |
+
half: false
|
| 43 |
+
dnn: false
|
| 44 |
+
plots: true
|
| 45 |
+
source: null
|
| 46 |
+
vid_stride: 1
|
| 47 |
+
stream_buffer: false
|
| 48 |
+
visualize: false
|
| 49 |
+
augment: false
|
| 50 |
+
agnostic_nms: false
|
| 51 |
+
classes: null
|
| 52 |
+
retina_masks: false
|
| 53 |
+
embed: null
|
| 54 |
+
show: false
|
| 55 |
+
save_frames: false
|
| 56 |
+
save_txt: false
|
| 57 |
+
save_conf: false
|
| 58 |
+
save_crop: false
|
| 59 |
+
show_labels: true
|
| 60 |
+
show_conf: true
|
| 61 |
+
show_boxes: true
|
| 62 |
+
line_width: null
|
| 63 |
+
format: torchscript
|
| 64 |
+
keras: false
|
| 65 |
+
optimize: false
|
| 66 |
+
int8: false
|
| 67 |
+
dynamic: false
|
| 68 |
+
simplify: true
|
| 69 |
+
opset: null
|
| 70 |
+
workspace: null
|
| 71 |
+
nms: false
|
| 72 |
+
lr0: 0.01
|
| 73 |
+
lrf: 0.01
|
| 74 |
+
momentum: 0.937
|
| 75 |
+
weight_decay: 0.0005
|
| 76 |
+
warmup_epochs: 3.0
|
| 77 |
+
warmup_momentum: 0.8
|
| 78 |
+
warmup_bias_lr: 0.1
|
| 79 |
+
box: 7.5
|
| 80 |
+
cls: 0.5
|
| 81 |
+
dfl: 1.5
|
| 82 |
+
pose: 12.0
|
| 83 |
+
kobj: 1.0
|
| 84 |
+
nbs: 64
|
| 85 |
+
hsv_h: 0.015
|
| 86 |
+
hsv_s: 0.7
|
| 87 |
+
hsv_v: 0.4
|
| 88 |
+
degrees: 0.0
|
| 89 |
+
translate: 0.1
|
| 90 |
+
scale: 0.5
|
| 91 |
+
shear: 0.0
|
| 92 |
+
perspective: 0.0
|
| 93 |
+
flipud: 0.0
|
| 94 |
+
fliplr: 0.5
|
| 95 |
+
bgr: 0.0
|
| 96 |
+
mosaic: 1.0
|
| 97 |
+
mixup: 0.0
|
| 98 |
+
cutmix: 0.0
|
| 99 |
+
copy_paste: 0.0
|
| 100 |
+
copy_paste_mode: flip
|
| 101 |
+
auto_augment: randaugment
|
| 102 |
+
erasing: 0.4
|
| 103 |
+
cfg: null
|
| 104 |
+
tracker: botsort.yaml
|
| 105 |
+
save_dir: /cluster/home/cachtari/runs/yolo11_1500_leadname_class
|
models/digitization_models/yolo11_lead/.ipynb_checkpoints/confusion_matrix-checkpoint.png
ADDED
|
Git LFS Details
|
models/digitization_models/yolo11_lead/.ipynb_checkpoints/confusion_matrix_normalized-checkpoint.png
ADDED
|
Git LFS Details
|
models/digitization_models/yolo11_lead/.ipynb_checkpoints/labels-checkpoint.jpg
ADDED
|
Git LFS Details
|
models/digitization_models/yolo11_lead/.ipynb_checkpoints/labels_correlogram-checkpoint.jpg
ADDED
|
Git LFS Details
|
models/digitization_models/yolo11_lead/.ipynb_checkpoints/results-checkpoint.csv
ADDED
|
@@ -0,0 +1,101 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
epoch,time,train/box_loss,train/cls_loss,train/dfl_loss,metrics/precision(B),metrics/recall(B),metrics/mAP50(B),metrics/mAP50-95(B),val/box_loss,val/cls_loss,val/dfl_loss,lr/pg0,lr/pg1,lr/pg2
|
| 2 |
+
1,886.991,0.93774,0.93982,0.89298,0.98053,0.95723,0.98816,0.78539,0.80694,0.38397,0.85301,0.00333207,0.00333207,0.00333207
|
| 3 |
+
2,1741.49,0.79997,0.45797,0.84868,0.9877,0.96985,0.99114,0.81433,0.7171,0.35068,0.84086,0.00659942,0.00659942,0.00659942
|
| 4 |
+
3,2587.75,0.79563,0.44381,0.84741,0.9897,0.97185,0.99174,0.82059,0.70896,0.34325,0.83936,0.00980076,0.00980076,0.00980076
|
| 5 |
+
4,3432.96,0.75251,0.41456,0.84102,0.98979,0.97597,0.99211,0.84977,0.62835,0.30703,0.82137,0.009703,0.009703,0.009703
|
| 6 |
+
5,4277.5,0.69831,0.37499,0.8299,0.99061,0.97892,0.99275,0.85163,0.62284,0.29359,0.81764,0.009604,0.009604,0.009604
|
| 7 |
+
6,5120.98,0.6707,0.35552,0.82548,0.99217,0.98008,0.99302,0.86206,0.60218,0.28213,0.813,0.009505,0.009505,0.009505
|
| 8 |
+
7,5963.46,0.65473,0.34407,0.82271,0.9925,0.98065,0.99334,0.86668,0.59208,0.27651,0.81129,0.009406,0.009406,0.009406
|
| 9 |
+
8,6809.02,0.64014,0.33376,0.82023,0.99242,0.98181,0.99349,0.87175,0.57154,0.26426,0.80764,0.009307,0.009307,0.009307
|
| 10 |
+
9,7654.95,0.6281,0.32268,0.81734,0.99407,0.98343,0.99377,0.87428,0.57133,0.25716,0.80706,0.009208,0.009208,0.009208
|
| 11 |
+
10,8498.47,0.61846,0.31632,0.81624,0.99353,0.98488,0.99378,0.87652,0.56228,0.25264,0.80488,0.009109,0.009109,0.009109
|
| 12 |
+
11,9341.04,0.6143,0.3126,0.81357,0.99451,0.98566,0.99391,0.87779,0.55769,0.24658,0.80416,0.00901,0.00901,0.00901
|
| 13 |
+
12,10184.5,0.60103,0.30526,0.81238,0.99462,0.98598,0.99398,0.88037,0.55125,0.24366,0.8027,0.008911,0.008911,0.008911
|
| 14 |
+
13,11031.8,0.59679,0.30283,0.81192,0.99405,0.98645,0.99409,0.87994,0.55084,0.24197,0.8026,0.008812,0.008812,0.008812
|
| 15 |
+
14,11876.7,0.59333,0.29933,0.80987,0.99453,0.98626,0.99402,0.88191,0.54564,0.23862,0.80157,0.008713,0.008713,0.008713
|
| 16 |
+
15,12722.1,0.58581,0.29494,0.80972,0.99523,0.98615,0.99412,0.88326,0.53836,0.23728,0.80025,0.008614,0.008614,0.008614
|
| 17 |
+
16,13565.6,0.58006,0.29061,0.80786,0.99542,0.98692,0.99403,0.88504,0.53513,0.23455,0.79935,0.008515,0.008515,0.008515
|
| 18 |
+
17,14411.3,0.57672,0.28795,0.80764,0.99561,0.98806,0.99427,0.88538,0.53908,0.23224,0.79983,0.008416,0.008416,0.008416
|
| 19 |
+
18,15258.7,0.57058,0.28694,0.80589,0.99513,0.98829,0.99422,0.88413,0.53386,0.23022,0.79868,0.008317,0.008317,0.008317
|
| 20 |
+
19,16103.3,0.56536,0.28123,0.80563,0.99493,0.98831,0.99425,0.8901,0.52304,0.22738,0.79724,0.008218,0.008218,0.008218
|
| 21 |
+
20,16946.7,0.56263,0.28067,0.80593,0.99553,0.98797,0.99426,0.88536,0.53237,0.22661,0.79834,0.008119,0.008119,0.008119
|
| 22 |
+
21,17789.4,0.55984,0.27929,0.80478,0.99557,0.98861,0.99427,0.88764,0.52399,0.226,0.7972,0.00802,0.00802,0.00802
|
| 23 |
+
22,18638.8,0.55793,0.27731,0.80402,0.99598,0.98838,0.99433,0.89168,0.51603,0.22054,0.79583,0.007921,0.007921,0.007921
|
| 24 |
+
23,19485,0.55206,0.27382,0.80311,0.99603,0.98875,0.99435,0.8932,0.51208,0.2185,0.79524,0.007822,0.007822,0.007822
|
| 25 |
+
24,20326.6,0.55068,0.27161,0.8036,0.99586,0.98887,0.99437,0.89249,0.51388,0.21833,0.79516,0.007723,0.007723,0.007723
|
| 26 |
+
25,21169.1,0.54745,0.27037,0.80235,0.99581,0.98928,0.99437,0.89536,0.50525,0.21598,0.79397,0.007624,0.007624,0.007624
|
| 27 |
+
26,22012.1,0.54547,0.26931,0.80247,0.99554,0.99001,0.99433,0.89836,0.49678,0.21205,0.79302,0.007525,0.007525,0.007525
|
| 28 |
+
27,22860.5,0.54133,0.26642,0.80207,0.99579,0.99015,0.99439,0.89703,0.4992,0.21184,0.79319,0.007426,0.007426,0.007426
|
| 29 |
+
28,23703.3,0.53816,0.26553,0.80118,0.99645,0.98994,0.99437,0.89788,0.49825,0.20991,0.79271,0.007327,0.007327,0.007327
|
| 30 |
+
29,24544.8,0.53313,0.2634,0.80064,0.99592,0.99047,0.99434,0.8979,0.49459,0.20961,0.79266,0.007228,0.007228,0.007228
|
| 31 |
+
30,25388.8,0.52973,0.26017,0.79985,0.99619,0.99031,0.99446,0.90008,0.48961,0.20792,0.79182,0.007129,0.007129,0.007129
|
| 32 |
+
31,26233.2,0.5269,0.25997,0.79837,0.99702,0.98979,0.9944,0.89964,0.49105,0.20738,0.79174,0.00703,0.00703,0.00703
|
| 33 |
+
32,27079.8,0.52688,0.25916,0.79927,0.99672,0.99078,0.9944,0.90069,0.48737,0.20495,0.79115,0.006931,0.006931,0.006931
|
| 34 |
+
33,27921.5,0.52129,0.25693,0.79811,0.99683,0.99049,0.99454,0.90165,0.48401,0.20341,0.79081,0.006832,0.006832,0.006832
|
| 35 |
+
34,28764.2,0.52114,0.25471,0.79824,0.99719,0.99018,0.99447,0.90317,0.47991,0.20193,0.79012,0.006733,0.006733,0.006733
|
| 36 |
+
35,29606.6,0.51987,0.25446,0.79883,0.9966,0.99087,0.99438,0.90382,0.47828,0.20192,0.79008,0.006634,0.006634,0.006634
|
| 37 |
+
36,30456.1,0.51496,0.25208,0.79743,0.99719,0.99037,0.99437,0.90508,0.47383,0.20009,0.78931,0.006535,0.006535,0.006535
|
| 38 |
+
37,31300.5,0.511,0.24972,0.79762,0.99658,0.99091,0.99437,0.90613,0.47076,0.19873,0.7887,0.006436,0.006436,0.006436
|
| 39 |
+
38,32143.9,0.51136,0.25042,0.79623,0.99692,0.99084,0.99445,0.90558,0.47097,0.19822,0.78883,0.006337,0.006337,0.006337
|
| 40 |
+
39,32985.8,0.50541,0.24745,0.79565,0.99724,0.99077,0.99448,0.90597,0.46946,0.19736,0.78834,0.006238,0.006238,0.006238
|
| 41 |
+
40,33830.3,0.50479,0.24695,0.79518,0.99724,0.99083,0.99452,0.90742,0.46542,0.19641,0.78804,0.006139,0.006139,0.006139
|
| 42 |
+
41,34677.4,0.5029,0.24614,0.79591,0.99696,0.99136,0.99454,0.90771,0.46354,0.19568,0.78762,0.00604,0.00604,0.00604
|
| 43 |
+
42,35520.5,0.49842,0.24446,0.79503,0.99725,0.99107,0.99452,0.90858,0.46129,0.19448,0.78728,0.005941,0.005941,0.005941
|
| 44 |
+
43,36364.6,0.49839,0.24419,0.7955,0.99732,0.99082,0.99444,0.90947,0.45884,0.19314,0.78682,0.005842,0.005842,0.005842
|
| 45 |
+
44,37206.5,0.49216,0.24015,0.79435,0.99722,0.99107,0.99449,0.91033,0.45611,0.19261,0.78638,0.005743,0.005743,0.005743
|
| 46 |
+
45,38049.1,0.49418,0.24203,0.79385,0.99703,0.99144,0.99451,0.91087,0.4543,0.19214,0.78613,0.005644,0.005644,0.005644
|
| 47 |
+
46,38897.8,0.48824,0.23884,0.79406,0.99717,0.99129,0.99453,0.91115,0.45407,0.19199,0.78622,0.005545,0.005545,0.005545
|
| 48 |
+
47,39741.3,0.48797,0.23914,0.79344,0.99707,0.99147,0.99454,0.91148,0.45203,0.19103,0.78606,0.005446,0.005446,0.005446
|
| 49 |
+
48,40584.9,0.48281,0.23633,0.79246,0.99712,0.9916,0.99454,0.91221,0.44899,0.19046,0.7859,0.005347,0.005347,0.005347
|
| 50 |
+
49,41428.6,0.4808,0.23643,0.79376,0.99757,0.99099,0.99447,0.91283,0.4469,0.19003,0.78564,0.005248,0.005248,0.005248
|
| 51 |
+
50,42275,0.4799,0.23503,0.7923,0.99698,0.99156,0.99457,0.91367,0.44459,0.18939,0.78536,0.005149,0.005149,0.005149
|
| 52 |
+
51,43117.5,0.47796,0.23351,0.79202,0.99712,0.99146,0.99457,0.91408,0.44334,0.18854,0.7851,0.00505,0.00505,0.00505
|
| 53 |
+
52,43961,0.47425,0.23329,0.79199,0.99745,0.99117,0.9946,0.91474,0.44112,0.18791,0.78484,0.004951,0.004951,0.004951
|
| 54 |
+
53,44804,0.46837,0.22882,0.79024,0.99725,0.99152,0.99463,0.91509,0.43921,0.1873,0.78449,0.004852,0.004852,0.004852
|
| 55 |
+
54,45648,0.46714,0.22827,0.79055,0.99736,0.99148,0.99464,0.91561,0.4377,0.18655,0.78424,0.004753,0.004753,0.004753
|
| 56 |
+
55,46495.3,0.46614,0.22808,0.79074,0.99736,0.99148,0.99463,0.91603,0.43667,0.18608,0.78411,0.004654,0.004654,0.004654
|
| 57 |
+
56,47339.5,0.4637,0.2255,0.78996,0.99686,0.99221,0.99465,0.91641,0.43518,0.18521,0.78393,0.004555,0.004555,0.004555
|
| 58 |
+
57,48184.1,0.46273,0.22496,0.78987,0.99748,0.99156,0.99458,0.91676,0.43384,0.18469,0.78368,0.004456,0.004456,0.004456
|
| 59 |
+
58,49029.3,0.45682,0.2228,0.79042,0.99749,0.99158,0.99462,0.91735,0.43217,0.18397,0.78336,0.004357,0.004357,0.004357
|
| 60 |
+
59,49877.6,0.45594,0.22277,0.79048,0.99748,0.99154,0.99461,0.9177,0.43028,0.18332,0.78314,0.004258,0.004258,0.004258
|
| 61 |
+
60,50721.8,0.45355,0.2194,0.78919,0.99744,0.99169,0.99461,0.91808,0.42898,0.18286,0.78299,0.004159,0.004159,0.004159
|
| 62 |
+
61,51566.9,0.44861,0.21824,0.7893,0.99742,0.99164,0.99466,0.91869,0.42741,0.18213,0.78286,0.00406,0.00406,0.00406
|
| 63 |
+
62,52412.3,0.44858,0.21778,0.78838,0.99736,0.99171,0.99467,0.91925,0.42557,0.18148,0.78267,0.003961,0.003961,0.003961
|
| 64 |
+
63,53257.7,0.44843,0.21688,0.78885,0.99726,0.99179,0.99467,0.91961,0.42442,0.18094,0.78249,0.003862,0.003862,0.003862
|
| 65 |
+
64,54107,0.44477,0.21446,0.78753,0.99715,0.99188,0.99467,0.91993,0.42315,0.18063,0.78229,0.003763,0.003763,0.003763
|
| 66 |
+
65,54953,0.44175,0.21357,0.7874,0.99734,0.99172,0.99467,0.92038,0.42182,0.18018,0.78208,0.003664,0.003664,0.003664
|
| 67 |
+
66,55800,0.43928,0.21206,0.78695,0.99751,0.99154,0.99469,0.92063,0.42075,0.17969,0.78192,0.003565,0.003565,0.003565
|
| 68 |
+
67,56645.5,0.43963,0.21207,0.78765,0.9974,0.99168,0.99469,0.92097,0.41954,0.17931,0.78178,0.003466,0.003466,0.003466
|
| 69 |
+
68,57496.2,0.43514,0.21014,0.78565,0.99751,0.99162,0.99468,0.92127,0.41848,0.17882,0.78167,0.003367,0.003367,0.003367
|
| 70 |
+
69,58350.2,0.42996,0.20736,0.78585,0.99755,0.9915,0.99462,0.92163,0.41739,0.17858,0.78156,0.003268,0.003268,0.003268
|
| 71 |
+
70,59199.3,0.42742,0.20551,0.78586,0.99759,0.9915,0.99461,0.9218,0.41656,0.1783,0.78142,0.003169,0.003169,0.003169
|
| 72 |
+
71,60048.6,0.42606,0.20497,0.78549,0.99755,0.99153,0.99469,0.92203,0.41563,0.17809,0.78128,0.00307,0.00307,0.00307
|
| 73 |
+
72,60897.2,0.42453,0.20353,0.78485,0.99738,0.99163,0.99471,0.92237,0.41471,0.17778,0.78115,0.002971,0.002971,0.002971
|
| 74 |
+
73,61751.8,0.42075,0.20147,0.78419,0.99736,0.99168,0.99471,0.92261,0.41389,0.17746,0.781,0.002872,0.002872,0.002872
|
| 75 |
+
74,62622.8,0.41807,0.20066,0.78485,0.99734,0.99175,0.99471,0.92299,0.41298,0.17724,0.7809,0.002773,0.002773,0.002773
|
| 76 |
+
75,63476.1,0.41766,0.20019,0.78364,0.99736,0.99163,0.99472,0.92325,0.41214,0.17706,0.78079,0.002674,0.002674,0.002674
|
| 77 |
+
76,64327.2,0.416,0.19909,0.78397,0.99737,0.99158,0.99473,0.92332,0.41149,0.17701,0.78071,0.002575,0.002575,0.002575
|
| 78 |
+
77,65176.8,0.41124,0.19658,0.78418,0.99716,0.99183,0.99473,0.92365,0.41066,0.17665,0.78061,0.002476,0.002476,0.002476
|
| 79 |
+
78,66035,0.40958,0.19567,0.78388,0.99712,0.99182,0.99473,0.92393,0.40989,0.17629,0.78052,0.002377,0.002377,0.002377
|
| 80 |
+
79,66886.2,0.40919,0.19443,0.78283,0.99715,0.9918,0.99473,0.92408,0.40921,0.1762,0.78042,0.002278,0.002278,0.002278
|
| 81 |
+
80,67738.1,0.40715,0.19381,0.78292,0.99708,0.99192,0.99472,0.92434,0.40846,0.17597,0.78031,0.002179,0.002179,0.002179
|
| 82 |
+
81,68589.2,0.40328,0.19213,0.78175,0.99709,0.99194,0.99472,0.9245,0.40775,0.17579,0.78023,0.00208,0.00208,0.00208
|
| 83 |
+
82,69443.4,0.40319,0.19144,0.78211,0.99719,0.99182,0.99472,0.92463,0.40726,0.17563,0.78012,0.001981,0.001981,0.001981
|
| 84 |
+
83,70289.6,0.39893,0.18903,0.78224,0.99705,0.99196,0.99471,0.92486,0.40658,0.17557,0.78004,0.001882,0.001882,0.001882
|
| 85 |
+
84,71137.2,0.39812,0.18907,0.78186,0.99678,0.99223,0.9947,0.92504,0.40588,0.17558,0.77995,0.001783,0.001783,0.001783
|
| 86 |
+
85,71985.2,0.39451,0.18653,0.78083,0.99717,0.9918,0.9947,0.92519,0.40524,0.17548,0.77985,0.001684,0.001684,0.001684
|
| 87 |
+
86,72835.6,0.39281,0.18653,0.78087,0.99708,0.99194,0.9947,0.92533,0.40454,0.17543,0.77977,0.001585,0.001585,0.001585
|
| 88 |
+
87,73688.6,0.38824,0.18342,0.78001,0.99712,0.99192,0.99469,0.92568,0.40373,0.17536,0.77965,0.001486,0.001486,0.001486
|
| 89 |
+
88,74539.5,0.38613,0.18234,0.78023,0.99715,0.99184,0.9947,0.92593,0.40293,0.17522,0.77955,0.001387,0.001387,0.001387
|
| 90 |
+
89,75389.8,0.38569,0.18154,0.78027,0.99716,0.99181,0.9947,0.92611,0.40243,0.17525,0.77946,0.001288,0.001288,0.001288
|
| 91 |
+
90,76239.2,0.38303,0.17992,0.7801,0.99713,0.99185,0.9947,0.92634,0.40173,0.17533,0.77937,0.001189,0.001189,0.001189
|
| 92 |
+
91,77086.7,0.37095,0.17454,0.77936,0.99716,0.99186,0.9947,0.92656,0.40097,0.17513,0.77928,0.00109,0.00109,0.00109
|
| 93 |
+
92,77938.8,0.36715,0.17267,0.77984,0.9971,0.99194,0.99469,0.9268,0.40027,0.17496,0.77918,0.000991,0.000991,0.000991
|
| 94 |
+
93,78809.8,0.36022,0.16862,0.77824,0.99715,0.99183,0.99469,0.92707,0.39938,0.17471,0.77907,0.000892,0.000892,0.000892
|
| 95 |
+
94,79725.6,0.36091,0.16907,0.77796,0.99718,0.99181,0.99469,0.92743,0.39847,0.17434,0.77896,0.000793,0.000793,0.000793
|
| 96 |
+
95,80641,0.35823,0.16708,0.77799,0.99714,0.99188,0.99469,0.92771,0.39764,0.17414,0.77886,0.000694,0.000694,0.000694
|
| 97 |
+
96,81570.1,0.35413,0.16608,0.77726,0.99715,0.9918,0.99469,0.92805,0.39675,0.17394,0.77875,0.000595,0.000595,0.000595
|
| 98 |
+
97,82477.9,0.35209,0.16355,0.77688,0.99703,0.9919,0.99469,0.92822,0.39607,0.17364,0.77865,0.000496,0.000496,0.000496
|
| 99 |
+
98,83378.5,0.34975,0.1626,0.77636,0.99695,0.99201,0.99469,0.92863,0.39533,0.17326,0.77857,0.000397,0.000397,0.000397
|
| 100 |
+
99,84261.4,0.34865,0.16152,0.77566,0.99699,0.99195,0.99469,0.92895,0.39459,0.1728,0.77846,0.000298,0.000298,0.000298
|
| 101 |
+
100,85165.5,0.34551,0.15998,0.77677,0.99728,0.99157,0.99469,0.92913,0.39388,0.17234,0.77837,0.000199,0.000199,0.000199
|
models/digitization_models/yolo11_lead/.ipynb_checkpoints/results-checkpoint.png
ADDED
|
Git LFS Details
|
models/digitization_models/yolo11_lead/.ipynb_checkpoints/val_batch0_labels-checkpoint.jpg
ADDED
|
Git LFS Details
|
models/digitization_models/yolo11_lead/.ipynb_checkpoints/val_batch0_pred-checkpoint.jpg
ADDED
|
Git LFS Details
|
models/digitization_models/yolo11_lead/.ipynb_checkpoints/val_batch1_labels-checkpoint.jpg
ADDED
|
Git LFS Details
|
models/digitization_models/yolo11_lead/F1_curve.png
ADDED
|
Git LFS Details
|