#!/usr/bin/env python3 """ Find optimal ensemble ratio between V10.1 model and Friend's Net3 model. Evaluates on random training samples with all variants. Uses pre-processed stage1 images from /data/ecg-digitization/stage1_data/train/ """ import os import sys import numpy as np import pandas as pd import cv2 import torch import torch.nn as nn import torch.nn.functional as F import torchvision.transforms as T import timm from pathlib import Path from tqdm import tqdm from scipy.signal import savgol_filter import scipy.signal import scipy.optimize import random from typing import Tuple # Add baseline path BASELINE_PATH = '/data/ecg-digitization/hengck23-submit-physionet/hengck23-submit-physionet' sys.path.insert(0, BASELINE_PATH) import stage2_common as s2c from stage2_model import MyCoordUnetDecoder, encode_with_resnet # 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 SOFT_ARGMAX_TEMP = 100.0 DEVICE = "cuda:0" VALID_VARIANTS = ['0001', '0003', '0004', '0005', '0006', '0009', '0010', '0011', '0012'] # ============================================================================= # V10.1 Model # ============================================================================= class CoordDecoderBlock(nn.Module): def __init__(self, in_ch, skip_ch, out_ch, scale=2): super().__init__() self.scale = scale self.conv = nn.Sequential( nn.Conv2d(in_ch + skip_ch + 2, out_ch, 3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), ) def forward(self, x, skip=None): x = F.interpolate(x, scale_factor=self.scale, mode='nearest') if skip is not None: x = torch.cat([x, skip], dim=1) b, c, h, w = x.shape cy, cx = torch.meshgrid( torch.linspace(-1, 1, h, device=x.device, dtype=x.dtype), torch.linspace(-1, 1, w, device=x.device, dtype=x.dtype), indexing='ij' ) coord = torch.stack([cx, cy]).unsqueeze(0).expand(b, -1, -1, -1) x = torch.cat([x, coord], dim=1) return self.conv(x) class ECGNetV10(nn.Module): def __init__(self, decoder_dims=[256, 128, 64, 32, 16]): super().__init__() self.encoder = timm.create_model( 'efficientnet_b4.ra2_in1k', pretrained=False, features_only=True, out_indices=(0, 1, 2, 3, 4) ) enc_dims = [24, 32, 56, 160, 448] self.dec_blocks = nn.ModuleList() in_ch = enc_dims[-1] skip_chs = enc_dims[:-1][::-1] + [0] while len(decoder_dims) < len(skip_chs): decoder_dims.append(decoder_dims[-1]) decoder_dims = decoder_dims[:len(skip_chs)] for skip_ch, out_ch in zip(skip_chs, decoder_dims): self.dec_blocks.append(CoordDecoderBlock(in_ch, skip_ch, out_ch)) in_ch = out_ch self.seg_head = nn.Conv2d(decoder_dims[-1], 4, 1) self.reg_head = nn.Sequential( nn.Conv2d(decoder_dims[-1], 64, 3, padding=1), nn.ReLU(inplace=True), nn.AdaptiveAvgPool2d((1, None)), ) self.reg_out = nn.Sequential( nn.Conv1d(64, 32, 3, padding=1), nn.ReLU(inplace=True), nn.Conv1d(32, 4, 1), nn.Sigmoid() ) def forward(self, x): input_size = x.shape[2:] enc = self.encoder(x) d = enc[-1] skips = enc[:-1][::-1] + [None] for block, skip in zip(self.dec_blocks, skips): d = block(d, skip) if d.shape[2:] != input_size: d = F.interpolate(d, size=input_size, mode='bilinear', align_corners=False) seg_logits = self.seg_head(d) reg_feat = self.reg_head(d).squeeze(2) reg_coords = self.reg_out(reg_feat) return seg_logits, reg_coords # ============================================================================= # Friend's Net3 Model # ============================================================================= class Net3(nn.Module): def __init__(self, pretrained=True): super().__init__() encoder_dim = [64, 128, 256, 512] decoder_dim = [128, 64, 32, 16] self.encoder = timm.create_model( model_name='resnet34.a3_in1k', pretrained=pretrained, in_chans=3, num_classes=0, global_pool='' ) self.decoder = MyCoordUnetDecoder( in_channel=encoder_dim[-1], skip_channel=encoder_dim[:-1][::-1] + [0], out_channel=decoder_dim, scale=[2, 2, 2, 2] ) self.pixel = nn.Conv2d(decoder_dim[-1], 4, 1) def forward(self, image): encode = encode_with_resnet(self.encoder, image) last, _ = self.decoder(feature=encode[-1], skip=encode[:-1][::-1] + [None]) return self.pixel(last) # ============================================================================= # Helper Functions # ============================================================================= def soft_argmax(heatmap, temperature=SOFT_ARGMAX_TEMP): B, C, H, W = heatmap.shape y_coords = torch.arange(H, device=heatmap.device, dtype=heatmap.dtype).view(1, 1, H, 1) weights = F.softmax(heatmap * temperature, dim=2) return (weights * y_coords).sum(dim=2) def resample_signal(signal, target_length): if len(signal) == target_length: return signal x_old = np.linspace(0, 1, len(signal)) x_new = np.linspace(0, 1, target_length) return np.interp(x_new, x_old, signal) # ============================================================================= # Competition-style SNR (with alignment) # ============================================================================= PERFECT_SCORE = 10**(384/10) # Maximum SNR ratio MAX_TIME_SHIFT = 0.2 # 200ms max shift (as fraction of second) # Lead layout for extraction ROW_LAYOUT = [ ['I', 'aVR', 'V1', 'V4'], ['II_short', 'aVL', 'V2', 'V5'], ['III', 'aVF', 'V3', 'V6'], ] def align_signals(label: np.ndarray, pred: np.ndarray, max_shift: float = float('inf')) -> np.ndarray: """Align prediction to label using cross-correlation (competition style).""" if np.any(~np.isfinite(label)): raise ValueError('values in label should all be finite') if np.sum(np.isfinite(pred)) == 0: raise ValueError('prediction can not all be infinite') label_arr = np.asarray(label, dtype=np.float64) pred_arr = np.asarray(pred, dtype=np.float64) label_mean = np.mean(label_arr) pred_mean = np.nanmean(pred_arr[np.isfinite(pred_arr)]) label_arr_centered = label_arr - label_mean pred_arr_centered = np.where(np.isfinite(pred_arr), pred_arr - pred_mean, 0) correlation = scipy.signal.correlate(label_arr_centered, pred_arr_centered, mode='full') n_label = np.size(label_arr) n_pred = np.size(pred_arr) lags = scipy.signal.correlation_lags(n_label, n_pred, mode='full') valid_lags_mask = (lags >= -max_shift) & (lags <= max_shift) max_correlation = np.nanmax(correlation[valid_lags_mask]) all_max_indices = np.flatnonzero(correlation == max_correlation) best_idx = min(all_max_indices, key=lambda i: abs(lags[i])) time_shift = lags[best_idx] start_padding_len = max(time_shift, 0) pred_slice_start = max(-time_shift, 0) pred_slice_end = min(n_label - time_shift, n_pred) end_padding_len = max(n_label - n_pred - time_shift, 0) aligned_pred = np.concatenate(( np.full(start_padding_len, np.nan), pred_arr[pred_slice_start:pred_slice_end], np.full(end_padding_len, np.nan) )) def objective_func(v_shift): return np.nansum((label_arr - (aligned_pred - v_shift)) ** 2) if np.any(np.isfinite(label_arr) & np.isfinite(aligned_pred)): results = scipy.optimize.minimize_scalar(objective_func, method='Brent') vertical_shift = results.x aligned_pred -= vertical_shift return aligned_pred def compute_power(label: np.ndarray, prediction: np.ndarray) -> Tuple[float, float]: """Compute signal and noise power (competition style).""" if label.ndim != 1 or prediction.ndim != 1: raise ValueError('Inputs must be 1-dimensional arrays.') finite_mask = np.isfinite(prediction) if not np.any(finite_mask): raise ValueError("The 'prediction' array contains no finite values.") prediction = prediction.copy() prediction[~np.isfinite(prediction)] = 0 noise = label - prediction p_signal = np.sum(label**2) p_noise = np.sum(noise**2) return p_signal, p_noise def series_to_leads(series_mv): """Convert 4-row mV series to 12-lead dictionary.""" leads = {} segment_width = series_mv.shape[1] // 4 for row_idx in range(3): for seg_idx, lead_name in enumerate(ROW_LAYOUT[row_idx]): actual_name = 'II_short' if lead_name == 'II_short' else lead_name start = seg_idx * segment_width end = (seg_idx + 1) * segment_width leads[actual_name] = series_mv[row_idx, start:end] leads['II'] = series_mv[3] # Full rhythm strip return leads def compute_competition_snr(pred_mv, gt_df, sig_len): """Compute SNR using competition-style scoring with alignment.""" # Get sampling frequency (assume 10-second recording) fs = sig_len // 10 if sig_len >= 10 else 500 max_shift = int(fs * MAX_TIME_SHIFT) # Convert pred to leads pred_leads = series_to_leads(pred_mv) # Get short lead length short_lead_len = len(gt_df['I'].dropna()) if 'I' in gt_df.columns else sig_len // 4 sum_signal = 0.0 sum_noise = 0.0 # Standard leads (short duration) short_leads = ['I', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6'] for lead in short_leads: if lead not in gt_df.columns or lead not in pred_leads: continue gt_lead = gt_df[lead].dropna().values if len(gt_lead) == 0: continue pred_lead = pred_leads[lead] pred_resampled = resample_signal(pred_lead, len(gt_lead)) try: aligned_pred = align_signals(gt_lead, pred_resampled, max_shift) p_signal, p_noise = compute_power(gt_lead, aligned_pred) sum_signal += p_signal sum_noise += p_noise except Exception: # Fallback - direct comparison p_signal = np.sum(gt_lead**2) p_noise = np.sum((gt_lead - pred_resampled)**2) sum_signal += p_signal sum_noise += p_noise # II_short (first 2.5s of II) if 'II' in gt_df.columns and 'II_short' in pred_leads: full_ii = gt_df['II'].dropna().values gt_ii_short = full_ii[:short_lead_len] pred_ii_short = pred_leads['II_short'] pred_resampled = resample_signal(pred_ii_short, len(gt_ii_short)) try: aligned_pred = align_signals(gt_ii_short, pred_resampled, max_shift) p_signal, p_noise = compute_power(gt_ii_short, aligned_pred) sum_signal += p_signal sum_noise += p_noise except Exception: p_signal = np.sum(gt_ii_short**2) p_noise = np.sum((gt_ii_short - pred_resampled)**2) sum_signal += p_signal sum_noise += p_noise # Full II (rhythm strip) - only count the portion NOT covered by II_short # Actually competition may score II separately. Let's include it. if 'II' in gt_df.columns and 'II' in pred_leads: gt_ii = gt_df['II'].dropna().values pred_ii = pred_leads['II'] pred_resampled = resample_signal(pred_ii, len(gt_ii)) try: aligned_pred = align_signals(gt_ii, pred_resampled, max_shift) p_signal, p_noise = compute_power(gt_ii, aligned_pred) sum_signal += p_signal sum_noise += p_noise except Exception: p_signal = np.sum(gt_ii**2) p_noise = np.sum((gt_ii - pred_resampled)**2) sum_signal += p_signal sum_noise += p_noise # Compute final SNR if sum_noise == 0: return PERFECT_SCORE elif sum_signal == 0: return 0 else: snr_ratio = min((sum_signal / sum_noise), PERFECT_SCORE) return snr_ratio # Return ratio, convert to dB later def compute_snr_4row(pred_mv, gt_mv): """Compute average SNR across all 4 rows (legacy, for quick comparison).""" snrs = [] for row in range(4): pred = resample_signal(pred_mv[row], len(gt_mv[row])) signal_power = np.mean(gt_mv[row] ** 2) noise_power = np.mean((pred - gt_mv[row]) ** 2) if noise_power > 1e-10 and signal_power > 1e-10: snr = 10 * np.log10(signal_power / noise_power) snrs.append(np.clip(snr, -10, 50)) return np.mean(snrs) if snrs else 0.0 # ============================================================================= # Inference Functions # ============================================================================= @torch.no_grad() def process_v10(model, image_bgr, device): """V10 inference on stage1 image. Returns 4-row mV signal.""" h, w = image_bgr.shape[:2] crop_h, crop_w = min(h, Y1), min(w, X1) image_cropped = image_bgr[:crop_h, :crop_w] image_resized = cv2.resize(image_cropped, (TARGET_WIDTH, TARGET_HEIGHT), interpolation=cv2.INTER_LINEAR) image_tensor = torch.from_numpy(image_resized.astype(np.float32) / 255.0) image_tensor = image_tensor.permute(2, 0, 1).unsqueeze(0).to(device) with torch.amp.autocast('cuda', dtype=torch.float32): seg_logits, _ = model(image_tensor) seg_logits = torch.nan_to_num(seg_logits, nan=0.0, posinf=0.0, neginf=0.0) seg_probs = torch.sigmoid(seg_logits.float()) signal_full = soft_argmax(seg_probs).cpu().numpy()[0] signal_pixel = signal_full[:, T0:T1] signal_mv = np.zeros_like(signal_pixel) for row_idx in range(4): signal_mv[row_idx] = (ZERO_MV[row_idx] - signal_pixel[row_idx]) / MV_TO_PIXEL return np.clip(signal_mv, -7.0, 7.0) @torch.no_grad() def process_net3(model, image_bgr, sig_len, device): """Net3 (Friend's) inference on stage1 image. Returns 4-row mV signal.""" resize = T.Resize((1696, 4352), interpolation=T.InterpolationMode.BILINEAR) # Use RGB for Net3 (friend's model expects RGB) image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB) img = image_rgb[Y0:Y1, X0:X1] / 255.0 batch = resize(torch.from_numpy(np.ascontiguousarray(img.transpose(2, 0, 1))).unsqueeze(0)) batch = batch.float().to(device) with torch.amp.autocast('cuda', dtype=torch.float32): output = model(batch) pixel = torch.sigmoid(output).float().cpu().numpy()[0] series_in_pixel = s2c.pixel_to_series(pixel[..., T0:T1], ZERO_MV.tolist(), sig_len) series_mv = (ZERO_MV.reshape(4, 1) - series_in_pixel) / MV_TO_PIXEL return np.clip(series_mv, -7.0, 7.0) def load_ground_truth(csv_path, sig_len): """Load ground truth from Kaggle CSV and convert to 4-row mV format.""" df = pd.read_csv(csv_path) segment_width = sig_len // 4 gt_4row = np.zeros((4, sig_len), dtype=np.float32) # Get short lead length (use I as reference) short_lead_len = len(df['I'].dropna()) if 'I' in df.columns else segment_width row_layout = [ ['I', 'aVR', 'V1', 'V4'], ['II_short', 'aVL', 'V2', 'V5'], # II_short = first part of II matching short lead length ['III', 'aVF', 'V3', 'V6'], ] for row_idx in range(3): for seg_idx, lead_name in enumerate(row_layout[row_idx]): # Handle II_short specially - use first portion of II if lead_name == 'II_short': if 'II' in df.columns: full_ii = df['II'].dropna().values # Take first short_lead_len samples (same duration as other short leads) ii_short = full_ii[:short_lead_len] resampled = resample_signal(ii_short, segment_width) start = seg_idx * segment_width end = (seg_idx + 1) * segment_width gt_4row[row_idx, start:end] = resampled elif lead_name in df.columns: signal = df[lead_name].dropna().values resampled = resample_signal(signal, segment_width) start = seg_idx * segment_width end = (seg_idx + 1) * segment_width gt_4row[row_idx, start:end] = resampled # Row 3: full II if 'II' in df.columns: gt_4row[3] = resample_signal(df['II'].dropna().values, sig_len) return gt_4row def ensemble_predictions(pred_v10, pred_friend, weight_v10): """Ensemble two 4-row predictions.""" weight_friend = 1.0 - weight_v10 target_len = max(pred_v10.shape[1], pred_friend.shape[1]) result = np.zeros((4, target_len), dtype=np.float32) for row in range(4): v10_row = resample_signal(pred_v10[row], target_len) friend_row = resample_signal(pred_friend[row], target_len) result[row] = weight_v10 * v10_row + weight_friend * friend_row return result def main(): print("="*70) print("Finding Optimal Ensemble Ratio") print("="*70) # Paths - use pre-processed stage1 images data_dir = Path("/data/ecg-digitization/stage1_data/train") # Model paths - downloaded from Kaggle v10_weights = "/data/ecg-digitization/checkpoints/ecg_v10_best.pth" net3_weights = "/data/ecg-digitization/checkpoints/net3_kaggle/iter_0004200.pt" # Check if weights exist, otherwise try alternative locations if not os.path.exists(v10_weights): alt_path = "/data/ecg-digitization/checkpoints/v10_efficientnet_b4_best_reg.pth" if os.path.exists(alt_path): v10_weights = alt_path else: print(f"ERROR: V10 weights not found at {v10_weights}") print("Please download from: https://www.kaggle.com/models/udaybhatia/ecg-v10-best/PyTorch/pytorch/4") return if not os.path.exists(net3_weights): alt_path = f"{BASELINE_PATH}/weight/stage2-00005810.checkpoint.pth" if os.path.exists(alt_path): net3_weights = alt_path else: print(f"ERROR: Net3 weights not found at {net3_weights}") print("Please download from: https://www.kaggle.com/models/wasupandceacar/physio-seg-public/PyTorch/net3_009_4200/1") return # Load V10 model print(f"\nLoading V10 model from {v10_weights}...") v10_model = ECGNetV10() checkpoint = torch.load(v10_weights, map_location='cpu', weights_only=False) state_dict = checkpoint.get('model', checkpoint) if isinstance(state_dict, dict) and list(state_dict.keys())[0].startswith('module.'): state_dict = {k[7:]: v for k, v in state_dict.items()} v10_model.load_state_dict(state_dict) v10_model.to(DEVICE).eval() print(f"V10 loaded - epoch {checkpoint.get('epoch', '?')}") # Load Net3 model (friend's) print(f"\nLoading Net3 model from {net3_weights}...") net3_model = Net3(pretrained=False).to(DEVICE).eval() st = torch.load(net3_weights, map_location="cpu") if isinstance(st, dict) and "state_dict" in st: st = st["state_dict"] net3_model.load_state_dict(st, strict=True) print("Net3 loaded.") # Collect samples - use pre-processed stage1 images all_samples = [] for sample_dir in sorted(data_dir.iterdir()): 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 variant in VALID_VARIANTS: img_path = sample_dir / f"{sample_dir.name}-{variant}.png" if img_path.exists(): all_samples.append((img_path, csv_path, sample_dir.name, variant)) print(f"\nTotal samples with all variants: {len(all_samples)}") # Random sample 1000 random.seed(42) if len(all_samples) > 1000: samples = random.sample(all_samples, 1000) else: samples = all_samples print(f"Evaluating on {len(samples)} samples...") # Collect predictions results = [] for img_path, csv_path, sample_id, variant in tqdm(samples, desc="Processing"): try: # Read pre-processed stage1 image (already rectified) img_bgr = cv2.imread(str(img_path)) if img_bgr is None: continue # Get ground truth gt_df = pd.read_csv(csv_path) sig_len = len(gt_df['II'].dropna()) if 'II' in gt_df.columns else 5000 gt_4row = load_ground_truth(csv_path, sig_len) # V10 prediction (expects BGR) pred_v10 = process_v10(v10_model, img_bgr, DEVICE) # Net3 prediction (friend's model) pred_net3 = process_net3(net3_model, img_bgr, sig_len, DEVICE) # Resample to same length target_len = max(pred_v10.shape[1], pred_net3.shape[1]) pred_v10_resampled = np.zeros((4, target_len), dtype=np.float32) pred_net3_resampled = np.zeros((4, target_len), dtype=np.float32) for row in range(4): pred_v10_resampled[row] = resample_signal(pred_v10[row], target_len) pred_net3_resampled[row] = resample_signal(pred_net3[row], target_len) results.append({ 'sample_id': sample_id, 'variant': variant, 'pred_v10': pred_v10_resampled, 'pred_friend': pred_net3_resampled, 'gt': gt_4row, 'gt_df': gt_df, 'sig_len': sig_len }) except Exception as e: print(f"Error {sample_id}-{variant}: {e}") continue print(f"\nSuccessfully processed {len(results)} samples") if len(results) == 0: print("ERROR: No samples processed successfully!") return # Evaluate different ratios print("\n" + "="*70) print("Evaluating ensemble ratios (competition-style with alignment)...") print("="*70) ratios = np.arange(0.0, 1.05, 0.05) ratio_results = [] for ratio in ratios: snr_ratios = [] # SNR as ratio, not dB for r in results: ensemble = ensemble_predictions(r['pred_v10'], r['pred_friend'], ratio) snr_ratio = compute_competition_snr(ensemble, r['gt_df'], r['sig_len']) snr_ratios.append(snr_ratio) # Competition final score: mean of ratios, then to dB mean_ratio = np.mean(snr_ratios) avg_snr_db = max(10 * np.log10(mean_ratio), -384) if mean_ratio > 0 else -384 ratio_results.append((ratio, avg_snr_db, mean_ratio)) print(f"V10 weight: {ratio:.2f}, Friend weight: {1-ratio:.2f} -> SNR: {avg_snr_db:.3f} dB") # Find best ratio best_ratio, best_snr, _ = max(ratio_results, key=lambda x: x[1]) print("\n" + "="*70) print(f"BEST RATIO: V10={best_ratio:.2f}, Friend={1-best_ratio:.2f}") print(f"Best SNR: {best_snr:.3f} dB") print("="*70) # Also show individual model performance with competition scoring v10_ratios = [compute_competition_snr(r['pred_v10'], r['gt_df'], r['sig_len']) for r in results] friend_ratios = [compute_competition_snr(r['pred_friend'], r['gt_df'], r['sig_len']) for r in results] v10_snr_db = max(10 * np.log10(np.mean(v10_ratios)), -384) friend_snr_db = max(10 * np.log10(np.mean(friend_ratios)), -384) print(f"\nV10 only: {v10_snr_db:.3f} dB") print(f"Friend only: {friend_snr_db:.3f} dB") print(f"Ensemble: {best_snr:.3f} dB (improvement: +{best_snr - max(v10_snr_db, friend_snr_db):.3f} dB)") # Breakdown by variant print("\n" + "="*70) print("SNR by variant (using best ensemble ratio):") print("="*70) for variant in VALID_VARIANTS: variant_results = [r for r in results if r['variant'] == variant] if variant_results: v10_r = [compute_competition_snr(r['pred_v10'], r['gt_df'], r['sig_len']) for r in variant_results] friend_r = [compute_competition_snr(r['pred_friend'], r['gt_df'], r['sig_len']) for r in variant_results] ensemble_r = [] for r in variant_results: ensemble = ensemble_predictions(r['pred_v10'], r['pred_friend'], best_ratio) ensemble_r.append(compute_competition_snr(ensemble, r['gt_df'], r['sig_len'])) v10_db = max(10 * np.log10(np.mean(v10_r)), -384) friend_db = max(10 * np.log10(np.mean(friend_r)), -384) ensemble_db = max(10 * np.log10(np.mean(ensemble_r)), -384) print(f" {variant}: V10={v10_db:.2f}, Friend={friend_db:.2f}, Ensemble={ensemble_db:.2f} dB (n={len(variant_results)})") return best_ratio if __name__ == "__main__": best_ratio = main()