code/scripts/generate_synthetic_v2.py

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