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