{ "cells": [ { "cell_type": "code", "execution_count": null, "id": "a45073f8", "metadata": {}, "outputs": [], "source": [ "# Install missing dependencies\n", "!pip install -q cc3d connected-components-3d" ] }, { "cell_type": "code", "execution_count": null, "id": "9f6dfb8d", "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", "import timm\n", "\n", "# Kaggle paths\n", "BASELINE_PATH = '/kaggle/input/hengck23-submit-physionet/hengck23-submit-physionet'\n", "WEIGHTS_PATH = '/kaggle/input/ecg-v9-weights'\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": "code", "execution_count": null, "id": "c3fde343", "metadata": {}, "outputs": [], "source": [ "# Constants\n", "TARGET_HEIGHT, TARGET_WIDTH = 1696, 4352\n", "ZERO_MV = np.array([703.5, 987.5, 1271.5, 1531.5])\n", "MV_TO_PIXEL = 78.5\n", "T0, T1 = 235, 4161\n", "X0, X1 = 0, 2176\n", "Y0, Y1 = 0, 1696\n", "OUTPUT_WIDTH = T1 - T0\n", "SOFT_ARGMAX_TEMP = 100.0 # Temperature for soft-argmax\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": "2f9f365c", "metadata": {}, "source": [ "## Model Architecture" ] }, { "cell_type": "code", "execution_count": null, "id": "82869781", "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 ECGNetV9(nn.Module):\n", " \"\"\"Network with 1-pixel line segmentation + regression head.\"\"\"\n", " def __init__(self, encoder='resnet34', decoder_dims=[128, 64, 32, 16]):\n", " super().__init__()\n", " enc_dims = [64, 128, 256, 512]\n", " \n", " self.encoder = timm.create_model(\n", " f'{encoder}.a3_in1k', pretrained=False, in_chans=3, num_classes=0, global_pool=''\n", " )\n", " \n", " self.dec_blocks = nn.ModuleList()\n", " in_ch = enc_dims[-1]\n", " skip_chs = enc_dims[:-1][::-1] + [0]\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", " self.seg_head = nn.Conv2d(decoder_dims[-1], 4, 1)\n", " \n", " self.reg_head = nn.Sequential(\n", " nn.Conv2d(decoder_dims[-1], 32, 3, padding=1),\n", " nn.ReLU(inplace=True),\n", " nn.AdaptiveAvgPool2d((1, None)),\n", " )\n", " self.reg_out = nn.Sequential(\n", " nn.Conv1d(32, 16, 3, padding=1),\n", " nn.ReLU(inplace=True),\n", " nn.Conv1d(16, 4, 1),\n", " nn.Sigmoid()\n", " )\n", " \n", " def encode(self, x):\n", " enc = []\n", " x = self.encoder.conv1(x)\n", " x = self.encoder.bn1(x)\n", " x = self.encoder.act1(x)\n", " x = self.encoder.layer1(x); enc.append(x)\n", " x = self.encoder.layer2(x); enc.append(x)\n", " x = self.encoder.layer3(x); enc.append(x)\n", " x = self.encoder.layer4(x); enc.append(x)\n", " return enc\n", " \n", " def forward(self, x):\n", " enc = self.encode(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", " \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": "aa7973f4", "metadata": {}, "source": [ "## Load Models" ] }, { "cell_type": "code", "execution_count": null, "id": "009a0fe6", "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 V9 model\n", "print(\"Loading V9 model...\")\n", "model = ECGNetV9(encoder='resnet34')\n", "checkpoint = torch.load(f'{WEIGHTS_PATH}/v9_best_reg.pth', map_location='cpu', weights_only=False)\n", "\n", "# Key is 'model' not 'model_state_dict'\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', '?')}, SNR: {checkpoint.get('snr_reg', '?')} dB\")\n", "print(\"All models loaded!\")" ] }, { "cell_type": "markdown", "id": "d5eddee6", "metadata": {}, "source": [ "## Inference Functions" ] }, { "cell_type": "code", "execution_count": null, "id": "b6526503", "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 V9 model.\n", " \n", " CRITICAL: Must match training preprocessing exactly:\n", " 1. Crop to [Y0:Y1, X0:X1] = [0:1696, 0:2176]\n", " 2. Resize to TARGET_WIDTH x TARGET_HEIGHT = 4352 x 1696\n", " 3. Model outputs [4, 4352] pixel coordinates\n", " 4. Extract signal region [T0:T1] = [235:4161] = 3926 pixels\n", " \"\"\"\n", " h, w = image.shape[:2]\n", " \n", " # Crop to left region - MATCH TRAINING EXACTLY\n", " crop_h = min(h, Y1)\n", " crop_w = min(w, X1)\n", " image_cropped = image[:crop_h, :crop_w]\n", " \n", " # Resize to model input size\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, 4352]\n", " else:\n", " signal_norm = reg_coords.float().cpu().numpy()[0]\n", " signal_full = signal_norm * (TARGET_HEIGHT - 1) # [4, 4352]\n", " \n", " # Extract signal region T0:T1 - THIS IS CRITICAL\n", " signal_pixel = signal_full[:, T0:T1] # [4, OUTPUT_WIDTH=3926]\n", " \n", " return signal_pixel\n", "\n", "\n", "def pixel_to_mv(signal_pixel):\n", " \"\"\"Convert pixel Y-coordinates 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", " return signal_mv\n", "\n", "\n", "def series_to_leads(series_mv):\n", " \"\"\"Convert 4-row series to 12-lead dictionary.\"\"\"\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", " leads['II'] = series_mv[3] # Full rhythm strip\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:\n", " normalized = process_stage0(image)\n", " except Exception as e:\n", " normalized = image\n", " \n", " try:\n", " rectified = process_stage1(normalized)\n", " except Exception as e:\n", " rectified = normalized\n", " \n", " signal_pixel = process_stage2(rectified, use_soft_argmax=True)\n", " signal_mv = pixel_to_mv(signal_pixel)\n", " \n", " return signal_mv\n" ] }, { "cell_type": "code", "execution_count": null, "id": "c55be7ad", "metadata": {}, "outputs": [], "source": [ "# DEBUG: Test soft-argmax vs regression\n", "\n", "@torch.no_grad()\n", "def process_stage2_debug(image, use_soft_argmax=True, crop_half=True):\n", " \"\"\"Stage 2 with debug options.\"\"\"\n", " h, w = image.shape[:2]\n", " print(f\" Input image: {w}x{h}\")\n", " \n", " if crop_half:\n", " crop_w = min(w, 2176)\n", " crop_h = min(h, 1696)\n", " image_cropped = image[:crop_h, :crop_w]\n", " print(f\" After crop: {crop_w}x{crop_h}\")\n", " else:\n", " image_cropped = image\n", " \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, 4352]\n", " print(f\" Using SOFT-ARGMAX (26.7 dB on val)\")\n", " else:\n", " signal_norm = reg_coords.float().cpu().numpy()[0]\n", " signal_full = signal_norm * (TARGET_HEIGHT - 1) # [4, 4352]\n", " print(f\" Using REGRESSION (20.2 dB on val)\")\n", " \n", " print(f\" Full signal shape: {signal_full.shape}\")\n", " print(f\" Pixel range: [{signal_full.min():.1f}, {signal_full.max():.1f}]\")\n", " print(f\" Expected zero mV positions: {ZERO_MV}\")\n", " \n", " for row in range(4):\n", " mean_pos = signal_full[row, T0:T1].mean()\n", " dist_from_zero = abs(mean_pos - ZERO_MV[row])\n", " print(f\" Row {row}: mean={mean_pos:.1f}, expected zero={ZERO_MV[row]:.1f}, diff={dist_from_zero:.1f}\")\n", " \n", " # CRITICAL FIX: Extract T0:T1 region, NOT interpolate full signal\n", " # The model outputs 4352 pixels, but signal is only valid in [T0:T1] = [235:4161] = 3926 pixels\n", " signal_out = signal_full[:, T0:T1] # [4, OUTPUT_WIDTH=3926]\n", " print(f\" Extracted signal region [{T0}:{T1}] = {signal_out.shape[1]} pixels\")\n", " \n", " return signal_out\n", "\n", "\n", "def process_image_debug(image_path, skip_stage01=False, use_soft_argmax=True):\n", " \"\"\"Full pipeline with debug.\"\"\"\n", " image = cv2.imread(str(image_path))\n", " image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n", " print(f\"Original: {image.shape[1]}x{image.shape[0]}\")\n", " \n", " if skip_stage01:\n", " rectified = image\n", " print(\"Skipping Stage 0/1\")\n", " else:\n", " try:\n", " normalized = process_stage0(image)\n", " print(f\"After Stage 0: {normalized.shape[1]}x{normalized.shape[0]}\")\n", " except Exception as e:\n", " print(f\"Stage 0 failed: {e}\")\n", " normalized = image\n", " \n", " try:\n", " rectified = process_stage1(normalized)\n", " print(f\"After Stage 1: {rectified.shape[1]}x{rectified.shape[0]}\")\n", " except Exception as e:\n", " print(f\"Stage 1 failed: {e}\")\n", " rectified = normalized\n", " \n", " signal_pixel = process_stage2_debug(rectified, use_soft_argmax=use_soft_argmax, crop_half=True)\n", " signal_mv = pixel_to_mv(signal_pixel)\n", " \n", " print(f\"Final mV range: [{signal_mv.min():.2f}, {signal_mv.max():.2f}]\")\n", " return signal_mv\n", "\n", "\n", "# Test both methods\n", "test_images = list(Path(f'{COMPETITION_PATH}/test').glob('*.png'))[:1]\n", "if test_images:\n", " test_image_path = test_images[0]\n", " print(f\"=== Testing: {test_image_path.name} ===\\n\")\n", " \n", " print(\"--- SOFT-ARGMAX (recommended, 26.7 dB) ---\")\n", " sig1 = process_image_debug(test_image_path, skip_stage01=False, use_soft_argmax=True)\n", " \n", " print(\"\\n--- REGRESSION (20.2 dB) ---\")\n", " sig2 = process_image_debug(test_image_path, skip_stage01=False, use_soft_argmax=False)\n", "else:\n", " print(\"No test images found\")" ] }, { "cell_type": "markdown", "id": "df121fe2", "metadata": {}, "source": [ "## Generate Submission" ] }, { "cell_type": "code", "execution_count": null, "id": "cca71652", "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...\")" ] }, { "cell_type": "code", "execution_count": null, "id": "d728431c", "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": "da89354d", "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 }