code/scripts/gen_0001_parallel.py

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

Spawns multiple subprocess workers in parallel to maximize CPU utilization.
Each worker calls ecg-image-kit directly - no imports needed in this script.

Usage:
    python gen_0001_parallel.py --n_samples 21799 --n_workers 32
"""

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

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

# Output specs
RAW_HEIGHT = 1700
RAW_WIDTH = 2200
T0 = 235
T1 = 4161
OUTPUT_GT_WIDTH = T1 - T0


def find_ptbxl_records(ptbxl_dir):
    """Find all PTB-XL 500Hz records."""
    records = []
    ptbxl_dir = Path(ptbxl_dir)
    records500 = ptbxl_dir / 'records500'
    if records500.exists():
        for subdir in sorted(records500.iterdir()):
            if subdir.is_dir():
                for hea in subdir.glob('*_hr.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 GT signal from record."""
    try:
        record_path = hea_file.replace('.hea', '')
        record = wfdb.rdrecord(record_path)
        signals = record.p_signal
        sig_names = record.sig_name
        fs = record.fs
        
        if fs != 500:
            num_samples = int(signals.shape[0] * 500 / fs)
            signals = scipy_signal.resample(signals, num_samples, axis=0)
        
        lead_signals = {n.upper(): signals[:, i] for i, n in enumerate(sig_names)}
        
        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(3):
            for col_idx in range(3):
                lead_name = row_leads[row_idx][col_idx]
                if lead_name in lead_signals:
                    lead_data = lead_signals[lead_name]
                    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
                    if len(segment) > 0:
                        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)
        
        if 'II' in lead_signals:
            lead_ii = lead_signals['II']
            x_old = np.linspace(0, 1, len(lead_ii))
            x_new = np.linspace(0, 1, output_width)
            gt_signal[3, :] = np.interp(x_new, x_old, lead_ii)
        
        return gt_signal
    except:
        return None


def generate_worker(args):
    """Worker function - spawns ecg-image-kit subprocess."""
    idx, hea_file, dat_file, output_dir, ecg_kit_path = args
    
    sample_id = f'syn_{idx:08d}'
    out_img = output_dir / 'stage1' / f'{sample_id}-0001.png'
    out_gt = output_dir / 'gt' / f'{sample_id}.npy'
    
    if out_img.exists() and out_gt.exists():
        return 'skip', idx
    
    temp_dir = output_dir / 'temp' / f't_{idx}'
    
    try:
        temp_dir.mkdir(parents=True, exist_ok=True)
        
        # Call ecg-image-kit as subprocess
        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',
            '--num_columns', '4',
            '--full_mode', 'II',
            '--standard_grid_color', '5',  # Red grid = 0001 style
        ]
        
        # Set environment to avoid matplotlib lock issues
        env = os.environ.copy()
        env['MPLBACKEND'] = 'Agg'
        env['MPLCONFIGDIR'] = str(temp_dir)  # Each worker uses own matplotlib config
        
        result = subprocess.run(
            cmd,
            cwd=str(ecg_kit_path),
            capture_output=True,
            text=True,
            timeout=180,
            env=env
        )
        
        if result.returncode != 0:
            shutil.rmtree(temp_dir, ignore_errors=True)
            return 'fail', idx
        
        # Find and process generated image
        gen_imgs = list(temp_dir.glob('*.png'))
        if not gen_imgs:
            shutil.rmtree(temp_dir, ignore_errors=True)
            return 'fail', idx
        
        img = cv2.imread(str(gen_imgs[0]))
        if img is None:
            shutil.rmtree(temp_dir, ignore_errors=True)
            return 'fail', idx
        
        # Resize to Kaggle format
        img_resized = cv2.resize(img, (RAW_WIDTH, RAW_HEIGHT), interpolation=cv2.INTER_LINEAR)
        cv2.imwrite(str(out_img), img_resized)
        
        # Extract GT
        gt_signal = extract_gt_signal(hea_file, OUTPUT_GT_WIDTH)
        if gt_signal is not None:
            np.save(out_gt, gt_signal)
        else:
            shutil.rmtree(temp_dir, ignore_errors=True)
            return 'fail', idx
        
        shutil.rmtree(temp_dir, ignore_errors=True)
        return 'ok', idx
        
    except subprocess.TimeoutExpired:
        shutil.rmtree(temp_dir, ignore_errors=True)
        return 'timeout', idx
    except Exception as e:
        shutil.rmtree(temp_dir, ignore_errors=True)
        return 'error', idx


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--ptbxl_dir', type=str, default=str(PTBXL_PATH))
    parser.add_argument('--output_dir', type=str, default='/data/ecg-digitization/synthetic_0001_parallel')
    parser.add_argument('--n_samples', type=int, default=None)
    parser.add_argument('--n_workers', type=int, default=32)
    parser.add_argument('--start_idx', type=int, default=0)
    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)
    (output_dir / 'temp').mkdir(exist_ok=True)
    
    print(f"Finding PTB-XL records...")
    records = find_ptbxl_records(Path(args.ptbxl_dir))
    print(f"Found {len(records)} records")
    
    n_samples = args.n_samples or len(records)
    n_samples = min(n_samples, len(records))
    
    print(f"\nGenerating {n_samples} 0001-type images with {args.n_workers} parallel workers")
    print(f"Output: {output_dir}")
    print()
    
    # Prepare tasks
    tasks = []
    for i in range(args.start_idx, args.start_idx + n_samples):
        record_idx = i % len(records)
        hea, dat = records[record_idx]
        tasks.append((i, hea, dat, output_dir, ECG_IMAGE_KIT))
    
    ok_count = 0
    skip_count = 0
    fail_count = 0
    
    # Run in parallel
    with ProcessPoolExecutor(max_workers=args.n_workers) as executor:
        futures = {executor.submit(generate_worker, task): task[0] for task in tasks}
        
        pbar = tqdm(as_completed(futures), total=len(futures), desc="Generating")
        for future in pbar:
            try:
                status, idx = future.result(timeout=300)
                if status == 'ok':
                    ok_count += 1
                elif status == 'skip':
                    skip_count += 1
                else:
                    fail_count += 1
            except Exception as e:
                fail_count += 1
            
            pbar.set_postfix({'ok': ok_count, 'skip': skip_count, 'fail': fail_count})
    
    print(f"\n\nDone! ok={ok_count}, skip={skip_count}, fail={fail_count}")
    
    # Create manifest
    manifest = []
    for img in sorted((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, 'image': 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")


if __name__ == '__main__':
    main()