#!/usr/bin/env python3 """ Generate ONLY 0001-type synthetic ECG images using ecg-image-kit. Based on Kaggle discussion: train only on 0001 type (clean red grid) and use heavy augmentation to simulate other variants. Usage: python generate_synthetic_0001_only.py --output_dir /data/ecg-digitization/synthetic_0001_all """ import os import sys import argparse import subprocess import numpy as np import cv2 import torch import random import json import traceback from pathlib import Path from tqdm import tqdm import shutil import wfdb from scipy import signal as scipy_signal # Paths PROJECT_ROOT = Path(__file__).parent.parent ECG_IMAGE_KIT = PROJECT_ROOT / 'ecg-image-kit' / 'codes' / 'ecg-image-generator' BASELINE_PATH = PROJECT_ROOT / 'data' / 'hengck23-submit-physionet' / 'hengck23-submit-physionet' PTBXL_PATH = Path('/data/ecg-digitization/ptbxl/physionet.org/files/ptb-xl/1.0.3') sys.path.insert(0, str(BASELINE_PATH)) # Target specifications - match Kaggle format exactly RAW_HEIGHT = 1700 RAW_WIDTH = 2200 # GT extraction parameters T0 = 235 T1 = 4161 OUTPUT_GT_WIDTH = T1 - T0 # 3926 def find_ptbxl_records(ptbxl_dir): """Find all PTB-XL record files (500Hz high-res).""" records = [] ptbxl_dir = Path(ptbxl_dir) # Use records500 (500Hz) for better quality records500 = ptbxl_dir / 'records500' if records500.exists(): for subdir in sorted(records500.iterdir()): if subdir.is_dir(): for hea in subdir.glob('*_hr.hea'): # High-res files only dat = hea.with_suffix('.dat') if dat.exists(): records.append((str(hea), str(dat))) return records def extract_gt_signal(hea_file, output_width=3926): """Extract ground truth signal from PTB-XL record.""" try: record_path = hea_file.replace('.hea', '') record = wfdb.rdrecord(record_path) signals = record.p_signal # [samples, leads] sig_names = 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) # Map lead names lead_signals = {} for i, name in enumerate(sig_names): lead_signals[name.upper()] = signals[:, i] # 4-row output: each row shows 3 leads, 2.5s each (standard 12-lead layout) # Row 0: I, aVR, V1, V4 # Row 1: II, aVL, V2, V5 # Row 2: III, aVF, V3, V6 # Row 3: Long II strip row_leads = [ ['I', 'AVR', 'V1', 'V4'], ['II', 'AVL', 'V2', 'V5'], ['III', 'AVF', 'V3', 'V6'], ] segment_len = output_width // 3 gt_signal = np.zeros((4, output_width), dtype=np.float32) # First 3 rows: 3 leads each, 2.5s segments for row_idx in range(3): for col_idx in range(3): lead_name = row_leads[row_idx][col_idx] if lead_name in lead_signals: lead_data = lead_signals[lead_name] # 2.5s at 500Hz = 1250 samples start = col_idx * 1250 end = min(start + 1250, len(lead_data)) segment = lead_data[start:end] col_start = col_idx * segment_len col_end = (col_idx + 1) * segment_len if col_idx < 2 else output_width if len(segment) > 0: x_old = np.linspace(0, 1, len(segment)) x_new = np.linspace(0, 1, col_end - col_start) gt_signal[row_idx, col_start:col_end] = np.interp(x_new, x_old, segment) # Row 4: Long Lead II strip (10s) if 'II' in lead_signals: lead_ii = lead_signals['II'] # Full 10 seconds x_old = np.linspace(0, 1, len(lead_ii)) x_new = np.linspace(0, 1, output_width) gt_signal[3, :] = np.interp(x_new, x_old, lead_ii) return gt_signal except Exception as e: print(f"Error extracting GT from {hea_file}: {e}") return None def generate_single_0001(hea_file, dat_file, output_dir, sample_idx, temp_base): """Generate a single 0001-type image (clean red grid, no distortion).""" sample_id = f'syn_{sample_idx:08d}' out_img = output_dir / 'stage1' / f'{sample_id}-0001.png' out_gt = output_dir / 'gt' / f'{sample_id}.npy' # Skip if exists if out_img.exists() and out_gt.exists(): return True, sample_idx, "exists" temp_dir = temp_base / f'temp_{sample_idx}' try: temp_dir.mkdir(parents=True, exist_ok=True) # Generate image with ecg-image-kit # 0001 = clean red grid, no augmentation, no wrinkles 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', # 200 DPI '--num_columns', '4', '--full_mode', 'II', '--standard_grid_color', '5', # Red grid (0001 style) # NO --augment, NO --wrinkles for clean baseline ] result = subprocess.run( cmd, cwd=str(ECG_IMAGE_KIT), capture_output=True, text=True, timeout=120 ) if result.returncode != 0: shutil.rmtree(temp_dir, ignore_errors=True) return False, sample_idx, f"ecg-kit error: {result.stderr[:200]}" # Find generated image gen_imgs = list(temp_dir.glob('*.png')) if not gen_imgs: shutil.rmtree(temp_dir, ignore_errors=True) return False, sample_idx, "no image generated" # Load image img = cv2.imread(str(gen_imgs[0])) if img is None: shutil.rmtree(temp_dir, ignore_errors=True) return False, sample_idx, "failed to read image" # Resize to match Kaggle format (1700x2200) img_resized = cv2.resize(img, (RAW_WIDTH, RAW_HEIGHT), interpolation=cv2.INTER_LINEAR) # Save image out_img.parent.mkdir(parents=True, exist_ok=True) cv2.imwrite(str(out_img), img_resized) # Extract and save GT signal gt_signal = extract_gt_signal(hea_file, OUTPUT_GT_WIDTH) if gt_signal is not None: out_gt.parent.mkdir(parents=True, exist_ok=True) np.save(out_gt, gt_signal) else: shutil.rmtree(temp_dir, ignore_errors=True) return False, sample_idx, "failed to extract GT" # Cleanup shutil.rmtree(temp_dir, ignore_errors=True) return True, sample_idx, "success" except subprocess.TimeoutExpired: shutil.rmtree(temp_dir, ignore_errors=True) return False, sample_idx, "timeout" except Exception as e: shutil.rmtree(temp_dir, ignore_errors=True) return False, sample_idx, str(e) def main(): parser = argparse.ArgumentParser() parser.add_argument('--ptbxl_dir', type=str, default=str(PTBXL_PATH)) parser.add_argument('--output_dir', type=str, default='/data/ecg-digitization/synthetic_0001_all') parser.add_argument('--start_idx', type=int, default=0) parser.add_argument('--max_samples', type=int, default=None, help='Max samples to generate (default: all)') args = parser.parse_args() output_dir = Path(args.output_dir) output_dir.mkdir(parents=True, exist_ok=True) (output_dir / 'stage1').mkdir(exist_ok=True) (output_dir / 'gt').mkdir(exist_ok=True) temp_base = output_dir / 'temp' temp_base.mkdir(exist_ok=True) # Find PTB-XL records print(f"Finding PTB-XL records in {args.ptbxl_dir}...") records = find_ptbxl_records(Path(args.ptbxl_dir)) print(f"Found {len(records)} PTB-XL records") if not records: print("No records found!") return # Determine how many to generate n_samples = len(records) if args.max_samples: n_samples = min(n_samples, args.max_samples) print(f"\nGenerating {n_samples} 0001-type synthetic images...") print(f" Output: {output_dir / 'stage1'}") print(f" GT: {output_dir / 'gt'}") print(f" Format: {RAW_WIDTH}x{RAW_HEIGHT} (matches Kaggle)") print() # Generate one by one (simple, reliable) success_count = 0 skip_count = 0 fail_count = 0 pbar = tqdm(range(args.start_idx, args.start_idx + n_samples), desc="Generating") for i in pbar: record_idx = i % len(records) hea, dat = records[record_idx] success, idx, msg = generate_single_0001(hea, dat, output_dir, i, temp_base) if success: if msg == "exists": skip_count += 1 else: success_count += 1 else: fail_count += 1 if fail_count <= 10: print(f"\n Failed {i}: {msg}") pbar.set_postfix({'ok': success_count, 'skip': skip_count, 'fail': fail_count}) print(f"\n\nGeneration complete!") print(f" Success: {success_count}") print(f" Skipped (existing): {skip_count}") print(f" Failed: {fail_count}") # Create manifest manifest = [] for img in sorted((output_dir / 'stage1').glob('*.png')): sample_id = img.stem.replace('-0001', '') gt_path = output_dir / 'gt' / f'{sample_id}.npy' if gt_path.exists(): manifest.append({ 'sample_id': sample_id, 'image': str(img), 'gt': str(gt_path), }) with open(output_dir / 'manifest.json', 'w') as f: json.dump(manifest, f, indent=2) print(f"\nManifest saved with {len(manifest)} samples") print(f"Output directory: {output_dir}") if __name__ == '__main__': main()