""" ECG Dataset V2 - With Multi-Version Support and Style-Specific Augmentation Based on expert recommendations: - Use Stage 1 recovered images for training consistency - Heavy augmentation based on different image types (0001-0012) - Exclude bad data samples - Train on all 9-12 image versions with augmentation Usage: from dataset_v2 import create_dataloaders_v2 """ import os import torch import numpy as np import cv2 import pandas as pd from pathlib import Path from torch.utils.data import Dataset, DataLoader, ConcatDataset import albumentations as A from albumentations.pytorch import ToTensorV2 import random import json # Target specifications TARGET_HEIGHT = 1696 TARGET_WIDTH = 4352 # Calibration constants ZERO_MV = [703.5, 987.5, 1271.5, 1531.5] MV_TO_PIXEL = 78.5 T0, T1 = 235, 4161 # Style-specific augmentation strengths (0001-0012) STYLE_AUGMENTATION = { '0001': {'noise': 0.3, 'blur': 0.2, 'color': 0.5, 'quality': 0.3}, # Clean print '0002': {'noise': 0.5, 'blur': 0.3, 'color': 0.6, 'quality': 0.4}, # Aged paper '0003': {'noise': 0.3, 'blur': 0.2, 'color': 0.5, 'quality': 0.3}, # Green grid '0004': {'noise': 0.5, 'blur': 0.3, 'color': 0.6, 'quality': 0.4}, # Green aged '0005': {'noise': 0.3, 'blur': 0.3, 'color': 0.7, 'quality': 0.3}, # Blue trace '0006': {'noise': 0.6, 'blur': 0.4, 'color': 0.7, 'quality': 0.5}, # Grey bg '0007': {'noise': 0.6, 'blur': 0.4, 'color': 0.7, 'quality': 0.5}, # Yellow aged '0008': {'noise': 0.4, 'blur': 0.3, 'color': 0.5, 'quality': 0.4}, # Faded green '0009': {'noise': 0.3, 'blur': 0.2, 'color': 0.6, 'quality': 0.3}, # Orange grid '0010': {'noise': 0.4, 'blur': 0.3, 'color': 0.6, 'quality': 0.4}, # High contrast '0011': {'noise': 0.5, 'blur': 0.3, 'color': 0.5, 'quality': 0.5}, # Low contrast '0012': {'noise': 0.3, 'blur': 0.2, 'color': 0.4, 'quality': 0.3}, # No grid } def get_style_transforms(style_key='0001', height=TARGET_HEIGHT, width=TARGET_WIDTH): """ Get style-specific augmentation transforms. Different image styles need different augmentation intensities. """ aug_params = STYLE_AUGMENTATION.get(style_key, STYLE_AUGMENTATION['0001']) transforms = [ A.Resize(height, width), ] # Color augmentations (intensity varies by style) color_p = aug_params['color'] transforms.append( A.OneOf([ A.ColorJitter( brightness=0.15 * color_p, contrast=0.15 * color_p, saturation=0.1 * color_p, hue=0.05 * color_p, p=0.8 ), A.RandomBrightnessContrast( brightness_limit=0.15 * color_p, contrast_limit=0.15 * color_p, p=0.8 ), A.HueSaturationValue( hue_shift_limit=int(10 * color_p), sat_shift_limit=int(20 * color_p), val_shift_limit=int(20 * color_p), p=0.8 ), ], p=color_p) ) # Noise augmentations noise_p = aug_params['noise'] transforms.append( A.OneOf([ A.GaussNoise(var_limit=(5, int(30 * noise_p)), p=0.5), A.ISONoise( color_shift=(0.01, 0.05 * noise_p), intensity=(0.1, 0.3 * noise_p), p=0.5 ), ], p=noise_p) ) # Blur augmentations blur_p = aug_params['blur'] transforms.append( A.OneOf([ A.GaussianBlur(blur_limit=(3, max(3, int(5 * blur_p))), p=0.5), A.MotionBlur(blur_limit=max(3, int(5 * blur_p)), p=0.3), ], p=blur_p) ) # Quality degradation quality_p = aug_params['quality'] transforms.append( A.OneOf([ A.ImageCompression( quality_lower=max(50, int(100 - 30 * quality_p)), quality_upper=100, p=0.3 ), A.Downscale( scale_min=max(0.7, 1.0 - 0.2 * quality_p), scale_max=0.95, p=0.3 ), ], p=quality_p) ) # Final transforms transforms.extend([ A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ToTensorV2(), ]) return A.Compose(transforms) def get_heavy_augmentation(height=TARGET_HEIGHT, width=TARGET_WIDTH): """ Heavy augmentation for training on all image versions. Used when we want to be robust across all styles. """ return A.Compose([ A.Resize(height, width), # Strong color augmentations A.OneOf([ A.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2, hue=0.1, p=1.0), A.RandomBrightnessContrast(brightness_limit=0.3, contrast_limit=0.3, p=1.0), A.HueSaturationValue(hue_shift_limit=15, sat_shift_limit=30, val_shift_limit=30, p=1.0), A.RGBShift(r_shift_limit=20, g_shift_limit=20, b_shift_limit=20, p=1.0), ], p=0.8), # Channel manipulations A.OneOf([ A.ChannelShuffle(p=0.3), A.ToGray(p=0.1), A.CLAHE(clip_limit=2.0, p=0.3), ], p=0.3), # Noise A.OneOf([ A.GaussNoise(var_limit=(10, 50), p=0.5), A.ISONoise(color_shift=(0.01, 0.1), intensity=(0.1, 0.5), p=0.5), A.MultiplicativeNoise(multiplier=(0.9, 1.1), p=0.3), ], p=0.5), # Blur A.OneOf([ A.GaussianBlur(blur_limit=(3, 7), p=0.4), A.MotionBlur(blur_limit=5, p=0.3), A.MedianBlur(blur_limit=5, p=0.2), ], p=0.3), # Quality A.OneOf([ A.ImageCompression(quality_lower=60, quality_upper=100, p=0.4), A.Downscale(scale_min=0.7, scale_max=0.95, p=0.3), ], p=0.3), # Shadows and lighting A.RandomShadow(shadow_roi=(0, 0, 1, 1), num_shadows_lower=1, num_shadows_upper=3, p=0.2), # Normalize and convert A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ToTensorV2(), ]) def get_val_transforms(height=TARGET_HEIGHT, width=TARGET_WIDTH): """Validation transforms - resize and normalize only.""" return A.Compose([ A.Resize(height, width), A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ToTensorV2(), ]) class SyntheticECGDatasetV2(Dataset): """ Synthetic ECG Dataset - loads Stage 1 processed images with ground truth. Directory structure (from generate_synthetic_v2.py): synthetic_dir/ stage1/ <- Main training images (Stage 1 rectified) syn_00000000-0001.png ... gt/ <- Ground truth signals syn_00000000.npy ... manifest.json <- Sample metadata """ def __init__(self, synthetic_dir, transform=None, max_samples=None, use_styles=None): self.synthetic_dir = Path(synthetic_dir) self.transform = transform or get_heavy_augmentation() # Load manifest if exists manifest_path = self.synthetic_dir / 'manifest.json' if manifest_path.exists(): with open(manifest_path) as f: self.manifest = json.load(f) else: self.manifest = None # Find stage1 images stage1_dir = self.synthetic_dir / 'stage1' if stage1_dir.exists(): self.image_files = sorted(stage1_dir.glob('*.png')) else: # Fallback to images directory image_dir = self.synthetic_dir / 'images' if image_dir.exists(): self.image_files = sorted(image_dir.glob('*.png')) else: self.image_files = [] # Filter by style if specified if use_styles: self.image_files = [ f for f in self.image_files if any(f'-{s}' in f.stem for s in use_styles) ] # Limit samples if specified if max_samples: self.image_files = self.image_files[:max_samples] print(f"SyntheticDatasetV2: {len(self.image_files)} images") def __len__(self): return len(self.image_files) def __getitem__(self, idx): image_path = self.image_files[idx] # Load image image = cv2.imread(str(image_path)) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # Determine style from filename (e.g., syn_00000000-0001.png) filename = image_path.stem style = '0001' # default for s in STYLE_AUGMENTATION.keys(): if f'-{s}' in filename: style = s break # Load ground truth # Extract sample_id (e.g., syn_00000000 from syn_00000000-0001) sample_id = filename.split('-')[0] if '-' in filename else filename gt_path = self.synthetic_dir / 'gt' / f'{sample_id}.npy' if gt_path.exists(): target = np.load(gt_path) # Shape: [4, width] else: # No GT, return zeros target = np.zeros((4, T1 - T0), dtype=np.float32) # Apply style-specific or heavy augmentation if self.transform: transformed = self.transform(image=image) image = transformed['image'] # Convert mV to normalized pixel coordinates target_pixel = np.zeros_like(target) for row_idx in range(4): target_pixel[row_idx] = ZERO_MV[row_idx] - target[row_idx] * MV_TO_PIXEL target_normalized = target_pixel / TARGET_HEIGHT target_normalized = np.clip(target_normalized, 0, 1) # Ensure correct output width output_width = T1 - T0 if target_normalized.shape[1] != output_width: target_resized = np.zeros((4, output_width), dtype=np.float32) for row_idx in range(4): x_old = np.linspace(0, 1, target_normalized.shape[1]) x_new = np.linspace(0, 1, output_width) target_resized[row_idx] = np.interp(x_new, x_old, target_normalized[row_idx]) target_normalized = target_resized return { 'image': image, 'target': torch.from_numpy(target_normalized.astype(np.float32)), 'id': filename, 'style': style, } class KaggleECGDatasetV2(Dataset): """ Kaggle ECG Dataset V2 - loads ALL image versions (0001-0012) with augmentation. Key features: - Uses Stage 1 rectified images for consistency - Loads all versions (0001-0012) of each image - Applies style-specific augmentation - Excludes bad data samples (configurable) Directory structure: kaggle_dir/ train/ {id}/ {id}-0001.png, {id}-0002.png, ..., {id}-0012.png {id}.csv (ground truth) stage1_output/ {id}-0001.png, {id}-0002.png, ... (rectified) """ def __init__( self, kaggle_dir, stage1_dir=None, df=None, transform=None, use_all_versions=True, versions=None, exclude_ids=None, ): self.kaggle_dir = Path(kaggle_dir) self.stage1_dir = Path(stage1_dir) if stage1_dir else None self.transform = transform or get_heavy_augmentation() self.use_all_versions = use_all_versions self.versions = versions or [f'{i:04d}' for i in range(1, 13)] # 0001-0012 self.exclude_ids = set(exclude_ids or []) # Get image IDs from train directory or df if df is not None: self.image_ids = [str(x) for x in df['id'].unique() if str(x) not in self.exclude_ids] else: train_dir = self.kaggle_dir / 'train' if train_dir.exists(): self.image_ids = [ d.name for d in train_dir.iterdir() if d.is_dir() and d.name not in self.exclude_ids ] else: self.image_ids = [] # Build sample list: (image_id, version) self.samples = [] for img_id in self.image_ids: if use_all_versions: for v in self.versions: self.samples.append((img_id, v)) else: # Just use first version self.samples.append((img_id, self.versions[0])) print(f"KaggleDatasetV2: {len(self.image_ids)} images, {len(self.samples)} samples (versions: {len(self.versions)})") def __len__(self): return len(self.samples) def __getitem__(self, idx): image_id, version = self.samples[idx] # Try Stage 1 rectified first image = None if self.stage1_dir: stage1_path = self.stage1_dir / f"{image_id}-{version}.png" if stage1_path.exists(): image = cv2.imread(str(stage1_path)) else: # Try without version suffix stage1_path = self.stage1_dir / f"{image_id}.png" if stage1_path.exists(): image = cv2.imread(str(stage1_path)) # Fall back to original if image is None: original_path = self.kaggle_dir / 'train' / image_id / f"{image_id}-{version}.png" if original_path.exists(): image = cv2.imread(str(original_path)) else: # Last resort - try first version original_path = self.kaggle_dir / 'train' / image_id / f"{image_id}-0001.png" if original_path.exists(): image = cv2.imread(str(original_path)) else: # Return dummy image image = np.zeros((TARGET_HEIGHT, TARGET_WIDTH, 3), dtype=np.uint8) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # Load ground truth csv_path = self.kaggle_dir / 'train' / image_id / f"{image_id}.csv" if csv_path.exists(): gt_df = pd.read_csv(csv_path) target = self._csv_to_target(gt_df) else: target = np.zeros((4, T1 - T0), dtype=np.float32) # Apply transforms (style-specific or heavy) if self.transform: transformed = self.transform(image=image) image = transformed['image'] return { 'image': image, 'target': torch.from_numpy(target.astype(np.float32)), 'id': image_id, 'version': version, } def _csv_to_target(self, df): """Convert ground truth CSV to normalized target array.""" lead_layout = [ ['I', 'aVR', 'V1', 'V4'], ['II', 'aVL', 'V2', 'V5'], ['III', 'aVF', 'V3', 'V6'], ] output_width = T1 - T0 target = np.zeros((4, output_width), dtype=np.float32) # Process rows 0-2 for row_idx in range(3): row_signals = [] quarter_width = output_width // 4 for lead in lead_layout[row_idx]: if lead in df.columns: signal = df[lead].dropna().values if len(signal) > 0: x_old = np.linspace(0, 1, len(signal)) x_new = np.linspace(0, 1, quarter_width) signal_resampled = np.interp(x_new, x_old, signal) else: signal_resampled = np.zeros(quarter_width) else: signal_resampled = np.zeros(quarter_width) row_signals.append(signal_resampled) target[row_idx] = np.concatenate(row_signals) # Row 3: Full II rhythm strip if 'II' in df.columns: signal_ii = df['II'].dropna().values if len(signal_ii) > 0: x_old = np.linspace(0, 1, len(signal_ii)) x_new = np.linspace(0, 1, output_width) target[3] = np.interp(x_new, x_old, signal_ii) # Convert to normalized pixel coordinates target_pixel = np.zeros_like(target) for row_idx in range(4): target_pixel[row_idx] = ZERO_MV[row_idx] - target[row_idx] * MV_TO_PIXEL target_normalized = target_pixel / TARGET_HEIGHT return np.clip(target_normalized, 0, 1) class MixedDatasetV2(Dataset): """ Combined synthetic + real dataset for training. """ def __init__(self, synthetic_dataset, kaggle_dataset, synthetic_ratio=0.3): self.synthetic_dataset = synthetic_dataset self.kaggle_dataset = kaggle_dataset n_kaggle = len(kaggle_dataset) n_synthetic_to_use = int(n_kaggle * synthetic_ratio / (1 - synthetic_ratio)) # Create index mapping self.indices = [] # Add all kaggle samples for i in range(n_kaggle): self.indices.append(('kaggle', i)) # Add synthetic samples (with wrap-around) for i in range(n_synthetic_to_use): syn_idx = i % len(synthetic_dataset) self.indices.append(('synthetic', syn_idx)) random.shuffle(self.indices) print(f"MixedDatasetV2: {n_kaggle} kaggle + {n_synthetic_to_use} synthetic = {len(self.indices)} total") def __len__(self): return len(self.indices) def __getitem__(self, idx): source, real_idx = self.indices[idx] if source == 'kaggle': return self.kaggle_dataset[real_idx] else: return self.synthetic_dataset[real_idx] def create_dataloaders_v2( synthetic_dir=None, kaggle_dir=None, stage1_dir=None, train_df=None, val_df=None, batch_size=8, num_workers=4, synthetic_ratio=0.3, use_all_versions=True, exclude_ids=None, distributed=False, ): """ Create DataLoaders V2 with multi-version support. Args: synthetic_dir: Path to synthetic data (with stage1/ subdirectory) kaggle_dir: Path to Kaggle competition data stage1_dir: Path to Stage 1 processed Kaggle images train_df: Training DataFrame (image IDs) val_df: Validation DataFrame batch_size: Batch size per GPU num_workers: Number of data loading workers synthetic_ratio: Ratio of synthetic data in mixed training use_all_versions: Whether to use all 12 image versions exclude_ids: List of image IDs to exclude (bad data) distributed: Whether using DDP Returns: train_loader, val_loader """ train_transform = get_heavy_augmentation() val_transform = get_val_transforms() # Create training dataset if synthetic_dir and kaggle_dir: # Mixed training synthetic_ds = SyntheticECGDatasetV2(synthetic_dir, train_transform) kaggle_ds = KaggleECGDatasetV2( kaggle_dir, stage1_dir, train_df, train_transform, use_all_versions=use_all_versions, exclude_ids=exclude_ids, ) train_dataset = MixedDatasetV2(synthetic_ds, kaggle_ds, synthetic_ratio) elif synthetic_dir: # Synthetic pre-training train_dataset = SyntheticECGDatasetV2(synthetic_dir, train_transform) elif kaggle_dir: # Kaggle fine-tuning train_dataset = KaggleECGDatasetV2( kaggle_dir, stage1_dir, train_df, train_transform, use_all_versions=use_all_versions, exclude_ids=exclude_ids, ) else: raise ValueError("Must provide synthetic_dir or kaggle_dir") # Validation dataset (use -0001 version only, no heavy augmentation) val_dataset = None if val_df is not None and kaggle_dir: val_dataset = KaggleECGDatasetV2( kaggle_dir, stage1_dir, val_df, val_transform, use_all_versions=False, # Only use 0001 for validation exclude_ids=exclude_ids, ) # Samplers for distributed training train_sampler = None val_sampler = None if distributed: train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset) if val_dataset: val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset, shuffle=False) # Create loaders train_loader = DataLoader( train_dataset, batch_size=batch_size, shuffle=(train_sampler is None), sampler=train_sampler, num_workers=num_workers, pin_memory=True, drop_last=True, ) val_loader = None if val_dataset: val_loader = DataLoader( val_dataset, batch_size=batch_size, shuffle=False, sampler=val_sampler, num_workers=num_workers, pin_memory=True, ) return train_loader, val_loader if __name__ == '__main__': # Test the datasets print("Testing Dataset V2 classes...") # Test transforms train_tf = get_heavy_augmentation(424, 1088) dummy_image = np.random.randint(0, 255, (424, 1088, 3), dtype=np.uint8) result = train_tf(image=dummy_image) print(f"Transform output shape: {result['image'].shape}") # Test style-specific transforms for style in ['0001', '0006', '0012']: tf = get_style_transforms(style, 424, 1088) result = tf(image=dummy_image) print(f"Style {style} transform: {result['image'].shape}") print("\nDataset V2 tests completed!")