#!/usr/bin/env python3 """ V19 Fine-tuning with Soft-DTW Loss on Kaggle Training Data Uses the 977 Kaggle training samples with ground truth signals. Fine-tunes from best V19 checkpoint (epoch 19) using: - Soft-DTW loss for temporal alignment tolerance - MSE loss for point-wise accuracy - Very low learning rate to preserve features """ import os import sys import numpy as np import torch import torch.nn as nn import torch.nn.functional as F import timm from torch.utils.data import DataLoader, Dataset from pathlib import Path import cv2 from tqdm import tqdm import random import pandas as pd # Configuration BASELINE_PATH = '/home/azureuser/tmp/hengck23/hengck23-submit-physionet' V19_CHECKPOINT = '/data/ecg-digitization/checkpoints/v19_enhanced_epoch019.pth' TRAIN_DIR = Path('/data/ecg-digitization/kaggle/train') OUTPUT_DIR = Path('/data/ecg-digitization/checkpoints') DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') LEARNING_RATE = 5e-6 # Very low for fine-tuning EPOCHS = 5 BATCH_SIZE = 2 DTW_GAMMA = 0.5 LAMBDA_DTW = 0.2 # 20% DTW, 80% MSE NUM_WORKERS = 4 # ECG constants TARGET_HEIGHT, TARGET_WIDTH = 1696, 4352 ZERO_MV = np.array([703.5, 987.5, 1271.5, 1531.5]) MV_TO_PIXEL = 78.5 T0, T1 = 235, 4161 Y0, Y1 = 0, 1696 CROP_HALF_HEIGHT = 250 ROW_HEIGHT = 500 OUTPUT_WIDTH = T1 - T0 ECG_MV_MIN, ECG_MV_MAX = -7.0, 7.0 ROW_LAYOUT = [['I', 'aVR', 'V1', 'V4'], ['II', 'aVL', 'V2', 'V5'], ['III', 'aVF', 'V3', 'V6']] try: from torchvision.ops import DeformConv2d HAS_DEFORM_CONV = True except ImportError: HAS_DEFORM_CONV = False # ======================= Soft-DTW Loss ======================= class SoftDTWLoss(nn.Module): """Efficient Soft-DTW for fine-tuning.""" def __init__(self, gamma=0.5): super().__init__() self.gamma = gamma def forward(self, pred, target): """Compute soft-DTW on downsampled sequences.""" B, T = pred.shape # Compute pairwise distances D = (pred.unsqueeze(2) - target.unsqueeze(1)) ** 2 # [B, T, T] # Simple DP with soft-min R = torch.zeros(B, T+1, T+1, device=pred.device) R[:, 0, 1:] = float('inf') R[:, 1:, 0] = float('inf') for i in range(1, T+1): for j in range(1, T+1): options = torch.stack([R[:, i-1, j-1], R[:, i-1, j], R[:, i, j-1]], dim=1) R[:, i, j] = -self.gamma * torch.logsumexp(-options / self.gamma, dim=1) + D[:, i-1, j-1] return R[:, T, T].mean() class CombinedLoss(nn.Module): """Combined MSE + Soft-DTW with gradient scaling.""" def __init__(self, lambda_dtw=0.2, gamma=0.5, downsample=16): super().__init__() self.lambda_dtw = lambda_dtw self.mse = nn.MSELoss() self.soft_dtw = SoftDTWLoss(gamma=gamma) self.downsample = downsample def forward(self, pred, target): # MSE on full resolution mse_loss = self.mse(pred, target) # DTW on heavily downsampled for efficiency pred_ds = pred[:, ::self.downsample] target_ds = target[:, ::self.downsample] # Clamp to prevent explosion with torch.amp.autocast('cuda', enabled=False): pred_ds = pred_ds.float() target_ds = target_ds.float() dtw_loss = self.soft_dtw(pred_ds, target_ds) dtw_loss = torch.clamp(dtw_loss, 0, 10) total = (1 - self.lambda_dtw) * mse_loss + self.lambda_dtw * dtw_loss * 0.01 # Scale DTW return total, mse_loss, dtw_loss # ======================= V19 Model (same as before) ======================= class DeformableConvBlock(nn.Module): def __init__(self, in_ch, out_ch, kernel_size=3, stride=1, padding=1): super().__init__() if HAS_DEFORM_CONV: self.offset_conv = nn.Sequential( nn.Conv2d(in_ch, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(True), nn.Conv2d(64, 2 * kernel_size * kernel_size, 3, padding=1)) self.deform_conv = DeformConv2d(in_ch, out_ch, kernel_size, stride=stride, padding=padding) else: self.conv = nn.Conv2d(in_ch, out_ch, kernel_size, stride=stride, padding=padding) self.norm = nn.BatchNorm2d(out_ch) self.act = nn.GELU() def forward(self, x): if HAS_DEFORM_CONV: out = self.deform_conv(x, self.offset_conv(x)) else: out = self.conv(x) return self.act(self.norm(out)) class BiLSTMHead(nn.Module): def __init__(self, input_dim, hidden_dim=128, num_layers=2, dropout=0.1): super().__init__() self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers, batch_first=True, bidirectional=True, dropout=dropout if num_layers > 1 else 0) self.output_proj = nn.Sequential(nn.Linear(hidden_dim * 2, hidden_dim), nn.LayerNorm(hidden_dim), nn.GELU()) self.output_dim = hidden_dim def forward(self, x): x = x.permute(0, 2, 1) lstm_out, _ = self.lstm(x) return self.output_proj(lstm_out) class AuxiliaryHeads(nn.Module): def __init__(self, feature_dim): super().__init__() self.grid_head = nn.Sequential(nn.Conv2d(32, 16, 3, padding=1), nn.BatchNorm2d(16), nn.ReLU(True), nn.Conv2d(16, 1, 1), nn.Sigmoid()) self.gradient_head = nn.Sequential(nn.Linear(feature_dim, 64), nn.GELU(), nn.Linear(64, 1), nn.Tanh()) self.uncertainty_head = nn.Sequential(nn.Linear(feature_dim, 64), nn.GELU(), nn.Linear(64, 1)) def forward(self, f2d, f1d): return self.grid_head(f2d), self.gradient_head(f1d).squeeze(-1), self.uncertainty_head(f1d).squeeze(-1) class CoordConv2d(nn.Module): def __init__(self, in_ch, out_ch, kernel_size, **kwargs): super().__init__() self.conv = nn.Conv2d(in_ch + 2, out_ch, kernel_size, **kwargs) def forward(self, x): B, C, H, W = x.shape yy = torch.linspace(-1, 1, H, device=x.device).view(1, 1, H, 1).expand(B, 1, H, W) xx = torch.linspace(-1, 1, W, device=x.device).view(1, 1, 1, W).expand(B, 1, H, W) return self.conv(torch.cat([x, yy, xx], dim=1)) class UNetDecoderBlockV19(nn.Module): def __init__(self, in_ch, skip_ch, out_ch, use_deform=False): super().__init__() if use_deform and HAS_DEFORM_CONV: self.conv1 = DeformableConvBlock(in_ch + skip_ch, out_ch) else: self.conv1 = nn.Sequential(nn.Conv2d(in_ch + skip_ch, out_ch, 3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.GELU()) self.conv2 = nn.Sequential(nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.GELU()) self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) def forward(self, x, skip=None): x = self.upsample(x) if skip is not None: if x.shape[2:] != skip.shape[2:]: x = F.interpolate(x, size=skip.shape[2:], mode='bilinear', align_corners=True) x = torch.cat([x, skip], dim=1) return self.conv2(self.conv1(x)) class PerLeadNetV19(nn.Module): def __init__(self, encoder_name='convnext_base.fb_in22k_ft_in1k', pretrained=True): super().__init__() self.encoder = timm.create_model(encoder_name, pretrained=pretrained, features_only=True, out_indices=(0, 1, 2, 3)) enc_channels = self.encoder.feature_info.channels() decoder_dims = [256, 128, 64, 32] self.dec_blocks = nn.ModuleList() in_ch = enc_channels[-1] for i, (skip_ch, out_ch) in enumerate(zip(enc_channels[:-1][::-1] + [0], decoder_dims)): self.dec_blocks.append(UNetDecoderBlockV19(in_ch, skip_ch, out_ch, use_deform=(i >= 2))) in_ch = out_ch self.final_up = nn.Sequential(nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True), nn.Conv2d(32, 32, 3, padding=1, bias=False), nn.BatchNorm2d(32), nn.GELU()) self.height_attention = nn.Sequential(CoordConv2d(32, 64, 3, padding=1), nn.BatchNorm2d(64), nn.GELU(), nn.Conv2d(64, 1, 1)) self.bilstm = BiLSTMHead(32, hidden_dim=128, num_layers=2, dropout=0.1) self.regression_head = nn.Sequential(nn.Linear(128, 64), nn.GELU(), nn.Linear(64, 1), nn.Sigmoid()) self.aux_heads = AuxiliaryHeads(128) def forward(self, x, return_aux=False): B, C, H, W = x.shape features = self.encoder(x) d = features[-1] for block, skip in zip(self.dec_blocks, features[:-1][::-1] + [None]): d = block(d, skip) features_2d = d d = self.final_up(d) if d.shape[3] != W: d = F.interpolate(d, size=(d.shape[2], W), mode='bilinear', align_corners=True) attn = F.softmax(self.height_attention(d), dim=2) pooled = (d * attn).sum(dim=2) temporal_features = self.bilstm(pooled) y_pred = self.regression_head(temporal_features).squeeze(-1) if return_aux: return y_pred, self.aux_heads(features_2d, temporal_features) return y_pred # ======================= Preprocessing ======================= sys.path.insert(0, BASELINE_PATH) import stage0_common as s0c import stage1_common as s1c from stage0_model import Net as Stage0Net from stage1_model import Net as Stage1Net def load_preprocessing_models(): stage0 = s0c.load_net(Stage0Net(pretrained=False), f'{BASELINE_PATH}/weight/stage0-last.checkpoint.pth').to(DEVICE).eval() stage1 = s1c.load_net(Stage1Net(pretrained=False), f'{BASELINE_PATH}/weight/stage1-last.checkpoint.pth').to(DEVICE).eval() return stage0, stage1 def change_color(image_rgb): hsv = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2HSV) h, s, v = cv2.split(hsv) v_denoised = cv2.fastNlMeansDenoising(v, h=5.46) std = np.std(v_denoised) clip_limit = max(1.0, min(3.5, 2.0 + std / 25)) clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=(8, 8)) v_enhanced = clahe.apply(v_denoised) return cv2.cvtColor(cv2.merge([h, s, v_enhanced]), cv2.COLOR_HSV2RGB) @torch.no_grad() def run_stage0(image_rgb, stage0_net): batch = s0c.image_to_batch(change_color(image_rgb)) output = stage0_net(batch) rotated, keypoint = s0c.output_to_predict(image_rgb, batch, output) normalized, _, _ = s0c.normalise_by_homography(rotated, keypoint) return normalized @torch.no_grad() def run_stage1(image_rgb, stage1_net): batch = {'image': torch.from_numpy(np.ascontiguousarray(image_rgb.transpose(2, 0, 1))).unsqueeze(0)} output = stage1_net(batch) gridpoint_xy, _ = s1c.output_to_predict(image_rgb, batch, output) return s1c.rectify_image(image_rgb, gridpoint_xy) # ======================= Dataset ======================= class KaggleTrainDataset(Dataset): """Dataset using Kaggle training data with ground truth signals.""" def __init__(self, train_dir, stage0_net, stage1_net, max_samples=None, cache_dir=None): self.train_dir = Path(train_dir) self.stage0_net = stage0_net self.stage1_net = stage1_net self.cache_dir = Path(cache_dir) if cache_dir else None # Find all training samples self.samples = [] for sample_dir in sorted(self.train_dir.iterdir()): if not sample_dir.is_dir(): continue gt_path = sample_dir / f"{sample_dir.name}.csv" if not gt_path.exists(): continue # Find all image variants for img_path in sorted(sample_dir.glob('*.png')): variant = img_path.stem.split('-')[-1] self.samples.append({ 'sample_id': sample_dir.name, 'image_path': img_path, 'gt_path': gt_path, 'variant': variant }) if max_samples and max_samples < len(self.samples): random.shuffle(self.samples) self.samples = self.samples[:max_samples] print(f"Loaded {len(self.samples)} training samples") def __len__(self): return len(self.samples) * 4 # 4 rows per image def __getitem__(self, idx): sample_idx = idx // 4 row_idx = idx % 4 sample = self.samples[sample_idx] # Check cache cache_key = f"{sample['sample_id']}_{sample['variant']}_{row_idx}" if self.cache_dir: cache_path = self.cache_dir / f"{cache_key}.npz" if cache_path.exists(): data = np.load(cache_path) return torch.from_numpy(data['image']).float(), torch.from_numpy(data['target']).float() # Load and preprocess image img_bgr = cv2.imread(str(sample['image_path'])) img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) try: normalized = run_stage0(img_rgb, self.stage0_net) rectified = run_stage1(normalized, self.stage1_net) except Exception as e: print(f"Preprocessing failed for {sample['image_path']}: {e}") # Return zeros return torch.zeros(3, ROW_HEIGHT, OUTPUT_WIDTH), torch.zeros(OUTPUT_WIDTH) # Resize and crop row rectified_bgr = cv2.cvtColor(rectified, cv2.COLOR_RGB2BGR) h, w = rectified_bgr.shape[:2] image_cropped = rectified_bgr[:min(h, Y1), :min(w, 2176)] image_resized = cv2.resize(image_cropped, (TARGET_WIDTH, TARGET_HEIGHT)) # Crop row baseline_y = int(ZERO_MV[row_idx]) y_start = max(0, baseline_y - CROP_HALF_HEIGHT) y_end = min(TARGET_HEIGHT, baseline_y + CROP_HALF_HEIGHT) row_crop = image_resized[y_start:y_end, T0:T1, :].copy() if row_crop.shape[0] < ROW_HEIGHT: pad_top = max(0, CROP_HALF_HEIGHT - baseline_y) pad_bottom = max(0, (baseline_y + CROP_HALF_HEIGHT) - TARGET_HEIGHT) row_crop = np.pad(row_crop, ((pad_top, pad_bottom), (0, 0), (0, 0)), mode='edge') # Load ground truth gt_df = pd.read_csv(sample['gt_path']) # Get leads for this row (row_idx=3 is the long Lead II strip) # Layout: each lead occupies 1/4 of the signal duration for rows 0-2 # Lead positions: 0-25%, 25-50%, 50-75%, 75-100% total_samples = len(gt_df) segment_samples = total_samples // 4 segment_width = OUTPUT_WIDTH // 4 if row_idx < 3: lead_names = ROW_LAYOUT[row_idx] y_target = np.zeros(OUTPUT_WIDTH) for seg_idx, lead_name in enumerate(lead_names): if lead_name in gt_df.columns: lead_signal = gt_df[lead_name].values # Extract the valid segment (each lead has data only in its 1/4 segment) start_idx = seg_idx * segment_samples end_idx = (seg_idx + 1) * segment_samples lead_segment = lead_signal[start_idx:end_idx] # Handle NaN values valid_mask = ~np.isnan(lead_segment) if valid_mask.sum() < 10: continue # Skip if too few valid samples lead_segment = lead_segment[valid_mask] resampled = np.interp( np.linspace(0, 1, segment_width), np.linspace(0, 1, len(lead_segment)), lead_segment ) y_target[seg_idx * segment_width:(seg_idx + 1) * segment_width] = resampled else: # Row 3 is the full Lead II strip if 'II' in gt_df.columns: lead_signal = gt_df['II'].values valid_mask = ~np.isnan(lead_signal) lead_signal = lead_signal[valid_mask] if valid_mask.sum() > 10 else lead_signal y_target = np.interp( np.linspace(0, 1, OUTPUT_WIDTH), np.linspace(0, 1, len(lead_signal)), lead_signal ) else: y_target = np.zeros(OUTPUT_WIDTH) # Convert mV to normalized pixel coordinates y_pixel = ZERO_MV[row_idx] - y_target * MV_TO_PIXEL # mV to pixel y y_crop_coord = y_pixel - (baseline_y - CROP_HALF_HEIGHT) # Relative to crop y_normalized = np.clip(y_crop_coord / ROW_HEIGHT, 0, 1) # Normalize to [0, 1] # Convert to tensors image_tensor = torch.from_numpy(row_crop.astype(np.float32) / 255.0).permute(2, 0, 1) target_tensor = torch.from_numpy(y_normalized.astype(np.float32)) # Save to cache if self.cache_dir: self.cache_dir.mkdir(parents=True, exist_ok=True) np.savez_compressed(cache_path, image=image_tensor.numpy(), target=target_tensor.numpy()) return image_tensor, target_tensor # ======================= Training ======================= def compute_snr(pred, target): """Compute SNR in dB.""" mse = ((pred - target) ** 2).mean(dim=1) signal_power = (target ** 2).mean(dim=1) snr = 10 * torch.log10(signal_power / (mse + 1e-10)) return snr.mean().item() def train_epoch(model, dataloader, optimizer, criterion, device, epoch, scaler): model.train() total_loss, total_mse, total_dtw = 0, 0, 0 total_snr = 0 n_batches = 0 pbar = tqdm(dataloader, desc=f'Epoch {epoch}') for images, targets in pbar: if images.sum() == 0: continue # Skip failed samples images = images.to(device) targets = targets.to(device) optimizer.zero_grad() with torch.amp.autocast('cuda', dtype=torch.float16): pred = model(images, return_aux=False) loss, mse_loss, dtw_loss = criterion(pred, targets) scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) scaler.step(optimizer) scaler.update() total_loss += loss.item() total_mse += mse_loss.item() total_dtw += dtw_loss.item() with torch.no_grad(): snr = compute_snr(pred, targets) total_snr += snr n_batches += 1 pbar.set_postfix({'loss': f'{loss.item():.4f}', 'snr': f'{snr:.2f}'}) return total_loss / n_batches, total_mse / n_batches, total_dtw / n_batches, total_snr / n_batches def validate(model, dataloader, device): model.eval() total_snr = 0 n_samples = 0 with torch.no_grad(): for images, targets in tqdm(dataloader, desc='Validation'): if images.sum() == 0: continue images = images.to(device) targets = targets.to(device) with torch.amp.autocast('cuda', dtype=torch.float16): pred = model(images, return_aux=False) mse = ((pred - targets) ** 2).mean(dim=1) signal_power = (targets ** 2).mean(dim=1) snr = 10 * torch.log10(signal_power / (mse + 1e-10)) total_snr += snr.sum().item() n_samples += images.shape[0] return total_snr / n_samples if n_samples > 0 else 0 def main(): print(f"Device: {DEVICE}") print(f"Deformable Conv: {HAS_DEFORM_CONV}") # Load preprocessing models print("Loading preprocessing models...") stage0_net, stage1_net = load_preprocessing_models() # Load V19 model print("Loading V19 model...") model = PerLeadNetV19(pretrained=False) checkpoint = torch.load(V19_CHECKPOINT, map_location='cpu', weights_only=False) model.load_state_dict(checkpoint['model'], strict=True) model.to(DEVICE) print(f"Loaded V19 epoch {checkpoint['epoch']}, SNR={checkpoint['snr']:.2f} dB") # Create dataset cache_dir = OUTPUT_DIR / 'training_cache' print("Creating dataset...") dataset = KaggleTrainDataset(TRAIN_DIR, stage0_net, stage1_net, max_samples=500, cache_dir=cache_dir) # Split train/val n_val = min(100, len(dataset) // 5) n_train = len(dataset) - n_val train_dataset, val_dataset = torch.utils.data.random_split(dataset, [n_train, n_val]) train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKERS, pin_memory=True) val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS, pin_memory=True) print(f"Train: {len(train_dataset)}, Val: {len(val_dataset)}") # Loss and optimizer criterion = CombinedLoss(lambda_dtw=LAMBDA_DTW, gamma=DTW_GAMMA) optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=0.01) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS, eta_min=LEARNING_RATE / 10) scaler = torch.amp.GradScaler('cuda') best_snr = checkpoint['snr'] for epoch in range(1, EPOCHS + 1): loss, mse, dtw, train_snr = train_epoch(model, train_loader, optimizer, criterion, DEVICE, epoch, scaler) val_snr = validate(model, val_loader, DEVICE) scheduler.step() print(f"Epoch {epoch}: Loss={loss:.4f}, MSE={mse:.4f}, DTW={dtw:.4f}, Train SNR={train_snr:.2f}, Val SNR={val_snr:.2f}") # Save checkpoint ckpt = {'epoch': checkpoint['epoch'] + epoch, 'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'snr': val_snr, 'loss': loss} torch.save(ckpt, OUTPUT_DIR / f'v19_softdtw_epoch{epoch:03d}.pth') if val_snr > best_snr: best_snr = val_snr torch.save(ckpt, OUTPUT_DIR / 'v19_softdtw_best.pth') print(f"New best! SNR={val_snr:.2f} dB") print(f"\nFine-tuning complete. Best SNR: {best_snr:.2f} dB") if __name__ == '__main__': main()