{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51894314",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import gc\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\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",
    "V10_WEIGHTS_PATH = '/kaggle/input/ecg-v10-best/pytorch/pytorch/4/ecg_v10_best.pth'\n",
    "V16_WEIGHTS_PATH = '/kaggle/input/ecg-v16-perlead/pytorch/v16-epoch10/1/v16_perlead_best.pth'\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}\")\n",
    "print(f\"PyTorch: {torch.__version__}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36b5e01d",
   "metadata": {},
   "source": [
    "## Constants"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "baedf2dc",
   "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",
    "# V16 per-row crop parameters\n",
    "CROP_HALF_HEIGHT = 250\n",
    "ROW_HEIGHT = 500\n",
    "\n",
    "# V10 soft-argmax temperature\n",
    "SOFT_ARGMAX_TEMP = 100.0\n",
    "\n",
    "# ECG amplitude limits (mV)\n",
    "ECG_MV_MIN, ECG_MV_MAX = -7.0, 7.0\n",
    "\n",
    "# Ensemble weight\n",
    "WEIGHT_V10 = 0.5  # 50% V10, 50% V16\n",
    "WEIGHT_V16 = 1.0 - WEIGHT_V10\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",
    "print(f\"Ensemble weights: V10={WEIGHT_V10:.0%}, V16={WEIGHT_V16:.0%}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "061e17f3",
   "metadata": {},
   "source": [
    "## V10 Model Architecture (EfficientNet-B4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4412c941",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CoordDecoderBlock(nn.Module):\n",
    "    \"\"\"Decoder block with coordinate channels.\"\"\"\n",
    "    def __init__(self, in_ch, skip_ch, out_ch, scale=2):\n",
    "        super().__init__()\n",
    "        self.scale = scale\n",
    "        self.conv = nn.Sequential(\n",
    "            nn.Conv2d(in_ch + skip_ch + 2, out_ch, 3, padding=1, bias=False),\n",
    "            nn.BatchNorm2d(out_ch),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False),\n",
    "            nn.BatchNorm2d(out_ch),\n",
    "            nn.ReLU(inplace=True),\n",
    "        )\n",
    "\n",
    "    def forward(self, x, skip=None):\n",
    "        x = F.interpolate(x, scale_factor=self.scale, mode='nearest')\n",
    "        if skip is not None:\n",
    "            x = torch.cat([x, skip], dim=1)\n",
    "        b, c, h, w = x.shape\n",
    "        cy, cx = torch.meshgrid(\n",
    "            torch.linspace(-1, 1, h, device=x.device, dtype=x.dtype),\n",
    "            torch.linspace(-1, 1, w, device=x.device, dtype=x.dtype),\n",
    "            indexing='ij'\n",
    "        )\n",
    "        coord = torch.stack([cx, cy]).unsqueeze(0).expand(b, -1, -1, -1)\n",
    "        x = torch.cat([x, coord], dim=1)\n",
    "        return self.conv(x)\n",
    "\n",
    "\n",
    "class ECGNetV10(nn.Module):\n",
    "    \"\"\"V10 with EfficientNet-B4 encoder.\"\"\"\n",
    "    def __init__(self, decoder_dims=[256, 128, 64, 32, 16]):\n",
    "        super().__init__()\n",
    "        self.encoder = timm.create_model(\n",
    "            'efficientnet_b4.ra2_in1k', pretrained=False, \n",
    "            features_only=True, out_indices=(0, 1, 2, 3, 4)\n",
    "        )\n",
    "        enc_dims = [24, 32, 56, 160, 448]\n",
    "        \n",
    "        self.dec_blocks = nn.ModuleList()\n",
    "        in_ch = enc_dims[-1]\n",
    "        skip_chs = enc_dims[:-1][::-1] + [0]\n",
    "        while len(decoder_dims) < len(skip_chs):\n",
    "            decoder_dims.append(decoder_dims[-1])\n",
    "        decoder_dims = decoder_dims[:len(skip_chs)]\n",
    "        \n",
    "        for skip_ch, out_ch in zip(skip_chs, decoder_dims):\n",
    "            self.dec_blocks.append(CoordDecoderBlock(in_ch, skip_ch, out_ch))\n",
    "            in_ch = out_ch\n",
    "        \n",
    "        self.seg_head = nn.Conv2d(decoder_dims[-1], 4, 1)\n",
    "        self.reg_head = nn.Sequential(\n",
    "            nn.Conv2d(decoder_dims[-1], 64, 3, padding=1),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.AdaptiveAvgPool2d((1, None)),\n",
    "        )\n",
    "        self.reg_out = nn.Sequential(\n",
    "            nn.Conv1d(64, 32, 3, padding=1),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Conv1d(32, 4, 1),\n",
    "            nn.Sigmoid()\n",
    "        )\n",
    "    \n",
    "    def forward(self, x):\n",
    "        input_size = x.shape[2:]\n",
    "        enc = self.encoder(x)\n",
    "        d = enc[-1]\n",
    "        skips = enc[:-1][::-1] + [None]\n",
    "        for block, skip in zip(self.dec_blocks, skips):\n",
    "            d = block(d, skip)\n",
    "        if d.shape[2:] != input_size:\n",
    "            d = F.interpolate(d, size=input_size, mode='bilinear', align_corners=False)\n",
    "        seg_logits = self.seg_head(d)\n",
    "        reg_feat = self.reg_head(d).squeeze(2)\n",
    "        reg_coords = self.reg_out(reg_feat)\n",
    "        return seg_logits, reg_coords\n",
    "\n",
    "print(\"V10 model defined.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "921fcd68",
   "metadata": {},
   "source": [
    "## V16 Model Architecture (ConvNeXt-Base)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ff9bc10",
   "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)\n",
    "\n",
    "print(\"V16 model defined.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6f174e8",
   "metadata": {},
   "source": [
    "## Load All Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c12fe3f6",
   "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 V10 model\n",
    "print(\"Loading V10 model...\")\n",
    "v10_model = ECGNetV10()\n",
    "checkpoint_v10 = torch.load(V10_WEIGHTS_PATH, map_location='cpu', weights_only=False)\n",
    "state_dict = checkpoint_v10['model']\n",
    "if list(state_dict.keys())[0].startswith('module.'):\n",
    "    state_dict = {k[7:]: v for k, v in state_dict.items()}\n",
    "v10_model.load_state_dict(state_dict)\n",
    "v10_model.to(device).eval()\n",
    "print(f\"V10 loaded - epoch {checkpoint_v10.get('epoch', '?')}, SNR: {checkpoint_v10.get('holdout_snr_soft', 0):.2f} dB\")\n",
    "\n",
    "# Load V16 model\n",
    "print(\"Loading V16 model...\")\n",
    "v16_model = PerLeadNet(encoder_name='convnext_base.fb_in22k_ft_in1k', pretrained=False)\n",
    "checkpoint_v16 = torch.load(V16_WEIGHTS_PATH, map_location='cpu', weights_only=False)\n",
    "state_dict = checkpoint_v16['model']\n",
    "if list(state_dict.keys())[0].startswith('module.'):\n",
    "    state_dict = {k[7:]: v for k, v in state_dict.items()}\n",
    "v16_model.load_state_dict(state_dict)\n",
    "v16_model.to(device).eval()\n",
    "print(f\"V16 loaded - epoch {checkpoint_v16.get('epoch', '?')}, SNR: {checkpoint_v16.get('snr', 0):.2f} dB\")\n",
    "\n",
    "print(\"\\nAll models loaded!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ba24e27",
   "metadata": {},
   "source": [
    "## Preprocessing Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16579c98",
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def process_stage0(image_rgb):\n",
    "    \"\"\"Stage 0: Orientation correction. Expects RGB input.\"\"\"\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. Expects RGB input, returns RGB.\"\"\"\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 preprocess_image(image_path):\n",
    "    \"\"\"Full preprocessing pipeline: Load -> Stage0 -> Stage1 -> Resize.\"\"\"\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",
    "    return image_resized\n",
    "\n",
    "print(\"Preprocessing functions defined.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd3a6df5",
   "metadata": {},
   "source": [
    "## V10 Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7183dce",
   "metadata": {},
   "outputs": [],
   "source": [
    "def soft_argmax(heatmap, temperature=SOFT_ARGMAX_TEMP):\n",
    "    \"\"\"Extract sub-pixel coordinates using soft-argmax.\"\"\"\n",
    "    B, C, H, W = heatmap.shape\n",
    "    y_coords = torch.arange(H, device=heatmap.device, dtype=heatmap.dtype).view(1, 1, H, 1)\n",
    "    weights = F.softmax(heatmap * temperature, dim=2)\n",
    "    return (weights * y_coords).sum(dim=2)\n",
    "\n",
    "\n",
    "def interpolate_nan(signal_1d):\n",
    "    \"\"\"Interpolate NaN values from valid neighbors.\"\"\"\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",
    "    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\n",
    "\n",
    "\n",
    "@torch.no_grad()\n",
    "def run_v10_inference(image_bgr):\n",
    "    \"\"\"Run V10 inference. Returns 4-row mV signal [4, OUTPUT_WIDTH].\"\"\"\n",
    "    # Prepare input tensor\n",
    "    image_tensor = torch.from_numpy(image_bgr.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.float32):\n",
    "        seg_logits, _ = v10_model(image_tensor)\n",
    "    \n",
    "    # Handle NaN/Inf\n",
    "    seg_logits = torch.nan_to_num(seg_logits, nan=0.0, posinf=0.0, neginf=0.0)\n",
    "    seg_probs = torch.sigmoid(seg_logits.float())\n",
    "    signal_full = soft_argmax(seg_probs).cpu().numpy()[0]  # [4, TARGET_WIDTH]\n",
    "    \n",
    "    # Handle any remaining NaN\n",
    "    for row_idx in range(4):\n",
    "        if not np.isfinite(signal_full[row_idx]).all():\n",
    "            signal_full[row_idx] = interpolate_nan(signal_full[row_idx].copy())\n",
    "            remaining_nan = ~np.isfinite(signal_full[row_idx])\n",
    "            if remaining_nan.any():\n",
    "                signal_full[row_idx, remaining_nan] = ZERO_MV[row_idx]\n",
    "    \n",
    "    # Extract signal region T0:T1\n",
    "    signal_pixel = signal_full[:, T0:T1]  # [4, OUTPUT_WIDTH]\n",
    "    \n",
    "    # Convert to mV\n",
    "    signal_mv = np.zeros_like(signal_pixel)\n",
    "    for row_idx in range(4):\n",
    "        signal_mv[row_idx] = (ZERO_MV[row_idx] - signal_pixel[row_idx]) / MV_TO_PIXEL\n",
    "    \n",
    "    return np.clip(signal_mv, ECG_MV_MIN, ECG_MV_MAX)\n",
    "\n",
    "print(\"V10 inference defined.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2de824a2",
   "metadata": {},
   "source": [
    "## V16 Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cba19ba8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def crop_row(image, row_idx):\n",
    "    \"\"\"Crop a single row centered on its baseline, signal region only (T0:T1).\"\"\"\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_row_v16(row_crop):\n",
    "    \"\"\"Run V16 inference on a single row crop.\"\"\"\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 = v16_model(image_tensor)\n",
    "    \n",
    "    # Convert from normalized [0, 1] to crop-relative pixels\n",
    "    pred_y_crop = output[0].cpu().numpy() * ROW_HEIGHT\n",
    "    return pred_y_crop\n",
    "\n",
    "\n",
    "def convert_crop_to_mv(pred_y_crop, row_idx):\n",
    "    \"\"\"Convert crop-relative y-coordinates to millivolts.\"\"\"\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",
    "    return pred_mv\n",
    "\n",
    "\n",
    "def run_v16_inference(image_bgr):\n",
    "    \"\"\"Run V16 inference. Returns 4-row mV signal [4, OUTPUT_WIDTH].\"\"\"\n",
    "    signal_mv = np.zeros((4, OUTPUT_WIDTH), dtype=np.float32)\n",
    "    \n",
    "    for row_idx in range(4):\n",
    "        row_crop = crop_row(image_bgr, row_idx)\n",
    "        pred_y_crop = predict_row_v16(row_crop)\n",
    "        pred_mv = convert_crop_to_mv(pred_y_crop, row_idx)\n",
    "        \n",
    "        # Handle NaN\n",
    "        pred_mv = interpolate_nan(pred_mv.copy())\n",
    "        \n",
    "        signal_mv[row_idx] = np.clip(pred_mv, ECG_MV_MIN, ECG_MV_MAX)\n",
    "    \n",
    "    return signal_mv\n",
    "\n",
    "print(\"V16 inference defined.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0adc23ff",
   "metadata": {},
   "source": [
    "## Post-Processing & Ensemble"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72647700",
   "metadata": {},
   "outputs": [],
   "source": [
    "def series_to_leads(series_mv):\n",
    "    \"\"\"Convert 4-row series to 12-lead dictionary.\n",
    "    \n",
    "    Also creates II_short for Einthoven correction.\n",
    "    \"\"\"\n",
    "    leads = {}\n",
    "    segment_width = series_mv.shape[1] // 4\n",
    "    \n",
    "    # Row layout with II_short for short segment\n",
    "    row_layout_with_short = [\n",
    "        ['I', 'aVR', 'V1', 'V4'],\n",
    "        ['II_short', 'aVL', 'V2', 'V5'],  # Use II_short for Einthoven\n",
    "        ['III', 'aVF', 'V3', 'V6'],\n",
    "    ]\n",
    "    \n",
    "    for row_idx in range(3):\n",
    "        for seg_idx, lead_name in enumerate(row_layout_with_short[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].copy()\n",
    "    \n",
    "    # Row 3 is the full rhythm strip (Lead II, 10s)\n",
    "    leads['II'] = series_mv[3].copy()\n",
    "    \n",
    "    return leads\n",
    "\n",
    "\n",
    "def apply_einthoven_correction(leads_dict, alpha=0.33):\n",
    "    \"\"\"Apply Einthoven's law correction on SHORT leads: II_short ≈ I + III.\"\"\"\n",
    "    if all(k in leads_dict for k in ['I', 'II_short', 'III']):\n",
    "        L1 = leads_dict['I']\n",
    "        L2s = leads_dict['II_short']\n",
    "        L3 = leads_dict['III']\n",
    "        \n",
    "        # Compute error: e = II_short - (I + III)\n",
    "        e = L2s - (L1 + L3)\n",
    "        \n",
    "        # Distribute error\n",
    "        leads_dict['I'] = L1 + alpha * e\n",
    "        leads_dict['III'] = L3 + alpha * e\n",
    "        leads_dict['II_short'] = L2s - alpha * e\n",
    "    \n",
    "    return leads_dict\n",
    "\n",
    "\n",
    "def apply_savgol_smoothing(leads_dict, window=7, polyorder=2):\n",
    "    \"\"\"Apply Savitzky-Golay smoothing to all leads.\"\"\"\n",
    "    smoothed = {}\n",
    "    for lead, signal in leads_dict.items():\n",
    "        if len(signal) >= window:\n",
    "            smoothed[lead] = savgol_filter(signal, window_length=window, polyorder=polyorder)\n",
    "        else:\n",
    "            smoothed[lead] = signal\n",
    "    return smoothed\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)\n",
    "\n",
    "\n",
    "def ensemble_leads(leads_v10, leads_v16, weight_v10=0.5):\n",
    "    \"\"\"Ensemble two lead dictionaries with given weights.\"\"\"\n",
    "    weight_v16 = 1.0 - weight_v10\n",
    "    ensemble = {}\n",
    "    \n",
    "    all_leads = set(leads_v10.keys()) | set(leads_v16.keys())\n",
    "    \n",
    "    for lead in all_leads:\n",
    "        sig_v10 = leads_v10.get(lead, None)\n",
    "        sig_v16 = leads_v16.get(lead, None)\n",
    "        \n",
    "        if sig_v10 is None and sig_v16 is None:\n",
    "            continue\n",
    "        elif sig_v10 is None:\n",
    "            ensemble[lead] = sig_v16\n",
    "        elif sig_v16 is None:\n",
    "            ensemble[lead] = sig_v10\n",
    "        else:\n",
    "            # Ensure same length\n",
    "            target_len = max(len(sig_v10), len(sig_v16))\n",
    "            sig_v10 = resample_signal(sig_v10, target_len)\n",
    "            sig_v16 = resample_signal(sig_v16, target_len)\n",
    "            ensemble[lead] = weight_v10 * sig_v10 + weight_v16 * sig_v16\n",
    "    \n",
    "    return ensemble\n",
    "\n",
    "print(\"Post-processing functions defined.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d8e605b",
   "metadata": {},
   "source": [
    "## Full Inference Pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98c1ee3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_image_ensemble(image_path, apply_smoothing=True, apply_einthoven=True):\n",
    "    \"\"\"Full ensemble inference pipeline for a single image.\n",
    "    \n",
    "    Flow:\n",
    "    1. Preprocess image (Stage0 -> Stage1 -> Resize)\n",
    "    2. Run V10 inference\n",
    "    3. Run V16 inference\n",
    "    4. Convert to lead dictionaries\n",
    "    5. Apply Einthoven correction to each\n",
    "    6. Ensemble the results\n",
    "    7. Apply smoothing\n",
    "    \n",
    "    Returns:\n",
    "        leads: dict mapping lead name to signal array in mV\n",
    "    \"\"\"\n",
    "    # Preprocess\n",
    "    image_bgr = preprocess_image(image_path)\n",
    "    \n",
    "    # Run both models\n",
    "    signal_v10 = run_v10_inference(image_bgr)  # [4, OUTPUT_WIDTH]\n",
    "    signal_v16 = run_v16_inference(image_bgr)  # [4, OUTPUT_WIDTH]\n",
    "    \n",
    "    # Convert to lead dictionaries\n",
    "    leads_v10 = series_to_leads(signal_v10)\n",
    "    leads_v16 = series_to_leads(signal_v16)\n",
    "    \n",
    "    # Apply Einthoven correction to each before ensembling\n",
    "    if apply_einthoven:\n",
    "        leads_v10 = apply_einthoven_correction(leads_v10, alpha=0.33)\n",
    "        leads_v16 = apply_einthoven_correction(leads_v16, alpha=0.33)\n",
    "    \n",
    "    # Ensemble\n",
    "    leads = ensemble_leads(leads_v10, leads_v16, weight_v10=WEIGHT_V10)\n",
    "    \n",
    "    # Apply smoothing after ensemble\n",
    "    if apply_smoothing:\n",
    "        leads = apply_savgol_smoothing(leads, window=7, polyorder=2)\n",
    "    \n",
    "    return leads\n",
    "\n",
    "print(\"Ensemble pipeline defined.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1f5d0cf",
   "metadata": {},
   "source": [
    "## Generate Submission"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dea260ef",
   "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\"Ensemble: V10={WEIGHT_V10:.0%} + V16={WEIGHT_V16:.0%}\")\n",
    "print(f\"\\nSample test.csv rows:\")\n",
    "print(test_df.head())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2218520a",
   "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_ensemble(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",
    "        # Map II_short to the short II segment if needed\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",
    "    # Memory cleanup\n",
    "    if len(all_rows) % 100000 == 0:\n",
    "        gc.collect()\n",
    "\n",
    "print(f\"\\nTotal rows: {len(all_rows)}\")\n",
    "print(f\"Failed images: {len(failed_images)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "339c3e4d",
   "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": "5f30d85f",
   "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} (V10+V16 Ensemble)')\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
}