added installation guide

.gitignore

@@ -1,3 +1,3 @@
-dataset
+dataset/ptb_xl/
 .venv
 runs

README.md

@@ -1,2 +1,58 @@
 # ecg_reconstruction
-Just Playing around for now.
+
+Self-supervised **12-lead ECG reconstruction** with a **masked autoencoder (MAE)** built using the idea of **[CoRe-ECG](https://arxiv.org/abs/2604.11359)**: spatio-temporal dual masking (STDM), a **visibility-restricted encoder**, and a full **decoder** that predicts masked patches. This repo implements the **reconstruction branch only** (no downstream classification head in the training loop).
+
+## Techniques at a glance
+
+
+| Idea                                    | What it does here                                                                                                                                                                                                                                                                                                                                             |
+| --------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| **Patch embedding**                     | Each lead is split into non-overlapping time patches (default **75 samples** per patch, **66 patches** → **4950** samples per lead at **500 Hz**). Patches are linearly embedded plus learned **2D position** over leads × time indices.                                                                                                                      |
+| **STDM (Spatio-Temporal Dual Masking)** | Per batch and per time index, either **full temporal masking** (all leads supervised at that index) or **partial masking**: **k** leads stay **visible** (encoder input), remaining leads are split into **reconstruction targets** vs **dropped** (no loss). Controlled by `p_time`, `p_lead`, and `num_visible_leads` (`k`) in `MAEConfig` / `mae/stdm.py`. |
+| **Visibility-restricted encoder**       | Standard self-attention would let masked positions “peek” at neighbors. Here, **additive attention bias** restricts visible tokens to attend to visible keys; non-visible positions use a **stabilized identity row** so the encoder does not mix information across the visibility boundary (`mae/encoder.py`).                                              |
+| **Decoder + mask token**                | Encoder outputs are projected; **visible** slots keep signal tokens, **non-visible** slots get a learned **mask token** plus decoder positions. A shallow **decoder stack** (full attention over all positions) predicts **patch pixels**; loss is **MSE only on STDM supervision mask `M`** (`mae/losses.py`).                                               |
+| **Preprocessing**                       | **Butterworth bandpass** 0.65–40 Hz, zero-phase `filtfilt`, per lead (`preprocessor.py`), aligned with common ECG MAE setups.                                                                                                                                                                                                                                 |
+| **Data**                                | **PTB-XL** high-rate records (`filename_hr` @ 500 Hz), official **stratified folds** 1–8 train, 9 val, 10 test (`ptb_xl_dataset.py`).                                                                                                                                                                                                                         |
+
+
+## Project layout
+
+- `train_ecg_mae.py` — training loop: **AdamW**, TensorBoard scalars and periodic **reconstruction figures**, checkpoints.
+- `mae/` — config, **ECGDataMAE** model, STDM sampling, encoder/decoder blocks, loss.
+- `preprocessor.py` — filtering, patch length / signal window constants shared with the model.
+- `ptb_xl_dataset.py` — CSV-driven PTB-XL loading via **WFDB**.
+- `inference.py` — `load_pipeline`, `reconstruct`, checkpoint I/O, plotting helpers for dashboards or notebooks.
+- `visualization.ipynb` — exploratory plots.
+
+## Setup
+
+Create a virtual environment:
+
+```bash
+python -m venv .venv
+source .venv/bin/activate
+```
+
+Install **PyTorch** separately (CUDA build as you need). Other dependencies:
+
+```bash
+pip install -r requirements.txt
+```
+
+Check this [GUIDE](./dataset/README.md) to install the PTB-XL dataset.
+
+## Train
+
+```bash
+python train_ecg_mae.py --data-root dataset/ptb_xl --log-dir runs/[experiment name] --epochs 80
+```
+
+Useful flags: `--batch-size`, `--lr`, `--weight-decay`, `--resume path/to/checkpoint.pt`, `--ckpt-every`, `--vis-every`. Logs and checkpoints go under `--log-dir`; that folder is listed in `.gitignore` so artifacts stay local.
+
+## Inference
+
+Load a saved checkpoint and run `reconstruct()` from `inference.py` on tensors shaped `(batch, 12, signal_length)` with the same preprocessing as training.
+
+## Note on patch length vs. paper
+
+The implementation keeps **75 samples per patch** (paper patch length in samples). At **500 Hz** the temporal extent per patch differs from a 250 Hz setup; the code comments in `preprocessor.py` describe this tradeoff explicitly.

dataset/README.md

@@ -0,0 +1,22 @@
+# PTB-XL dataset (local install)
+
+Training and evaluation use **PTB-XL v1.0.3** at **500 Hz**: waveforms are read from paths in the `filename_hr` column of `ptbxl_database.csv` (under `records500/`). See `ptb_xl_dataset.py` for how splits use `strat_fold`.
+
+`dataset/ptb_xl/` is listed in the repo `.gitignore`, so downloaded files stay on your machine and are not committed.
+
+## Prerequisites
+
+- **AWS CLI v2** with the `aws` command available ([install guide](https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html)).
+- Enough **disk space** for the release (full **1.0.3** sync is on the order of **many gigabytes**; exact size depends on PhysioNet’s current layout).
+- Network access to the public **PhysioNet** S3 bucket (no AWS account required for this bucket).
+
+## Download
+
+From the **repository root** (parent of `dataset/`):
+
+```bash
+cd dataset
+aws s3 sync --no-sign-request s3://physionet-open/ptb-xl/1.0.3/ ptb_xl
+```
+
+This creates `dataset/ptb_xl/` with metadata and WFDB files, including `ptbxl_database.csv` and the `records500/` tree used for high-rate records.