code/scripts/infer_v16_visualize_tta.py

#!/usr/bin/env python3
"""
V16 Inference with TTA and Visualization

Same as infer_v16_visualize.py but with Test-Time Augmentation (TTA).
TTA averages predictions from multiple augmented versions of each row crop.

Usage:
  python infer_v16_visualize_tta.py --use_holdout
  python infer_v16_visualize_tta.py --no_tta  # Disable TTA for comparison
"""

import os
import sys
import argparse
import random
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


# =============================================================================
# Constants (from train_v16_perlead.py)
# =============================================================================
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 - V16 input width

# Per-row crop parameters (V16 uses larger crop)
CROP_HALF_HEIGHT = 250
ROW_HEIGHT = 500

# ECG amplitude limits (mV) - values beyond this are likely errors
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 - same as training)
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 (same as training)
LEAD_LAYOUT = [
    ['I', 'aVR', 'V1', 'V4'],
    ['II', 'aVL', 'V2', 'V5'],
    ['III', 'aVF', 'V3', 'V6'],
]


# =============================================================================
# TTA Configuration
# =============================================================================
# Horizontal flip is the most effective augmentation for this task
# It averages out directional biases in the model
TTA_AUGMENTATIONS = [
    {'name': 'original', 'hflip': False, 'brightness': 1.0, 'contrast': 1.0},
    {'name': 'hflip', 'hflip': True, 'brightness': 1.0, 'contrast': 1.0},
    # Optional: add brightness/contrast variations (less impactful than hflip)
    # {'name': 'bright+', 'hflip': False, 'brightness': 1.08, 'contrast': 1.0},
    # {'name': 'bright-', 'hflip': False, 'brightness': 0.92, 'contrast': 1.0},
]


def apply_tta_augmentation(image, hflip=False, brightness=1.0, contrast=1.0):
    """Apply augmentation to image.
    
    Args:
        image: BGR image [H, W, 3]
        hflip: Horizontal flip
        brightness: Multiplicative brightness factor
        contrast: Multiplicative contrast factor  
    
    Returns:
        Augmented image
    """
    img = image.copy()
    
    # Horizontal flip
    if hflip:
        img = np.flip(img, axis=1).copy()
    
    # Brightness/contrast (optional, less impactful)
    if brightness != 1.0 or contrast != 1.0:
        img = img.astype(np.float32)
        if brightness != 1.0:
            img = img * brightness
        if contrast != 1.0:
            mean = img.mean()
            img = (img - mean) * contrast + mean
        img = np.clip(img, 0, 255).astype(np.uint8)
    
    return img


# =============================================================================
# Model Architecture (copy from train_v16_perlead.py)
# =============================================================================
class CoordConv2d(nn.Module):
    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 UNetDecoderBlock(nn.Module):
    def __init__(self, in_ch, skip_ch, out_ch):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch + skip_ch, out_ch, 3, padding=1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.GELU(),
            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.conv(x)


class PerLeadNet(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]
        skip_channels = enc_channels[:-1][::-1] + [0]
        
        for skip_ch, out_ch in zip(skip_channels, decoder_dims):
            self.dec_blocks.append(UNetDecoderBlock(in_ch, skip_ch, out_ch))
            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.regression_head = nn.Sequential(
            nn.Conv1d(decoder_dims[-1], 128, 7, padding=3),
            nn.BatchNorm1d(128),
            nn.GELU(),
            nn.Conv1d(128, 64, 5, padding=2),
            nn.BatchNorm1d(64),
            nn.GELU(),
            nn.Conv1d(64, 1, 1),
        )
    
    def forward(self, x):
        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)
        
        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)
        
        d = (d * attn).sum(dim=2)
        
        out = self.regression_head(d)
        out = torch.sigmoid(out)
        
        return out.squeeze(1)


# =============================================================================
# Inference Enhancement Functions
# =============================================================================
def apply_savgol_smoothing(signal_mv, window=7, polyorder=2):
    """Apply Savitzky-Golay smoothing to remove high-frequency noise."""
    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 on short lead segments."""
    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 from valid neighbors."""
    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


# =============================================================================
# Inference and Visualization
# =============================================================================
def load_model(checkpoint_path, device):
    """Load trained V16 model."""
    model = PerLeadNet(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 checkpoint from epoch {checkpoint['epoch']}")
    print(f"  SNR: {checkpoint.get('snr', 'N/A'):.2f} dB")
    print(f"  MAE: {checkpoint.get('mae', 'N/A'):.2f} px")
    
    return model


def crop_row(image, row_idx):
    """Crop a single row centered on its baseline, signal region only (T0:T1)."""
    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 (no TTA)."""
    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.cuda.amp.autocast():
            output = model(image_tensor)
    
    pred_y_crop = output[0].cpu().numpy() * ROW_HEIGHT
    
    return pred_y_crop


def predict_row_tta(model, row_crop, device):
    """Run inference with TTA on a single row crop.
    
    For horizontal flip:
    1. Flip input image horizontally
    2. Run inference
    3. Reverse the x-axis of predictions (flip back)
    4. Average with original predictions
    
    This averages out directional biases in the model.
    """
    all_preds = []
    
    for aug in TTA_AUGMENTATIONS:
        # Apply augmentation
        aug_crop = apply_tta_augmentation(
            row_crop, 
            hflip=aug.get('hflip', False),
            brightness=aug.get('brightness', 1.0),
            contrast=aug.get('contrast', 1.0),
        )
        
        # Convert to tensor
        image_tensor = torch.from_numpy(aug_crop.astype(np.float32) / 255.0)
        image_tensor = image_tensor.permute(2, 0, 1).unsqueeze(0).to(device)
        
        with torch.no_grad():
            with torch.cuda.amp.autocast():
                output = model(image_tensor)
        
        pred_y_crop = output[0].cpu().numpy() * ROW_HEIGHT
        
        # If horizontally flipped, reverse the x-axis of predictions
        if aug.get('hflip', False):
            pred_y_crop = pred_y_crop[::-1].copy()
        
        all_preds.append(pred_y_crop)
    
    # Average all predictions
    avg_pred = np.mean(all_preds, axis=0)
    
    return avg_pred


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 as dots on the image."""
    vis_image = image[:, T0:T1, :].copy()
    
    colors = [
        (255, 0, 0),    # Blue for row 0
        (0, 255, 0),    # Green for row 1
        (255, 0, 255),  # Magenta for row 2
        (0, 165, 255),  # Orange for row 3 (rhythm)
    ]
    
    for row_idx, pred_y in enumerate(predictions):
        color = colors[row_idx]
        
        for x in range(len(pred_y)):
            y = int(np.clip(pred_y[x], 0, TARGET_HEIGHT - 1))
            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
            
            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:
                    quarter_len = len(gt_mv) // 4
                    gt_mv = gt_mv[:quarter_len]
            
            seg_start = seg_idx * segment_width
            seg_end = (seg_idx + 1) * segment_width
            pred_y_seg = pred_y_full[seg_start:seg_end]
            
            pred_mv_pixels = (baseline_y - pred_y_seg) / MV_TO_PIXEL
            
            x_pred = np.linspace(0, 1, len(pred_mv_pixels))
            x_gt = np.linspace(0, 1, len(gt_mv))
            pred_mv_resampled = np.interp(x_gt, x_pred, pred_mv_pixels)
            
            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:
                snr = 10 * np.log10(signal_power / (noise_power + epsilon))
                lead_snrs[lead_name] = float(snr)
    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_pixels = (baseline_y - pred_y_full) / MV_TO_PIXEL
        
        x_pred = np.linspace(0, 1, len(pred_mv_pixels))
        x_gt = np.linspace(0, 1, len(gt_mv))
        pred_mv_resampled = np.interp(x_gt, x_pred, pred_mv_pixels)
        
        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:
            snr = 10 * np.log10(signal_power / (noise_power + epsilon))
            lead_snrs['II_rhythm'] = float(snr)
    
    return lead_snrs


def process_image(model, image_path, csv_path, output_dir, device, 
                  apply_smoothing=True, apply_einthoven=True, use_tta=True, negative_dir=None):
    """Process a single image with optional TTA."""
    # Load image
    image = cv2.imread(str(image_path), cv2.IMREAD_COLOR)
    if image is None:
        print(f"Failed to load: {image_path}")
        return None, None, None
    
    # Preprocess
    image = image[Y0:Y1, X0:X1]
    image = cv2.resize(image, (TARGET_WIDTH, TARGET_HEIGHT), interpolation=cv2.INTER_LINEAR)
    
    # Load ground truth
    df = pd.read_csv(csv_path)
    
    # Predict each row
    predictions = []
    pred_mv_rows = {}
    
    for row_idx in range(4):
        row_crop = crop_row(image, row_idx)
        
        # Use TTA or regular prediction
        if use_tta:
            pred_y_crop = predict_row_tta(model, row_crop, device)
        else:
            pred_y_crop = predict_row(model, row_crop, device)
        
        pred_y_full = convert_crop_to_full(pred_y_crop, row_idx)
        
        # Convert to mV for post-processing
        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)
    
    # Apply Einthoven correction
    if apply_einthoven:
        pred_mv_rows = apply_einthoven_correction(pred_mv_rows, alpha=0.33)
    
    # Convert corrected mV back to pixel coordinates for visualization
    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)
    
    # Compute SNR using corrected predictions
    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)
    
    # Find minimum SNR for this image
    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
    
    # Visualize
    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)
    
    # Save to negative folder if SNR is low
    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='/data/ecg-digitization/checkpoints/v16_perlead_best_snr.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/v16_tta'))
    parser.add_argument('--num_samples', type=int, default=None,
                        help='Number of samples (default: all holdout samples)')
    parser.add_argument('--use_holdout', action='store_true', default=True,
                        help='Use holdout/validation set instead of random samples')
    parser.add_argument('--no_tta', action='store_true',
                        help='Disable Test-Time Augmentation')
    parser.add_argument('--no_smoothing', action='store_true',
                        help='Disable Savitzky-Golay smoothing')
    parser.add_argument('--no_einthoven', action='store_true',
                        help='Disable Einthoven law correction')
    parser.add_argument('--seed', type=int, default=42)
    args = parser.parse_args()
    
    # Setup
    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/v16_tta/negpreds'))
    
    # Delete old predictions
    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)
    
    use_tta = not args.no_tta
    
    print(f"{'='*70}")
    print(f"V16 Inference with TTA")
    print(f"{'='*70}")
    print(f"Checkpoint: {args.checkpoint}")
    print(f"Output: {output_dir}")
    print(f"Device: {device}")
    print(f"TTA: {'ON (' + str(len(TTA_AUGMENTATIONS)) + ' augmentations)' if use_tta else 'OFF'}")
    if use_tta:
        print(f"  Augmentations: {[a['name'] for a in TTA_AUGMENTATIONS]}")
    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"Low SNR threshold: {LOW_SNR_THRESHOLD} dB")
    print(f"{'='*70}")
    
    # Load model
    model = load_model(args.checkpoint, device)
    
    # Find all valid images with their CSV files
    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))
    
    set_type = "holdout" if args.use_holdout else "all"
    print(f"Found {len(all_samples)} valid images in {set_type} set")
    
    # Select samples
    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...")
    
    # Collect per-lead SNRs
    all_snrs = {lead: [] for lead in ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', 
                                       'V1', 'V2', 'V3', 'V4', 'V5', 'V6', 'II_rhythm']}
    low_snr_samples = []
    
    # Process each image
    for img_path, csv_path in tqdm(selected):
        try:
            lead_snrs, min_snr, output_path = process_image(
                model, img_path, csv_path, output_dir, device,
                apply_smoothing=not args.no_smoothing,
                apply_einthoven=not args.no_einthoven,
                use_tta=use_tta,
                negative_dir=negative_dir
            )
            if lead_snrs:
                for lead, snr in lead_snrs.items():
                    if snr is not None:
                        all_snrs[lead].append(snr)
                
                if min_snr is not None and min_snr < LOW_SNR_THRESHOLD:
                    worst_lead = min(lead_snrs, key=lambda k: lead_snrs[k] if lead_snrs[k] is not None else float('inf'))
                    low_snr_samples.append({
                        'sample': str(img_path.parent.name),
                        'variant': img_path.stem.split('-')[-1] if '-' in img_path.stem else '0000',
                        'min_snr': min_snr,
                        'worst_lead': worst_lead
                    })
        except Exception as e:
            print(f"Error processing {img_path}: {e}")
    
    # Print per-lead SNR statistics
    print(f"\n{'='*70}")
    print(f"Per-Lead SNR Statistics (dB) - TTA {'ON' if use_tta else 'OFF'}")
    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:
            mean_snr = np.mean(snrs)
            std_snr = np.std(snrs)
            min_snr = np.min(snrs)
            max_snr = np.max(snrs)
            print(f"{lead:<12} {mean_snr:>8.2f} {std_snr:>8.2f} {min_snr:>8.2f} {max_snr:>8.2f} {len(snrs):>6}")
            total_snrs.extend(snrs)
        else:
            print(f"{lead:<12} {'N/A':>8} {'N/A':>8} {'N/A':>8} {'N/A':>8} {0:>6}")
    
    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}")
    
    # Save SNR report
    snr_report_path = output_dir / 'snr_report.csv'
    with open(snr_report_path, 'w') as f:
        f.write("Lead,Mean_SNR,Std_SNR,Min_SNR,Max_SNR,Count\n")
        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:
                f.write(f"{lead},{np.mean(snrs):.2f},{np.std(snrs):.2f},{np.min(snrs):.2f},{np.max(snrs):.2f},{len(snrs)}\n")
            else:
                f.write(f"{lead},N/A,N/A,N/A,N/A,0\n")
        if len(total_snrs) > 0:
            f.write(f"OVERALL,{np.mean(total_snrs):.2f},{np.std(total_snrs):.2f},{np.min(total_snrs):.2f},{np.max(total_snrs):.2f},{len(total_snrs)}\n")
    
    print(f"\nSNR report saved to: {snr_report_path}")
    print(f"Done! Saved {len(selected)} visualizations to {output_dir}")
    
    # Print low-SNR summary
    if low_snr_samples:
        print(f"\n{'='*70}")
        print(f"Low SNR Samples (< {LOW_SNR_THRESHOLD} dB)")
        print(f"{'='*70}")
        low_snr_samples.sort(key=lambda x: x['min_snr'])
        
        for item in low_snr_samples[:20]:
            print(f"{item['sample']:<15} {item['variant']:<8} {item['min_snr']:>10.2f} {item['worst_lead']:<12}")
        
        if len(low_snr_samples) > 20:
            print(f"... and {len(low_snr_samples) - 20} more")
        
        print(f"\nTotal low-SNR samples: {len(low_snr_samples)}")
    
    print(f"{'='*70}")


if __name__ == '__main__':
    main()