| |
| """ |
| Compute per-segment baseline offsets for V19 model. |
| |
| V19: ConvNeXt-Base + BiLSTM + Deformable Conv (per-row, no cross-row attention) |
| """ |
|
|
| 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 |
|
|
| try: |
| from torchvision.ops import DeformConv2d |
| HAS_DEFORM_CONV = True |
| except ImportError: |
| HAS_DEFORM_CONV = False |
|
|
|
|
| |
| |
| |
| 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 |
|
|
| CROP_HALF_HEIGHT = 250 |
| ROW_HEIGHT = 500 |
|
|
| 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'] |
|
|
|
|
| |
| |
| |
| class DeformableConvBlock(nn.Module): |
| def __init__(self, in_ch, out_ch, kernel_size=3, stride=1, padding=1): |
| super().__init__() |
| self.kernel_size = kernel_size |
| 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), |
| ) |
| 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) |
| return self.act(self.norm(out)) |
|
|
|
|
| class BiLSTMHead(nn.Module): |
| 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) |
| lstm_out, _ = self.lstm(x) |
| return self.output_proj(lstm_out) |
|
|
|
|
| class AuxiliaryHeads(nn.Module): |
| 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): |
| return self.grid_head(features_2d), self.gradient_head(features_1d).squeeze(-1), self.uncertainty_head(features_1d).squeeze(-1) |
|
|
|
|
| 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) |
| return self.conv(torch.cat([x, yy, xx], dim=1)) |
|
|
|
|
| class UNetDecoderBlockV19(nn.Module): |
| 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) |
| return self.conv2(self.conv1(x)) |
|
|
|
|
| class PerLeadNetV19(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 i, (skip_ch, out_ch) in enumerate(zip(skip_channels, decoder_dims)): |
| self.dec_blocks.append(UNetDecoderBlockV19(in_ch, skip_ch, out_ch, use_deform=(i >= 2))) |
| 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(decoder_dims[-1], 128, 2, 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 = F.softmax(self.height_attention(d), 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: |
| return y_pred, self.aux_heads(features_2d, temporal_features) |
| return y_pred |
|
|
|
|
| |
| |
| |
| 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 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 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) |
| return pred_y_crop - pad_top + y_start |
|
|
|
|
| def load_model(checkpoint_path, device): |
| model = PerLeadNetV19(encoder_name='convnext_base.fb_in22k_ft_in1k', pretrained=False) |
| 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 V19 from epoch {checkpoint.get('epoch', 'N/A')}, SNR: {checkpoint.get('snr', checkpoint.get('best_snr', 'N/A')):.2f} dB") |
| return model |
|
|
|
|
| def compute_baseline_offset(pred_mv_segment, gt_mv): |
| 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): |
| segment_width = OUTPUT_WIDTH // 4 |
| pred_mv_rows = {} |
| |
| for row_idx in range(3): |
| row_crop = crop_row(image, row_idx) |
| 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'): |
| pred = model(image_tensor, return_aux=False) |
| |
| pred_y_crop = pred[0].cpu().numpy() * ROW_HEIGHT |
| 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 |
| |
| pred_mv_rows = apply_einthoven_correction(pred_mv_rows, alpha=0.33) |
| |
| 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 |
| 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] |
| offsets[lead_name] = compute_baseline_offset(pred_segment, gt_mv) |
| |
| return offsets |
|
|
|
|
| def describe_distribution(values, lead_name): |
| values = np.array([v for v in values if not np.isnan(v)]) |
| if len(values) == 0: |
| return None |
| binned = np.round(values * 100) / 100 |
| mode_result = stats.mode(binned, keepdims=True) |
| return { |
| 'lead': lead_name, 'count': len(values), 'mean': float(np.mean(values)), |
| 'median': float(np.median(values)), 'mode': float(mode_result.mode[0]), |
| 'mode_count': int(mode_result.count[0]), 'std': float(np.std(values)), |
| 'min': float(np.min(values)), 'max': float(np.max(values)), |
| 'range': float(np.max(values) - np.min(values)), |
| 'p5': float(np.percentile(values, 5)), 'p25': float(np.percentile(values, 25)), |
| 'p75': float(np.percentile(values, 75)), 'p95': float(np.percentile(values, 95)), |
| 'iqr': float(np.percentile(values, 75) - np.percentile(values, 25)), |
| } |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description='Compute per-lead baseline offsets for V19') |
| 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=4000) |
| parser.add_argument('--output_dir', type=str, default='/home/azureuser/tmp/baseline_analysis_v19') |
| 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 V19") |
| print(f"{'='*80}") |
| print(f"Checkpoint: {args.checkpoint}") |
| print(f"Num samples: {args.num_samples}") |
| print(f"Device: {device}") |
| print(f"{'='*80}") |
| |
| model = load_model(args.checkpoint, device) |
| |
| kaggle_dir = Path(args.kaggle_data) |
| sample_paths = [] |
| 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 |
| if len(sample_paths) >= args.num_samples: |
| break |
| |
| print(f"Found {len(sample_paths)} samples") |
| |
| all_offsets = {lead: [] for lead in ALL_LEADS} |
| raw_data = [] |
| |
| 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) |
| 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: {e}") |
| |
| raw_df = pd.DataFrame(raw_data) |
| raw_df.to_csv(output_dir / 'baseline_offsets_raw.csv', index=False) |
| |
| print(f"\n{'='*80}") |
| print(f"BASELINE OFFSET STATISTICS (mV) - V19") |
| print(f"{'='*80}") |
| print(f"{'Lead':<8} {'Count':>6} {'Mean':>8} {'Median':>8} {'Std':>8}") |
| print("-" * 50) |
| |
| stats_data = [] |
| for lead in ALL_LEADS: |
| s = describe_distribution(all_offsets[lead], lead) |
| if s: |
| stats_data.append(s) |
| print(f"{lead:<8} {s['count']:>6} {s['mean']:>8.4f} {s['median']:>8.4f} {s['std']:>8.4f}") |
| |
| print(f"\n{'='*80}") |
| print(f"RECOMMENDED HARDCODED BASELINE OFFSETS (mV) - V19") |
| print(f"{'='*80}\n") |
| |
| recommended = {} |
| print("BASELINE_OFFSETS_V19 = {") |
| for lead in ALL_LEADS: |
| if all_offsets[lead]: |
| median_offset = float(np.median(all_offsets[lead])) |
| recommended[lead] = median_offset |
| print(f" '{lead}': {median_offset:.4f},") |
| print("}") |
| |
| with open(output_dir / 'baseline_stats.json', '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 to: {output_dir}") |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
