# ECG Digitization Experiments

This repository contains archived PyTorch `.pth` checkpoints from ECG digitization research experiments (Kaggle competition).

## About

These are model weights for ECG image-to-signal digitization - extracting ECG waveform values from paper ECG images.

**Competition:** [ECG Image Digitization Challenge](https://www.kaggle.com/competitions/ecg-image-digitization)

## Experiments Overview

| Experiment | Architecture | Description | Best SNR |
|------------|--------------|-------------|----------|
| v7 | EfficientNet-B4 | Early baseline | ~15 dB |
| v9 | EfficientNet-B4 | Improved training | ~16 dB |
| v10 | EfficientNet-B4 | Multi-scale features | ~17 dB |
| v10_1 | EfficientNet-B4 | Refinements | ~17 dB |
| v11 | EfficientNet-B4 | 1.5x scale | ~17 dB |
| v14 | ConvNeXt + SimDR | SimDR heatmap approach | ~18 dB |
| v15 | Per-Lead CNN | Per-lead extraction | ~18 dB |
| v16 | Per-Lead + BiLSTM | Temporal coherence | ~19 dB |
| v18 | Refiner Network | Post-processing refiner | ~19 dB |
| v19 | Augraphy Augmentation | Paper degradation aug | ~20 dB |
| v20 | Integral Regression | Integral loss | ~20 dB |
| v21 | GRU Refiner | GRU-based refinement | ~20 dB |
| v22 | ConvNeXt-Base + U-Net | Multi-scale fusion + Height Attention + BiLSTM | ~22 dB |
| **v23** | **V22 + DSNT** | **DSNT Sub-Pixel Head (Rank 3 technique)** | **~22.35 dB** |
| mixed | Various | Early mixed training | ~16 dB |
| mixed_v4 | Various | Mixed v4 | ~17 dB |
| mixed_v5 | Various | Mixed v5 | ~17 dB |

## V23 Architecture (Best Model)

```
① ConvNeXt-Base Encoder → L3 + L4 multi-scale features
② Feature Fusion (fine details + semantics)
③ U-Net Decoder with skip connections
④ Height Attention (learns vertical regions)
⑤ BiLSTM (temporal coherence)
⑥ Conv1D + Linear → Sigmoid (V22 head)
⑦ DSNT Sub-Pixel Head (Rank 3 winner technique)
```

## Directory Structure

```
├── README.md
├── v7/                    # EfficientNet-B4 baseline
├── v9/                    # Improved training
├── v10/                   # Multi-scale features
├── v10_1/                 # Refinements
├── v11/                   # 1.5x scale
├── v14/                   # ConvNeXt + SimDR
├── v15/                   # Per-lead CNN
├── v16/                   # Per-lead + BiLSTM
├── v18/                   # Refiner network
├── v19/                   # Augraphy augmentation
├── v20/                   # Integral regression
├── v21/                   # GRU refiner
├── v22/                   # ConvNeXt-Base + U-Net
├── v23/                   # V22 + DSNT (BEST)
├── mixed/                 # Mixed training v1
├── mixed_v4/              # Mixed training v4
├── mixed_v5/              # Mixed training v5
└── code/
    ├── scripts/           # All training scripts
    └── notebooks/         # Kaggle inference notebooks
```

## Checkpoint Format

Each `.pth` file contains:
- `model`: Model state dict
- `opt`: Optimizer state dict
- `epoch`: Training epoch
- `snr`: Validation SNR (dB)
- `mae`: Mean Absolute Error (pixels)

## Usage

```python
import torch

# Load checkpoint
checkpoint = torch.load("v23/v23_epoch038.pth", map_location="cpu")

# Load model state
model.load_state_dict(checkpoint['model'])

# Check metrics
print(f"Epoch: {checkpoint['epoch']}")
print(f"SNR: {checkpoint['snr']:.2f} dB")
```

## Key Files

- `v23/v23_epoch038.pth` - Best V23 checkpoint (22.35 dB)
- `v22/v22_best_snr.pth` - Best V22 checkpoint
- `code/scripts/train_v23.py` - V23 training script
- `code/scripts/inference_v23.py` - V23 inference script
- `code/notebooks/kaggle_v23_inference.ipynb` - Kaggle submission notebook

## License

Research use only. Please cite if you use these weights.

## Framework

- PyTorch 2.x
- timm (ConvNeXt-Base encoder)
- albumentations (augmentation)