code/scripts/infer_v19_visualize.py

#!/usr/bin/env python3
"""
V19 Inference with Visualization

Runs V19 per-lead model (BiLSTM + Deformable Conv) on Kaggle images.
Automatically SCPs the latest checkpoint from remote training VM.
"""

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']

# Validation sample IDs (holdout set)
# VAL_SAMPLE_IDS = [
#     '1006427285', '1006867983', '1012423188', '10140238', '1015663939',
#     '102150619', '1026034238', '1041099777', '104573050', '1048962695',
#     '1052007218', '1053922973', '1059602762', '1063816858', '106482869',
#     '1067371646', '1067975047', '1068062585', '1072767337', '1079294623',
#     '1084993373', '108599929'
# ]

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

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

# 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


# =============================================================================
# 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'):
            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)
            # Lead II short strip is in row 1, segment 0 (first column)
            # Training uses first 25% of GT (0-2.5s), so inference must match
            if lead_name == 'II':
                ref_len = len(df['I'].dropna().values) if 'I' in df.columns else len(gt_mv) // 4
                if len(gt_mv) > ref_len * 2:  # Lead II is 10s data
                    quarter_len = len(gt_mv) // 4
                    gt_mv = gt_mv[:quarter_len]  # First quarter (0-2.5s)
            
            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, 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)
    
    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


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--checkpoint', type=str, default=None,
                        help='Local checkpoint path (if not provided, will SCP from remote)')
    parser.add_argument('--remote_host', type=str, default=REMOTE_HOST,
                        help='Remote host for SCP')
    parser.add_argument('--remote_dir', type=str, default=REMOTE_CHECKPOINT_DIR,
                        help='Remote checkpoint directory')
    parser.add_argument('--local_cache', type=str, 
                        default='/tmp/v19_checkpoints',
                        help='Local cache for downloaded checkpoints')
    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'))
    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_scp', action='store_true',
                        help='Skip SCP, use local checkpoint only')
    parser.add_argument('--seed', type=int, default=42)
    args = parser.parse_args()
    
    random.seed(args.seed)
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    output_dir = Path(args.output_dir)
    negative_dir = Path(os.path.expanduser('~/tmp/pred/v19/negpreds'))
    
    print(f"{'='*70}")
    print(f"V19 Inference: BiLSTM + Deformable Conv")
    print(f"{'='*70}")
    
    # Get checkpoint
    checkpoint = args.checkpoint
    local_cache = Path(args.local_cache)
    
    if not args.no_scp and checkpoint is None:
        print(f"\nFetching checkpoint from {args.remote_host}...")
        local_cache.mkdir(parents=True, exist_ok=True)
        
        # Try latest first
        remote_latest = f"{args.remote_dir}/v19_enhanced_latest.pth"
        checkpoint = scp_checkpoint(args.remote_host, remote_latest, 
                                    local_cache / 'v19_enhanced_latest.pth')
        
        if checkpoint is None:
            print("Failed to download checkpoint. Use --checkpoint to specify local path.")
            sys.exit(1)
    elif checkpoint is None:
        # Look for local checkpoint
        checkpoint = str(local_cache / 'v19_enhanced_latest.pth')
        if not Path(checkpoint).exists():
            checkpoint = '/data/ecg-digitization/checkpoints/v19_enhanced_latest.pth'
    
    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"{'='*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...")
    
    all_snrs = {lead: [] for lead in ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', 
                                       'V1', 'V2', 'V3', 'V4', 'V5', 'V6', 'II_rhythm']}
    
    for img_path, csv_path in tqdm(selected):
        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,
                negative_dir=negative_dir
            )
            if lead_snrs:
                for lead, snr in lead_snrs.items():
                    if snr is not None:
                        all_snrs[lead].append(snr)
        except Exception as e:
            print(f"Error processing {img_path}: {e}")
    
    # Print results
    print(f"\n{'='*70}")
    print(f"Per-Lead SNR Statistics (dB)")
    print(f"{'='*70}")
    print(f"{'Lead':<12} {'Mean':>8} {'Std':>8} {'Min':>8} {'Max':>8} {'Count':>6}")
    print(f"{'-'*70}")
    
    total_snrs = []
    for lead in ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6', 'II_rhythm']:
        snrs = all_snrs[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} {len(snrs):>6}")
            total_snrs.extend(snrs)
    
    print(f"{'-'*70}")
    if len(total_snrs) > 0:
        print(f"{'OVERALL':<12} {np.mean(total_snrs):>8.2f} {np.std(total_snrs):>8.2f} "
              f"{np.min(total_snrs):>8.2f} {np.max(total_snrs):>8.2f} {len(total_snrs):>6}")
    print(f"{'='*70}")
    
    print(f"\nDone! Saved {len(selected)} visualizations to {output_dir}")


if __name__ == '__main__':
    main()