{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9553bd45",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install connected-components-3d from baseline\n",
    "!uv pip install --no-deps --system --no-index --find-links='/kaggle/input/hengck23-submit-physionet/hengck23-submit-physionet/setup' 'connected-components-3d'\n",
    "\n",
    "# Imports\n",
    "import os\n",
    "import sys\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from pathlib import Path\n",
    "from tqdm import tqdm\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import cv2\n",
    "import timm\n",
    "\n",
    "# =============================================================================\n",
    "# Kaggle Paths\n",
    "# =============================================================================\n",
    "BASELINE_PATH = '/kaggle/input/hengck23-submit-physionet/hengck23-submit-physionet'\n",
    "WEIGHTS_PATH = '/kaggle/input/ecg-digitization-v23-weights'\n",
    "COMPETITION_PATH = '/kaggle/input/physionet-ecg-image-digitization'\n",
    "\n",
    "# Add baseline to path for Stage0/Stage1\n",
    "sys.path.insert(0, BASELINE_PATH)\n",
    "\n",
    "# Import baseline preprocessing (root-level modules)\n",
    "import stage0_common as s0c\n",
    "import stage1_common as s1c\n",
    "from stage0_model import Net as Stage0Net\n",
    "from stage1_model import Net as Stage1Net\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "print(f\"Device: {device}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "980c9ec8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# =============================================================================\n",
    "# V23 Constants\n",
    "# =============================================================================\n",
    "TARGET_HEIGHT = 1696\n",
    "TARGET_WIDTH = 4352\n",
    "\n",
    "# Crop region\n",
    "X0, X1 = 0, 2176\n",
    "Y0, Y1 = 0, 1696\n",
    "\n",
    "# Signal extraction region\n",
    "T0, T1 = 235, 4161\n",
    "OUTPUT_WIDTH = T1 - T0  # 3926 samples\n",
    "\n",
    "# ECG calibration\n",
    "ZERO_MV = np.array([703.5, 987.5, 1271.5, 1531.5])  # Zero line Y for each row\n",
    "MV_TO_PIXEL = 78.5\n",
    "\n",
    "# Per-row crop parameters\n",
    "CROP_HALF_HEIGHT = 250\n",
    "ROW_HEIGHT = 500\n",
    "\n",
    "# ECG amplitude limits\n",
    "ECG_MV_MIN, ECG_MV_MAX = -10.0, 10.0\n",
    "\n",
    "# Lead layout (12-lead ECG standard)\n",
    "LEAD_NAMES = ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6']\n",
    "LEAD_LAYOUT = [\n",
    "    ['I', 'aVR', 'V1', 'V4'],    # Row 0\n",
    "    ['II', 'aVL', 'V2', 'V5'],   # Row 1\n",
    "    ['III', 'aVF', 'V3', 'V6'],  # Row 2\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5b17a54",
   "metadata": {},
   "source": [
    "## V23 Model Architecture (V22 + DSNT Sub-Pixel Refinement)\n",
    "\n",
    "V23 builds on V22 by adding DSNT (Differentiable Spatial to Numerical Transform) for sub-pixel Y-coordinate prediction (Rank 3 winner technique):\n",
    "\n",
    "1. **ConvNeXt-Base Encoder** with multi-scale feature extraction (L3 + L4)\n",
    "2. **Feature Fusion** combining fine details and semantic features\n",
    "3. **U-Net Decoder** with skip connections\n",
    "4. **Height Attention** to learn which vertical regions matter\n",
    "5. **BiLSTM** for temporal coherence\n",
    "6. **Dual Heads**: V22 head + DSNT sub-pixel head with learnable blend weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b61017b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# =============================================================================\n",
    "# V23 Model Architecture Components\n",
    "# =============================================================================\n",
    "\n",
    "class MultiScaleFeatureFusion(nn.Module):\n",
    "    \"\"\"Fuse L3 and L4 encoder features.\"\"\"\n",
    "    def __init__(self, l3_channels, l4_channels, out_channels=256):\n",
    "        super().__init__()\n",
    "        self.l3_proj = nn.Sequential(\n",
    "            nn.Conv2d(l3_channels, out_channels, 1, bias=False),\n",
    "            nn.BatchNorm2d(out_channels),\n",
    "            nn.GELU(),\n",
    "        )\n",
    "        self.l4_proj = nn.Sequential(\n",
    "            nn.Conv2d(l4_channels, out_channels, 1, bias=False),\n",
    "            nn.BatchNorm2d(out_channels),\n",
    "            nn.GELU(),\n",
    "        )\n",
    "        self.fusion = nn.Sequential(\n",
    "            nn.Conv2d(out_channels * 2, out_channels, 3, padding=1, bias=False),\n",
    "            nn.BatchNorm2d(out_channels),\n",
    "            nn.GELU(),\n",
    "            nn.Conv2d(out_channels, out_channels, 3, padding=1, bias=False),\n",
    "            nn.BatchNorm2d(out_channels),\n",
    "            nn.GELU(),\n",
    "        )\n",
    "        self.out_channels = out_channels\n",
    "    \n",
    "    def forward(self, l3_feat, l4_feat):\n",
    "        l3_proj = self.l3_proj(l3_feat)\n",
    "        l4_proj = self.l4_proj(l4_feat)\n",
    "        l4_up = F.interpolate(l4_proj, size=l3_proj.shape[2:], mode='bilinear', align_corners=True)\n",
    "        fused = torch.cat([l3_proj, l4_up], dim=1)\n",
    "        fused = self.fusion(fused)\n",
    "        return fused\n",
    "\n",
    "\n",
    "class HeightAttention(nn.Module):\n",
    "    \"\"\"Height-wise attention with coordinate encoding.\"\"\"\n",
    "    def __init__(self, in_channels):\n",
    "        super().__init__()\n",
    "        self.attention = nn.Sequential(\n",
    "            nn.Conv2d(in_channels + 2, 64, 3, padding=1),\n",
    "            nn.BatchNorm2d(64),\n",
    "            nn.GELU(),\n",
    "            nn.Conv2d(64, 32, 3, padding=1),\n",
    "            nn.BatchNorm2d(32),\n",
    "            nn.GELU(),\n",
    "            nn.Conv2d(32, 1, 1),\n",
    "        )\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_coord = torch.cat([x, yy, xx], dim=1)\n",
    "        attn = self.attention(x_coord)\n",
    "        attn = F.softmax(attn, dim=2)\n",
    "        pooled = (x * attn).sum(dim=2)\n",
    "        return pooled\n",
    "\n",
    "\n",
    "class BiLSTMTemporal(nn.Module):\n",
    "    \"\"\"Bidirectional LSTM for temporal modeling.\"\"\"\n",
    "    def __init__(self, input_dim, hidden_dim=128, num_layers=2, dropout=0.1):\n",
    "        super().__init__()\n",
    "        self.lstm = nn.LSTM(\n",
    "            input_size=input_dim, hidden_size=hidden_dim,\n",
    "            num_layers=num_layers, batch_first=True,\n",
    "            bidirectional=True, dropout=dropout if num_layers > 1 else 0,\n",
    "        )\n",
    "        self.output_proj = nn.Sequential(\n",
    "            nn.Linear(hidden_dim * 2, hidden_dim),\n",
    "            nn.LayerNorm(hidden_dim),\n",
    "            nn.GELU(),\n",
    "        )\n",
    "        self.output_dim = hidden_dim\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = x.permute(0, 2, 1)  # [B, W, C]\n",
    "        lstm_out, _ = self.lstm(x)\n",
    "        out = self.output_proj(lstm_out)\n",
    "        return out\n",
    "\n",
    "\n",
    "class FinalHead(nn.Module):\n",
    "    \"\"\"V22 regression head.\"\"\"\n",
    "    def __init__(self, input_dim, hidden_dim=64):\n",
    "        super().__init__()\n",
    "        self.conv1d = nn.Sequential(\n",
    "            nn.Conv1d(input_dim, hidden_dim, kernel_size=5, padding=2),\n",
    "            nn.BatchNorm1d(hidden_dim),\n",
    "            nn.GELU(),\n",
    "            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm1d(hidden_dim),\n",
    "            nn.GELU(),\n",
    "        )\n",
    "        self.linear = nn.Sequential(\n",
    "            nn.Linear(hidden_dim, 32),\n",
    "            nn.GELU(),\n",
    "            nn.Linear(32, 1),\n",
    "            nn.Sigmoid(),\n",
    "        )\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = x.permute(0, 2, 1)  # [B, C, W]\n",
    "        x = self.conv1d(x)\n",
    "        x = x.permute(0, 2, 1)  # [B, W, C]\n",
    "        out = self.linear(x).squeeze(-1)\n",
    "        return out\n",
    "\n",
    "\n",
    "class ColumnDSNT(nn.Module):\n",
    "    \"\"\"DSNT Sub-Pixel Head for Y-coordinate prediction (Rank 3 technique).\"\"\"\n",
    "    def __init__(self, in_channels, hidden_channels=32):\n",
    "        super().__init__()\n",
    "        self.conv = nn.Sequential(\n",
    "            nn.Conv2d(in_channels, hidden_channels, 3, padding=1),\n",
    "            nn.BatchNorm2d(hidden_channels),\n",
    "            nn.GELU(),\n",
    "            nn.Conv2d(hidden_channels, hidden_channels, 3, padding=1),\n",
    "            nn.BatchNorm2d(hidden_channels),\n",
    "            nn.GELU(),\n",
    "            nn.Conv2d(hidden_channels, 1, 1),\n",
    "        )\n",
    "        \n",
    "    def forward(self, x):\n",
    "        B, C, H, W = x.shape\n",
    "        logits = self.conv(x)  # [B, 1, H, W]\n",
    "        probs = F.softmax(logits, dim=2)  # softmax over height\n",
    "        y_coords_normalized = torch.linspace(0, 1, H, device=x.device)\n",
    "        y_coords_normalized = y_coords_normalized.view(1, 1, H, 1)\n",
    "        y_expected = (probs * y_coords_normalized).sum(dim=2)  # weighted expectation\n",
    "        y_expected = y_expected.squeeze(1)  # [B, W]\n",
    "        return y_expected, probs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9827310",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PerLeadNetV23(nn.Module):\n",
    "    \"\"\"V23 = V22 + DSNT Sub-Pixel Refinement (Rank 3 Winner Technique)\"\"\"\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, pretrained=pretrained,\n",
    "            features_only=True, out_indices=(0, 1, 2, 3),\n",
    "        )\n",
    "        enc_channels = self.encoder.feature_info.channels()\n",
    "        \n",
    "        # Multi-scale feature fusion (L3 + L4)\n",
    "        self.ms_fusion = MultiScaleFeatureFusion(\n",
    "            l3_channels=enc_channels[2], l4_channels=enc_channels[3], out_channels=256,\n",
    "        )\n",
    "        \n",
    "        # U-Net decoder blocks\n",
    "        self.dec_blocks = nn.ModuleList()\n",
    "        self.dec_blocks.append(self._make_dec_block(256, enc_channels[1], 128))\n",
    "        self.dec_blocks.append(self._make_dec_block(128, enc_channels[0], 64))\n",
    "        self.dec_blocks.append(self._make_dec_block(64, 0, 32))\n",
    "        self.dec_blocks.append(self._make_dec_block(32, 0, 32))\n",
    "        \n",
    "        self.final_up = nn.Sequential(\n",
    "            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),\n",
    "            nn.Conv2d(32, 32, 3, padding=1, bias=False),\n",
    "            nn.BatchNorm2d(32),\n",
    "            nn.GELU(),\n",
    "        )\n",
    "        \n",
    "        # Height attention and BiLSTM (V22 components)\n",
    "        self.height_attention = HeightAttention(in_channels=32)\n",
    "        self.bilstm = BiLSTMTemporal(input_dim=32, hidden_dim=128, num_layers=2, dropout=0.1)\n",
    "        self.head = FinalHead(input_dim=128, hidden_dim=64)\n",
    "        \n",
    "        # DSNT Sub-Pixel Head (NEW in V23)\n",
    "        self.dsnt_head = ColumnDSNT(in_channels=32, hidden_channels=32)\n",
    "        \n",
    "        # Learnable blend weight between V22 and DSNT heads\n",
    "        self.dsnt_weight = nn.Parameter(torch.tensor(0.3))\n",
    "    \n",
    "    def _make_dec_block(self, in_ch, skip_ch, out_ch):\n",
    "        return nn.Sequential(\n",
    "            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),\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",
    "    \n",
    "    def forward(self, x):\n",
    "        B, C, H, W = x.shape\n",
    "        features = self.encoder(x)\n",
    "        l1, l2, l3, l4 = features\n",
    "        \n",
    "        # Multi-scale fusion of L3 and L4\n",
    "        fused = self.ms_fusion(l3, l4)\n",
    "        \n",
    "        # Decoder with skip connections\n",
    "        d = self.dec_blocks[0][0](fused)  # Upsample\n",
    "        if d.shape[2:] != l2.shape[2:]:\n",
    "            d = F.interpolate(d, size=l2.shape[2:], mode='bilinear', align_corners=True)\n",
    "        d = torch.cat([d, l2], dim=1)\n",
    "        d = self.dec_blocks[0][1:](d)\n",
    "        \n",
    "        d = self.dec_blocks[1][0](d)  # Upsample\n",
    "        if d.shape[2:] != l1.shape[2:]:\n",
    "            d = F.interpolate(d, size=l1.shape[2:], mode='bilinear', align_corners=True)\n",
    "        d = torch.cat([d, l1], dim=1)\n",
    "        d = self.dec_blocks[1][1:](d)\n",
    "        \n",
    "        d = self.dec_blocks[2](d)\n",
    "        d = self.dec_blocks[3](d)\n",
    "        d = self.final_up(d)\n",
    "        \n",
    "        # Match width\n",
    "        if d.shape[3] != W:\n",
    "            d = F.interpolate(d, size=(d.shape[2], W), mode='bilinear', align_corners=True)\n",
    "        \n",
    "        # DSNT branch (sub-pixel prediction)\n",
    "        y_dsnt, heatmap = self.dsnt_head(d)\n",
    "        \n",
    "        # V22 branch (height attention + BiLSTM)\n",
    "        pooled = self.height_attention(d)\n",
    "        temporal = self.bilstm(pooled)\n",
    "        y_v22 = self.head(temporal)\n",
    "        \n",
    "        # Blend outputs with learned weight\n",
    "        w = torch.sigmoid(self.dsnt_weight)\n",
    "        y_pred = (1 - w) * y_v22 + w * y_dsnt\n",
    "        \n",
    "        return y_pred"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "947023e0",
   "metadata": {},
   "source": [
    "## Load Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f22050d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load Stage 0 (orientation correction)\n",
    "print(\"Loading Stage 0...\")\n",
    "stage0_net = s0c.load_net(Stage0Net(pretrained=False), f'{BASELINE_PATH}/weight/stage0-last.checkpoint.pth')\n",
    "stage0_net = stage0_net.to(device).eval()\n",
    "\n",
    "# Load Stage 1 (grid rectification)\n",
    "print(\"Loading Stage 1...\")\n",
    "stage1_net = s1c.load_net(Stage1Net(pretrained=False), f'{BASELINE_PATH}/weight/stage1-last.checkpoint.pth')\n",
    "stage1_net = stage1_net.to(device).eval()\n",
    "\n",
    "# Load V23 model\n",
    "print(\"Loading V23 model...\")\n",
    "model = PerLeadNetV23(encoder_name='convnext_base.fb_in22k_ft_in1k', pretrained=False)\n",
    "checkpoint = torch.load(f'{WEIGHTS_PATH}/v23_epoch036.pth', map_location='cpu', weights_only=False)\n",
    "\n",
    "state_dict = checkpoint['model']\n",
    "if any(k.startswith('module.') for k in state_dict.keys()):\n",
    "    state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}\n",
    "\n",
    "model.load_state_dict(state_dict)\n",
    "model.to(device).eval()\n",
    "\n",
    "epoch = checkpoint.get('epoch', '?')\n",
    "snr = checkpoint.get('snr', checkpoint.get('best_snr', 0))\n",
    "dsnt_weight = torch.sigmoid(model.dsnt_weight).item()\n",
    "\n",
    "print(f\"Loaded V23 epoch {epoch}\")\n",
    "print(f\"  Training SNR: {snr:.2f} dB\")\n",
    "print(f\"  DSNT blend weight: {dsnt_weight:.3f} (V22: {1-dsnt_weight:.3f})\")\n",
    "print(\"All models loaded!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac4602e8",
   "metadata": {},
   "source": [
    "## Inference Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8efa672a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def change_color(image_rgb):\n",
    "    \"\"\"CLAHE + denoising - applied inside Stage0.\"\"\"\n",
    "    hsv = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2HSV)\n",
    "    h, s, v = cv2.split(hsv)\n",
    "    v_denoised = cv2.fastNlMeansDenoising(v, h=5.46)\n",
    "    std = np.std(v_denoised)\n",
    "    clip_limit = max(1.0, min(3.5, 2.0 + std / 25))\n",
    "    clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=(8, 8))\n",
    "    v_enhanced = clahe.apply(v_denoised)\n",
    "    hsv_enhanced = cv2.merge([h, s, v_enhanced])\n",
    "    return cv2.cvtColor(hsv_enhanced, cv2.COLOR_HSV2RGB)\n",
    "\n",
    "\n",
    "@torch.no_grad()\n",
    "def process_stage0(img_bgr):\n",
    "    \"\"\"Stage 0: Orientation correction.\"\"\"\n",
    "    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)\n",
    "    img_for_model = change_color(img_rgb)\n",
    "    batch = s0c.image_to_batch(img_for_model)\n",
    "    with torch.amp.autocast(device.type, dtype=torch.float32):\n",
    "        output = stage0_net(batch)\n",
    "    rotated, keypoint = s0c.output_to_predict(img_rgb, batch, output)\n",
    "    normalised, _, _ = s0c.normalise_by_homography(rotated, keypoint)\n",
    "    return normalised\n",
    "\n",
    "\n",
    "@torch.no_grad()\n",
    "def process_stage1(stage0_img_rgb):\n",
    "    \"\"\"Stage 1: Grid rectification.\"\"\"\n",
    "    batch = {'image': torch.from_numpy(np.ascontiguousarray(stage0_img_rgb.transpose(2, 0, 1))).unsqueeze(0)}\n",
    "    with torch.amp.autocast(device.type, dtype=torch.float32):\n",
    "        output = stage1_net(batch)\n",
    "    gridpoint_xy, _ = s1c.output_to_predict(stage0_img_rgb, batch, output)\n",
    "    return s1c.rectify_image(stage0_img_rgb, gridpoint_xy)\n",
    "\n",
    "\n",
    "def crop_row(image, row_idx):\n",
    "    \"\"\"Crop a single row centered on its baseline.\"\"\"\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",
    "    row_crop = image[y_start:y_end, T0:T1, :].copy()\n",
    "    \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(row_crop):\n",
    "    \"\"\"Run V23 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'):\n",
    "        output = model(image_tensor)\n",
    "    \n",
    "    pred_y_crop = output[0].cpu().numpy() * ROW_HEIGHT\n",
    "    return pred_y_crop\n",
    "\n",
    "\n",
    "def convert_crop_to_full(pred_y_crop, row_idx):\n",
    "    \"\"\"Convert crop-relative Y to full image coordinates.\"\"\"\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",
    "    pred_y_full = pred_y_crop - pad_top + y_start\n",
    "    return pred_y_full\n",
    "\n",
    "\n",
    "def convert_y_to_mv(pred_y_full, row_idx):\n",
    "    \"\"\"Convert Y-coordinates to mV.\"\"\"\n",
    "    baseline_y = ZERO_MV[row_idx]\n",
    "    return (baseline_y - pred_y_full) / MV_TO_PIXEL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a9f0a00",
   "metadata": {},
   "outputs": [],
   "source": [
    "# =============================================================================\n",
    "# Output Conversion (NO post-processing - best validation results)\n",
    "# =============================================================================\n",
    "\n",
    "def series_to_leads(pred_mv_rows):\n",
    "    \"\"\"Convert 4-row predictions to 12-lead dictionary.\"\"\"\n",
    "    leads = {}\n",
    "    segment_width = len(pred_mv_rows[0]) // 4\n",
    "    \n",
    "    for row_idx in range(3):\n",
    "        lead_names = LEAD_LAYOUT[row_idx]\n",
    "        for seg_idx, lead_name in enumerate(lead_names):\n",
    "            seg_start = seg_idx * segment_width\n",
    "            seg_end = (seg_idx + 1) * segment_width\n",
    "            leads[lead_name] = pred_mv_rows[row_idx][seg_start:seg_end]\n",
    "    \n",
    "    # Row 3: Full Lead II rhythm strip (10s)\n",
    "    leads['II'] = pred_mv_rows[3]\n",
    "    \n",
    "    return leads"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af2fbd03",
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_image(image_path):\n",
    "    \"\"\"Full V23 inference pipeline for single image.\n",
    "    \n",
    "    Flow:\n",
    "    1. Load image (BGR)\n",
    "    2. Stage0: orientation correction (takes BGR, returns RGB)\n",
    "    3. Stage1: grid rectification (takes RGB, returns RGB)\n",
    "    4. Convert to BGR, crop, resize\n",
    "    5. V23: per-row inference\n",
    "    6. NO post-processing (best local validation: 19.75 dB)\n",
    "    \n",
    "    Returns:\n",
    "        pred_mv_rows: dict of {row_idx: mV signal array}\n",
    "    \"\"\"\n",
    "    # Load as BGR\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",
    "    # Stage 0: Orientation correction (takes BGR, returns RGB)\n",
    "    try:\n",
    "        normalized_rgb = process_stage0(image_bgr)\n",
    "    except Exception:\n",
    "        normalized_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)\n",
    "    \n",
    "    # Stage 1: Grid rectification (takes RGB, returns RGB)\n",
    "    try:\n",
    "        rectified_rgb = process_stage1(normalized_rgb)\n",
    "    except Exception:\n",
    "        rectified_rgb = normalized_rgb\n",
    "    \n",
    "    # Convert to BGR, crop, resize\n",
    "    rectified_bgr = cv2.cvtColor(rectified_rgb, cv2.COLOR_RGB2BGR)\n",
    "    image = rectified_bgr[Y0:Y1, X0:X1]\n",
    "    image = cv2.resize(image, (TARGET_WIDTH, TARGET_HEIGHT), interpolation=cv2.INTER_LINEAR)\n",
    "    \n",
    "    # V23: Process all 4 rows\n",
    "    pred_mv_rows = {}\n",
    "    \n",
    "    for row_idx in range(4):\n",
    "        row_crop = crop_row(image, row_idx)\n",
    "        pred_y_crop = predict_row(row_crop)\n",
    "        pred_y_full = convert_crop_to_full(pred_y_crop, row_idx)\n",
    "        pred_mv = convert_y_to_mv(pred_y_full, row_idx)\n",
    "        \n",
    "        # NO smoothing, NO Einthoven correction, NO baseline correction\n",
    "        # (Best local validation: 19.75 dB without post-processing)\n",
    "        \n",
    "        # Just clamp to reasonable ECG range\n",
    "        pred_mv = np.clip(pred_mv, ECG_MV_MIN, ECG_MV_MAX)\n",
    "        \n",
    "        pred_mv_rows[row_idx] = pred_mv\n",
    "    \n",
    "    return pred_mv_rows"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0351ca60",
   "metadata": {},
   "source": [
    "## Generate Submission"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7603eaae",
   "metadata": {},
   "outputs": [],
   "source": [
    "def resample_signal(signal, target_length):\n",
    "    \"\"\"Resample signal to target length.\"\"\"\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",
    "# 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} (T0={T0} to T1={T1})\")\n",
    "print(f\"Post-processing: DISABLED (best local validation)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69fead24",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_rows = []\n",
    "\n",
    "for img_id in tqdm(image_ids):\n",
    "    img_path = test_dir / f\"{img_id}.png\"\n",
    "    \n",
    "    if not img_path.exists():\n",
    "        print(f\"Missing: {img_path}\")\n",
    "        # Still need to fill in zeros for missing images\n",
    "        img_df = test_df[test_df['id'] == img_id]\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",
    "    img_df = test_df[test_df['id'] == img_id]\n",
    "    \n",
    "    try:\n",
    "        pred_mv_rows = process_image(img_path)\n",
    "        leads = series_to_leads(pred_mv_rows)\n",
    "    except Exception as e:\n",
    "        print(f\"Error {img_id}: {e}\")\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",
    "    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 in range(num_samples):\n",
    "                val = float(signal[i])\n",
    "                # Handle NaN/Inf values\n",
    "                if not np.isfinite(val):\n",
    "                    val = 0.0\n",
    "                all_rows.append({'id': f\"{img_id}_{i}_{lead}\", 'value': val})\n",
    "        else:\n",
    "            for i in range(num_samples):\n",
    "                all_rows.append({'id': f\"{img_id}_{i}_{lead}\", 'value': 0.0})\n",
    "\n",
    "print(f\"Total rows: {len(all_rows)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97f7deff",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create submission directly (no merge with sample_submission)\n",
    "submission_df = pd.DataFrame(all_rows)\n",
    "\n",
    "# Verify no NaN/Inf values\n",
    "nan_count = submission_df['value'].isna().sum()\n",
    "inf_count = np.isinf(submission_df['value']).sum()\n",
    "print(f\"NaN values: {nan_count}, Inf values: {inf_count}\")\n",
    "\n",
    "if nan_count > 0 or inf_count > 0:\n",
    "    print(\"WARNING: Replacing NaN/Inf 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, \"Found NaN values!\"\n",
    "assert np.isinf(submission_df['value']).sum() == 0, \"Found Inf values!\"\n",
    "\n",
    "# Save as CSV\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(submission_df.head(10))\n",
    "print(f\"\\nValue stats:\")\n",
    "print(submission_df['value'].describe())"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}