{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe9ab3b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# V11 2x Scale Constants\n",
    "TARGET_HEIGHT, TARGET_WIDTH = 3392, 8704\n",
    "ZERO_MV = np.array([1407.0, 1975.0, 2543.0, 3063.0])\n",
    "MV_TO_PIXEL = 157.0\n",
    "T0, T1 = 470, 8322\n",
    "OUTPUT_WIDTH = T1 - T0  # 7852 pixels\n",
    "SOFT_ARGMAX_TEMP = 100.0\n",
    "\n",
    "# BASE crop coordinates (at 1x resolution) - CRITICAL for matching training\n",
    "BASE_X0, BASE_X1 = 0, 2176\n",
    "BASE_Y0, BASE_Y1 = 0, 1696\n",
    "\n",
    "# Lead layout\n",
    "LEAD_NAMES = ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6']\n",
    "ROW_LAYOUT = [\n",
    "    ['I', 'aVR', 'V1', 'V4'],\n",
    "    ['II', 'aVL', 'V2', 'V5'],\n",
    "    ['III', 'aVF', 'V3', 'V6'],\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa0fdffc",
   "metadata": {},
   "source": [
    "## V11 Model Architecture (EfficientNet-B4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9b8c53e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CoordDecoderBlock(nn.Module):\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",
    "        \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 ECGNetV11(nn.Module):\n",
    "    \"\"\"V11 with EfficientNet-B4 encoder (same architecture as V10).\"\"\"\n",
    "    def __init__(self, encoder='efficientnet_b4', decoder_dims=[256, 128, 64, 32, 16]):\n",
    "        super().__init__()\n",
    "        \n",
    "        # EfficientNet-B4 with features_only\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.enc_dims = enc_dims\n",
    "        \n",
    "        # Decoder blocks\n",
    "        self.dec_blocks = nn.ModuleList()\n",
    "        in_ch = enc_dims[-1]\n",
    "        skip_chs = enc_dims[:-1][::-1] + [0]\n",
    "        \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 i, (skip_ch, out_ch) in enumerate(zip(skip_chs, decoder_dims)):\n",
    "            self.dec_blocks.append(CoordDecoderBlock(in_ch, skip_ch, out_ch))\n",
    "            in_ch = out_ch\n",
    "        \n",
    "        # Segmentation head\n",
    "        self.seg_head = nn.Conv2d(decoder_dims[-1], 4, 1)\n",
    "        \n",
    "        # Regression head\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",
    "        \n",
    "        enc = self.encoder(x)\n",
    "        \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",
    "        \n",
    "        # Ensure output matches input size\n",
    "        if d.shape[2:] != input_size:\n",
    "            d = F.interpolate(d, size=input_size, mode='bilinear', align_corners=False)\n",
    "        \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",
    "        \n",
    "        return seg_logits, reg_coords"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "546b011a",
   "metadata": {},
   "source": [
    "## Load Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6298545",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load Stage 0\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\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 V11 model\n",
    "print(\"Loading V11 2x model...\")\n",
    "model = ECGNetV11(encoder='efficientnet_b4')\n",
    "checkpoint = torch.load(f'{WEIGHTS_PATH}/v11_efficientnet_b4_scale2x_best.pth', map_location='cpu', weights_only=False)\n",
    "\n",
    "state_dict = 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",
    "print(f\"Loaded epoch {checkpoint.get('epoch', '?')}, Holdout SNR: {checkpoint.get('holdout_snr_soft', checkpoint.get('snr_soft', '?')):.2f} dB\")\n",
    "print(\"All models loaded!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "324b0a50",
   "metadata": {},
   "source": [
    "## Inference Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3160d19f",
   "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",
    "@torch.no_grad()\n",
    "def process_stage0(image):\n",
    "    \"\"\"Stage 0: Orientation correction.\"\"\"\n",
    "    batch = image_to_batch(image)\n",
    "    with torch.amp.autocast('cuda', dtype=torch.float32):\n",
    "        output = stage0_net(batch)\n",
    "    rotated, keypoint = output_to_predict(image, batch, output)\n",
    "    normalized, _, _ = normalise_by_homography(rotated, keypoint)\n",
    "    return normalized\n",
    "\n",
    "\n",
    "@torch.no_grad()\n",
    "def process_stage1(image):\n",
    "    \"\"\"Stage 1: Grid rectification.\"\"\"\n",
    "    batch = {'image': torch.from_numpy(np.ascontiguousarray(image.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, batch, output)\n",
    "    rectified = rectify_image(image, gridpoint_xy)\n",
    "    return rectified\n",
    "\n",
    "\n",
    "@torch.no_grad()\n",
    "def process_stage2(image, use_soft_argmax=True):\n",
    "    \"\"\"Stage 2: Signal extraction using V11 2x model.\n",
    "    \n",
    "    CRITICAL: Must match training preprocessing exactly:\n",
    "    1. Crop at BASE resolution [0:1696, 0:2176] (1x scale)\n",
    "    2. Resize to TARGET_WIDTH x TARGET_HEIGHT = 8704 x 3392 (2x scale)\n",
    "    3. Model outputs [4, 8704] pixel coordinates\n",
    "    4. Extract signal region [T0:T1] = [470:8322] = 7852 pixels\n",
    "    \"\"\"\n",
    "    h, w = image.shape[:2]\n",
    "    \n",
    "    # Crop at BASE resolution first (1x scale) - MATCH TRAINING EXACTLY\n",
    "    crop_h = min(h, BASE_Y1)\n",
    "    crop_w = min(w, BASE_X1)\n",
    "    image_cropped = image[:crop_h, :crop_w]\n",
    "    \n",
    "    # Resize to V11 2x target resolution\n",
    "    image_resized = cv2.resize(image_cropped, (TARGET_WIDTH, TARGET_HEIGHT), interpolation=cv2.INTER_LINEAR)\n",
    "    \n",
    "    image_tensor = torch.from_numpy(image_resized.astype(np.float32) / 255.0).permute(2, 0, 1).unsqueeze(0)\n",
    "    image_tensor = image_tensor.to(device)\n",
    "    \n",
    "    with torch.amp.autocast('cuda', dtype=torch.float16):\n",
    "        seg_logits, reg_coords = model(image_tensor)\n",
    "    \n",
    "    if use_soft_argmax:\n",
    "        seg_probs = torch.sigmoid(seg_logits.float())\n",
    "        signal_full = soft_argmax(seg_probs).cpu().numpy()[0]  # [4, 8704]\n",
    "    else:\n",
    "        signal_norm = reg_coords.float().cpu().numpy()[0]\n",
    "        signal_full = signal_norm * (TARGET_HEIGHT - 1)  # [4, 8704]\n",
    "    \n",
    "    # Extract signal region T0:T1\n",
    "    signal_pixel = signal_full[:, T0:T1]  # [4, OUTPUT_WIDTH=7852]\n",
    "    \n",
    "    return signal_pixel\n",
    "\n",
    "\n",
    "def pixel_to_mv(signal_pixel):\n",
    "    \"\"\"Convert pixel Y-coordinates to mV using V11 2x constants.\"\"\"\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",
    "    return signal_mv\n",
    "\n",
    "\n",
    "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)\n",
    "    - Row 1: II, aVL, V2, V5 (4 segments)  \n",
    "    - Row 2: III, aVF, V3, V6 (4 segments)\n",
    "    - Row 3: II full rhythm strip\n",
    "    \"\"\"\n",
    "    leads = {}\n",
    "    segment_width = series_mv.shape[1] // 4\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)\n",
    "    leads['II'] = series_mv[3]\n",
    "    return leads\n",
    "\n",
    "\n",
    "def process_image(image_path):\n",
    "    \"\"\"Full pipeline for single image.\"\"\"\n",
    "    image = cv2.imread(str(image_path))\n",
    "    if image is None:\n",
    "        raise ValueError(f\"Failed to load image: {image_path}\")\n",
    "    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "    \n",
    "    # Try Stage 0 (orientation correction)\n",
    "    try:\n",
    "        normalized = process_stage0(image)\n",
    "    except Exception as e:\n",
    "        normalized = image\n",
    "    \n",
    "    # Try Stage 1 (grid rectification)\n",
    "    try:\n",
    "        rectified = process_stage1(normalized)\n",
    "    except Exception as e:\n",
    "        rectified = normalized\n",
    "    \n",
    "    # Stage 2: Signal extraction with V11 2x\n",
    "    signal_pixel = process_stage2(rectified, use_soft_argmax=True)\n",
    "    signal_mv = pixel_to_mv(signal_pixel)\n",
    "    \n",
    "    return signal_mv"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23039435",
   "metadata": {},
   "source": [
    "## Generate Submission"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e04980b",
   "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\"V11 2x: {TARGET_HEIGHT}x{TARGET_WIDTH}, signal region T0={T0} to T1={T1}\")\n",
    "print(f\"Signal output width: {OUTPUT_WIDTH} pixels\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "087391f2",
   "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",
    "        continue\n",
    "    \n",
    "    img_df = test_df[test_df['id'] == img_id]\n",
    "    \n",
    "    try:\n",
    "        signal_mv = process_image(img_path)\n",
    "        leads = series_to_leads(signal_mv)\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, val in enumerate(signal):\n",
    "                all_rows.append({'id': f\"{img_id}_{i}_{lead}\", 'value': float(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": "a6906f03",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create submission\n",
    "submission_df = pd.DataFrame(all_rows)\n",
    "submission_df.to_parquet('/kaggle/working/submission.parquet', index=False)\n",
    "\n",
    "print(f\"Submission 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
}