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