{ "cells": [ { "cell_type": "code", "execution_count": null, "id": "b1335e66", "metadata": {}, "outputs": [], "source": [ "import os\n", "import sys\n", "import numpy as np\n", "import pandas as pd\n", "import cv2\n", "import torch\n", "import torch.nn as nn\n", "import torch.nn.functional as F\n", "from pathlib import Path\n", "from tqdm.auto import tqdm\n", "from scipy.signal import savgol_filter\n", "import timm\n", "\n", "# Kaggle paths - UPDATE THESE FOR YOUR SETUP\n", "BASELINE_PATH = '/kaggle/input/hengck23-submit-physionet/hengck23-submit-physionet'\n", "V16_WEIGHTS_PATH = '/kaggle/input/ecg-v16-perlead'\n", "V18_WEIGHTS_PATH = '/kaggle/input/ecg-v18-refiner'\n", "COMPETITION_PATH = '/kaggle/input/physionet-ecg-image-digitization'\n", "\n", "sys.path.insert(0, BASELINE_PATH)\n", "\n", "# Import baseline stage 0/1\n", "from stage0_model import Net as Stage0Net\n", "from stage0_common import image_to_batch, output_to_predict, normalise_by_homography, load_net\n", "from stage1_model import Net as Stage1Net\n", "from stage1_common import output_to_predict as stage1_output_to_predict, rectify_image\n", "\n", "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n", "print(f\"Device: {device}\")" ] }, { "cell_type": "markdown", "id": "7983851a", "metadata": {}, "source": [ "## Constants" ] }, { "cell_type": "code", "execution_count": null, "id": "f7eb2667", "metadata": {}, "outputs": [], "source": [ "# Image dimensions after preprocessing\n", "TARGET_HEIGHT, TARGET_WIDTH = 1696, 4352\n", "\n", "# Baseline Y-coordinates for each row (0mV line)\n", "ZERO_MV = np.array([703.5, 987.5, 1271.5, 1531.5])\n", "\n", "# Conversion factor: pixels to millivolts\n", "MV_TO_PIXEL = 78.5\n", "\n", "# Signal region boundaries (excludes left/right margins)\n", "T0, T1 = 235, 4161\n", "OUTPUT_WIDTH = T1 - T0 # 3926 pixels\n", "\n", "# Crop region (left half of image used for training)\n", "X0, X1 = 0, 2176\n", "Y0, Y1 = 0, 1696\n", "\n", "# Per-row crop parameters (V16 uses larger crop)\n", "CROP_HALF_HEIGHT = 250\n", "ROW_HEIGHT = 500\n", "\n", "# ECG amplitude limits (mV) - values beyond this are errors\n", "ECG_MV_MIN, ECG_MV_MAX = -7.0, 7.0\n", "\n", "# Lead layout (standard 12-lead ECG)\n", "LEAD_NAMES = ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6']\n", "ROW_LAYOUT = [\n", " ['I', 'aVR', 'V1', 'V4'], # Row 0\n", " ['II', 'aVL', 'V2', 'V5'], # Row 1 \n", " ['III', 'aVF', 'V3', 'V6'], # Row 2\n", "]\n", "\n", "# Hardcoded baseline offsets (mV) - computed from 977 images\n", "# These are median(prediction - ground_truth) values\n", "# Positive means model predicts slightly higher than GT\n", "BASELINE_OFFSETS = {\n", " 'I': 0.0055, 'II': 0.0095, 'III': 0.0058,\n", " 'aVR': 0.0067, 'aVL': 0.0081, 'aVF': 0.0076,\n", " 'V1': 0.0069, 'V2': 0.0077, 'V3': 0.0076,\n", " 'V4': 0.0070, 'V5': 0.0080, 'V6': 0.0077,\n", "}" ] }, { "cell_type": "markdown", "id": "2dc622ce", "metadata": {}, "source": [ "## V16 Model Architecture\n", "\n", "ConvNeXt-Base encoder with U-Net decoder and height attention mechanism." ] }, { "cell_type": "code", "execution_count": null, "id": "1b9dd8da", "metadata": {}, "outputs": [], "source": [ "class CoordConv2d(nn.Module):\n", " \"\"\"Conv2d with coordinate channels for positional awareness.\"\"\"\n", " def __init__(self, in_channels, out_channels, kernel_size, **kwargs):\n", " super().__init__()\n", " self.conv = nn.Conv2d(in_channels + 2, out_channels, kernel_size, **kwargs)\n", " \n", " def forward(self, x):\n", " B, C, H, W = x.shape\n", " yy = torch.linspace(-1, 1, H, device=x.device).view(1, 1, H, 1).expand(B, 1, H, W)\n", " xx = torch.linspace(-1, 1, W, device=x.device).view(1, 1, 1, W).expand(B, 1, H, W)\n", " x = torch.cat([x, yy, xx], dim=1)\n", " return self.conv(x)\n", "\n", "\n", "class UNetDecoderBlock(nn.Module):\n", " \"\"\"U-Net style decoder block with skip connections.\"\"\"\n", " def __init__(self, in_ch, skip_ch, out_ch):\n", " super().__init__()\n", " self.conv = nn.Sequential(\n", " nn.Conv2d(in_ch + skip_ch, out_ch, 3, padding=1, bias=False),\n", " nn.BatchNorm2d(out_ch),\n", " nn.GELU(),\n", " nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False),\n", " nn.BatchNorm2d(out_ch),\n", " nn.GELU(),\n", " )\n", " self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)\n", " \n", " def forward(self, x, skip=None):\n", " x = self.upsample(x)\n", " if skip is not None:\n", " if x.shape[2:] != skip.shape[2:]:\n", " x = F.interpolate(x, size=skip.shape[2:], mode='bilinear', align_corners=True)\n", " x = torch.cat([x, skip], dim=1)\n", " return self.conv(x)\n", "\n", "\n", "class PerLeadNet(nn.Module):\n", " \"\"\"V16 Per-Lead ECG Extraction Network.\n", " \n", " Architecture:\n", " - ConvNeXt-Base encoder (pretrained on ImageNet-22k)\n", " - U-Net decoder with skip connections\n", " - Height attention for sub-pixel y-coordinate extraction\n", " - 1D regression head for final signal output\n", " \n", " Input: [B, 3, 500, 3926] - single row crop\n", " Output: [B, 3926] - normalized y-coordinates [0, 1]\n", " \"\"\"\n", " def __init__(self, encoder_name='convnext_base.fb_in22k_ft_in1k', pretrained=True):\n", " super().__init__()\n", " \n", " self.encoder = timm.create_model(\n", " encoder_name,\n", " pretrained=pretrained,\n", " features_only=True,\n", " out_indices=(0, 1, 2, 3),\n", " )\n", " enc_channels = self.encoder.feature_info.channels()\n", " \n", " decoder_dims = [256, 128, 64, 32]\n", " self.dec_blocks = nn.ModuleList()\n", " in_ch = enc_channels[-1]\n", " skip_channels = enc_channels[:-1][::-1] + [0]\n", " \n", " for skip_ch, out_ch in zip(skip_channels, decoder_dims):\n", " self.dec_blocks.append(UNetDecoderBlock(in_ch, skip_ch, out_ch))\n", " in_ch = out_ch\n", " \n", " self.final_up = nn.Sequential(\n", " nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),\n", " nn.Conv2d(decoder_dims[-1], decoder_dims[-1], 3, padding=1, bias=False),\n", " nn.BatchNorm2d(decoder_dims[-1]),\n", " nn.GELU(),\n", " )\n", " \n", " self.height_attention = nn.Sequential(\n", " CoordConv2d(decoder_dims[-1], 64, 3, padding=1),\n", " nn.BatchNorm2d(64),\n", " nn.GELU(),\n", " nn.Conv2d(64, 1, 1),\n", " )\n", " \n", " self.regression_head = nn.Sequential(\n", " nn.Conv1d(decoder_dims[-1], 128, 7, padding=3),\n", " nn.BatchNorm1d(128),\n", " nn.GELU(),\n", " nn.Conv1d(128, 64, 5, padding=2),\n", " nn.BatchNorm1d(64),\n", " nn.GELU(),\n", " nn.Conv1d(64, 1, 1),\n", " )\n", " \n", " def forward(self, x):\n", " B, C, H, W = x.shape\n", " \n", " features = self.encoder(x)\n", " \n", " d = features[-1]\n", " skips = features[:-1][::-1] + [None]\n", " \n", " for block, skip in zip(self.dec_blocks, skips):\n", " d = block(d, skip)\n", " \n", " d = self.final_up(d)\n", " \n", " if d.shape[3] != W:\n", " d = F.interpolate(d, size=(d.shape[2], W), mode='bilinear', align_corners=True)\n", " \n", " attn = self.height_attention(d)\n", " attn = F.softmax(attn, dim=2)\n", " \n", " d = (d * attn).sum(dim=2)\n", " \n", " out = self.regression_head(d)\n", " out = torch.sigmoid(out)\n", " \n", " return out.squeeze(1)" ] }, { "cell_type": "markdown", "id": "9d2f62dc", "metadata": {}, "source": [ "## V18 Refiner Model Architecture\n", "\n", "Cross-Row Attention refiner that takes V16 predictions and refines them by learning inter-row consistency." ] }, { "cell_type": "code", "execution_count": null, "id": "0bc41fe7", "metadata": {}, "outputs": [], "source": [ "class CrossRowAttention(nn.Module):\n", " \"\"\"Cross-Row Multi-Head Self-Attention for V18 refiner.\"\"\"\n", " def __init__(self, embed_dim, num_heads=8, dropout=0.1):\n", " super().__init__()\n", " self.num_heads = num_heads\n", " self.head_dim = embed_dim // num_heads\n", " self.scale = self.head_dim ** -0.5\n", " \n", " self.qkv = nn.Linear(embed_dim, embed_dim * 3)\n", " self.proj = nn.Linear(embed_dim, embed_dim)\n", " self.dropout = nn.Dropout(dropout)\n", " self.norm = nn.LayerNorm(embed_dim)\n", " \n", " def forward(self, x):\n", " B, num_rows, C, W = x.shape\n", " x = x.permute(0, 3, 1, 2).reshape(B * W, num_rows, C)\n", " \n", " residual = x\n", " x = self.norm(x)\n", " \n", " qkv = self.qkv(x).reshape(B * W, num_rows, 3, self.num_heads, self.head_dim)\n", " qkv = qkv.permute(2, 0, 3, 1, 4)\n", " q, k, v = qkv[0], qkv[1], qkv[2]\n", " \n", " attn = (q @ k.transpose(-2, -1)) * self.scale\n", " attn = attn.softmax(dim=-1)\n", " attn = self.dropout(attn)\n", " \n", " out = (attn @ v).transpose(1, 2).reshape(B * W, num_rows, C)\n", " out = self.proj(out)\n", " out = self.dropout(out)\n", " out = out + residual\n", " \n", " out = out.reshape(B, W, num_rows, C).permute(0, 2, 3, 1)\n", " return out\n", "\n", "\n", "class CrossRowTransformerBlock(nn.Module):\n", " \"\"\"Full transformer block with cross-row attention and FFN.\"\"\"\n", " def __init__(self, embed_dim, num_heads=8, mlp_ratio=4.0, dropout=0.1):\n", " super().__init__()\n", " \n", " self.attn = CrossRowAttention(embed_dim, num_heads, dropout)\n", " \n", " self.norm = nn.LayerNorm(embed_dim)\n", " hidden_dim = int(embed_dim * mlp_ratio)\n", " self.ffn = nn.Sequential(\n", " nn.Linear(embed_dim, hidden_dim),\n", " nn.GELU(),\n", " nn.Dropout(dropout),\n", " nn.Linear(hidden_dim, embed_dim),\n", " nn.Dropout(dropout),\n", " )\n", " \n", " def forward(self, x):\n", " x = self.attn(x)\n", " \n", " B, num_rows, C, W = x.shape\n", " residual = x\n", " x = x.permute(0, 3, 1, 2).reshape(B * W, num_rows, C)\n", " x = self.norm(x)\n", " x = self.ffn(x) + residual.permute(0, 3, 1, 2).reshape(B * W, num_rows, C)\n", " x = x.reshape(B, W, num_rows, C).permute(0, 2, 3, 1)\n", " \n", " return x\n", "\n", "\n", "class RefinerEncoder(nn.Module):\n", " \"\"\"Lightweight encoder for V18 refinement (4 input channels: RGB + guide).\"\"\"\n", " def __init__(self, encoder_name='efficientnet_b0', pretrained=True):\n", " super().__init__()\n", " \n", " self.encoder = timm.create_model(\n", " encoder_name,\n", " pretrained=pretrained,\n", " features_only=True,\n", " out_indices=(1, 2, 3),\n", " in_chans=4,\n", " )\n", " self.channels = self.encoder.feature_info.channels()\n", " \n", " def forward(self, x):\n", " return self.encoder(x)\n", "\n", "\n", "class V18RefinerNet(nn.Module):\n", " \"\"\"V18 Refiner Network: Takes V16 predictions and refines them with cross-row attention.\"\"\"\n", " \n", " def __init__(self, \n", " encoder_name='efficientnet_b0',\n", " pretrained=True,\n", " cross_row_layers=3,\n", " cross_row_dim=128,\n", " num_heads=4):\n", " super().__init__()\n", " \n", " self.row_encoder = RefinerEncoder(encoder_name, pretrained)\n", " enc_channels = self.row_encoder.channels\n", " \n", " self.feature_proj = nn.Sequential(\n", " nn.AdaptiveAvgPool2d((1, None)),\n", " nn.Flatten(1, 2),\n", " )\n", " self.channel_proj = nn.Conv1d(enc_channels[-1], cross_row_dim, 1)\n", " \n", " self.cross_row_blocks = nn.ModuleList([\n", " CrossRowTransformerBlock(\n", " embed_dim=cross_row_dim,\n", " num_heads=num_heads,\n", " mlp_ratio=2.0,\n", " dropout=0.1\n", " )\n", " for _ in range(cross_row_layers)\n", " ])\n", " \n", " self.residual_head = nn.Sequential(\n", " nn.Conv1d(cross_row_dim, 64, 5, padding=2),\n", " nn.BatchNorm1d(64),\n", " nn.GELU(),\n", " nn.Conv1d(64, 32, 3, padding=1),\n", " nn.BatchNorm1d(32),\n", " nn.GELU(),\n", " nn.Conv1d(32, 1, 1),\n", " nn.Tanh(),\n", " )\n", " \n", " self.residual_scale = nn.Parameter(torch.tensor(0.1))\n", " \n", " def create_guide_channel(self, v16_pred, height, sigma=15.0):\n", " \"\"\"Create a Gaussian guide channel from V16 prediction.\"\"\"\n", " B, W = v16_pred.shape\n", " device = v16_pred.device\n", " \n", " y_pred = v16_pred * height\n", " y_grid = torch.arange(height, device=device, dtype=torch.float32)\n", " y_grid = y_grid.view(1, height, 1)\n", " y_pred = y_pred.unsqueeze(1)\n", " \n", " guide = torch.exp(-0.5 * ((y_grid - y_pred) / sigma) ** 2)\n", " guide = guide.unsqueeze(1)\n", " \n", " return guide\n", " \n", " def forward(self, images, v16_preds):\n", " \"\"\"\n", " Args:\n", " images: [B, num_rows, C, H, W] - row crops\n", " v16_preds: [B, num_rows, W] - V16 normalized predictions\n", " \n", " Returns:\n", " refined: [B, num_rows, W] - refined predictions\n", " residuals: [B, num_rows, W] - learned residuals\n", " \"\"\"\n", " B, num_rows, C, H, W = images.shape\n", " \n", " all_features = []\n", " \n", " for row_idx in range(num_rows):\n", " row_image = images[:, row_idx]\n", " row_v16 = v16_preds[:, row_idx]\n", " \n", " guide = self.create_guide_channel(row_v16, H)\n", " row_input = torch.cat([row_image, guide], dim=1)\n", " \n", " enc_features = self.row_encoder(row_input)\n", " row_feat = enc_features[-1]\n", " \n", " row_feat = self.feature_proj(row_feat)\n", " row_feat = F.interpolate(row_feat, size=W, mode='linear', align_corners=True)\n", " row_feat = self.channel_proj(row_feat)\n", " \n", " all_features.append(row_feat)\n", " \n", " features = torch.stack(all_features, dim=1)\n", " \n", " for block in self.cross_row_blocks:\n", " features = block(features)\n", " \n", " residuals = []\n", " for row_idx in range(num_rows):\n", " row_feat = features[:, row_idx]\n", " residual = self.residual_head(row_feat)\n", " residuals.append(residual.squeeze(1))\n", " \n", " residuals = torch.stack(residuals, dim=1)\n", " scaled_residuals = residuals * self.residual_scale * 0.1\n", " \n", " refined = torch.clamp(v16_preds + scaled_residuals, 0, 1)\n", " \n", " return refined, scaled_residuals\n", "\n", "\n", "print(\"V18 Refiner architecture loaded!\")" ] }, { "cell_type": "markdown", "id": "c18d66df", "metadata": {}, "source": [ "## Load Models" ] }, { "cell_type": "code", "execution_count": null, "id": "eb609b15", "metadata": {}, "outputs": [], "source": [ "# Load Stage 0 (orientation correction)\n", "print(\"Loading Stage 0...\")\n", "stage0_net = Stage0Net(pretrained=False)\n", "stage0_net = load_net(stage0_net, f'{BASELINE_PATH}/weight/stage0-last.checkpoint.pth')\n", "stage0_net.to(device).eval()\n", "\n", "# Load Stage 1 (grid rectification)\n", "print(\"Loading Stage 1...\")\n", "stage1_net = Stage1Net(pretrained=False)\n", "stage1_net = load_net(stage1_net, f'{BASELINE_PATH}/weight/stage1-last.checkpoint.pth')\n", "stage1_net.to(device).eval()\n", "\n", "# Load V16 model\n", "print(\"Loading V16 model...\")\n", "model = PerLeadNet(encoder_name='convnext_base.fb_in22k_ft_in1k', pretrained=False)\n", "v16_checkpoint = torch.load(f'{V16_WEIGHTS_PATH}/v16_perlead_best_snr.pth', map_location='cpu', weights_only=False)\n", "\n", "# Handle DataParallel state dict\n", "state_dict = v16_checkpoint['model']\n", "if list(state_dict.keys())[0].startswith('module.'):\n", " state_dict = {k[7:]: v for k, v in state_dict.items()}\n", "model.load_state_dict(state_dict)\n", "model.to(device).eval()\n", "\n", "v16_epoch = v16_checkpoint.get('epoch', '?')\n", "v16_snr = v16_checkpoint.get('snr', 0)\n", "print(f\"Loaded V16 epoch {v16_epoch}, SNR: {v16_snr:.2f} dB\")\n", "\n", "# Load V18 refiner\n", "print(\"Loading V18 refiner...\")\n", "refiner = V18RefinerNet(encoder_name='efficientnet_b0', pretrained=False)\n", "v18_checkpoint = torch.load(f'{V18_WEIGHTS_PATH}/v18_refiner_best.pth', map_location='cpu', weights_only=False)\n", "refiner.load_state_dict(v18_checkpoint['refiner'])\n", "refiner.to(device).eval()\n", "\n", "v18_epoch = v18_checkpoint.get('epoch', '?')\n", "v18_snr = v18_checkpoint.get('snr', 0)\n", "print(f\"Loaded V18 epoch {v18_epoch}, Refined SNR: {v18_snr:.2f} dB\")\n", "print(f\"Residual scale: {refiner.residual_scale.item():.4f}\")\n", "\n", "print(\"\\nAll models loaded!\")" ] }, { "cell_type": "markdown", "id": "c1f47d8f", "metadata": {}, "source": [ "## Post-Processing Functions\n", "\n", "These enhancements improve signal quality and enforce physical constraints." ] }, { "cell_type": "code", "execution_count": null, "id": "b553f182", "metadata": {}, "outputs": [], "source": [ "def apply_savgol_smoothing(signal_mv, window=7, polyorder=2):\n", " \"\"\"Apply Savitzky-Golay smoothing to remove high-frequency noise.\n", " \n", " This preserves signal shape better than simple moving average.\n", " \"\"\"\n", " if len(signal_mv) >= window:\n", " return savgol_filter(signal_mv, window_length=window, polyorder=polyorder)\n", " return signal_mv\n", "\n", "\n", "def apply_einthoven_correction(pred_mv_rows, alpha=0.33):\n", " \"\"\"\n", " Apply Einthoven's law correction on short lead segments.\n", " \n", " Einthoven's Law: II = I + III (in mV)\n", " \n", " For each time point, if there's a violation e = II - (I + III), distribute error:\n", " - I' = I + α*e\n", " - III'= III + α*e \n", " - II' = II - α*e (skipped since we use rhythm strip for II)\n", " \n", " Args:\n", " pred_mv_rows: dict with row predictions in mV {0: array, 1: array, 2: array, 3: array}\n", " alpha: correction factor (0.33 distributes error equally)\n", " \"\"\"\n", " segment_width = len(pred_mv_rows[0]) // 4\n", " \n", " # Row 0 segment 0 = Lead I\n", " # Row 1 segment 0 = Lead II short (0-2.5s)\n", " # Row 2 segment 0 = Lead III\n", " \n", " lead_I = pred_mv_rows[0][:segment_width].copy()\n", " lead_II_short = pred_mv_rows[1][:segment_width].copy()\n", " lead_III = pred_mv_rows[2][:segment_width].copy()\n", " \n", " # Compute Einthoven violation: e = II - (I + III)\n", " derived_II = lead_I + lead_III\n", " error = lead_II_short - derived_II\n", " \n", " # Correct - distribute error\n", " lead_I_corrected = lead_I + alpha * error\n", " lead_III_corrected = lead_III + alpha * error\n", " \n", " # Update rows (only segment 0 for leads I and III)\n", " pred_mv_rows[0][:segment_width] = lead_I_corrected\n", " pred_mv_rows[2][:segment_width] = lead_III_corrected\n", " \n", " return pred_mv_rows\n", "\n", "\n", "def apply_hardcoded_baseline_correction(pred_mv_rows):\n", " \"\"\"\n", " Apply hardcoded baseline correction for each lead.\n", " \n", " Uses pre-computed median offsets from 977 images. These are tiny\n", " (~0.007 mV) so the effect is minimal, but included for completeness.\n", " \n", " This is production-ready - no GT required.\n", " \"\"\"\n", " segment_width = len(pred_mv_rows[0]) // 4\n", " \n", " for row_idx in range(3): # Only for rows 0-2 (not rhythm strip)\n", " lead_names = ROW_LAYOUT[row_idx]\n", " for seg_idx, lead_name in enumerate(lead_names):\n", " offset = BASELINE_OFFSETS.get(lead_name, 0.0)\n", " \n", " seg_start = seg_idx * segment_width\n", " seg_end = (seg_idx + 1) * segment_width\n", " pred_mv_rows[row_idx][seg_start:seg_end] -= offset\n", " \n", " # Rhythm strip uses Lead II offset\n", " pred_mv_rows[3] -= BASELINE_OFFSETS.get('II', 0.0)\n", " \n", " return pred_mv_rows\n", "\n", "\n", "def clamp_ecg_amplitude(signal_mv):\n", " \"\"\"Clamp signal to reasonable ECG range (±7mV).\"\"\"\n", " return np.clip(signal_mv, ECG_MV_MIN, ECG_MV_MAX)\n", "\n", "\n", "def interpolate_nan(signal_1d):\n", " \"\"\"Interpolate NaN values from valid neighbors. Falls back to 0 if all NaN.\"\"\"\n", " valid_mask = np.isfinite(signal_1d)\n", " if valid_mask.all():\n", " return signal_1d\n", " if not valid_mask.any():\n", " return np.zeros_like(signal_1d)\n", " \n", " x = np.arange(len(signal_1d))\n", " signal_1d[~valid_mask] = np.interp(x[~valid_mask], x[valid_mask], signal_1d[valid_mask])\n", " return signal_1d" ] }, { "cell_type": "markdown", "id": "5dc61330", "metadata": {}, "source": [ "## Inference Functions" ] }, { "cell_type": "code", "execution_count": null, "id": "6ee9e776", "metadata": {}, "outputs": [], "source": [ "@torch.no_grad()\n", "def process_stage0(image_rgb):\n", " \"\"\"Stage 0: Orientation correction.\n", " \n", " Detects corner keypoints and corrects rotation/perspective.\n", " NOTE: Baseline expects RGB input.\n", " \"\"\"\n", " batch = image_to_batch(image_rgb)\n", " with torch.amp.autocast('cuda', dtype=torch.float32):\n", " output = stage0_net(batch)\n", " rotated, keypoint = output_to_predict(image_rgb, batch, output)\n", " normalized, _, _ = normalise_by_homography(rotated, keypoint)\n", " return normalized\n", "\n", "\n", "@torch.no_grad()\n", "def process_stage1(image_rgb):\n", " \"\"\"Stage 1: Grid rectification.\n", " \n", " Detects grid lines and applies dewarping.\n", " NOTE: Baseline expects RGB input, returns RGB.\n", " \"\"\"\n", " batch = {'image': torch.from_numpy(np.ascontiguousarray(image_rgb.transpose(2, 0, 1))).unsqueeze(0)}\n", " with torch.amp.autocast('cuda', dtype=torch.float32):\n", " output = stage1_net(batch)\n", " gridpoint_xy, _ = stage1_output_to_predict(image_rgb, batch, output)\n", " rectified = rectify_image(image_rgb, gridpoint_xy)\n", " return rectified\n", "\n", "\n", "def crop_row(image, row_idx):\n", " \"\"\"Crop a single row centered on its baseline, signal region only (T0:T1).\n", " \n", " Args:\n", " image: Full preprocessed image [1696, 4352, 3]\n", " row_idx: Row index 0-3\n", " \n", " Returns:\n", " Row crop [500, 3926, 3]\n", " \"\"\"\n", " baseline_y = int(ZERO_MV[row_idx])\n", " y_start = max(0, baseline_y - CROP_HALF_HEIGHT)\n", " y_end = min(TARGET_HEIGHT, baseline_y + CROP_HALF_HEIGHT)\n", " \n", " # Crop x to signal region only (T0:T1)\n", " row_crop = image[y_start:y_end, T0:T1, :].copy()\n", " \n", " # Pad if necessary\n", " if row_crop.shape[0] < ROW_HEIGHT:\n", " pad_top = max(0, CROP_HALF_HEIGHT - baseline_y)\n", " pad_bottom = max(0, (baseline_y + CROP_HALF_HEIGHT) - TARGET_HEIGHT)\n", " row_crop = np.pad(row_crop, ((pad_top, pad_bottom), (0, 0), (0, 0)), mode='edge')\n", " \n", " return row_crop\n", "\n", "\n", "@torch.no_grad()\n", "def predict_v16_row(row_crop):\n", " \"\"\"Run V16 inference on a single row crop.\n", " \n", " Args:\n", " row_crop: [500, 3926, 3] BGR image\n", " \n", " Returns:\n", " pred_normalized: [3926] predictions normalized [0, 1]\n", " \"\"\"\n", " image_tensor = torch.from_numpy(row_crop.astype(np.float32) / 255.0)\n", " image_tensor = image_tensor.permute(2, 0, 1).unsqueeze(0).to(device)\n", " \n", " with torch.amp.autocast('cuda', dtype=torch.float16):\n", " output = model(image_tensor)\n", " \n", " return output[0] # Keep as tensor, normalized [0, 1]\n", "\n", "\n", "@torch.no_grad()\n", "def predict_with_v18_refiner(row_crops):\n", " \"\"\"Run V16 + V18 refiner on all 4 row crops.\n", " \n", " Args:\n", " row_crops: list of 4 row crops, each [500, 3926, 3] BGR\n", " \n", " Returns:\n", " refined_preds: list of 4 arrays [3926] in crop-relative pixels\n", " \"\"\"\n", " # Stack row crops: [1, 4, 3, H, W]\n", " row_tensors = []\n", " for crop in row_crops:\n", " t = torch.from_numpy(crop.astype(np.float32) / 255.0)\n", " t = t.permute(2, 0, 1) # [3, H, W]\n", " row_tensors.append(t)\n", " \n", " images = torch.stack(row_tensors, dim=0).unsqueeze(0).to(device) # [1, 4, 3, H, W]\n", " \n", " # Get V16 predictions for all rows\n", " with torch.amp.autocast('cuda', dtype=torch.float16):\n", " v16_preds = []\n", " for row_idx in range(4):\n", " row_pred = model(images[:, row_idx]) # [1, W] normalized\n", " v16_preds.append(row_pred)\n", " v16_preds = torch.stack(v16_preds, dim=1) # [1, 4, W]\n", " \n", " # Run V18 refiner\n", " refined, residuals = refiner(images, v16_preds) # [1, 4, W]\n", " \n", " # Convert to crop-relative pixels\n", " refined_np = refined[0].cpu().numpy() # [4, W]\n", " results = []\n", " for row_idx in range(4):\n", " pred_y_crop = refined_np[row_idx] * ROW_HEIGHT # Convert to pixels\n", " results.append(pred_y_crop)\n", " \n", " return results\n", "\n", "\n", "def convert_crop_to_mv(pred_y_crop, row_idx):\n", " \"\"\"Convert crop-relative y-coordinates to millivolts.\n", " \n", " Args:\n", " pred_y_crop: [3926] predictions in crop-relative pixels\n", " row_idx: Row index 0-3\n", " \n", " Returns:\n", " pred_mv: [3926] predictions in millivolts\n", " \"\"\"\n", " baseline_y = int(ZERO_MV[row_idx])\n", " y_start = max(0, baseline_y - CROP_HALF_HEIGHT)\n", " pad_top = max(0, CROP_HALF_HEIGHT - baseline_y)\n", " \n", " # Convert to full image y\n", " pred_y_full = pred_y_crop - pad_top + y_start\n", " \n", " # Convert to mV (positive y is down, positive mV is up)\n", " pred_mv = (ZERO_MV[row_idx] - pred_y_full) / MV_TO_PIXEL\n", " \n", " return pred_mv" ] }, { "cell_type": "code", "execution_count": null, "id": "e62be051", "metadata": {}, "outputs": [], "source": [ "def series_to_leads(series_mv):\n", " \"\"\"Convert 4-row series to 12-lead dictionary.\n", " \n", " Layout:\n", " - Row 0: I, aVR, V1, V4 (4 segments of 2.5s each)\n", " - Row 1: II, aVL, V2, V5 (4 segments)\n", " - Row 2: III, aVF, V3, V6 (4 segments)\n", " - Row 3: Lead II full rhythm strip (10s)\n", " \n", " Args:\n", " series_mv: dict {0: [3926], 1: [3926], 2: [3926], 3: [3926]}\n", " \n", " Returns:\n", " leads: dict mapping lead name to signal array\n", " \"\"\"\n", " leads = {}\n", " segment_width = OUTPUT_WIDTH // 4 # ~981 pixels per 2.5s segment\n", " \n", " for row_idx in range(3):\n", " for seg_idx, lead_name in enumerate(ROW_LAYOUT[row_idx]):\n", " start = seg_idx * segment_width\n", " end = (seg_idx + 1) * segment_width\n", " leads[lead_name] = series_mv[row_idx][start:end]\n", " \n", " # Row 3 is the full rhythm strip (Lead II, 10s)\n", " leads['II'] = series_mv[3]\n", " \n", " return leads\n", "\n", "\n", "def resample_signal(signal, target_length):\n", " \"\"\"Resample signal to target length using linear interpolation.\"\"\"\n", " if len(signal) == target_length:\n", " return signal\n", " x_old = np.linspace(0, 1, len(signal))\n", " x_new = np.linspace(0, 1, target_length)\n", " return np.interp(x_new, x_old, signal)" ] }, { "cell_type": "code", "execution_count": null, "id": "c7ade6e9", "metadata": {}, "outputs": [], "source": [ "def process_image(image_path, apply_smoothing=True, apply_einthoven=True, use_v18_refiner=True):\n", " \"\"\"Full inference pipeline for a single image.\n", " \n", " Flow:\n", " 1. Load image (BGR)\n", " 2. Convert to RGB for stage0/stage1 (baseline expects RGB)\n", " 3. Stage0: orientation correction\n", " 4. Stage1: grid rectification\n", " 5. Convert to BGR, resize to standard dimensions\n", " 6. For each row: crop\n", " 7. Run V16 + V18 refiner (if enabled) \n", " 8. Apply post-processing (smoothing, Einthoven, clamping)\n", " 9. Convert to 12-lead dictionary\n", " \n", " Args:\n", " image_path: Path to input image\n", " apply_smoothing: Apply Savitzky-Golay smoothing\n", " apply_einthoven: Apply Einthoven's law correction\n", " use_v18_refiner: Use V18 refiner on top of V16 (default: True)\n", " \n", " Returns:\n", " leads: dict mapping lead name to signal array in mV\n", " \"\"\"\n", " # Load image\n", " image_bgr = cv2.imread(str(image_path))\n", " if image_bgr is None:\n", " raise ValueError(f\"Failed to load image: {image_path}\")\n", " \n", " # Convert to RGB for baseline stage0/stage1\n", " image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)\n", " \n", " # Stage 0: Orientation correction\n", " try:\n", " normalized = process_stage0(image_rgb)\n", " except Exception as e:\n", " normalized = image_rgb\n", " \n", " # Stage 1: Grid rectification\n", " try:\n", " rectified = process_stage1(normalized)\n", " except Exception as e:\n", " rectified = normalized\n", " \n", " # Convert back to BGR and resize\n", " rectified_bgr = cv2.cvtColor(rectified, cv2.COLOR_RGB2BGR)\n", " \n", " # Crop to left region and resize to target dimensions\n", " h, w = rectified_bgr.shape[:2]\n", " crop_h = min(h, Y1)\n", " crop_w = min(w, X1)\n", " image_cropped = rectified_bgr[:crop_h, :crop_w]\n", " image_resized = cv2.resize(image_cropped, (TARGET_WIDTH, TARGET_HEIGHT), interpolation=cv2.INTER_LINEAR)\n", " \n", " # Crop all 4 rows\n", " row_crops = [crop_row(image_resized, row_idx) for row_idx in range(4)]\n", " \n", " # Predict with V18 refiner (V16 + cross-row attention) or V16 only\n", " if use_v18_refiner:\n", " row_preds = predict_with_v18_refiner(row_crops)\n", " else:\n", " # V16 only (legacy path)\n", " row_preds = []\n", " for row_idx, row_crop in enumerate(row_crops):\n", " pred_normalized = predict_v16_row(row_crop)\n", " pred_y_crop = pred_normalized.cpu().numpy() * ROW_HEIGHT\n", " row_preds.append(pred_y_crop)\n", " \n", " # Convert to mV and apply post-processing\n", " pred_mv_rows = {}\n", " for row_idx in range(4):\n", " pred_mv = convert_crop_to_mv(row_preds[row_idx], row_idx)\n", " \n", " # Apply smoothing\n", " if apply_smoothing:\n", " pred_mv = apply_savgol_smoothing(pred_mv, window=7, polyorder=2)\n", " \n", " # Clamp to reasonable range\n", " pred_mv = clamp_ecg_amplitude(pred_mv)\n", " \n", " # Handle NaN values\n", " pred_mv = interpolate_nan(pred_mv.copy())\n", " \n", " pred_mv_rows[row_idx] = pred_mv\n", " \n", " # Apply Einthoven correction\n", " if apply_einthoven:\n", " pred_mv_rows = apply_einthoven_correction(pred_mv_rows, alpha=0.33)\n", " \n", " # Apply hardcoded baseline correction (tiny ~0.007 mV offsets from 977-image analysis)\n", " pred_mv_rows = apply_hardcoded_baseline_correction(pred_mv_rows)\n", " \n", " # Convert to 12-lead format\n", " leads = series_to_leads(pred_mv_rows)\n", " \n", " return leads\n", "\n", "print(\"Pipeline: V16 → V18 Refiner → Post-processing (with baseline correction)\")" ] }, { "cell_type": "markdown", "id": "19bf4566", "metadata": {}, "source": [ "## Generate Submission" ] }, { "cell_type": "code", "execution_count": null, "id": "f9606825", "metadata": {}, "outputs": [], "source": [ "# Load test metadata\n", "test_df = pd.read_csv(f'{COMPETITION_PATH}/test.csv')\n", "test_dir = Path(f'{COMPETITION_PATH}/test')\n", "image_ids = test_df['id'].unique()\n", "\n", "print(f\"Processing {len(image_ids)} images...\")\n", "print(f\"Signal output width: {OUTPUT_WIDTH} pixels\")\n", "print(f\"Test metadata columns: {list(test_df.columns)}\")\n", "print(f\"\\nSample test.csv rows:\")\n", "print(test_df.head())" ] }, { "cell_type": "code", "execution_count": null, "id": "84c6485c", "metadata": {}, "outputs": [], "source": [ "all_rows = []\n", "failed_images = []\n", "\n", "for img_id in tqdm(image_ids, desc=\"Processing images\"):\n", " img_path = test_dir / f\"{img_id}.png\"\n", " \n", " if not img_path.exists():\n", " print(f\"Missing: {img_path}\")\n", " failed_images.append(img_id)\n", " continue\n", " \n", " img_df = test_df[test_df['id'] == img_id]\n", " \n", " try:\n", " leads = process_image(img_path, apply_smoothing=True, apply_einthoven=True)\n", " except Exception as e:\n", " print(f\"Error {img_id}: {e}\")\n", " failed_images.append(img_id)\n", " # Fallback to zeros\n", " for _, row in img_df.iterrows():\n", " for i in range(row['number_of_rows']):\n", " all_rows.append({'id': f\"{img_id}_{i}_{row['lead']}\", 'value': 0.0})\n", " continue\n", " \n", " # Generate submission rows for each lead\n", " for _, row in img_df.iterrows():\n", " lead = row['lead']\n", " num_samples = row['number_of_rows']\n", " \n", " if lead in leads:\n", " signal = resample_signal(leads[lead], num_samples)\n", " for i, val in enumerate(signal):\n", " all_rows.append({'id': f\"{img_id}_{i}_{lead}\", 'value': float(val)})\n", " else:\n", " # Unknown lead - fall back to zeros\n", " for i in range(num_samples):\n", " all_rows.append({'id': f\"{img_id}_{i}_{lead}\", 'value': 0.0})\n", "\n", "print(f\"\\nTotal rows: {len(all_rows)}\")\n", "print(f\"Failed images: {len(failed_images)}\")" ] }, { "cell_type": "code", "execution_count": null, "id": "a87ce82c", "metadata": {}, "outputs": [], "source": [ "# Create submission DataFrame\n", "submission_df = pd.DataFrame(all_rows)\n", "\n", "# CRITICAL: Check for NaN/Inf values - these cause submission errors\n", "nan_count = submission_df['value'].isna().sum()\n", "inf_count = (~np.isfinite(submission_df['value'])).sum() - nan_count\n", "print(f\"NaN values: {nan_count}\")\n", "print(f\"Inf values: {inf_count}\")\n", "\n", "if nan_count > 0 or inf_count > 0:\n", " print(\"WARNING: Found NaN/Inf values! Replacing with 0.0\")\n", " submission_df['value'] = submission_df['value'].replace([np.inf, -np.inf], 0.0)\n", " submission_df['value'] = submission_df['value'].fillna(0.0)\n", "\n", "# Final verification\n", "assert submission_df['value'].isna().sum() == 0, \"Still have NaN values!\"\n", "assert np.isfinite(submission_df['value']).all(), \"Still have Inf values!\"\n", "\n", "# Save submission\n", "submission_df.to_csv('/kaggle/working/submission.csv', index=False)\n", "\n", "print(f\"\\nSubmission saved!\")\n", "print(f\"Shape: {submission_df.shape}\")\n", "print(f\"\\nFirst 10 rows:\")\n", "print(submission_df.head(10))\n", "print(f\"\\nValue statistics:\")\n", "print(submission_df['value'].describe())" ] }, { "cell_type": "code", "execution_count": null, "id": "29f64eed", "metadata": {}, "outputs": [], "source": [ "# Sanity check: Plot a sample prediction\n", "import matplotlib.pyplot as plt\n", "\n", "# Get a sample image for visualization\n", "sample_id = str(image_ids[0])\n", "sample_rows = submission_df[submission_df['id'].str.startswith(sample_id)]\n", "\n", "# Group by lead\n", "for lead in ['I', 'II', 'III', 'V1']:\n", " lead_rows = sample_rows[sample_rows['id'].str.endswith(f'_{lead}')]\n", " if len(lead_rows) > 0:\n", " values = lead_rows['value'].values\n", " plt.figure(figsize=(12, 2))\n", " plt.plot(values)\n", " plt.title(f'{sample_id} - Lead {lead}')\n", " plt.ylabel('mV')\n", " plt.xlabel('Sample')\n", " plt.grid(True, alpha=0.3)\n", " plt.show()" ] } ], "metadata": { "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 5 }