code/scripts/generate_all_variants.py

#!/usr/bin/env python3
"""
Generate all 12 style variants for each synthetic ECG record.

This script takes existing PTB-XL records and generates 12 different style variants
(0001-0012) for each, similar to the Kaggle dataset format.

Usage:
    python generate_all_variants.py --n_workers 8
"""

import os
import sys
import subprocess
import numpy as np
import cv2
import json
import shutil
from pathlib import Path
from multiprocessing import Pool, cpu_count
from tqdm import tqdm
import argparse

# Paths
PROJECT_ROOT = Path(__file__).parent.parent
ECG_IMAGE_KIT = PROJECT_ROOT / 'ecg-image-kit' / 'codes' / 'ecg-image-generator'
PTBXL_PATH = Path('/data/ecg-digitization/ptbxl/physionet.org/files/ptb-xl/1.0.3')
OUTPUT_DIR = Path('/data/ecg-digitization/synthetic_ecgkit')

# Target specifications - match Kaggle format
TARGET_HEIGHT = 1700
TARGET_WIDTH = 2200

# 12 style configurations matching Kaggle variants 0001-0012
STYLE_CONFIGS = {
    1: {'grid': True, 'grid_color': 5, 'augment': False, 'wrinkles': False, 'hw_text': False},   # Clean red grid
    2: {'grid': True, 'grid_color': 5, 'augment': True, 'wrinkles': True, 'hw_text': False},    # Aged red  
    3: {'grid': True, 'grid_color': 2, 'augment': False, 'wrinkles': False, 'hw_text': True},   # Green grid + handwriting
    4: {'grid': True, 'grid_color': 2, 'augment': True, 'wrinkles': True, 'hw_text': False},    # Aged green
    5: {'grid': True, 'grid_color': 1, 'augment': False, 'wrinkles': False, 'hw_text': False},  # Blue grid
    6: {'grid': True, 'grid_color': 3, 'augment': True, 'wrinkles': False, 'hw_text': True},    # Grey + handwriting
    7: {'grid': True, 'grid_color': 4, 'augment': True, 'wrinkles': True, 'hw_text': False},    # Yellow aged
    8: {'grid': True, 'grid_color': 2, 'augment': True, 'wrinkles': False, 'hw_text': False},   # Faded green
    9: {'grid': True, 'grid_color': 6, 'augment': False, 'wrinkles': False, 'hw_text': True},   # Orange + handwriting
    10: {'grid': True, 'grid_color': 5, 'augment': True, 'wrinkles': False, 'hw_text': False},  # High contrast red
    11: {'grid': True, 'grid_color': 5, 'augment': True, 'wrinkles': True, 'hw_text': True},    # Low contrast + text
    12: {'grid': False, 'grid_color': 0, 'augment': False, 'wrinkles': False, 'hw_text': False}, # No grid
}


def find_ptbxl_records(ptbxl_dir, max_records=None):
    """Find all PTB-XL record files."""
    records = []
    
    records500 = ptbxl_dir / 'records500'
    if records500.exists():
        for subdir in sorted(records500.iterdir()):
            if subdir.is_dir():
                for hea in sorted(subdir.glob('*.hea')):
                    dat = hea.with_suffix('.dat')
                    if dat.exists():
                        records.append((str(hea), str(dat)))
                        if max_records and len(records) >= max_records:
                            return records
    
    return records


def generate_variant(args):
    """Generate a single style variant for one record."""
    hea_file, dat_file, sample_idx, style_idx, output_dir = args
    
    try:
        style = STYLE_CONFIGS[style_idx]
        sample_id = f'syn_{sample_idx:08d}'
        style_str = f'{style_idx:04d}'
        
        # Create temp directory
        temp_dir = output_dir / 'temp' / f'{sample_id}_{style_str}'
        temp_dir.mkdir(parents=True, exist_ok=True)
        
        # Build ecg-image-kit command
        cmd = [
            'python', 'gen_ecg_image_from_data.py',
            '-i', dat_file,
            '-hea', hea_file,
            '-o', str(temp_dir),
            '-st', str(sample_idx),
            '-se', str(sample_idx),
            '-r', '200',
            '--num_columns', '4',
            '--full_mode', 'II',
            '--store_config', '1',
        ]
        
        if style['grid']:
            cmd.extend(['--standard_grid_color', str(style['grid_color'])])
        else:
            cmd.extend(['--random_grid_present', '0'])
        
        if style['augment']:
            cmd.append('--augment')
        
        if style['wrinkles']:
            cmd.append('--wrinkles')
        
        if style['hw_text']:
            cmd.append('--print_txt')
        
        # Run with isolated matplotlib config
        env = os.environ.copy()
        env['MPLBACKEND'] = 'Agg'
        env['MPLCONFIGDIR'] = str(temp_dir)
        
        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 None
        
        # Find generated image
        generated_images = list(temp_dir.glob('*.png'))
        if not generated_images:
            shutil.rmtree(temp_dir, ignore_errors=True)
            return None
        
        gen_img_path = generated_images[0]
        
        # Read and resize
        img = cv2.imread(str(gen_img_path))
        if img is None:
            shutil.rmtree(temp_dir, ignore_errors=True)
            return None
        
        img = cv2.resize(img, (TARGET_WIDTH, TARGET_HEIGHT))
        
        # Save raw image
        raw_dir = output_dir / 'raw'
        raw_path = raw_dir / f'{sample_id}-{style_str}.png'
        cv2.imwrite(str(raw_path), img)
        
        # Clean up
        shutil.rmtree(temp_dir, ignore_errors=True)
        
        return {
            'sample_id': sample_id,
            'style': style_str,
            'raw_path': str(raw_path),
            'source_record': hea_file,
        }
        
    except Exception as e:
        try:
            shutil.rmtree(temp_dir, ignore_errors=True)
        except:
            pass
        return None


def process_through_stage01(raw_dir, output_dir, batch_size=50):
    """Process raw images through Stage 0 and Stage 1."""
    sys.path.insert(0, str(PROJECT_ROOT / 'data' / 'hengck23-submit-physionet' / 'hengck23-submit-physionet'))
    
    from stage0_model import Net as Stage0Net
    from stage1_model import Net as Stage1Net
    from stage0_common import image_to_batch, output_to_predict, normalise_by_homography, load_net
    from stage1_common import output_to_predict as stage1_output_to_predict, rectify_image
    
    raw_dir = Path(raw_dir)
    output_dir = Path(output_dir)
    
    stage1_dir = output_dir / 'stage1'
    stage1_dir.mkdir(parents=True, exist_ok=True)
    
    # Load models
    weight_dir = PROJECT_ROOT / 'data' / 'hengck23-submit-physionet' / 'hengck23-submit-physionet' / 'weight'
    device = 'cuda:0'
    
    stage0_net = Stage0Net(pretrained=False)
    stage0_net = load_net(stage0_net, str(weight_dir / 'stage0-last.checkpoint.pth'))
    stage0_net.to(device).eval()
    
    stage1_net = Stage1Net(pretrained=False)
    stage1_net = load_net(stage1_net, str(weight_dir / 'stage1-last.checkpoint.pth'))
    stage1_net.to(device).eval()
    
    raw_images = sorted(raw_dir.glob('*.png'))
    
    # Filter to only process new images (not already in stage1)
    existing = set(p.stem for p in stage1_dir.glob('*.png'))
    to_process = [p for p in raw_images if p.stem not in existing]
    
    print(f"Processing {len(to_process)} images through Stage 0/1 (skipping {len(existing)} existing)")
    
    for raw_path in tqdm(to_process, desc="Stage 0/1"):
        try:
            image = cv2.imread(str(raw_path))
            if image is None:
                continue
            
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            
            # Stage 0
            with torch.no_grad():
                batch = image_to_batch(image)
                output = stage0_net(batch)
            rotated, keypoint = output_to_predict(image, batch, output)
            normalized, _, _ = normalise_by_homography(rotated, keypoint)
            
            # Stage 1
            with torch.no_grad():
                batch = {'image': torch.from_numpy(normalized.transpose(2, 0, 1)).unsqueeze(0)}
                output = stage1_net(batch)
            gridpoint_xy, _ = stage1_output_to_predict(normalized, batch, output)
            rectified = rectify_image(normalized, gridpoint_xy)
            
            # Save
            rectified_bgr = cv2.cvtColor(rectified, cv2.COLOR_RGB2BGR)
            cv2.imwrite(str(stage1_dir / raw_path.name), rectified_bgr)
            
        except Exception as e:
            continue
    
    print(f"Stage 0/1 processing complete. {len(list(stage1_dir.glob('*.png')))} images in stage1/")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--n_samples', type=int, default=None, 
                       help='Number of PTB-XL records to use (default: all ~9000)')
    parser.add_argument('--n_workers', type=int, default=8)
    parser.add_argument('--output_dir', type=str, default=str(OUTPUT_DIR))
    parser.add_argument('--skip_generation', action='store_true',
                       help='Skip generation, only run Stage 0/1')
    parser.add_argument('--skip_stage01', action='store_true',
                       help='Skip Stage 0/1 processing')
    parser.add_argument('--styles', type=str, default='1-12',
                       help='Style range to generate, e.g., "1-12" or "2,4,6,8"')
    args = parser.parse_args()
    
    output_dir = Path(args.output_dir)
    raw_dir = output_dir / 'raw'
    raw_dir.mkdir(parents=True, exist_ok=True)
    
    # Parse styles
    if '-' in args.styles:
        start, end = map(int, args.styles.split('-'))
        styles = list(range(start, end + 1))
    else:
        styles = [int(s) for s in args.styles.split(',')]
    
    print(f"Generating styles: {styles}")
    
    if not args.skip_generation:
        # Find PTB-XL records
        records = find_ptbxl_records(PTBXL_PATH, args.n_samples)
        print(f"Found {len(records)} PTB-XL records")
        
        # Check existing raw images to avoid regenerating
        existing_raw = set()
        for p in raw_dir.glob('*.png'):
            # syn_00000000-0001.png -> (0, 1)
            parts = p.stem.split('-')
            if len(parts) == 2:
                sample_num = int(parts[0].replace('syn_', ''))
                style_num = int(parts[1])
                existing_raw.add((sample_num, style_num))
        
        print(f"Found {len(existing_raw)} existing raw images")
        
        # Generate all variants
        tasks = []
        for sample_idx, (hea, dat) in enumerate(records):
            for style_idx in styles:
                if (sample_idx, style_idx) not in existing_raw:
                    tasks.append((hea, dat, sample_idx, style_idx, output_dir))
        
        print(f"Generating {len(tasks)} new images...")
        
        if tasks:
            with Pool(args.n_workers) as pool:
                results = list(tqdm(
                    pool.imap_unordered(generate_variant, tasks),
                    total=len(tasks),
                    desc="Generating"
                ))
            
            successful = sum(1 for r in results if r is not None)
            print(f"Generated {successful}/{len(tasks)} images")
    
    # Process through Stage 0/1
    if not args.skip_stage01:
        import torch
        process_through_stage01(raw_dir, output_dir)
    
    # Generate CSV files for new stage1 images
    print("\nConverting GT to CSV format...")
    stage1_dir = output_dir / 'stage1'
    gt_dir = output_dir / 'gt'
    
    # Count new stage1 images without CSV
    stage1_images = set(p.stem for p in stage1_dir.glob('*.png'))
    existing_csv = set(p.stem for p in gt_dir.glob('*.csv'))
    new_images = stage1_images - existing_csv
    
    if new_images:
        print(f"Need to create {len(new_images)} new CSV files")
        
        # For each new image, find the base NPY and create CSV
        for img_stem in tqdm(new_images, desc="Creating CSVs"):
            # syn_00000000-0003 -> syn_00000000
            base_id = img_stem.rsplit('-', 1)[0]
            npy_path = gt_dir / f'{base_id}.npy'
            
            if npy_path.exists():
                # Load NPY and convert to CSV
                mv_signal = np.load(npy_path)
                
                # Convert mV to pixels
                ZERO_MV = np.array([703.5, 987.5, 1271.5, 1531.5])
                MV_TO_PIXEL = 78.5
                pixel_signal = np.zeros_like(mv_signal)
                for row in range(4):
                    pixel_signal[row] = ZERO_MV[row] - mv_signal[row] * MV_TO_PIXEL
                pixel_signal = np.clip(pixel_signal, 0, 1695)
                
                # Save CSV
                import pandas as pd
                df = pd.DataFrame(pixel_signal.T, columns=['row0', 'row1', 'row2', 'row3'])
                df.to_csv(gt_dir / f'{img_stem}.csv', index=False)
    
    # Summary
    print("\n" + "="*60)
    print("Summary:")
    print(f"  Raw images:    {len(list(raw_dir.glob('*.png')))}")
    print(f"  Stage1 images: {len(list(stage1_dir.glob('*.png')))}")
    print(f"  GT CSV files:  {len(list(gt_dir.glob('*.csv')))}")
    print("="*60)


if __name__ == '__main__':
    main()