code/scripts/compute_baseline_offsets.py

#!/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()