#!/usr/bin/env python3 """ V16 Inference with Visualization Runs V16 per-lead model on Kaggle images and plots predictions as dots on the strips. V16 uses signal-region-only input (T0:T1) and 500px crop height. """ 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', '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 (same as training) LEAD_LAYOUT = [ ['I', 'aVR', 'V1', 'V4'], ['II', 'aVL', 'V2', 'V5'], ['III', 'aVF', 'V3', 'V6'], ] # ============================================================================= # 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 (from V10 high-scoring solution) # ============================================================================= 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) For each time point, if there's a violation e = II - (I + III), distribute error: - I' = I + α*e - III'= III + α*e - II' = II - α*e (but we use rhythm strip, so skip) pred_mv_rows: dict with row predictions in mV """ segment_width = len(pred_mv_rows[0]) // 4 # Row 0 segment 0 = Lead I # Row 1 segment 0 = Lead II short (0-2.5s) # Row 2 segment 0 = Lead III 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() # Compute Einthoven violation: e = II - (I + III) derived_II = lead_I + lead_III error = lead_II_short - derived_II # Correct - distribute error lead_I_corrected = lead_I + alpha * error lead_III_corrected = lead_III + alpha * error # Update rows (only segment 0 for leads I and III) 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. Falls back to 0 if all NaN.""" 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 V15 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) # V16: Crop x to signal region only (T0:T1) row_crop = image[y_start:y_end, T0:T1, :].copy() # Pad if necessary 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.""" # 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 crop-relative pixels 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. V16: Output is OUTPUT_WIDTH (3926) values, maps to x positions T0:T1. """ baseline_y = int(ZERO_MV[row_idx]) y_start = max(0, baseline_y - CROP_HALF_HEIGHT) # Adjust for padding pad_top = max(0, CROP_HALF_HEIGHT - baseline_y) # Convert to full image 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. predictions: list of 4 arrays, one per row, each [OUTPUT_WIDTH] in full image y-coords V16: predictions are already OUTPUT_WIDTH (3926) values for x positions 0 to OUTPUT_WIDTH-1. Output is cropped to signal region. Plots EVERY predicted point. """ # 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] # V16: pred_y has OUTPUT_WIDTH values, directly maps to vis_image x coords for x in range(len(pred_y)): y = int(np.clip(pred_y[x], 0, TARGET_HEIGHT - 1)) # Use radius 1 for dense dots 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. IMPORTANT: Resample predictions UP to match GT sample rate (typically 500 Hz). Kaggle evaluates SNR at the original sample rate, not pixel resolution. pred_y_full: [OUTPUT_WIDTH] predictions in full image y-coords 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 // 4 lead_snrs = {} if row_idx < 3: # Rows 0-2: 4 leads each 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 # Get original GT signal in mV 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 # Lead II GT is typically 4x longer than other short leads (10s vs 2.5s) if lead_name == 'II': # Lead II in CSV contains 10s rhythm strip data # For short strip (2.5s), use first quarter # Compare with other short lead lengths to detect 10s data 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 at least 2x longer = it's 10s data quarter_len = len(gt_mv) // 4 gt_mv = gt_mv[:quarter_len] # First quarter (0-2.5s) - matches training # Get prediction segment in pixels seg_start = seg_idx * segment_width seg_end = (seg_idx + 1) * segment_width pred_y_seg = pred_y_full[seg_start:seg_end] # Convert prediction from pixels to mV pred_mv_pixels = (baseline_y - pred_y_seg) / MV_TO_PIXEL # Resample prediction UP to match GT length 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) # Compute SNR 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: # Row 3: Full Lead II rhythm strip (10s) 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 # Convert full prediction from pixels to mV pred_mv_pixels = (baseline_y - pred_y_full) / MV_TO_PIXEL # Resample prediction UP to match GT length 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) # Compute SNR 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. Args: apply_smoothing: Apply Savitzky-Golay smoothing to predictions apply_einthoven: Apply Einthoven's law correction negative_dir: If provided, save low-SNR predictions here Returns: Tuple of (all_lead_snrs dict, min_snr value, output_path) """ # 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 = {} # Store mV predictions for Einthoven correction 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) # 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) # Clamp to reasonable range pred_mv = clamp_ecg_amplitude(pred_mv) # Handle any NaN values pred_mv = interpolate_nan(pred_mv.copy()) pred_mv_rows[row_idx] = pred_mv predictions.append(pred_y_full) # Apply Einthoven correction (only for short leads in rows 0-2) 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_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/v16')) 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('--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/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) print(f"{'='*70}") print(f"V16 Inference Visualization (with enhancements)") print(f"{'='*70}") print(f"Checkpoint: {args.checkpoint}") print(f"Output: {output_dir}") print(f"Negative predictions: {negative_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"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 # Filter by holdout set if requested if args.use_holdout and sample_dir.name not in val_sample_set: continue # Find CSV file 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 (holdout = all by default, random otherwise) 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 and track low-SNR samples all_snrs = {lead: [] for lead in ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6', 'II_rhythm']} low_snr_samples = [] # Track samples with low SNR # 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, negative_dir=negative_dir ) if lead_snrs: for lead, snr in lead_snrs.items(): if snr is not None: all_snrs[lead].append(snr) # Track low SNR samples if min_snr is not None and min_snr < LOW_SNR_THRESHOLD: # Find which lead has the lowest SNR 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 to file 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}") # Sort by SNR (worst first) low_snr_samples.sort(key=lambda x: x['min_snr']) for item in low_snr_samples[:20]: # Show top 20 worst 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)}") # Save low-SNR report 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()