#!/usr/bin/env python3 """ Synthetic ECG Image Generator V2 - With Stage 0/1 Processing Based on expert recommendations: - Generate raw images using ecg-image-kit with various styles (similar to 0001-0012 versions) - Process through Stage 0 (orientation) and Stage 1 (rectification) for consistency - Keep high-resolution ground truth from ecg-image-kit - Heavy augmentation based on image type Usage: python generate_synthetic_v2.py --ptbxl_dir ../data/ptbxl --output_dir ../data/synthetic --n_samples 200000 --n_workers 16 """ import os import sys import argparse import numpy as np import pandas as pd import cv2 import wfdb from pathlib import Path from multiprocessing import Pool, cpu_count from tqdm import tqdm import json import random import traceback import tempfile import shutil from PIL import Image import torch import torch.nn.functional as F # Add paths PROJECT_ROOT = os.path.dirname(os.path.dirname(__file__)) ECG_IMAGE_KIT_PATH = os.path.join(PROJECT_ROOT, 'ecg-image-kit') BASELINE_PATH = os.path.join(PROJECT_ROOT, 'data', 'hengck23-submit-physionet', 'hengck23-submit-physionet') sys.path.insert(0, ECG_IMAGE_KIT_PATH) sys.path.insert(0, BASELINE_PATH) # Target specifications (matching competition Stage 1 output) TARGET_HEIGHT = 1696 TARGET_WIDTH = 4352 LEAD_NAMES = ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6'] # Calibration constants (from competition baseline) ZERO_MV = [703.5, 987.5, 1271.5, 1531.5] # Y positions for 4 rows MV_TO_PIXEL = 78.5 # mV to pixel conversion T0, T1 = 235, 4161 # Time crop boundaries # Standard ECG row layout (4 leads per row, 2.5s each) ROW_LAYOUT = [ ['I', 'aVR', 'V1', 'V4'], ['II', 'aVL', 'V2', 'V5'], ['III', 'aVF', 'V3', 'V6'], ] # Image style presets (similar to competition 0001-0012 versions) IMAGE_STYLES = { '0001': { 'grid_color': 'red', 'background': 'white', 'line_color': 'black', 'grid_major_color': (255, 200, 200), 'grid_minor_color': (255, 230, 230), 'scan_effect': False, 'paper_texture': False, }, '0002': { 'grid_color': 'red', 'background': 'cream', 'line_color': 'black', 'grid_major_color': (255, 180, 180), 'grid_minor_color': (255, 210, 210), 'scan_effect': True, 'paper_texture': True, }, '0003': { 'grid_color': 'green', 'background': 'white', 'line_color': 'black', 'grid_major_color': (200, 255, 200), 'grid_minor_color': (230, 255, 230), 'scan_effect': False, 'paper_texture': False, }, '0004': { 'grid_color': 'green', 'background': 'cream', 'line_color': 'black', 'grid_major_color': (180, 255, 180), 'grid_minor_color': (210, 255, 210), 'scan_effect': True, 'paper_texture': True, }, '0005': { 'grid_color': 'red', 'background': 'white', 'line_color': 'blue', 'grid_major_color': (255, 200, 200), 'grid_minor_color': (255, 230, 230), 'scan_effect': False, 'paper_texture': False, }, '0006': { 'grid_color': 'red', 'background': 'grey', 'line_color': 'black', 'grid_major_color': (200, 150, 150), 'grid_minor_color': (200, 180, 180), 'scan_effect': True, 'paper_texture': True, }, '0007': { 'grid_color': 'faded_red', 'background': 'yellow', 'line_color': 'black', 'grid_major_color': (255, 220, 200), 'grid_minor_color': (255, 240, 220), 'scan_effect': True, 'paper_texture': True, }, '0008': { 'grid_color': 'faded_green', 'background': 'white', 'line_color': 'black', 'grid_major_color': (220, 255, 220), 'grid_minor_color': (240, 255, 240), 'scan_effect': False, 'paper_texture': True, }, '0009': { 'grid_color': 'orange', 'background': 'white', 'line_color': 'black', 'grid_major_color': (255, 220, 180), 'grid_minor_color': (255, 240, 210), 'scan_effect': False, 'paper_texture': False, }, '0010': { 'grid_color': 'red', 'background': 'white', 'line_color': 'black', 'grid_major_color': (255, 190, 190), 'grid_minor_color': (255, 220, 220), 'scan_effect': True, 'paper_texture': True, 'contrast': 'high', }, '0011': { 'grid_color': 'red', 'background': 'white', 'line_color': 'black', 'grid_major_color': (255, 210, 210), 'grid_minor_color': (255, 235, 235), 'scan_effect': True, 'paper_texture': True, 'contrast': 'low', }, '0012': { 'grid_color': 'none', 'background': 'white', 'line_color': 'black', 'grid_major_color': (240, 240, 240), 'grid_minor_color': (250, 250, 250), 'scan_effect': False, 'paper_texture': False, }, } def load_ptbxl_record(record_path): """Load a PTB-XL record and return 12-lead signals""" try: record = wfdb.rdrecord(str(record_path)) signals = record.p_signal fs = record.fs signal_dict = {} for i, lead in enumerate(record.sig_name): lead_normalized = lead.upper().replace(' ', '') if lead_normalized in LEAD_NAMES: signal_dict[lead_normalized] = signals[:, i] return signal_dict, int(fs) except Exception as e: return None, None def create_ecg_image_matplotlib(signal_dict, fs, style_config, random_state=None): """ Create a synthetic ECG image using matplotlib with various styles Returns: np.ndarray: RGB image array np.ndarray: Ground truth signal array (4, width) """ import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt from matplotlib.patches import Rectangle if random_state is not None: np.random.seed(random_state) random.seed(random_state) # Figure setup fig_width = 43.52 fig_height = 16.96 dpi = 100 fig, ax = plt.subplots(figsize=(fig_width, fig_height), dpi=dpi) # Background color bg_color = style_config.get('background', 'white') if bg_color == 'cream': bg_rgb = '#FFF8E7' elif bg_color == 'grey': bg_rgb = '#E8E8E8' elif bg_color == 'yellow': bg_rgb = '#FFFDE7' else: bg_rgb = 'white' ax.set_facecolor(bg_rgb) fig.patch.set_facecolor(bg_rgb) # Draw grid grid_color = style_config.get('grid_color', 'red') if grid_color != 'none': major_color = tuple(c/255 for c in style_config.get('grid_major_color', (255, 200, 200))) minor_color = tuple(c/255 for c in style_config.get('grid_minor_color', (255, 230, 230))) # Draw grid lines ax.grid(True, which='major', color=major_color, linewidth=0.5, alpha=0.8) ax.grid(True, which='minor', color=minor_color, linewidth=0.2, alpha=0.5) ax.minorticks_on() # ECG line color line_color = style_config.get('line_color', 'black') if line_color == 'blue': line_rgb = '#0000AA' else: line_rgb = 'black' # Time parameters total_time = 10.0 # 10 seconds total time_per_lead = 2.5 # 2.5 seconds per lead column # Prepare ground truth output_width = T1 - T0 # 3926 pixels gt_signals = np.zeros((4, output_width), dtype=np.float32) # Y positions for 4 rows (normalized to figure coordinates) row_heights = [0.75, 0.50, 0.25, 0.05] # Normalized positions row_y_centers = [3.5, 2.5, 1.5, 0.5] # For plotting # Plot each row for row_idx in range(3): leads_in_row = ROW_LAYOUT[row_idx] row_signals = [] for lead_idx, lead in enumerate(leads_in_row): if lead in signal_dict: # Get signal segment (2.5 seconds) samples_needed = int(time_per_lead * fs) sig = signal_dict[lead] # Random offset within the signal max_offset = max(0, len(sig) - samples_needed) offset = random.randint(0, max_offset) if max_offset > 0 else 0 segment = sig[offset:offset + samples_needed] if len(segment) < samples_needed: segment = np.pad(segment, (0, samples_needed - len(segment))) row_signals.append(segment) else: row_signals.append(np.zeros(int(time_per_lead * fs))) # Concatenate and resample for ground truth full_row = np.concatenate(row_signals) # Resample to output width gt_row = np.interp( np.linspace(0, len(full_row)-1, output_width), np.arange(len(full_row)), full_row ) gt_signals[row_idx] = gt_row # Plot t = np.linspace(0, total_time, len(full_row)) y = full_row * 0.5 + row_y_centers[row_idx] # Scale and offset ax.plot(t, y, color=line_rgb, linewidth=0.8) # Row 4: Full II rhythm strip (10 seconds) if 'II' in signal_dict: samples_needed = int(total_time * fs) sig = signal_dict['II'] max_offset = max(0, len(sig) - samples_needed) offset = random.randint(0, max_offset) if max_offset > 0 else 0 rhythm_signal = sig[offset:offset + samples_needed] if len(rhythm_signal) < samples_needed: rhythm_signal = np.pad(rhythm_signal, (0, samples_needed - len(rhythm_signal))) # Resample for ground truth gt_signals[3] = np.interp( np.linspace(0, len(rhythm_signal)-1, output_width), np.arange(len(rhythm_signal)), rhythm_signal ) # Plot t = np.linspace(0, total_time, len(rhythm_signal)) y = rhythm_signal * 0.5 + 0.5 ax.plot(t, y, color=line_rgb, linewidth=0.8) # Set axis limits ax.set_xlim(0, total_time) ax.set_ylim(-0.5, 4.5) ax.axis('off') # Convert to image fig.canvas.draw() # Use buffer_rgba() for newer matplotlib versions buf = fig.canvas.buffer_rgba() img = np.asarray(buf)[:, :, :3] # RGBA -> RGB plt.close(fig) # Resize to target size img = cv2.resize(img, (TARGET_WIDTH, TARGET_HEIGHT), interpolation=cv2.INTER_LINEAR) return img, gt_signals def apply_scan_effects(img, style_config, random_state=None): """Apply realistic scan/photo effects to image""" if random_state is not None: np.random.seed(random_state) random.seed(random_state) img = img.astype(np.float32) # Add Gaussian noise if style_config.get('scan_effect', False): noise_level = random.uniform(0.5, 3.0) noise = np.random.normal(0, noise_level, img.shape) img = img + noise # Add paper texture if style_config.get('paper_texture', False): texture = np.random.normal(1.0, 0.01, img.shape[:2]) texture = cv2.GaussianBlur(texture, (5, 5), 0) img = img * texture[:, :, np.newaxis] # Contrast adjustment contrast = style_config.get('contrast', 'normal') if contrast == 'high': img = np.clip((img - 128) * 1.2 + 128, 0, 255) elif contrast == 'low': img = np.clip((img - 128) * 0.8 + 128, 0, 255) # Random slight rotation (to be corrected by Stage 0) if random.random() > 0.5: angle = random.uniform(-3, 3) M = cv2.getRotationMatrix2D((img.shape[1]//2, img.shape[0]//2), angle, 1.0) img = cv2.warpAffine(img, M, (img.shape[1], img.shape[0]), borderMode=cv2.BORDER_REPLICATE) # Random perspective transform (to be corrected by Stage 0/1) if random.random() > 0.7: h, w = img.shape[:2] # Small perspective distortion d = random.randint(10, 30) src = np.float32([[0, 0], [w, 0], [w, h], [0, h]]) dst = np.float32([ [random.randint(0, d), random.randint(0, d)], [w - random.randint(0, d), random.randint(0, d)], [w - random.randint(0, d), h - random.randint(0, d)], [random.randint(0, d), h - random.randint(0, d)] ]) M = cv2.getPerspectiveTransform(src, dst) img = cv2.warpPerspective(img, M, (w, h), borderMode=cv2.BORDER_REPLICATE) return np.clip(img, 0, 255).astype(np.uint8) def process_single_sample(args): """Process a single sample - generate image and process through Stage 0/1""" record_path, output_dir, sample_idx, style_key, use_stage01, stage0_net, stage1_net = args try: # Load PTB-XL record signal_dict, fs = load_ptbxl_record(record_path) if signal_dict is None or fs is None: return None # Get style configuration style_config = IMAGE_STYLES.get(style_key, IMAGE_STYLES['0001']) # Generate raw synthetic image raw_img, gt_signals = create_ecg_image_matplotlib( signal_dict, fs, style_config, random_state=sample_idx ) # Apply scan effects raw_img = apply_scan_effects(raw_img, style_config, random_state=sample_idx) # Output paths sample_id = f"syn_{sample_idx:08d}" output_dir = Path(output_dir) # Save raw image raw_path = output_dir / 'raw' / f'{sample_id}-{style_key}.png' raw_path.parent.mkdir(parents=True, exist_ok=True) cv2.imwrite(str(raw_path), cv2.cvtColor(raw_img, cv2.COLOR_RGB2BGR)) # Save ground truth signals (high-res from ecg-image-kit generation) gt_path = output_dir / 'gt' / f'{sample_id}.npy' gt_path.parent.mkdir(parents=True, exist_ok=True) np.save(str(gt_path), gt_signals) final_img = raw_img # Process through Stage 0 and Stage 1 if requested if use_stage01 and stage0_net is not None and stage1_net is not None: try: # Stage 0: Orientation correction from stage0_common import image_to_batch, output_to_predict, normalise_by_homography batch = image_to_batch(raw_img) with torch.no_grad(): output = stage0_net(batch) rotated, keypoint = output_to_predict(raw_img, batch, output) normalized, _, _ = normalise_by_homography(rotated, keypoint) # Save Stage 0 output stage0_path = output_dir / 'stage0' / f'{sample_id}-{style_key}.png' stage0_path.parent.mkdir(parents=True, exist_ok=True) cv2.imwrite(str(stage0_path), cv2.cvtColor(normalized, cv2.COLOR_RGB2BGR)) # Stage 1: Rectification from stage1_common import output_to_predict as stage1_output_to_predict, rectify_image batch = {'image': torch.from_numpy( np.ascontiguousarray(normalized.transpose(2, 0, 1)) ).unsqueeze(0)} with torch.no_grad(): output = stage1_net(batch) gridpoint_xy, _ = stage1_output_to_predict(normalized, batch, output) rectified = rectify_image(normalized, gridpoint_xy) # Save Stage 1 output (this is the main training image) stage1_path = output_dir / 'stage1' / f'{sample_id}-{style_key}.png' stage1_path.parent.mkdir(parents=True, exist_ok=True) cv2.imwrite(str(stage1_path), cv2.cvtColor(rectified, cv2.COLOR_RGB2BGR)) final_img = rectified except Exception as e: print(f"Stage 0/1 processing failed for {sample_id}: {e}") # Fall back to raw image stage1_path = output_dir / 'stage1' / f'{sample_id}-{style_key}.png' stage1_path.parent.mkdir(parents=True, exist_ok=True) cv2.imwrite(str(stage1_path), cv2.cvtColor(raw_img, cv2.COLOR_RGB2BGR)) else: # Just copy raw to stage1 directory stage1_path = output_dir / 'stage1' / f'{sample_id}-{style_key}.png' stage1_path.parent.mkdir(parents=True, exist_ok=True) cv2.imwrite(str(stage1_path), cv2.cvtColor(raw_img, cv2.COLOR_RGB2BGR)) return { 'sample_id': sample_id, 'style': style_key, 'record': str(record_path), 'gt_path': str(gt_path), 'stage1_path': str(stage1_path), } except Exception as e: print(f"Error processing sample {sample_idx}: {e}") traceback.print_exc() return None class SyntheticGenerator: """Main class to generate synthetic ECG training data""" def __init__(self, ptbxl_dir, output_dir, device='cuda:0', use_stage01=True): self.ptbxl_dir = Path(ptbxl_dir) self.output_dir = Path(output_dir) self.device = torch.device(device) self.use_stage01 = use_stage01 # Find PTB-XL records self.records = self._find_records() print(f"Found {len(self.records)} PTB-XL records") # Load Stage 0/1 models if needed self.stage0_net = None self.stage1_net = None if use_stage01: self._load_models() def _find_records(self): """Find all PTB-XL record paths""" records = [] # Look for records500 (high quality 500Hz) records500_dir = self.ptbxl_dir / 'physionet.org' / 'files' / 'ptb-xl' / '1.0.3' / 'records500' if records500_dir.exists(): for subdir in records500_dir.iterdir(): if subdir.is_dir(): for hea_file in subdir.glob('*.hea'): record_path = str(hea_file)[:-4] # Remove .hea extension records.append(record_path) # Fallback to records100 if not records: records100_dir = self.ptbxl_dir / 'physionet.org' / 'files' / 'ptb-xl' / '1.0.3' / 'records100' if records100_dir.exists(): for subdir in records100_dir.iterdir(): if subdir.is_dir(): for hea_file in subdir.glob('*.hea'): record_path = str(hea_file)[:-4] records.append(record_path) return records def _load_models(self): """Load Stage 0 and Stage 1 models""" weight_dir = Path(BASELINE_PATH) / 'weight' if not weight_dir.exists(): print(f"Warning: Weight directory not found at {weight_dir}") print("Stage 0/1 processing will be disabled") self.use_stage01 = False return try: print("Loading Stage 0 model...") from stage0_model import Net as Stage0Net from stage0_common import load_net self.stage0_net = Stage0Net(pretrained=False) self.stage0_net = load_net(self.stage0_net, str(weight_dir / 'stage0-last.checkpoint.pth')) self.stage0_net.to(self.device).eval() print("Loading Stage 1 model...") from stage1_model import Net as Stage1Net self.stage1_net = Stage1Net(pretrained=False) self.stage1_net = load_net(self.stage1_net, str(weight_dir / 'stage1-last.checkpoint.pth')) self.stage1_net.to(self.device).eval() print("Models loaded successfully!") except Exception as e: print(f"Error loading models: {e}") traceback.print_exc() self.use_stage01 = False def generate(self, n_samples, n_workers=1, styles=None): """ Generate n_samples synthetic ECG images Args: n_samples: Total number of samples to generate n_workers: Number of parallel workers (only for non-GPU tasks) styles: List of style keys to use (default: all 12 styles) """ if styles is None: styles = list(IMAGE_STYLES.keys()) # Create output directories for subdir in ['raw', 'stage0', 'stage1', 'gt']: (self.output_dir / subdir).mkdir(parents=True, exist_ok=True) # Prepare sample list samples = [] for i in range(n_samples): record_path = random.choice(self.records) style_key = random.choice(styles) samples.append((record_path, str(self.output_dir), i, style_key)) # Process samples results = [] # For GPU processing, do sequential (Stage 0/1 need GPU) if self.use_stage01: print(f"Processing {n_samples} samples with Stage 0/1...") for record_path, output_dir, idx, style_key in tqdm(samples): result = self._process_single_gpu(record_path, output_dir, idx, style_key) if result: results.append(result) else: # Use multiprocessing for CPU-only generation print(f"Processing {n_samples} samples (CPU only)...") args_list = [(r, o, i, s, False, None, None) for r, o, i, s in samples] with Pool(n_workers) as pool: for result in tqdm(pool.imap(process_single_sample, args_list), total=len(args_list)): if result: results.append(result) # Save manifest manifest_path = self.output_dir / 'manifest.json' with open(manifest_path, 'w') as f: json.dump(results, f, indent=2) print(f"\nGeneration complete!") print(f"Total samples: {len(results)}") print(f"Manifest saved to: {manifest_path}") return results def _process_single_gpu(self, record_path, output_dir, sample_idx, style_key): """Process a single sample with GPU for Stage 0/1""" try: # Load PTB-XL record signal_dict, fs = load_ptbxl_record(record_path) if signal_dict is None: return None # Get style configuration style_config = IMAGE_STYLES.get(style_key, IMAGE_STYLES['0001']) # Generate raw synthetic image raw_img, gt_signals = create_ecg_image_matplotlib( signal_dict, fs, style_config, random_state=sample_idx ) # Apply scan effects raw_img = apply_scan_effects(raw_img, style_config, random_state=sample_idx) # Output paths sample_id = f"syn_{sample_idx:08d}" output_dir = Path(output_dir) # Save raw image raw_path = output_dir / 'raw' / f'{sample_id}-{style_key}.png' raw_path.parent.mkdir(parents=True, exist_ok=True) cv2.imwrite(str(raw_path), cv2.cvtColor(raw_img, cv2.COLOR_RGB2BGR)) # Save ground truth gt_path = output_dir / 'gt' / f'{sample_id}.npy' gt_path.parent.mkdir(parents=True, exist_ok=True) np.save(str(gt_path), gt_signals) # Stage 0: Orientation correction from stage0_common import image_to_batch, output_to_predict, normalise_by_homography batch = image_to_batch(raw_img) with torch.no_grad(), torch.amp.autocast('cuda', dtype=torch.float32): output = self.stage0_net(batch) rotated, keypoint = output_to_predict(raw_img, batch, output) normalized, _, _ = normalise_by_homography(rotated, keypoint) # Save Stage 0 output stage0_path = output_dir / 'stage0' / f'{sample_id}-{style_key}.png' stage0_path.parent.mkdir(parents=True, exist_ok=True) cv2.imwrite(str(stage0_path), cv2.cvtColor(normalized, cv2.COLOR_RGB2BGR)) # Stage 1: Rectification from stage1_common import output_to_predict as stage1_output_to_predict, rectify_image h, w = normalized.shape[:2] normalized_tensor = torch.from_numpy( np.ascontiguousarray(normalized.transpose(2, 0, 1)).astype(np.float32) / 255.0 ).unsqueeze(0).to(self.device) batch = {'image': normalized_tensor} with torch.no_grad(), torch.amp.autocast('cuda', dtype=torch.float32): output = self.stage1_net(batch) gridpoint_xy, _ = stage1_output_to_predict(normalized, batch, output) rectified = rectify_image(normalized, gridpoint_xy) # Save Stage 1 output (main training image) stage1_path = output_dir / 'stage1' / f'{sample_id}-{style_key}.png' stage1_path.parent.mkdir(parents=True, exist_ok=True) cv2.imwrite(str(stage1_path), cv2.cvtColor(rectified, cv2.COLOR_RGB2BGR)) return { 'sample_id': sample_id, 'style': style_key, 'record': str(record_path), 'gt_path': str(gt_path), 'stage1_path': str(stage1_path), } except Exception as e: print(f"Error processing sample {sample_idx}: {e}") traceback.print_exc() return None def main(): parser = argparse.ArgumentParser(description='Generate synthetic ECG training data') parser.add_argument('--ptbxl_dir', type=str, required=True, help='Path to PTB-XL dataset') parser.add_argument('--output_dir', type=str, required=True, help='Output directory for synthetic data') parser.add_argument('--n_samples', type=int, default=200000, help='Number of samples to generate') parser.add_argument('--n_workers', type=int, default=8, help='Number of parallel workers') parser.add_argument('--device', type=str, default='cuda:0', help='CUDA device') parser.add_argument('--no_stage01', action='store_true', help='Disable Stage 0/1 processing') parser.add_argument('--styles', type=str, nargs='+', default=None, help='Specific styles to use (e.g., 0001 0002)') args = parser.parse_args() generator = SyntheticGenerator( ptbxl_dir=args.ptbxl_dir, output_dir=args.output_dir, device=args.device, use_stage01=not args.no_stage01, ) generator.generate( n_samples=args.n_samples, n_workers=args.n_workers, styles=args.styles, ) if __name__ == '__main__': main()