#!/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()