code/scripts/generate_and_process_all.py

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