# mygithth27/af-risk-prediction-by-ecg-dnn architecture

> Architecture diagram of the GitHub repository mygithth27/af-risk-prediction-by-ecg-dnn, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-30). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/mygithth27/af-risk-prediction-by-ecg-dnn
- Repository: https://github.com/mygithth27/af-risk-prediction-by-ecg-dnn

## Overview

This repository supports research workflows for three-class atrial-fibrillation risk prediction from 12-lead ECGs, plus survival-model analysis. The main workflow formulates patient/exam labels and splits, trains a residual network from ECG batches, and evaluates a saved model to produce class probabilities. It also includes ECG-subset preparation and survival-analysis notebooks. The training and evaluation links to their core model components are visible in the excerpts; survival-model internals were not sampled, so that branch is documented but not expanded.

Main entry points: [`evaluate.py`](https://github.com/mygithth27/af-risk-prediction-by-ecg-dnn/blob/main/evaluate.py) (Model evaluation), [`train.py`](https://github.com/mygithth27/af-risk-prediction-by-ecg-dnn/blob/main/train.py) (Model training), [`generate_subhdf.py`](https://github.com/mygithth27/af-risk-prediction-by-ecg-dnn/blob/main/generate_subhdf.py) (ECG subset preparation), [`af_prediction_formulation.py`](https://github.com/mygithth27/af-risk-prediction-by-ecg-dnn/blob/main/af_prediction_formulation.py) (AF labels and splits)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_formulation["Label preparation"]
  node_af_labels["AF labels and splits"]
  node_split_helper["Age-based splits"]
end

subgraph group_prediction["Model prediction"]
  node_training["Model training<br/>[train.py]"]
  node_batch_loader["Batch loader<br/>[dataloader.py]"]
  node_resnet["Residual network<br/>[resnet.py]"]
  node_evaluation["Model evaluation<br/>[evaluate.py]"]
end

subgraph group_data["ECG data"]
  node_subset_prep["ECG subset preparation<br/>[generate_subhdf.py]"]
  node_ecg_store[("ECG HDF5 data")]
  node_metadata["Exam metadata CSV"]
  node_prediction_csv["Prediction CSV"]
end

subgraph group_analysis["Risk analysis"]
  node_survival["Survival analysis"]
  node_survival_comorbidities["Comorbidity analysis"]
end

node_researcher(("Researcher"))

node_researcher -->|"formulates labels"| node_af_labels
node_af_labels -->|"writes metadata"| node_metadata
node_researcher -->|"creates splits"| node_split_helper
node_split_helper -.->|"provides splits"| node_training
node_researcher -->|"trains model"| node_training
node_training -->|"batches data"| node_batch_loader
node_training -->|"trains network"| node_resnet
node_researcher -->|"evaluates model"| node_evaluation
node_evaluation -->|"runs inference"| node_resnet
node_evaluation -->|"reads traces"| node_ecg_store
node_evaluation -->|"writes probabilities"| node_prediction_csv
node_researcher -->|"prepares subset"| node_subset_prep
node_subset_prep -->|"reads traces"| node_ecg_store
node_subset_prep -->|"reads labels"| node_metadata
node_researcher -->|"explores survival"| node_survival
node_researcher -->|"explores comorbidities"| node_survival_comorbidities

click node_af_labels "https://github.com/mygithth27/af-risk-prediction-by-ecg-dnn/blob/main/af_prediction_formulation.py"
click node_split_helper "https://github.com/mygithth27/af-risk-prediction-by-ecg-dnn/blob/main/archives/formulate_problem.py"
click node_training "https://github.com/mygithth27/af-risk-prediction-by-ecg-dnn/blob/main/train.py"
click node_batch_loader "https://github.com/mygithth27/af-risk-prediction-by-ecg-dnn/blob/main/dataloader.py"
click node_resnet "https://github.com/mygithth27/af-risk-prediction-by-ecg-dnn/blob/main/resnet.py"
click node_evaluation "https://github.com/mygithth27/af-risk-prediction-by-ecg-dnn/blob/main/evaluate.py"
click node_subset_prep "https://github.com/mygithth27/af-risk-prediction-by-ecg-dnn/blob/main/generate_subhdf.py"
click node_survival "https://github.com/mygithth27/af-risk-prediction-by-ecg-dnn/blob/main/notebooks/survival_model.ipynb"
click node_survival_comorbidities "https://github.com/mygithth27/af-risk-prediction-by-ecg-dnn/blob/main/notebooks/survival_model_comorbidities.ipynb"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_af_labels,node_split_helper toneBlue
class node_training,node_batch_loader,node_resnet,node_evaluation toneAmber
class node_subset_prep,node_ecg_store,node_metadata,node_prediction_csv toneMint
class node_survival,node_survival_comorbidities toneRose
class node_researcher toneIndigo
```

## Components

### Label preparation

- [AF labels and splits](https://github.com/mygithth27/af-risk-prediction-by-ecg-dnn/blob/main/af_prediction_formulation.py): `af_prediction_formulation.py`
- [Age-based splits](https://github.com/mygithth27/af-risk-prediction-by-ecg-dnn/blob/main/archives/formulate_problem.py): `archives/formulate_problem.py`

### Model prediction

- [Model training](https://github.com/mygithth27/af-risk-prediction-by-ecg-dnn/blob/main/train.py): `train.py`
- [Batch loader](https://github.com/mygithth27/af-risk-prediction-by-ecg-dnn/blob/main/dataloader.py): `dataloader.py`
- [Residual network](https://github.com/mygithth27/af-risk-prediction-by-ecg-dnn/blob/main/resnet.py): `resnet.py`
- [Model evaluation](https://github.com/mygithth27/af-risk-prediction-by-ecg-dnn/blob/main/evaluate.py): `evaluate.py`

### ECG data

- [ECG subset preparation](https://github.com/mygithth27/af-risk-prediction-by-ecg-dnn/blob/main/generate_subhdf.py): `generate_subhdf.py`
- **ECG HDF5 data**
- **Exam metadata CSV**
- **Prediction CSV**

### Risk analysis

- [Survival analysis](https://github.com/mygithth27/af-risk-prediction-by-ecg-dnn/blob/main/notebooks/survival_model.ipynb): `notebooks/survival_model.ipynb`
- [Comorbidity analysis](https://github.com/mygithth27/af-risk-prediction-by-ecg-dnn/blob/main/notebooks/survival_model_comorbidities.ipynb): `notebooks/survival_model_comorbidities.ipynb`

### Other components

- **Researcher**

## Connections

- Researcher → AF labels and splits: formulates labels (evidence: `README.md`)
- AF labels and splits → Exam metadata CSV: writes metadata (evidence: `af_prediction_formulation.py`)
- Researcher → Age-based splits: creates splits (evidence: `README.md`)
- Age-based splits → Model training: provides splits
- Researcher → Model training: trains model (evidence: `README.md`)
- Model training → Batch loader: batches data (evidence: `train.py`)
- Model training → Residual network: trains network (evidence: `train.py`)
- Researcher → Model evaluation: evaluates model (evidence: `README.md`)
- Model evaluation → Residual network: runs inference (evidence: `evaluate.py`)
- Model evaluation → ECG HDF5 data: reads traces (evidence: `evaluate.py`)
- Model evaluation → Prediction CSV: writes probabilities (evidence: `evaluate.py`)
- Researcher → ECG subset preparation: prepares subset (evidence: `generate_subhdf.py`)
- ECG subset preparation → ECG HDF5 data: reads traces (evidence: `generate_subhdf.py`)
- ECG subset preparation → Exam metadata CSV: reads labels (evidence: `generate_subhdf.py`)
- Researcher → Survival analysis: explores survival (evidence: `README.md`)
- Researcher → Comorbidity analysis: explores comorbidities (evidence: `README.md`)