code/scripts/finetune_v19_kaggle.py

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