code/scripts/generate_synthetic_0001.py

#!/usr/bin/env python3
"""
Generate 0001-type synthetic ECG images using ecg-image-kit.

Uses PTB-XL dataset records to generate clean baseline images 
that match the competition format.

This runs in parallel with training to generate more data.

Usage:
    python generate_synthetic_0001.py --num_samples 10000 --output_dir /data/ecg-digitization/synthetic_0001_v2
"""

import os
import sys
import argparse
import json
import random
import subprocess
from pathlib import Path
from concurrent.futures import ProcessPoolExecutor, as_completed
from tqdm import tqdm
import numpy as np
import cv2
import wfdb
from scipy import signal as scipy_signal
import shutil

# Paths
PROJECT_ROOT = Path(__file__).parent.parent
PTBXL_DIR = Path('/data/ecg-digitization/ptbxl/physionet.org/files/ptb-xl/1.0.3')
ECG_IMAGE_KIT = PROJECT_ROOT / 'ecg-image-kit' / 'codes' / 'ecg-image-generator'
BASELINE_PATH = PROJECT_ROOT / 'data' / 'hengck23-submit-physionet' / 'hengck23-submit-physionet'
OUTPUT_DIR = '/data/ecg-digitization/synthetic_0001_v2'

# Target dimensions (match competition Stage 1 output)
TARGET_WIDTH = 4352
TARGET_HEIGHT = 1696

# GT parameters matching v9 training
T0 = 235
T1 = 4161
OUTPUT_GT_WIDTH = T1 - T0  # 3926


def find_ptbxl_records(ptbxl_dir):
    """Find all PTB-XL record files."""
    records = []
    ptbxl_dir = Path(ptbxl_dir)
    
    # Use records500 (500Hz) - higher quality for generation
    records500 = ptbxl_dir / 'records500'
    if records500.exists():
        for subdir in sorted(records500.iterdir()):
            if subdir.is_dir():
                for hea in subdir.glob('*_hr.hea'):  # Only high-res files
                    dat = hea.with_suffix('.dat')
                    if dat.exists():
                        records.append((str(hea), str(dat)))
    
    # Fallback to records100 if needed
    if len(records) == 0:
        records100 = ptbxl_dir / 'records100'
        if records100.exists():
            for subdir in sorted(records100.iterdir()):
                if subdir.is_dir():
                    for hea in subdir.glob('*_lr.hea'):
                        dat = hea.with_suffix('.dat')
                        if dat.exists():
                            records.append((str(hea), str(dat)))
    
    return records


def extract_gt_signal(hea_file, output_width=3926):
    """Extract ground truth signal from PTB-XL record."""
    try:
        # Load WFDB record
        record_path = hea_file.replace('.hea', '')
        record = wfdb.rdrecord(record_path)
        
        signals = record.p_signal  # [samples, leads]
        sig_names = record.sig_name
        fs = record.fs
        
        # Resample to 500Hz if needed
        if fs != 500:
            num_samples = int(signals.shape[0] * 500 / fs)
            signals = scipy_signal.resample(signals, num_samples, axis=0)
        
        # Map lead names
        lead_signals = {}
        for i, name in enumerate(sig_names):
            lead_signals[name.upper()] = signals[:, i]
        
        # 4-row output: each row shows 3 leads, 2.5s each
        row_leads = [
            ['I', 'AVR', 'V1', 'V4'],
            ['II', 'AVL', 'V2', 'V5'],
            ['III', 'AVF', 'V3', 'V6'],
        ]
        
        segment_len = output_width // 3
        gt_signal = np.zeros((4, output_width), dtype=np.float32)
        
        for row_idx in range(4):
            leads_for_row = [row_leads[col][row_idx] for col in range(3)]
            
            for col_idx, lead_name in enumerate(leads_for_row):
                if lead_name in lead_signals:
                    lead_data = lead_signals[lead_name]
                    # 2.5s at 500Hz = 1250 samples
                    start = col_idx * 1250
                    end = min(start + 1250, len(lead_data))
                    segment = lead_data[start:end]
                    
                    col_start = col_idx * segment_len
                    col_end = (col_idx + 1) * segment_len if col_idx < 2 else output_width
                    
                    x_old = np.linspace(0, 1, len(segment))
                    x_new = np.linspace(0, 1, col_end - col_start)
                    gt_signal[row_idx, col_start:col_end] = np.interp(x_new, x_old, segment)
        
        return gt_signal
        
    except Exception as e:
        return None


def generate_single_image(args):
    """Generate a single 0001-type synthetic ECG image."""
    idx, hea_file, dat_file, output_dir, temp_base = args
    
    sample_id = f'syn2_{idx:08d}'
    out_img = output_dir / 'stage1' / f'{sample_id}-0001.png'
    out_gt = output_dir / 'gt' / f'{sample_id}.npy'
    
    # Skip if exists
    if out_img.exists() and out_gt.exists():
        return True, idx, "exists"
    
    temp_dir = temp_base / f'temp_{idx}'
    
    try:
        temp_dir.mkdir(parents=True, exist_ok=True)
        
        # Generate image with ecg-image-kit (0001 = clean red grid)
        cmd = [
            'python', 'gen_ecg_image_from_data.py',
            '-i', dat_file,
            '-hea', hea_file,
            '-o', str(temp_dir),
            '-st', str(idx),
            '-se', str(idx),
            '-r', '200',  # DPI
            '--num_columns', '4',
            '--full_mode', 'II',
            '--standard_grid_color', '5',  # Red grid (0001 style)
            # No augmentation, no wrinkles for clean baseline
        ]
        
        result = subprocess.run(
            cmd,
            cwd=str(ECG_IMAGE_KIT),
            capture_output=True,
            text=True,
            timeout=60
        )
        
        if result.returncode != 0:
            shutil.rmtree(temp_dir, ignore_errors=True)
            return False, idx, f"ecg-kit error: {result.stderr[:100]}"
        
        # Find generated image
        gen_imgs = list(temp_dir.glob('*.png'))
        if not gen_imgs:
            shutil.rmtree(temp_dir, ignore_errors=True)
            return False, idx, "no image generated"
        
        # Load and resize to target
        img = cv2.imread(str(gen_imgs[0]))
        if img is None:
            shutil.rmtree(temp_dir, ignore_errors=True)
            return False, idx, "failed to read image"
        
        img_resized = cv2.resize(img, (TARGET_WIDTH, TARGET_HEIGHT), interpolation=cv2.INTER_LINEAR)
        
        # Save image
        out_img.parent.mkdir(parents=True, exist_ok=True)
        cv2.imwrite(str(out_img), img_resized)
        
        # Extract GT signal
        gt_signal = extract_gt_signal(hea_file, OUTPUT_GT_WIDTH)
        if gt_signal is not None:
            out_gt.parent.mkdir(parents=True, exist_ok=True)
            np.save(out_gt, gt_signal)
        
        # Cleanup
        shutil.rmtree(temp_dir, ignore_errors=True)
        
        return True, idx, "success"
        
    except Exception as e:
        shutil.rmtree(temp_dir, ignore_errors=True)
        return False, idx, str(e)[:100]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--num_samples', type=int, default=10000)
    parser.add_argument('--output_dir', type=str, default=OUTPUT_DIR)
    parser.add_argument('--workers', type=int, default=8)
    args = parser.parse_args()
    
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    (output_dir / 'stage1').mkdir(exist_ok=True)
    (output_dir / 'gt').mkdir(exist_ok=True)
    
    temp_base = output_dir / 'temp'
    temp_base.mkdir(exist_ok=True)
    
    print(f"Finding PTB-XL records in {PTBXL_DIR}...")
    records = find_ptbxl_records(PTBXL_DIR)
    print(f"Found {len(records)} records")
    
    if len(records) == 0:
        print("No records found!")
        return
    
    # Sample records if needed
    if len(records) > args.num_samples:
        records = random.sample(records, args.num_samples)
    elif len(records) < args.num_samples:
        # Duplicate records if we need more samples
        while len(records) < args.num_samples:
            records.extend(random.sample(records[:len(records)//2], 
                                        min(len(records)//2, args.num_samples - len(records))))
    
    print(f"Generating {len(records)} 0001-type synthetic images...")
    
    # Prepare tasks
    tasks = [
        (i, hea, dat, output_dir, temp_base)
        for i, (hea, dat) in enumerate(records)
    ]
    
    success = 0
    failed = 0
    
    with ProcessPoolExecutor(max_workers=args.workers) as executor:
        futures = {executor.submit(generate_single_image, task): task for task in tasks}
        
        for future in tqdm(as_completed(futures), total=len(futures), desc="Generating"):
            ok, idx, msg = future.result()
            if ok:
                success += 1
            else:
                failed += 1
                if failed <= 5:
                    print(f"\n  Failed {idx}: {msg}")
    
    print(f"\n✓ Done! Success: {success}, Failed: {failed}")
    print(f"  Images: {output_dir / 'stage1'}")
    print(f"  GT: {output_dir / 'gt'}")
    
    # Create manifest
    manifest = []
    for img in (output_dir / 'stage1').glob('*.png'):
        sample_id = img.stem.replace('-0001', '')
        gt_path = output_dir / 'gt' / f'{sample_id}.npy'
        if gt_path.exists():
            manifest.append({
                'sample_id': sample_id,
                'style': '0001',
                'path': str(img),
                'gt': str(gt_path)
            })
    
    with open(output_dir / 'manifest.json', 'w') as f:
        json.dump(manifest, f, indent=2)
    
    print(f"  Manifest: {len(manifest)} samples")
    
    # Cleanup
    shutil.rmtree(temp_base, ignore_errors=True)


if __name__ == '__main__':
    main()