code/scripts/infer_v19_augraphy.py

#!/usr/bin/env python3
"""
V19 Augraphy Inference

Runs V19 model fine-tuned with Augraphy augmentation (Layer B).
This model is more robust to paper degradation, stains, and noise.

Default checkpoint: /data/ecg-digitization/checkpoints/v19_augraphy_latest.pth
"""

import os
import sys
import argparse
import random
import subprocess
import numpy as np
import pandas as pd
from pathlib import Path
from tqdm import tqdm
from scipy.signal import savgol_filter

import torch
import torch.nn as nn
import torch.nn.functional as F

import cv2
import timm

# Try to import torchvision deformable conv
try:
    from torchvision.ops import DeformConv2d
    HAS_DEFORM_CONV = True
except ImportError:
    HAS_DEFORM_CONV = False


# =============================================================================
# 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
X0, X1 = 0, 2176
Y0, Y1 = 0, 1696
OUTPUT_WIDTH = T1 - T0  # 3926

# Per-row crop parameters
CROP_HALF_HEIGHT = 250
ROW_HEIGHT = 500

# ECG amplitude limits (mV)
ECG_MV_MIN, ECG_MV_MAX = -10.0, 10.0

# SNR threshold for "bad" predictions
LOW_SNR_THRESHOLD = 5.0  # dB

VALID_VARIANTS = ['0001', '0003', '0004', '0005', '0006', '0009', '0010', '0011', '0012']

VAL_SAMPLE_IDS = [
    '1006427285', '1006867983', '1012423188', '10140238', '1015663939',
    '102150619', '1026034238', '1041099777', '104573050', '1048962695',
    '1052007218', '1053922973', '1059602762', '1063816858',
    '1067371646', '1067975047', '1068062585', '1072767337',
    '1084993373', '108599929'
]

LEAD_LAYOUT = [
    ['I', 'aVR', 'V1', 'V4'],
    ['II', 'aVL', 'V2', 'V5'],
    ['III', 'aVF', 'V3', 'V6'],
]

# Hardcoded baseline offsets (mV) - computed from 977 images using V19 Augraphy epoch 34
# These are median(prediction - ground_truth) values
# Positive means model predicts slightly higher than GT
# Augraphy model has ~0.006 mV avg offset
BASELINE_OFFSETS = {
    'I': 0.0042, 'II': 0.0064, 'III': 0.0041,
    'aVR': 0.0061, 'aVL': 0.0060, 'aVF': 0.0059,
    'V1': 0.0061, 'V2': 0.0061, 'V3': 0.0060,
    'V4': 0.0062, 'V5': 0.0062, 'V6': 0.0061,
}

# Remote VM configuration
REMOTE_HOST = os.environ.get('REMOTE_TRAIN_HOST', 'azureuser@172.212.222.231')
REMOTE_CHECKPOINT_DIR = '/data/ecg-digitization/checkpoints'


# =============================================================================
# Model Architecture (from train_v19_enhanced.py)
# =============================================================================
class DeformableConvBlock(nn.Module):
    """Deformable convolution block."""
    def __init__(self, in_ch, out_ch, kernel_size=3, stride=1, padding=1):
        super().__init__()
        self.kernel_size = kernel_size
        self.padding = padding
        self.stride = stride
        
        if HAS_DEFORM_CONV:
            self.offset_conv = nn.Sequential(
                nn.Conv2d(in_ch, 64, 3, padding=1),
                nn.BatchNorm2d(64),
                nn.ReLU(inplace=True),
                nn.Conv2d(64, 2 * kernel_size * kernel_size, 3, padding=1),
            )
            nn.init.zeros_(self.offset_conv[-1].weight)
            nn.init.zeros_(self.offset_conv[-1].bias)
            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:
            offset = self.offset_conv(x)
            out = self.deform_conv(x, offset)
        else:
            out = self.conv(x)
        out = self.norm(out)
        out = self.act(out)
        return out


class BiLSTMHead(nn.Module):
    """Bidirectional LSTM for temporal modeling."""
    def __init__(self, input_dim, hidden_dim=128, num_layers=2, dropout=0.1):
        super().__init__()
        self.lstm = nn.LSTM(
            input_size=input_dim,
            hidden_size=hidden_dim,
            num_layers=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)  # [B, W, C]
        lstm_out, _ = self.lstm(x)
        out = self.output_proj(lstm_out)
        return out


class AuxiliaryHeads(nn.Module):
    """Auxiliary prediction heads."""
    def __init__(self, feature_dim):
        super().__init__()
        self.grid_head = nn.Sequential(
            nn.Conv2d(32, 16, 3, padding=1),
            nn.BatchNorm2d(16),
            nn.ReLU(inplace=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, features_2d, features_1d):
        grid_pred = self.grid_head(features_2d)
        gradient_pred = self.gradient_head(features_1d).squeeze(-1)
        log_var = self.uncertainty_head(features_1d).squeeze(-1)
        return grid_pred, gradient_pred, log_var


class CoordConv2d(nn.Module):
    """Conv2d with coordinate channels."""
    def __init__(self, in_channels, out_channels, kernel_size, **kwargs):
        super().__init__()
        self.conv = nn.Conv2d(in_channels + 2, out_channels, 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)
        x = torch.cat([x, yy, xx], dim=1)
        return self.conv(x)


class UNetDecoderBlockV19(nn.Module):
    """U-Net decoder block with optional deformable convolution."""
    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)
        x = self.conv1(x)
        x = self.conv2(x)
        return x


class PerLeadNetV19(nn.Module):
    """V19 Per-Lead ECG Network with BiLSTM + Deformable Conv."""
    
    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]
        skip_channels = enc_channels[:-1][::-1] + [0]
        
        for i, (skip_ch, out_ch) in enumerate(zip(skip_channels, decoder_dims)):
            use_deform = (i >= 2)
            self.dec_blocks.append(UNetDecoderBlockV19(in_ch, skip_ch, out_ch, use_deform))
            in_ch = out_ch
        
        self.final_up = nn.Sequential(
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
            nn.Conv2d(decoder_dims[-1], decoder_dims[-1], 3, padding=1, bias=False),
            nn.BatchNorm2d(decoder_dims[-1]),
            nn.GELU(),
        )
        
        self.height_attention = nn.Sequential(
            CoordConv2d(decoder_dims[-1], 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.GELU(),
            nn.Conv2d(64, 1, 1),
        )
        
        self.bilstm = BiLSTMHead(
            input_dim=decoder_dims[-1],
            hidden_dim=128,
            num_layers=2,
            dropout=0.1,
        )
        
        self.regression_head = nn.Sequential(
            nn.Linear(self.bilstm.output_dim, 64),
            nn.GELU(),
            nn.Linear(64, 1),
            nn.Sigmoid(),
        )
        
        self.aux_heads = AuxiliaryHeads(self.bilstm.output_dim)
    
    def forward(self, x, return_aux=False):
        B, C, H, W = x.shape
        
        features = self.encoder(x)
        
        d = features[-1]
        skips = features[:-1][::-1] + [None]
        
        for block, skip in zip(self.dec_blocks, skips):
            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 = self.height_attention(d)
        attn = F.softmax(attn, 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:
            grid_pred, gradient_pred, log_var = self.aux_heads(features_2d, temporal_features)
            return y_pred, {'grid': grid_pred, 'gradient': gradient_pred, 'log_var': log_var}
        
        return y_pred


# =============================================================================
# Post-Processing Functions
# =============================================================================
def apply_savgol_smoothing(signal_mv, window=7, polyorder=2):
    """Apply Savitzky-Golay smoothing."""
    if len(signal_mv) >= window:
        return savgol_filter(signal_mv, window_length=window, polyorder=polyorder)
    return signal_mv


def apply_einthoven_correction(pred_mv_rows, alpha=0.33):
    """Apply Einthoven's law correction."""
    segment_width = len(pred_mv_rows[0]) // 4
    
    lead_I = pred_mv_rows[0][:segment_width].copy()
    lead_II_short = pred_mv_rows[1][:segment_width].copy()
    lead_III = pred_mv_rows[2][:segment_width].copy()
    
    derived_II = lead_I + lead_III
    error = lead_II_short - derived_II
    
    lead_I_corrected = lead_I + alpha * error
    lead_III_corrected = lead_III + alpha * error
    
    pred_mv_rows[0][:segment_width] = lead_I_corrected
    pred_mv_rows[2][:segment_width] = lead_III_corrected
    
    return pred_mv_rows


def clamp_ecg_amplitude(signal_mv):
    """Clamp signal to reasonable ECG range."""
    return np.clip(signal_mv, ECG_MV_MIN, ECG_MV_MAX)


def interpolate_nan(signal_1d):
    """Interpolate NaN values."""
    valid_mask = np.isfinite(signal_1d)
    if valid_mask.all():
        return signal_1d
    if not valid_mask.any():
        return np.zeros_like(signal_1d)
    x = np.arange(len(signal_1d))
    signal_1d[~valid_mask] = np.interp(x[~valid_mask], x[valid_mask], signal_1d[valid_mask])
    return signal_1d


def apply_hardcoded_baseline_correction(pred_mv_rows):
    """
    Apply hardcoded baseline correction for each lead.
    
    Uses pre-computed median offsets from 977 images. These are tiny
    (~0.007 mV) so the effect is minimal, but included for completeness.
    
    This is production-ready - no GT required.
    """
    segment_width = len(pred_mv_rows[0]) // 4
    
    for row_idx in range(3):  # Only for rows 0-2 (not rhythm strip)
        lead_names = LEAD_LAYOUT[row_idx]
        for seg_idx, lead_name in enumerate(lead_names):
            offset = BASELINE_OFFSETS.get(lead_name, 0.0)
            
            seg_start = seg_idx * segment_width
            seg_end = (seg_idx + 1) * segment_width
            pred_mv_rows[row_idx][seg_start:seg_end] -= offset
    
    # Rhythm strip uses Lead II offset
    pred_mv_rows[3] -= BASELINE_OFFSETS.get('II', 0.0)
    
    return pred_mv_rows


# =============================================================================
# SCP Checkpoint Helper
# =============================================================================
def scp_checkpoint(remote_host, remote_path, local_path):
    """SCP a checkpoint from remote to local."""
    local_path = Path(local_path)
    local_path.parent.mkdir(parents=True, exist_ok=True)
    
    cmd = ['scp', '-o', 'StrictHostKeyChecking=no', 
           f'{remote_host}:{remote_path}', str(local_path)]
    print(f"  Downloading: {remote_path}")
    result = subprocess.run(cmd, capture_output=True, text=True)
    
    if result.returncode != 0:
        print(f"  SCP failed: {result.stderr}")
        return None
    return str(local_path)


# =============================================================================
# Inference Functions
# =============================================================================
def load_model(checkpoint_path, device):
    """Load V19 model."""
    model = PerLeadNetV19(encoder_name='convnext_base.fb_in22k_ft_in1k', pretrained=False)
    
    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
    model.load_state_dict(checkpoint['model'])
    model = model.to(device)
    model.eval()
    
    print(f"Loaded V19 from epoch {checkpoint.get('epoch', 'N/A')}")
    print(f"  SNR: {checkpoint.get('snr', checkpoint.get('best_snr', 'N/A')):.2f} dB")
    
    return model


def crop_row(image, row_idx):
    """Crop a single row centered on its baseline."""
    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[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')
    
    return row_crop


def predict_row(model, row_crop, device):
    """Run inference on a single row crop."""
    image_tensor = torch.from_numpy(row_crop.astype(np.float32) / 255.0)
    image_tensor = image_tensor.permute(2, 0, 1).unsqueeze(0).to(device)
    
    with torch.no_grad():
        with torch.amp.autocast('cuda:1'):
            output = model(image_tensor, return_aux=False)
    
    pred_y_crop = output[0].cpu().numpy() * ROW_HEIGHT
    return pred_y_crop


def convert_crop_to_full(pred_y_crop, row_idx):
    """Convert crop-relative y-coordinates to full image coordinates."""
    baseline_y = int(ZERO_MV[row_idx])
    y_start = max(0, baseline_y - CROP_HALF_HEIGHT)
    pad_top = max(0, CROP_HALF_HEIGHT - baseline_y)
    pred_y_full = pred_y_crop - pad_top + y_start
    return pred_y_full


def visualize_predictions(image, predictions, output_path):
    """Draw predictions on the image."""
    vis_image = image[:, T0:T1, :].copy()
    colors = [(255, 0, 0), (0, 255, 0), (255, 0, 255), (0, 165, 255)]
    
    for row_idx, pred_y in enumerate(predictions):
        color = colors[row_idx]
        for x in range(len(pred_y)):
            y = int(round(pred_y[x]))
            if 0 <= y < vis_image.shape[0]:
                cv2.circle(vis_image, (x, y), 1, color, -1)
    
    cv2.imwrite(str(output_path), vis_image)


def compute_snr_per_lead(pred_y_full, df, row_idx, epsilon=1e-10):
    """Compute SNR in dB for each lead in a row."""
    baseline_y = ZERO_MV[row_idx]
    segment_width = OUTPUT_WIDTH // 4
    lead_snrs = {}
    
    if row_idx < 3:
        lead_names = LEAD_LAYOUT[row_idx]
        for seg_idx, lead_name in enumerate(lead_names):
            if lead_name not in df.columns:
                lead_snrs[lead_name] = None
                continue
            
            gt_mv = df[lead_name].dropna().values
            if len(gt_mv) == 0:
                lead_snrs[lead_name] = None
                continue
            
            # Special handling for Lead II (10s data in GT, but short strip shows 2.5s)
            # Use Lead I length as reference to ensure exact alignment
            if lead_name == 'II':
                ref_lead = 'I' if 'I' in df.columns else 'III'
                if ref_lead in df.columns:
                    ref_len = len(df[ref_lead].dropna().values)
                    if len(gt_mv) > ref_len:
                        gt_mv = gt_mv[:ref_len]
            
            seg_start = seg_idx * segment_width
            seg_end = (seg_idx + 1) * segment_width
            pred_segment = pred_y_full[seg_start:seg_end]
            pred_mv = (baseline_y - pred_segment) / MV_TO_PIXEL
            
            x_pred = np.linspace(0, 1, len(pred_mv))
            x_gt = np.linspace(0, 1, len(gt_mv))
            pred_mv_resampled = np.interp(x_gt, x_pred, pred_mv)
            
            signal_power = (gt_mv ** 2).mean()
            noise_power = ((pred_mv_resampled - gt_mv) ** 2).mean()
            
            if noise_power < epsilon:
                lead_snrs[lead_name] = 50.0
            else:
                lead_snrs[lead_name] = 10 * np.log10(signal_power / noise_power)
    else:
        if 'II' not in df.columns:
            lead_snrs['II_rhythm'] = None
            return lead_snrs
        
        gt_mv = df['II'].dropna().values
        if len(gt_mv) == 0:
            lead_snrs['II_rhythm'] = None
            return lead_snrs
        
        pred_mv = (baseline_y - pred_y_full) / MV_TO_PIXEL
        x_pred = np.linspace(0, 1, len(pred_mv))
        x_gt = np.linspace(0, 1, len(gt_mv))
        pred_mv_resampled = np.interp(x_gt, x_pred, pred_mv)
        
        signal_power = (gt_mv ** 2).mean()
        noise_power = ((pred_mv_resampled - gt_mv) ** 2).mean()
        
        if noise_power < epsilon:
            lead_snrs['II_rhythm'] = 50.0
        else:
            lead_snrs['II_rhythm'] = 10 * np.log10(signal_power / noise_power)
    
    return lead_snrs


def process_image(model, image_path, csv_path, output_dir, device,
                  apply_smoothing=True, apply_einthoven=True, apply_baseline_fix=True,
                  negative_dir=None):
    """Process a single image and compute SNR."""
    image = cv2.imread(str(image_path), cv2.IMREAD_COLOR)
    if image is None:
        print(f"Failed to load: {image_path}")
        return None, None, None
    
    image = image[Y0:Y1, X0:X1]
    image = cv2.resize(image, (TARGET_WIDTH, TARGET_HEIGHT), interpolation=cv2.INTER_LINEAR)
    
    df = pd.read_csv(csv_path)
    
    predictions = []
    pred_mv_rows = {}
    
    for row_idx in range(4):
        row_crop = crop_row(image, row_idx)
        pred_y_crop = predict_row(model, row_crop, device)
        pred_y_full = convert_crop_to_full(pred_y_crop, row_idx)
        
        baseline_y = ZERO_MV[row_idx]
        pred_mv = (baseline_y - pred_y_full) / MV_TO_PIXEL
        
        if apply_smoothing:
            pred_mv = apply_savgol_smoothing(pred_mv, window=7, polyorder=2)
        
        pred_mv = clamp_ecg_amplitude(pred_mv)
        pred_mv = interpolate_nan(pred_mv.copy())
        
        pred_mv_rows[row_idx] = pred_mv
        predictions.append(pred_y_full)
    
    if apply_einthoven:
        pred_mv_rows = apply_einthoven_correction(pred_mv_rows, alpha=0.33)
    
    # Apply hardcoded baseline correction (tiny effect, ~0.007 mV)
    if apply_baseline_fix:
        pred_mv_rows = apply_hardcoded_baseline_correction(pred_mv_rows)
    
    corrected_predictions = []
    for row_idx in range(4):
        baseline_y = ZERO_MV[row_idx]
        pred_y_corrected = baseline_y - pred_mv_rows[row_idx] * MV_TO_PIXEL
        corrected_predictions.append(pred_y_corrected)
    
    all_lead_snrs = {}
    for row_idx in range(4):
        lead_snrs = compute_snr_per_lead(corrected_predictions[row_idx], df, row_idx)
        all_lead_snrs.update(lead_snrs)
    
    valid_snrs = [v for v in all_lead_snrs.values() if v is not None]
    min_snr = min(valid_snrs) if valid_snrs else 0.0
    
    sample_id = image_path.parent.name
    variant = image_path.stem.split('-')[-1] if '-' in image_path.stem else '0000'
    output_path = output_dir / f"{sample_id}_{variant}.png"
    
    visualize_predictions(image, corrected_predictions, output_path)
    
    if negative_dir is not None and min_snr < LOW_SNR_THRESHOLD:
        neg_output_path = negative_dir / f"{sample_id}_{variant}_snr{min_snr:.1f}.png"
        visualize_predictions(image, corrected_predictions, neg_output_path)
    
    return all_lead_snrs, min_snr, output_path


# Default augraphy checkpoint
DEFAULT_AUGRAPHY_CHECKPOINT = '/data/ecg-digitization/checkpoints/v19_augraphy_latest.pth'


def main():
    parser = argparse.ArgumentParser(description='V19 Augraphy Inference (Layer B)')
    parser.add_argument('--checkpoint', type=str, default=DEFAULT_AUGRAPHY_CHECKPOINT,
                        help='Checkpoint path (default: v19_augraphy_latest.pth)')
    parser.add_argument('--kaggle_data', type=str, 
                        default='/data/ecg-digitization/stage1_data/train')
    parser.add_argument('--output_dir', type=str, 
                        default=os.path.expanduser('~/tmp/pred/v19_augraphy'))
    parser.add_argument('--num_samples', type=int, default=None)
    parser.add_argument('--use_holdout', action='store_true', default=True)
    parser.add_argument('--no_smoothing', action='store_true')
    parser.add_argument('--no_einthoven', action='store_true')
    parser.add_argument('--no_baseline_fix', action='store_true',
                        help='Disable hardcoded baseline correction')
    parser.add_argument('--seed', type=int, default=42)
    args = parser.parse_args()
    
    random.seed(args.seed)
    device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')
    output_dir = Path(args.output_dir)
    negative_dir = Path(os.path.expanduser('~/tmp/pred/v19_augraphy/negpreds'))
    
    print(f"{'='*70}")
    print(f"V19 Augraphy Inference (Layer B)")
    print(f"{'='*70}")
    
    # Get checkpoint - use default augraphy checkpoint
    checkpoint = args.checkpoint
    
    if not Path(checkpoint).exists():
        print(f"ERROR: Checkpoint not found: {checkpoint}")
        print("Make sure augraphy training has saved a checkpoint.")
        sys.exit(1)
    
    if output_dir.exists():
        import shutil
        shutil.rmtree(output_dir)
        print(f"Deleted old predictions in {output_dir}")
    output_dir.mkdir(parents=True, exist_ok=True)
    negative_dir.mkdir(parents=True, exist_ok=True)
    
    print(f"\nCheckpoint: {checkpoint}")
    print(f"Output: {output_dir}")
    print(f"Device: {device}")
    print(f"Smoothing: {'OFF' if args.no_smoothing else 'ON (Savgol w=7)'}")
    print(f"Einthoven correction: {'OFF' if args.no_einthoven else 'ON (alpha=0.33)'}")
    print(f"Baseline correction: {'OFF' if args.no_baseline_fix else 'ON (hardcoded offsets)'}")
    print(f"{'='*70}")
    
    # Load model
    model = load_model(checkpoint, device)
    
    # Find images
    kaggle_dir = Path(args.kaggle_data)
    val_sample_set = set(VAL_SAMPLE_IDS)
    all_samples = []
    
    for sample_dir in kaggle_dir.iterdir():
        if not sample_dir.is_dir():
            continue
        if args.use_holdout and sample_dir.name not in val_sample_set:
            continue
        
        csv_files = list(sample_dir.glob('*.csv'))
        if len(csv_files) != 1:
            continue
        csv_path = csv_files[0]
        
        for img_path in sample_dir.glob('*.png'):
            variant = img_path.stem.split('-')[-1] if '-' in img_path.stem else '0000'
            if variant in VALID_VARIANTS:
                all_samples.append((img_path, csv_path))
    
    print(f"Found {len(all_samples)} valid images in holdout set")
    
    if args.num_samples is not None and len(all_samples) > args.num_samples:
        selected = random.sample(all_samples, args.num_samples)
    else:
        selected = all_samples
    
    print(f"Processing {len(selected)} images...")
    
    # Run inference TWICE: once without baseline fix, once with
    all_leads = ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6', 'II_rhythm']
    
    results_no_offset = {lead: [] for lead in all_leads}
    results_with_offset = {lead: [] for lead in all_leads}
    
    print(f"\n[Pass 1/2] Computing SNR WITHOUT baseline offsets...")
    for img_path, csv_path in tqdm(selected, desc="No offsets"):
        try:
            lead_snrs, _, _ = process_image(
                model, img_path, csv_path, output_dir, device,
                apply_smoothing=not args.no_smoothing,
                apply_einthoven=not args.no_einthoven,
                apply_baseline_fix=False,  # NO baseline fix
                negative_dir=None
            )
            if lead_snrs:
                for lead, snr in lead_snrs.items():
                    if snr is not None:
                        results_no_offset[lead].append(snr)
        except Exception as e:
            print(f"Error processing {img_path}: {e}")
    
    print(f"\n[Pass 2/2] Computing SNR WITH baseline offsets...")
    for img_path, csv_path in tqdm(selected, desc="With offsets"):
        try:
            lead_snrs, min_snr, _ = process_image(
                model, img_path, csv_path, output_dir, device,
                apply_smoothing=not args.no_smoothing,
                apply_einthoven=not args.no_einthoven,
                apply_baseline_fix=True,  # WITH baseline fix
                negative_dir=negative_dir
            )
            if lead_snrs:
                for lead, snr in lead_snrs.items():
                    if snr is not None:
                        results_with_offset[lead].append(snr)
        except Exception as e:
            print(f"Error processing {img_path}: {e}")
    
    # Print comparison table
    print(f"\n{'='*90}")
    print(f"Per-Lead SNR Comparison: WITHOUT vs WITH Baseline Offsets (dB)")
    print(f"{'='*90}")
    print(f"{'Lead':<12} {'No Offset':>12} {'With Offset':>12} {'Δ':>8} {'Count':>6}")
    print(f"{'-'*90}")
    
    total_no_offset = []
    total_with_offset = []
    
    for lead in all_leads:
        snrs_no = results_no_offset[lead]
        snrs_with = results_with_offset[lead]
        if len(snrs_no) > 0 and len(snrs_with) > 0:
            mean_no = np.mean(snrs_no)
            mean_with = np.mean(snrs_with)
            delta = mean_with - mean_no
            print(f"{lead:<12} {mean_no:>12.2f} {mean_with:>12.2f} {delta:>+8.3f} {len(snrs_no):>6}")
            total_no_offset.extend(snrs_no)
            total_with_offset.extend(snrs_with)
    
    print(f"{'-'*90}")
    if len(total_no_offset) > 0:
        overall_no = np.mean(total_no_offset)
        overall_with = np.mean(total_with_offset)
        overall_delta = overall_with - overall_no
        print(f"{'OVERALL':<12} {overall_no:>12.2f} {overall_with:>12.2f} {overall_delta:>+8.3f} {len(total_no_offset):>6}")
    print(f"{'='*90}")
    
    # Also print detailed stats for each mode
    for mode_name, mode_results in [("WITHOUT Baseline Offsets", results_no_offset), 
                                     ("WITH Baseline Offsets", results_with_offset)]:
        print(f"\n{'='*70}")
        print(f"Detailed Stats: {mode_name}")
        print(f"{'='*70}")
        print(f"{'Lead':<12} {'Mean':>8} {'Std':>8} {'Min':>8} {'Max':>8}")
        print(f"{'-'*70}")
        for lead in all_leads:
            snrs = mode_results[lead]
            if len(snrs) > 0:
                print(f"{lead:<12} {np.mean(snrs):>8.2f} {np.std(snrs):>8.2f} "
                      f"{np.min(snrs):>8.2f} {np.max(snrs):>8.2f}")
    
    print(f"\nDone! Saved {len(selected)} visualizations to {output_dir}")


if __name__ == '__main__':
    main()