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