#!/usr/bin/env python3 """ V16.1 Inference with Visualization Runs V16.1 per-lead model (2x upscaled) on Kaggle images and plots predictions. Automatically SCPs the latest checkpoint from remote training VM before inference. Key differences from V16: - Input upscaled 2x: 1000 × 7852 (from 500 × 3926) - Lanczos4 interpolation for crisp lines - Predictions are downsampled back to base resolution for visualization/SNR """ 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 # ============================================================================= # Constants (from train_v16_1_perlead.py - 2x upscaled) # ============================================================================= 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_BASE = T1 - T0 # 3926 - base signal width # V16.1: 2x upscaling UPSCALE_FACTOR = 2 OUTPUT_WIDTH = OUTPUT_WIDTH_BASE * UPSCALE_FACTOR # 7852 # Per-row crop parameters (at base resolution, before upscale) CROP_HALF_HEIGHT_BASE = 250 ROW_HEIGHT_BASE = 500 # After 2x upscale CROP_HALF_HEIGHT = CROP_HALF_HEIGHT_BASE * UPSCALE_FACTOR # 500 ROW_HEIGHT = ROW_HEIGHT_BASE * UPSCALE_FACTOR # 1000 # Model input dimensions (after 2x upscale) INPUT_HEIGHT = ROW_HEIGHT # 1000 INPUT_WIDTH = OUTPUT_WIDTH # 7852 # MV_TO_PIXEL also scales 2x MV_TO_PIXEL_UPSCALED = MV_TO_PIXEL * UPSCALE_FACTOR # 157 # 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 LEAD_LAYOUT = [ ['I', 'aVR', 'V1', 'V4'], ['II', 'aVL', 'V2', 'V5'], ['III', 'aVF', 'V3', 'V6'], ] # Remote VM configuration for SCP REMOTE_HOST = os.environ.get('REMOTE_TRAIN_HOST', 'azureuser@172.212.222.231') REMOTE_CHECKPOINT_DIR = '/data/ecg-digitization/checkpoints' REMOTE_CHECKPOINT_PATTERN = 'v16_1_perlead_best_snr' # Fetch best SNR checkpoint # ============================================================================= # SCP Checkpoint from Remote VM # ============================================================================= def scp_latest_checkpoint(remote_host, remote_dir, pattern, local_dir): """ SCP the latest checkpoint matching pattern from remote VM. Args: remote_host: SSH host (can be alias from ~/.ssh/config or user@ip) remote_dir: Remote directory containing checkpoints pattern: Pattern to match checkpoint files local_dir: Local directory to save checkpoint Returns: Path to local checkpoint file, or None if failed """ local_dir = Path(local_dir) local_dir.mkdir(parents=True, exist_ok=True) print(f"{'='*70}") print(f"Fetching latest checkpoint from remote VM...") print(f" Remote: {remote_host}:{remote_dir}") print(f" Pattern: {pattern}") print(f"{'='*70}") # Find latest checkpoint on remote # List files matching pattern and sort by modification time find_cmd = f"ssh {remote_host} 'ls -t {remote_dir}/{pattern}*.pth 2>/dev/null | head -1'" try: result = subprocess.run(find_cmd, shell=True, capture_output=True, text=True, timeout=30) if result.returncode != 0 or not result.stdout.strip(): print(f"ERROR: Could not find checkpoint matching '{pattern}' on remote") print(f" stderr: {result.stderr}") return None remote_path = result.stdout.strip() filename = os.path.basename(remote_path) local_path = local_dir / filename print(f" Found: {remote_path}") # Check if we already have this file (by size comparison) if local_path.exists(): # Get remote file size size_cmd = f"ssh {remote_host} 'stat -c %s {remote_path}'" size_result = subprocess.run(size_cmd, shell=True, capture_output=True, text=True, timeout=10) if size_result.returncode == 0: remote_size = int(size_result.stdout.strip()) local_size = local_path.stat().st_size if remote_size == local_size: print(f" Local copy already up-to-date: {local_path}") return local_path # SCP the file print(f" Downloading to: {local_path}") scp_cmd = f"scp {remote_host}:{remote_path} {local_path}" result = subprocess.run(scp_cmd, shell=True, capture_output=True, text=True, timeout=300) if result.returncode != 0: print(f"ERROR: SCP failed") print(f" stderr: {result.stderr}") return None print(f" ✓ Downloaded successfully ({local_path.stat().st_size / 1024 / 1024:.1f} MB)") return local_path except subprocess.TimeoutExpired: print("ERROR: SSH/SCP command timed out") return None except Exception as e: print(f"ERROR: {e}") return None # ============================================================================= # Model Architecture (from train_v16_1_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): """ Per-lead regression network for 2x upscaled input. V16.1 Input: (B, 3, 1000, 7852) Output: (B, 7852) - y-coordinate for each x position [0, 1] """ 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. Einthoven's Law: II = I + III (in mV) """ 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.1 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 (base resolution)") return model def crop_row(image, row_idx): """ Crop a single row centered on its baseline, signal region only (T0:T1). Returns 2x upscaled crop using Lanczos interpolation. """ baseline_y = int(ZERO_MV[row_idx]) y_start = max(0, baseline_y - CROP_HALF_HEIGHT_BASE) y_end = min(TARGET_HEIGHT, baseline_y + CROP_HALF_HEIGHT_BASE) # Crop x to signal region (T0:T1) at base resolution row_crop = image[y_start:y_end, T0:T1, :].copy() # Pad if necessary (at base resolution) if row_crop.shape[0] < ROW_HEIGHT_BASE: pad_top = max(0, CROP_HALF_HEIGHT_BASE - baseline_y) pad_bottom = max(0, (baseline_y + CROP_HALF_HEIGHT_BASE) - TARGET_HEIGHT) row_crop = np.pad(row_crop, ((pad_top, pad_bottom), (0, 0), (0, 0)), mode='edge') # 2x Upscale with Lanczos for crisp lines row_crop = cv2.resize(row_crop, (OUTPUT_WIDTH, ROW_HEIGHT), interpolation=cv2.INTER_LANCZOS4) return row_crop def predict_row(model, row_crop, device): """Run inference on a single 2x upscaled row crop.""" # Prepare input 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) # Convert from normalized [0, 1] to 2x crop-relative pixels pred_y_crop_2x = output[0].cpu().numpy() * ROW_HEIGHT return pred_y_crop_2x def convert_crop_to_full(pred_y_crop_2x, row_idx): """ Convert 2x crop-relative y-coordinates to base resolution full image coordinates. pred_y_crop_2x: [OUTPUT_WIDTH] predictions at 2x resolution Returns: [OUTPUT_WIDTH_BASE] predictions in full image y-coords (base resolution) """ baseline_y = int(ZERO_MV[row_idx]) y_start = max(0, baseline_y - CROP_HALF_HEIGHT_BASE) # Adjust for padding pad_top = max(0, CROP_HALF_HEIGHT_BASE - baseline_y) # Convert from 2x crop-relative to 2x full image y pred_y_full_2x = (pred_y_crop_2x - pad_top * UPSCALE_FACTOR) / UPSCALE_FACTOR + y_start # Downsample from 2x resolution to base resolution x_2x = np.linspace(0, 1, len(pred_y_full_2x)) x_base = np.linspace(0, 1, OUTPUT_WIDTH_BASE) pred_y_full_base = np.interp(x_base, x_2x, pred_y_full_2x) return pred_y_full_base def visualize_predictions(image, predictions, output_path): """ Draw predictions as dots on the image. predictions: list of 4 arrays, one per row, each [OUTPUT_WIDTH_BASE] in full image y-coords Output is cropped to signal region. """ # Crop image to signal region (x: T0 to T1) vis_image = image[:, T0:T1, :].copy() # Colors for each row (BGR) 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. pred_y_full: [OUTPUT_WIDTH_BASE] predictions in full image y-coords (base resolution) df: DataFrame with ground truth signals in mV row_idx: 0-2 for short leads (4 segments), 3 for rhythm strip Returns dict of lead_name -> SNR in dB """ baseline_y = ZERO_MV[row_idx] segment_width = OUTPUT_WIDTH_BASE // 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 # Lead II handling (10s data in GT) 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, negative_dir=None): """Process a single image, save visualization, and compute per-lead SNR.""" # 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 (with 2x upscaling internally) predictions = [] pred_mv_rows = {} for row_idx in range(4): row_crop = crop_row(image, row_idx) # Returns 2x upscaled crop pred_y_crop_2x = predict_row(model, row_crop, device) # Predict at 2x pred_y_full = convert_crop_to_full(pred_y_crop_2x, row_idx) # Convert to base resolution # Convert to mV for post-processing baseline_y = ZERO_MV[row_idx] pred_mv = (baseline_y - pred_y_full) / MV_TO_PIXEL # Apply smoothing 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 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 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 # 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) 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 SSH host for SCP (default: gpu-vm)') parser.add_argument('--remote_dir', type=str, default=REMOTE_CHECKPOINT_DIR, help='Remote checkpoint directory') parser.add_argument('--local_cache', type=str, default='/tmp/v16_1_checkpoints', help='Local directory to cache 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/v16_1')) 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_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('--no_scp', action='store_true', help='Skip SCP, use local checkpoint only') 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_1/negpreds')) # Get checkpoint checkpoint_path = None if args.checkpoint: checkpoint_path = Path(args.checkpoint) if not checkpoint_path.exists(): print(f"ERROR: Checkpoint not found: {checkpoint_path}") return elif not args.no_scp: # SCP latest from remote checkpoint_path = scp_latest_checkpoint( args.remote_host, args.remote_dir, REMOTE_CHECKPOINT_PATTERN, args.local_cache ) if checkpoint_path is None: print("ERROR: Failed to fetch checkpoint from remote") return else: # Look for local checkpoint local_cache = Path(args.local_cache) if local_cache.exists(): checkpoints = sorted(local_cache.glob(f'{REMOTE_CHECKPOINT_PATTERN}*.pth')) if checkpoints: checkpoint_path = checkpoints[-1] print(f"Using local checkpoint: {checkpoint_path}") if checkpoint_path is None: print("ERROR: No checkpoint found. Provide --checkpoint or remove --no_scp") return # 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) print(f"{'='*70}") print(f"V16.1 Inference Visualization (2x Upscaled)") print(f"{'='*70}") print(f"Checkpoint: {checkpoint_path}") print(f"Output: {output_dir}") print(f"Negative predictions: {negative_dir}") print(f"Device: {device}") print(f"Input resolution: {INPUT_HEIGHT} x {INPUT_WIDTH} (2x upscaled)") 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(checkpoint_path, 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") 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 = [] 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, 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)") 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) - Saved to {negative_dir}") print(f"{'='*70}") print(f"{'Sample':<15} {'Variant':<8} {'Min SNR':>10} {'Worst Lead':<12}") 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)}") low_snr_report_path = negative_dir / 'low_snr_report.csv' with open(low_snr_report_path, 'w') as f: f.write("sample,variant,min_snr,worst_lead\n") for item in low_snr_samples: f.write(f"{item['sample']},{item['variant']},{item['min_snr']:.2f},{item['worst_lead']}\n") print(f"Low-SNR report saved to: {low_snr_report_path}") else: print(f"\nNo samples with SNR < {LOW_SNR_THRESHOLD} dB") print(f"{'='*70}") if __name__ == '__main__': main()