#!/usr/bin/env python3 """ Compute per-segment baseline offsets for all leads across many images. This analyzes what the typical baseline offset is between model predictions and ground truth for each lead. If the offsets are consistent, we can use hardcoded values for production inference without needing ground truth. """ 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 from scipy import stats import json import torch import torch.nn as nn import torch.nn.functional as F import cv2 import timm # ============================================================================= # 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 CROP_HALF_HEIGHT = 250 ROW_HEIGHT = 500 NUM_BINS = 128 ECG_MV_MIN, ECG_MV_MAX = -10.0, 10.0 VALID_VARIANTS = ['0001', '0003', '0004', '0005', '0006', '0009', '0010', '0011', '0012'] LEAD_LAYOUT = [ ['I', 'aVR', 'V1', 'V4'], ['II', 'aVL', 'V2', 'V5'], ['III', 'aVF', 'V3', 'V6'], ] ALL_LEADS = ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6'] # ============================================================================= # Model Architecture (same as V20) # ============================================================================= 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 TCNBlock(nn.Module): def __init__(self, in_channels, out_channels, kernel_size=3, dilation=1, dropout=0.1): super().__init__() padding = (kernel_size - 1) * dilation // 2 self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size, padding=padding, dilation=dilation) self.bn1 = nn.BatchNorm1d(out_channels) self.conv2 = nn.Conv1d(out_channels, out_channels, kernel_size, padding=padding, dilation=dilation) self.bn2 = nn.BatchNorm1d(out_channels) self.dropout = nn.Dropout(dropout) self.activation = nn.GELU() self.residual = nn.Conv1d(in_channels, out_channels, 1) if in_channels != out_channels else nn.Identity() def forward(self, x): residual = self.residual(x) out = self.dropout(self.activation(self.bn1(self.conv1(x)))) out = self.activation(self.bn2(self.conv2(out)) + residual) return out class TCNRefiner(nn.Module): def __init__(self, in_channels, hidden_channels=64, num_layers=4, kernel_size=5, dropout=0.1): super().__init__() layers = [] dilations = [2**i for i in range(num_layers)] layers.append(TCNBlock(in_channels, hidden_channels, kernel_size, dilations[0], dropout)) for dilation in dilations[1:]: layers.append(TCNBlock(hidden_channels, hidden_channels, kernel_size, dilation, dropout)) self.tcn = nn.Sequential(*layers) self.output_proj = nn.Conv1d(hidden_channels, in_channels, 1) def forward(self, x): residual = x out = self.tcn(x) out = self.output_proj(out) return out + residual 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 IntegralRegressionHead(nn.Module): def __init__(self, in_channels, num_bins=NUM_BINS, temperature=1.0): super().__init__() self.num_bins = num_bins self.temperature = temperature bin_centers = torch.linspace(0, 1, num_bins) self.register_buffer('bin_centers', bin_centers) self.heatmap_conv = nn.Sequential( CoordConv2d(in_channels, 128, 3, padding=1), nn.BatchNorm2d(128), nn.GELU(), nn.Conv2d(128, 64, 3, padding=1), nn.BatchNorm2d(64), nn.GELU(), nn.Conv2d(64, num_bins, 1), ) self.height_attention = nn.Sequential( nn.Conv2d(in_channels, 64, 3, padding=1), nn.BatchNorm2d(64), nn.GELU(), nn.Conv2d(64, 1, 1), ) def forward(self, x, temperature=None): B, C, H, W = x.shape temp = temperature if temperature is not None else self.temperature heatmap_2d = self.heatmap_conv(x) attn = F.softmax(self.height_attention(x), dim=2) heatmap = (heatmap_2d * attn).sum(dim=2) heatmap_prob = F.softmax(heatmap / temp, dim=1) coords = (heatmap_prob * self.bin_centers.view(1, -1, 1)).sum(dim=1) return coords, heatmap_prob class IntegralRegressionNet(nn.Module): def __init__(self, encoder_name='convnext_base.fb_in22k_ft_in1k', pretrained=True, num_bins=NUM_BINS, temperature=1.0, use_tcn=True, tcn_layers=4): super().__init__() self.num_bins = num_bins self.use_tcn = use_tcn 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_pool = nn.Sequential( nn.Conv2d(decoder_dims[-1], 64, 3, padding=1), nn.BatchNorm2d(64), nn.GELU(), nn.Conv2d(64, 1, 1), ) if use_tcn: self.tcn_refiner = TCNRefiner(decoder_dims[-1], 64, tcn_layers, 5, 0.1) self.integral_head = IntegralRegressionHead(decoder_dims[-1], num_bins, temperature) def forward(self, x, temperature=None): 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) if self.use_tcn: attn = F.softmax(self.height_pool(d), dim=2) d_1d = (d * attn).sum(dim=2) d_1d = self.tcn_refiner(d_1d) d = d_1d.unsqueeze(2).expand(-1, -1, d.shape[2], -1) coords, heatmap = self.integral_head(d, temperature) return coords, heatmap # ============================================================================= # Helper Functions # ============================================================================= def apply_savgol_smoothing(signal_mv, window=7, polyorder=2): 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): 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 pred_mv_rows[0][:segment_width] = lead_I + alpha * error pred_mv_rows[2][:segment_width] = lead_III + alpha * error return pred_mv_rows def clamp_ecg_amplitude(signal_mv): return np.clip(signal_mv, ECG_MV_MIN, ECG_MV_MAX) def interpolate_nan(signal_1d): 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 load_model(checkpoint_path, device): model = IntegralRegressionNet( encoder_name='convnext_base.fb_in22k_ft_in1k', pretrained=False, num_bins=NUM_BINS, temperature=1.0, use_tcn=True, tcn_layers=4 ) checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False) state_dict = checkpoint['model'] if any(k.startswith('module.') for k in state_dict.keys()): state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()} model.load_state_dict(state_dict) model = model.to(device) model.eval() print(f"Loaded model from epoch {checkpoint.get('epoch', 'N/A')}, SNR: {checkpoint.get('snr', 'N/A'):.2f} dB") return model def crop_row(image, row_idx): 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, temperature=1.0): 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'): coords, heatmap = model(image_tensor, temperature=temperature) pred_y_crop = coords[0].cpu().numpy() * ROW_HEIGHT return pred_y_crop def convert_crop_to_full(pred_y_crop, row_idx): 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 compute_baseline_offset(pred_mv_segment, gt_mv): """Compute baseline offset between prediction and GT.""" x_pred = np.linspace(0, 1, len(pred_mv_segment)) x_gt = np.linspace(0, 1, len(gt_mv)) pred_resampled = np.interp(x_gt, x_pred, pred_mv_segment) return np.median(pred_resampled - gt_mv) def process_image(model, image, df, device, temperature=1.0): """ Process a single image and return baseline offsets for all 12 leads. """ pred_mv_rows = {} segment_width = OUTPUT_WIDTH // 4 for row_idx in range(3): # Only rows 0-2 (not rhythm strip) row_crop = crop_row(image, row_idx) pred_y_crop = predict_row(model, row_crop, device, temperature=temperature) 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 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 # Apply Einthoven correction (this affects I, III) pred_mv_rows = apply_einthoven_correction(pred_mv_rows, alpha=0.33) # Compute baseline offset for each lead offsets = {} for row_idx in range(3): lead_names = LEAD_LAYOUT[row_idx] for seg_idx, lead_name in enumerate(lead_names): if lead_name not in df.columns: offsets[lead_name] = np.nan continue gt_mv = df[lead_name].dropna().values if len(gt_mv) == 0: offsets[lead_name] = np.nan continue # Handle Lead II short strip (use Lead I length as reference) 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_mv_rows[row_idx][seg_start:seg_end] offset = compute_baseline_offset(pred_segment, gt_mv) offsets[lead_name] = offset return offsets def describe_distribution(values, lead_name): """Compute descriptive statistics for a distribution.""" values = np.array([v for v in values if not np.isnan(v)]) if len(values) == 0: return None # Mode calculation (discretize to 0.01 mV bins) binned = np.round(values * 100) / 100 mode_result = stats.mode(binned, keepdims=True) mode_val = mode_result.mode[0] mode_count = mode_result.count[0] return { 'lead': lead_name, 'count': len(values), 'mean': np.mean(values), 'median': np.median(values), 'mode': mode_val, 'mode_count': mode_count, 'std': np.std(values), 'min': np.min(values), 'max': np.max(values), 'range': np.max(values) - np.min(values), 'p5': np.percentile(values, 5), 'p25': np.percentile(values, 25), 'p75': np.percentile(values, 75), 'p95': np.percentile(values, 95), 'iqr': np.percentile(values, 75) - np.percentile(values, 25), } def main(): parser = argparse.ArgumentParser(description='Compute per-lead baseline offsets') parser.add_argument('--checkpoint', type=str, required=True) parser.add_argument('--kaggle_data', type=str, default='/data/ecg-digitization/stage1_data/train') parser.add_argument('--num_samples', type=int, default=3000) parser.add_argument('--output_dir', type=str, default='/home/azureuser/tmp/baseline_analysis') parser.add_argument('--seed', type=int, default=42) args = parser.parse_args() random.seed(args.seed) np.random.seed(args.seed) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') output_dir = Path(args.output_dir) output_dir.mkdir(parents=True, exist_ok=True) print(f"{'='*80}") print(f"Baseline Offset Analysis for V20") print(f"{'='*80}") print(f"Checkpoint: {args.checkpoint}") print(f"Num samples: {args.num_samples}") print(f"Output dir: {output_dir}") print(f"Device: {device}") print(f"{'='*80}") # Load model model = load_model(args.checkpoint, device) # Find samples kaggle_dir = Path(args.kaggle_data) sample_paths = [] print("Finding samples...") sample_dirs = list(kaggle_dir.iterdir()) random.shuffle(sample_dirs) for sample_dir in sample_dirs: if not sample_dir.is_dir(): 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: sample_paths.append((img_path, csv_path)) break # One image per sample directory for diversity if len(sample_paths) >= args.num_samples: break print(f"Found {len(sample_paths)} samples") # Process all samples all_offsets = {lead: [] for lead in ALL_LEADS} raw_data = [] # For saving individual values for img_path, csv_path in tqdm(sample_paths, desc="Computing offsets"): try: image = cv2.imread(str(img_path), cv2.IMREAD_COLOR) if image is None: continue image = image[Y0:Y1, X0:X1] image = cv2.resize(image, (TARGET_WIDTH, TARGET_HEIGHT), interpolation=cv2.INTER_LINEAR) df = pd.read_csv(csv_path) offsets = process_image(model, image, df, device) # Store raw values row = {'image': img_path.stem} for lead in ALL_LEADS: offset = offsets.get(lead, np.nan) row[lead] = offset if not np.isnan(offset): all_offsets[lead].append(offset) raw_data.append(row) except Exception as e: print(f"Error processing {img_path}: {e}") continue # Save raw data to CSV raw_df = pd.DataFrame(raw_data) raw_csv_path = output_dir / 'baseline_offsets_raw.csv' raw_df.to_csv(raw_csv_path, index=False) print(f"\nSaved raw data to: {raw_csv_path}") # Compute and print statistics print(f"\n{'='*80}") print(f"BASELINE OFFSET STATISTICS (mV)") print(f"{'='*80}") print(f"Offset = median(prediction - ground_truth)") print(f"Positive offset means model predicts HIGHER than GT") print(f"Negative offset means model predicts LOWER than GT") print(f"{'='*80}\n") stats_data = [] # Print header print(f"{'Lead':<8} {'Count':>6} {'Mean':>8} {'Median':>8} {'Mode':>8} {'Std':>8} {'Min':>8} {'Max':>8} {'Range':>8} {'IQR':>8}") print("-" * 90) for lead in ALL_LEADS: stats = describe_distribution(all_offsets[lead], lead) if stats: stats_data.append(stats) print(f"{lead:<8} {stats['count']:>6} {stats['mean']:>8.4f} {stats['median']:>8.4f} " f"{stats['mode']:>8.4f} {stats['std']:>8.4f} {stats['min']:>8.4f} " f"{stats['max']:>8.4f} {stats['range']:>8.4f} {stats['iqr']:>8.4f}") # Print percentiles print(f"\n{'='*80}") print(f"PERCENTILES (mV)") print(f"{'='*80}") print(f"{'Lead':<8} {'P5':>10} {'P25':>10} {'P50':>10} {'P75':>10} {'P95':>10}") print("-" * 60) for stats in stats_data: print(f"{stats['lead']:<8} {stats['p5']:>10.4f} {stats['p25']:>10.4f} " f"{stats['median']:>10.4f} {stats['p75']:>10.4f} {stats['p95']:>10.4f}") # Compute recommended hardcoded values print(f"\n{'='*80}") print(f"RECOMMENDED HARDCODED BASELINE OFFSETS (mV)") print(f"{'='*80}") print("Use these values to subtract from predictions in production:") print() recommended = {} print("BASELINE_OFFSETS = {") for lead in ALL_LEADS: values = all_offsets[lead] if values: median_offset = np.median(values) recommended[lead] = median_offset print(f" '{lead}': {median_offset:.4f},") print("}") # Save stats to JSON stats_json_path = output_dir / 'baseline_stats.json' with open(stats_json_path, 'w') as f: json.dump({ 'stats': stats_data, 'recommended_offsets': recommended, 'num_samples': len(sample_paths), 'checkpoint': args.checkpoint, }, f, indent=2) print(f"\nSaved stats to: {stats_json_path}") # Analysis: Are offsets consistent enough to use hardcoded values? print(f"\n{'='*80}") print(f"ANALYSIS: Should we use hardcoded offsets?") print(f"{'='*80}") # A good hardcoded offset should have low std relative to the offset magnitude print("\nConsistency check (lower is better for hardcoding):") print(f"{'Lead':<8} {'|Median|':>10} {'Std':>10} {'Std/|Med|':>12} {'Verdict':>15}") print("-" * 60) for lead in ALL_LEADS: values = all_offsets[lead] if values: median = np.median(values) std = np.std(values) if abs(median) > 0.01: ratio = std / abs(median) if ratio < 1.0: verdict = "✓ Good" elif ratio < 2.0: verdict = "~ Marginal" else: verdict = "✗ Too variable" else: verdict = "≈ Zero offset" ratio = 0 print(f"{lead:<8} {abs(median):>10.4f} {std:>10.4f} {ratio:>12.2f} {verdict:>15}") print(f"\n{'='*80}") print("Done!") if __name__ == '__main__': main()