#!/usr/bin/env python3 """ Generate synthetic ECG images with GT and Stage 0/1 processing. This script: 1. Generates 0001-style images using ecg-image-kit 2. Extracts GT values (y-pixel coordinates) from WFDB records 3. Runs Stage 0/1 processing on generated images 4. Saves everything in PTB-XL-compatible format Output structure: synthetic_0001_{source}/raw/*.png - Raw generated images (2200x1700) synthetic_0001_{source}/gt/*.csv - Ground truth (3926 x 4) synthetic_0001_{source}_processed/raw/*.png - Stage 0/1 processed (2200x1700) synthetic_0001_{source}_processed/gt/*.csv - GT copied from raw Usage: python generate_and_process_all.py --source georgia --workers 48 python generate_and_process_all.py --source all --workers 80 """ import os import sys import argparse import subprocess import shutil from pathlib import Path from concurrent.futures import ProcessPoolExecutor, as_completed from multiprocessing import cpu_count import numpy as np import cv2 import traceback # Paths PROJECT_ROOT = Path('/home/azureuser/ecg-digitization') ECG_IMAGE_KIT = PROJECT_ROOT / 'ecg-image-kit' / 'codes' / 'ecg-image-generator' BASELINE_PATH = PROJECT_ROOT / 'data' / 'hengck23-submit-physionet' / 'hengck23-submit-physionet' VENV_ECGKIT = '/home/azureuser/ecgkit_venv' VENV_TORCH = '/home/azureuser/.venv' # Add baseline to path for Stage 0/1 sys.path.insert(0, str(BASELINE_PATH)) # Data sources DATA_ROOT = Path('/data/ecg-digitization') DATA_SOURCES = { 'georgia': DATA_ROOT / 'georgia', 'chapman': DATA_ROOT / 'chapman', 'cpsc': DATA_ROOT / 'cpsc', 'ningbo': DATA_ROOT / 'ningbo', } # Image dimensions RAW_WIDTH = 2200 RAW_HEIGHT = 1700 # GT parameters (match PTB-XL format) 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, for GT) ZERO_MV = np.array([703.5, 987.5, 1271.5, 1531.5]) MV_TO_PIXEL = 78.5 # pixels per mV def find_records(source, data_dir): """Find all ECG record .hea files for a given source.""" records = [] data_dir = Path(data_dir) for hea in data_dir.rglob('*.hea'): mat = hea.with_suffix('.mat') if mat.exists(): records.append(str(hea)) return sorted(records) def extract_gt_from_wfdb(hea_file): """Extract ground truth signal from WFDB record in PTB-XL format.""" 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) fs = 500 # Map lead names lead_signals = {} for i, name in enumerate(sig_names): if name == 'AVR': name = 'aVR' elif name == 'AVL': name = 'aVL' elif name == 'AVF': name = 'aVF' lead_signals[name] = signals[:, i] # Layout: 4 rows, each showing 3 segments of 2.5s # Row 0: I (col0), aVR (col1), V1 (col2) # Row 1: II (col0), aVL (col1), V2 (col2) # Row 2: III (col0), aVF (col1), V3 (col2) # Row 3: V4 (col0), V5 (col1), V6 (col2) - OR full lead II row_leads = [ ['I', 'aVR', 'V1'], ['II', 'aVL', 'V2'], ['III', 'aVF', 'V3'], ['V4', 'V5', 'V6'], # Actually this row is Lead II full in 0001 format ] segment_len = OUTPUT_GT_WIDTH // 3 # ~1308 pixels per segment gt_data = np.zeros((OUTPUT_GT_WIDTH, 4), dtype=np.float32) samples_per_segment = int(2.5 * fs) # 1250 samples at 500Hz for row_idx in range(4): for col_idx in range(3): if row_idx < 3: lead_name = row_leads[row_idx][col_idx] else: # Row 3 is full Lead II lead_name = 'II' # Get signal 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: col_start = col_idx * segment_len col_end = (col_idx + 1) * segment_len if col_idx < 2 else OUTPUT_GT_WIDTH gt_data[col_start:col_end, row_idx] = ZERO_MV[row_idx] continue # Time segment if row_idx < 3: start = col_idx * samples_per_segment else: # Full Lead II uses continuous segments start = col_idx * samples_per_segment end = min(start + samples_per_segment, len(signal)) segment = signal[start:end] if end > start else np.zeros(1) if len(segment) == 0: segment = np.zeros(samples_per_segment) # GT pixel positions col_start = col_idx * segment_len col_end = (col_idx + 1) * segment_len if col_idx < 2 else OUTPUT_GT_WIDTH num_gt_pixels = col_end - col_start # Resample signal to GT resolution if len(segment) > 1: 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) else: signal_resampled = np.zeros(num_gt_pixels) # Convert to pixel coordinates (baseline - signal * scale) 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: return None def generate_single_sample(args): """Generate a single synthetic sample with GT.""" idx, hea_file, source, output_dir, temp_base = args base_name = Path(hea_file).stem sample_id = f'syn_{source}_{base_name}' 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, sample_id, "exists" temp_dir = temp_base / f'temp_{idx}' mat_file = hea_file.replace('.hea', '.mat') try: temp_dir.mkdir(parents=True, exist_ok=True) # Set matplotlib config for this worker env = os.environ.copy() env['MPLBACKEND'] = 'Agg' env['MPLCONFIGDIR'] = str(temp_dir) # Generate image with ecg-image-kit cmd = [ f'{VENV_ECGKIT}/bin/python', 'gen_ecg_image_from_data.py', '-i', mat_file, '-hea', hea_file, '-o', str(temp_dir), '-st', '0', '-se', '0', '-r', '200', '--num_columns', '4', '--full_mode', 'II', '--standard_grid_color', '5', ] result = subprocess.run( cmd, cwd=str(ECG_IMAGE_KIT), capture_output=True, text=True, timeout=120, env=env ) if result.returncode != 0: shutil.rmtree(temp_dir, ignore_errors=True) return False, idx, sample_id, "ecg-kit error" # Find generated image gen_imgs = list(temp_dir.glob('*.png')) if not gen_imgs: shutil.rmtree(temp_dir, ignore_errors=True) return False, idx, sample_id, "no image" # Load image img = cv2.imread(str(gen_imgs[0])) if img is None: shutil.rmtree(temp_dir, ignore_errors=True) return False, idx, sample_id, "read error" # Resize to exact target if needed if img.shape[:2] != (RAW_HEIGHT, RAW_WIDTH): img = cv2.resize(img, (RAW_WIDTH, RAW_HEIGHT)) # Save raw image raw_dir.mkdir(parents=True, exist_ok=True) cv2.imwrite(str(out_img), img) # Extract and save GT gt_data = extract_gt_from_wfdb(hea_file) 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, sample_id, "success" except Exception as e: shutil.rmtree(temp_dir, ignore_errors=True) return False, idx, sample_id, str(e)[:50] # Stage 0/1 processing functions (loaded lazily) _stage0 = None _stage1 = None _device = None def load_stage_models(device='cuda:0'): """Load Stage 0/1 models.""" global _stage0, _stage1, _device if _stage0 is not None: return _stage0, _stage1 import torch from stage0_model import Net as Stage0Net from stage1_model import Net as Stage1Net _device = device # Load Stage 0 _stage0 = Stage0Net(pretrained=False).to(device) ckpt0 = torch.load(BASELINE_PATH / 'weight' / 'stage0-last.checkpoint.pth', map_location=device) _stage0.load_state_dict(ckpt0['state_dict'], strict=True) _stage0.eval() # Load Stage 1 _stage1 = Stage1Net(pretrained=False).to(device) ckpt1 = torch.load(BASELINE_PATH / 'weight' / 'stage1-last.checkpoint.pth', map_location=device) _stage1.load_state_dict(ckpt1['state_dict'], strict=True) _stage1.eval() return _stage0, _stage1 def process_single_stage01(args): """Process a single image through Stage 0/1.""" sample_id, raw_path, processed_dir, gt_path, device = args out_img = processed_dir / 'raw' / f'{sample_id}.png' out_gt = processed_dir / 'gt' / f'{sample_id}.csv' # Skip if exists if out_img.exists() and out_gt.exists(): return True, sample_id, "exists" try: import torch from stage0_common import image_to_batch from stage0_common import output_to_predict as stage0_output from stage0_common import normalise_by_homography from stage1_common import output_to_predict as stage1_output from stage1_common import rectify_image # Load models (cached) stage0, stage1 = load_stage_models(device) # Load image image = cv2.imread(str(raw_path)) if image is None: return False, sample_id, "read error" image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # Stage 0: orientation + keypoint detection batch = image_to_batch(image) batch['image'] = batch['image'].to(device) with torch.no_grad(): with torch.cuda.amp.autocast(): output0 = stage0(batch) # output_to_predict returns (rotated, keypoint) rotated, keypoint = stage0_output(image, batch, output0) # Normalize by homography normalised, keypoint_with_match, homo = normalise_by_homography(rotated, keypoint) # Stage 1: grid detection + rectification batch1 = image_to_batch(normalised) batch1['image'] = batch1['image'].to(device) with torch.no_grad(): with torch.cuda.amp.autocast(): output1 = stage1(batch1) # output_to_predict returns (gridpoint_xy, more) gridpoint_xy, more = stage1_output(normalised, batch1, output1) # Rectify image using grid rectified = rectify_image(normalised, gridpoint_xy) # Output is 2200x1700 processed_image = rectified # Resize to raw dimensions if needed if processed_image.shape[:2] != (RAW_HEIGHT, RAW_WIDTH): processed_image = cv2.resize(processed_image, (RAW_WIDTH, RAW_HEIGHT)) # Save (processed_dir / 'raw').mkdir(parents=True, exist_ok=True) cv2.imwrite(str(out_img), cv2.cvtColor(processed_image, cv2.COLOR_RGB2BGR)) # Copy GT if gt_path and gt_path.exists(): (processed_dir / 'gt').mkdir(parents=True, exist_ok=True) shutil.copy(gt_path, out_gt) return True, sample_id, "success" except Exception as e: # Fallback: copy original try: image = cv2.imread(str(raw_path)) if image is not None: (processed_dir / 'raw').mkdir(parents=True, exist_ok=True) cv2.imwrite(str(out_img), image) if gt_path and gt_path.exists(): (processed_dir / 'gt').mkdir(parents=True, exist_ok=True) shutil.copy(gt_path, out_gt) return True, sample_id, "error-fallback" except: pass return False, sample_id, str(e)[:50] def process_source(source, data_dir, output_base, workers, device): """Process a single data source: generate + Stage 0/1.""" from tqdm import tqdm output_dir = output_base / f'synthetic_0001_{source}' processed_dir = output_base / f'synthetic_0001_{source}_processed' temp_base = Path(f'/tmp/ecg_gen_{source}') temp_base.mkdir(parents=True, exist_ok=True) print(f"\n{'='*60}") print(f"Processing: {source}") print(f"{'='*60}") # Find records print("Finding records...") records = find_records(source, data_dir) print(f"Found {len(records)} records") if len(records) == 0: return # ======================================== # Phase 1: Generate raw images + GT # ======================================== print(f"\n--- Phase 1: Generate raw images + GT ---") work_args = [(i, rec, source, output_dir, temp_base) for i, rec in enumerate(records)] existing = sum(1 for a in work_args if (output_dir / 'raw' / f'syn_{source}_{Path(a[1]).stem}-0001.png').exists()) print(f"Already generated: {existing}/{len(records)}") if existing < len(records): success = existing failed = 0 with ProcessPoolExecutor(max_workers=workers) as executor: futures = {executor.submit(generate_single_sample, arg): arg for arg in work_args} with tqdm(total=len(records), initial=existing, desc=f"Gen {source}") as pbar: for future in as_completed(futures): ok, idx, sample_id, msg = future.result() if ok: if msg != "exists": success += 1 else: failed += 1 pbar.update(1) print(f"Generation: {success} success, {failed} failed") # ======================================== # Phase 2: Stage 0/1 processing (GPU) # ======================================== print(f"\n--- Phase 2: Stage 0/1 processing ---") # Find all generated raw images raw_dir = output_dir / 'raw' gt_dir = output_dir / 'gt' if not raw_dir.exists(): print("No raw images to process") return raw_images = sorted(raw_dir.glob('*.png')) print(f"Found {len(raw_images)} raw images to process") # Check existing processed existing_processed = sum(1 for img in raw_images if (processed_dir / 'raw' / img.name).exists()) print(f"Already processed: {existing_processed}/{len(raw_images)}") if existing_processed < len(raw_images): # Stage 0/1 is GPU-bound, process sequentially with batch loading import torch # Pre-load models print(f"Loading Stage 0/1 models on {device}...") load_stage_models(device) success = existing_processed failed = 0 with tqdm(total=len(raw_images), initial=existing_processed, desc=f"S01 {source}") as pbar: for img_path in raw_images: sample_id = img_path.stem gt_path = gt_dir / f'{sample_id}.csv' out_img = processed_dir / 'raw' / f'{sample_id}.png' if out_img.exists(): continue ok, _, msg = process_single_stage01((sample_id, img_path, processed_dir, gt_path, device)) if ok: success += 1 else: failed += 1 pbar.update(1) print(f"Stage 0/1: {success} success, {failed} failed") # Cleanup temp shutil.rmtree(temp_base, ignore_errors=True) def main(): parser = argparse.ArgumentParser(description='Generate and process synthetic ECG images') parser.add_argument('--source', type=str, required=True, choices=['georgia', 'chapman', 'cpsc', 'ningbo', 'all']) parser.add_argument('--workers', type=int, default=None, help='Number of workers (default: all CPUs)') parser.add_argument('--device', type=str, default='cuda:0') args = parser.parse_args() workers = args.workers or cpu_count() print(f"Using {workers} workers for generation") print(f"Using {args.device} for Stage 0/1") if args.source == 'all': sources = ['georgia', 'cpsc', 'chapman', 'ningbo'] else: sources = [args.source] for source in sources: data_dir = DATA_SOURCES.get(source) if data_dir is None or not data_dir.exists(): print(f"Data not found: {data_dir}") continue process_source(source, data_dir, DATA_ROOT, workers, args.device) print("\n=== ALL COMPLETE ===") if __name__ == '__main__': main()