{
 "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",
    "    NOTE: Baseline expects RGB input.\n",
    "    \"\"\"\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",
    "    NOTE: Baseline expects RGB input, returns RGB.\n",
    "    \"\"\"\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_bgr, use_soft_argmax=True):\n",
    "    \"\"\"Stage 2: Signal extraction using V9 model.\n",
    "    \n",
    "    CRITICAL: V9 was trained on BGR images (cv2.imread without conversion).\n",
    "    Input should be BGR format.\n",
    "    \n",
    "    Processing:\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_bgr.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_bgr[: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",
    "    # Normalize to [0, 1] - training uses BGR directly\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",
    "    \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",
    "    \n",
    "    Flow:\n",
    "    1. Load image (BGR)\n",
    "    2. Convert to RGB for stage0/stage1 (baseline expects RGB)\n",
    "    3. Stage0: orientation correction (RGB->RGB)\n",
    "    4. Stage1: grid rectification (RGB->RGB)\n",
    "    5. Convert back to BGR for V9 model (trained on BGR)\n",
    "    6. Stage2: signal extraction\n",
    "    \"\"\"\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",
    "    # Try Stage 0 (orientation correction) - expects RGB\n",
    "    try:\n",
    "        normalized = process_stage0(image_rgb)\n",
    "    except Exception as e:\n",
    "        normalized = image_rgb\n",
    "    \n",
    "    # Try Stage 1 (grid rectification) - expects RGB\n",
    "    try:\n",
    "        rectified = process_stage1(normalized)\n",
    "    except Exception as e:\n",
    "        rectified = normalized\n",
    "    \n",
    "    # Convert back to BGR for V9 model (trained on BGR)\n",
    "    rectified_bgr = cv2.cvtColor(rectified, cv2.COLOR_RGB2BGR)\n",
    "    \n",
    "    # Stage 2: Signal extraction - expects BGR\n",
    "    signal_pixel = process_stage2(rectified_bgr, use_soft_argmax=True)\n",
    "    signal_mv = pixel_to_mv(signal_pixel)\n",
    "    \n",
    "    return signal_mv"
   ]
  },
  {
   "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_bgr, use_soft_argmax=True, crop_half=True):\n",
    "    \"\"\"Stage 2 with debug options. Expects BGR input.\"\"\"\n",
    "    h, w = image_bgr.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_bgr[:crop_h, :crop_w]\n",
    "        print(f\"  After crop: {crop_w}x{crop_h}\")\n",
    "    else:\n",
    "        image_cropped = image_bgr\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_bgr = cv2.imread(str(image_path))\n",
    "    print(f\"Original: {image_bgr.shape[1]}x{image_bgr.shape[0]} (BGR)\")\n",
    "    \n",
    "    # Convert to RGB for stage0/stage1 baseline\n",
    "    image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)\n",
    "    \n",
    "    if skip_stage01:\n",
    "        rectified_rgb = image_rgb\n",
    "        print(\"Skipping Stage 0/1\")\n",
    "    else:\n",
    "        try:\n",
    "            normalized = process_stage0(image_rgb)\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_rgb\n",
    "        \n",
    "        try:\n",
    "            rectified_rgb = process_stage1(normalized)\n",
    "            print(f\"After Stage 1: {rectified_rgb.shape[1]}x{rectified_rgb.shape[0]}\")\n",
    "        except Exception as e:\n",
    "            print(f\"Stage 1 failed: {e}\")\n",
    "            rectified_rgb = normalized\n",
    "    \n",
    "    # Convert back to BGR for V9 model (trained on BGR)\n",
    "    rectified_bgr = cv2.cvtColor(rectified_rgb, cv2.COLOR_RGB2BGR)\n",
    "    \n",
    "    signal_pixel = process_stage2_debug(rectified_bgr, 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
}