code/scripts/generate_synthetic_ecgkit_v2.py

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

Uses ecg-image-kit to generate proper PTB-XL-style images with:
- Correct 4-column layout (3 short + 1 full lead II)
- Lead labels (I, II, III, aVR, aVL, aVF, V1-V6)
- Red grid background
- Scale markers

IMPORTANT: Requires the patched ecg-image-kit (write_wfdb_file commented out)
IMPORTANT: Requires ecgkit_venv environment

Usage:
    source ~/ecgkit_venv/bin/activate
    python generate_synthetic_ecgkit_v2.py --source georgia --workers 8
    python generate_synthetic_ecgkit_v2.py --source all --workers 8
"""

import os
import sys
import argparse
import subprocess
import shutil
from pathlib import Path
from concurrent.futures import ProcessPoolExecutor, as_completed
from tqdm import tqdm
import numpy as np
import cv2

# Add ecg-image-kit to path
ECG_IMAGE_KIT = Path('/home/azureuser/ecg-digitization/ecg-image-kit/codes/ecg-image-generator')

# =============================================================================
# Configuration
# =============================================================================
DATA_ROOT = Path('/data/ecg-digitization')

DATA_SOURCES = {
    'chapman': DATA_ROOT / 'chapman',
    'georgia': DATA_ROOT / 'georgia',
    'cpsc': DATA_ROOT / 'cpsc',
    'ningbo': DATA_ROOT / 'ningbo',
}

# Output dimensions to match PTB-XL synthetic
RAW_WIDTH = 2200
RAW_HEIGHT = 1700

# GT parameters
T0 = 235
T1 = 4161
OUTPUT_GT_WIDTH = T1 - T0  # 3926

# Target dimensions for Stage 1
TARGET_WIDTH = 4352
TARGET_HEIGHT = 1696

# Baseline y-positions (in 1696-height image)
ZERO_MV = np.array([703.5, 987.5, 1271.5, 1531.5])
MV_TO_PIXEL = 78.5  # pixels per mV


# =============================================================================
# Record Discovery
# =============================================================================
def find_records(source, data_dir):
    """Find all ECG record files for a given source."""
    records = []
    data_dir = Path(data_dir)
    
    if source == 'chapman':
        for hea in data_dir.rglob('JS*.hea'):
            if hea.with_suffix('.mat').exists():
                records.append(str(hea))
    elif source == 'georgia':
        for hea in data_dir.rglob('E*.hea'):
            if hea.with_suffix('.mat').exists():
                records.append(str(hea))
    elif source == 'cpsc':
        for hea in data_dir.rglob('A*.hea'):
            if hea.with_suffix('.mat').exists():
                records.append(str(hea))
    elif source == 'ningbo':
        for hea in data_dir.rglob('*.hea'):
            if hea.with_suffix('.mat').exists():
                records.append(str(hea))
    
    return sorted(records)


# =============================================================================
# GT Extraction
# =============================================================================
def extract_gt_from_wfdb(hea_file, output_width=3926):
    """Extract ground truth signal from WFDB record."""
    try:
        import wfdb
        from scipy import signal as scipy_signal
        
        record_path = hea_file.replace('.hea', '')
        record = wfdb.rdrecord(record_path)
        
        signals = record.p_signal  # [samples, leads]
        sig_names = [name.upper() for name in 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):
            # Normalize names
            if name == 'AVR':
                name = 'aVR'
            elif name == 'AVL':
                name = 'aVL'
            elif name == 'AVF':
                name = 'aVF'
            lead_signals[name] = signals[:, i]
        
        # 4-row output layout (matching competition format)
        # Row 0: I (col0), aVR (col1), V1 (col2), V4 (col3)
        # Row 1: II (col0), aVL (col1), V2 (col2), V5 (col3)
        # Row 2: III (col0), aVF (col1), V3 (col2), V6 (col3)
        # Row 3: Full lead II (cols 0-3)
        
        row_leads = [
            ['I', 'aVR', 'V1', 'V4'],
            ['II', 'aVL', 'V2', 'V5'],
            ['III', 'aVF', 'V3', 'V6'],
        ]
        
        segment_len = output_width // 3  # ~1308 pixels per 2.5s segment
        gt_data = np.zeros((output_width, 4), dtype=np.float32)
        
        for row_idx in range(4):
            if row_idx < 3:
                # Standard 3-column layout
                leads_for_row = row_leads[0][row_idx], row_leads[1][row_idx], row_leads[2][row_idx]
            else:
                # Row 3 is full lead II (but we use the same 3-segment structure)
                leads_for_row = ['II', 'II', 'II']
            
            for col_idx in range(3):
                lead_name = leads_for_row[col_idx] if row_idx < 3 else 'II'
                
                # Get signal (handle various name formats)
                signal = None
                for variant in [lead_name, lead_name.upper(), lead_name.lower()]:
                    if variant in lead_signals:
                        signal = lead_signals[variant]
                        break
                
                if signal is None:
                    # Use baseline for missing leads
                    col_start = col_idx * segment_len
                    col_end = (col_idx + 1) * segment_len if col_idx < 2 else output_width
                    gt_data[col_start:col_end, row_idx] = ZERO_MV[row_idx]
                    continue
                
                # 2.5s at 500Hz = 1250 samples per segment
                samples_per_segment = 1250
                
                if row_idx < 3:
                    # Standard leads: segment based on column
                    start = col_idx * samples_per_segment
                else:
                    # Full lead II: use different time segments
                    start = col_idx * samples_per_segment * 3 // 3  # Adjust for full 10s
                
                end = min(start + samples_per_segment, len(signal))
                segment = signal[start:end]
                
                if len(segment) == 0:
                    col_start = col_idx * segment_len
                    col_end = (col_idx + 1) * segment_len if col_idx < 2 else output_width
                    gt_data[col_start:col_end, row_idx] = ZERO_MV[row_idx]
                    continue
                
                # GT pixel positions
                col_start = col_idx * segment_len
                col_end = (col_idx + 1) * segment_len if col_idx < 2 else output_width
                num_gt_pixels = col_end - col_start
                
                # Resample signal to GT resolution
                x_old = np.linspace(0, 1, len(segment))
                x_new = np.linspace(0, 1, num_gt_pixels)
                signal_resampled = np.interp(x_new, x_old, segment)
                
                # Convert to pixel coordinates
                y_pixels = ZERO_MV[row_idx] - signal_resampled * MV_TO_PIXEL
                y_pixels = np.clip(y_pixels, 0, TARGET_HEIGHT - 1)
                gt_data[col_start:col_end, row_idx] = y_pixels
        
        return gt_data
        
    except Exception as e:
        print(f"GT extraction error: {e}")
        return None


# =============================================================================
# Image Generation
# =============================================================================
def generate_single_image(args):
    """Generate a single ECG image using ecg-image-kit."""
    idx, hea_file, source, output_dir, temp_base = args
    
    sample_id = f'syn_{source}_{idx:08d}'
    style = '0001'
    
    raw_dir = output_dir / 'raw'
    gt_dir = output_dir / 'gt'
    
    out_img = raw_dir / f'{sample_id}-{style}.png'
    out_gt = gt_dir / f'{sample_id}-{style}.csv'
    
    # Skip if both exist
    if out_img.exists() and out_gt.exists():
        return True, idx, "exists"
    
    temp_dir = temp_base / f'temp_{idx}'
    dat_file = hea_file.replace('.hea', '.mat')
    
    try:
        temp_dir.mkdir(parents=True, exist_ok=True)
        
        # Generate image with ecg-image-kit
        cmd = [
            'python', 'gen_ecg_image_from_data.py',
            '-i', dat_file,
            '-hea', hea_file,
            '-o', str(temp_dir),
            '-st', '0',
            '-r', '200',  # DPI for 2200x1700
            '--num_columns', '4',
            '--full_mode', 'II',
            '--standard_grid_color', '5',  # Red grid
        ]
        
        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 save image
        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"
        
        # Resize to exact target dimensions if needed
        if img.shape[:2] != (RAW_HEIGHT, RAW_WIDTH):
            img = cv2.resize(img, (RAW_WIDTH, RAW_HEIGHT), interpolation=cv2.INTER_LINEAR)
        
        # Save image
        raw_dir.mkdir(parents=True, exist_ok=True)
        cv2.imwrite(str(out_img), img)
        
        # Extract GT
        gt_data = extract_gt_from_wfdb(hea_file, OUTPUT_GT_WIDTH)
        if gt_data is not None:
            gt_dir.mkdir(parents=True, exist_ok=True)
            np.savetxt(out_gt, gt_data, delimiter=',', fmt='%.2f')
        
        # 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, f"error: {str(e)[:80]}"


# =============================================================================
# Main
# =============================================================================
def main():
    parser = argparse.ArgumentParser(description='Generate 0001-style synthetic ECGs using ecg-image-kit')
    parser.add_argument('--source', type=str, required=True,
                       choices=['chapman', 'georgia', 'cpsc', 'ningbo', 'all'])
    parser.add_argument('--workers', type=int, default=8)
    parser.add_argument('--limit', type=int, default=None)
    parser.add_argument('--output_base', type=str, default='/data/ecg-digitization')
    args = parser.parse_args()
    
    output_base = Path(args.output_base)
    temp_base = Path('/tmp/ecg_gen')
    temp_base.mkdir(parents=True, exist_ok=True)
    
    if args.source == 'all':
        sources = ['chapman', 'georgia', 'cpsc', 'ningbo']
    else:
        sources = [args.source]
    
    for source in sources:
        print(f"\n{'='*60}")
        print(f"Processing: {source}")
        print(f"{'='*60}")
        
        data_dir = DATA_SOURCES.get(source)
        if data_dir is None or not data_dir.exists():
            print(f"Data directory not found: {data_dir}")
            continue
        
        output_dir = output_base / f'synthetic_0001_{source}'
        output_dir.mkdir(parents=True, exist_ok=True)
        
        print(f"Finding records...")
        records = find_records(source, data_dir)
        print(f"Found {len(records)} records")
        
        if len(records) == 0:
            continue
        
        if args.limit:
            records = records[:args.limit]
            print(f"Limited to {len(records)} records")
        
        # Prepare arguments
        work_args = [
            (i, rec, source, output_dir, temp_base)
            for i, rec in enumerate(records)
        ]
        
        # Check existing
        existing = sum(1 for args in work_args 
                      if (output_dir / 'raw' / f'syn_{source}_{args[0]:08d}-0001.png').exists())
        print(f"Already generated: {existing}")
        
        # Process in parallel
        success = existing
        failed = 0
        
        with ProcessPoolExecutor(max_workers=args.workers) as executor:
            futures = {executor.submit(generate_single_image, arg): arg for arg in work_args}
            
            with tqdm(total=len(records), initial=existing, desc=source) as pbar:
                for future in as_completed(futures):
                    ok, idx, msg = future.result()
                    if ok:
                        if msg != "exists":
                            success += 1
                            pbar.update(1)
                    else:
                        failed += 1
                        if failed <= 10:
                            tqdm.write(f"  Failed {idx}: {msg}")
                        pbar.update(1)
        
        print(f"\nCompleted: {success} success, {failed} failed")
    
    # Cleanup temp
    shutil.rmtree(temp_base, ignore_errors=True)
    print("\nDone!")


if __name__ == '__main__':
    main()