{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "684b46fa",
   "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.notebook import tqdm\n",
    "import matplotlib.pyplot as plt\n",
    "import timm\n",
    "\n",
    "# Paths\n",
    "BASELINE_PATH = '/data/ecg-digitization/hengck23-submit-physionet/hengck23-submit-physionet'\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",
    "print(\"Imports loaded!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7fc4588",
   "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",
    "\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",
    "]\n",
    "\n",
    "# Checkpoint path\n",
    "CHECKPOINT_DIR = '/data/ecg-digitization/checkpoints'\n",
    "V9_CHECKPOINT = f'{CHECKPOINT_DIR}/v9_best_reg.pth'  # Best regression checkpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8955ee61",
   "metadata": {},
   "source": [
    "## Model Architecture (from train_v9_optimal.py)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e2a1fbb",
   "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 output + 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=True, in_chans=3, num_classes=0, global_pool=''\n",
    "        )\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",
    "        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 (4 channels for 4 rows)\n",
    "        self.seg_head = nn.Conv2d(decoder_dims[-1], 4, 1)\n",
    "        \n",
    "        # Regression head (direct y-coordinate prediction)\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)),  # Pool H to 1\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()  # Output in [0, 1]\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",
    "        \n",
    "        # Decode\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",
    "        # Segmentation output\n",
    "        seg_logits = self.seg_head(d)  # [B, 4, H, W]\n",
    "        \n",
    "        # Regression output\n",
    "        reg_feat = self.reg_head(d).squeeze(2)  # [B, 32, W]\n",
    "        reg_coords = self.reg_out(reg_feat)  # [B, 4, W]\n",
    "        \n",
    "        return seg_logits, reg_coords\n",
    "\n",
    "print(\"Model architecture defined!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2713302d",
   "metadata": {},
   "source": [
    "## Load Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96d6a8b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
    "print(f\"Using device: {device}\")\n",
    "\n",
    "# Load Stage 0 (orientation)\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 (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 v9 model\n",
    "print(\"Loading v9 model...\")\n",
    "model = ECGNetV9(encoder='resnet34')\n",
    "checkpoint = torch.load(V9_CHECKPOINT, map_location='cpu')\n",
    "\n",
    "# Handle DDP state dict\n",
    "state_dict = checkpoint['model_state_dict']\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 checkpoint from epoch {checkpoint.get('epoch', 'unknown')}\")\n",
    "print(f\"Best SNR: {checkpoint.get('best_snr', checkpoint.get('snr', 'unknown'))} dB\")\n",
    "print(\"All models loaded!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6756e28",
   "metadata": {},
   "source": [
    "## Inference Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cf06dfa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def soft_argmax(heatmap, temperature=100.0):\n",
    "    \"\"\"Soft-argmax for sub-pixel Y-coordinate extraction.\"\"\"\n",
    "    B, C, H, W = heatmap.shape\n",
    "    \n",
    "    # Apply softmax along H dimension\n",
    "    weights = F.softmax(heatmap * temperature, dim=2)\n",
    "    \n",
    "    # Create Y-coordinate grid\n",
    "    y_coords = torch.linspace(0, 1, H, device=heatmap.device, dtype=heatmap.dtype)\n",
    "    y_coords = y_coords.view(1, 1, H, 1).expand(B, C, -1, W)\n",
    "    \n",
    "    # Weighted sum\n",
    "    expected_y = (weights * y_coords).sum(dim=2)  # [B, C, W]\n",
    "    \n",
    "    return expected_y\n",
    "\n",
    "\n",
    "@torch.no_grad()\n",
    "def process_stage0(image):\n",
    "    \"\"\"Stage 0: Orientation correction.\"\"\"\n",
    "    batch = image_to_batch(image)\n",
    "    \n",
    "    with torch.amp.autocast('cuda', dtype=torch.float32):\n",
    "        output = stage0_net(batch)\n",
    "    \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",
    "    \n",
    "    with torch.amp.autocast('cuda', dtype=torch.float32):\n",
    "        output = stage1_net(batch)\n",
    "    \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_regression=True):\n",
    "    \"\"\"\n",
    "    Stage 2: Signal extraction using v9 model.\n",
    "    \n",
    "    Args:\n",
    "        image: Rectified RGB image\n",
    "        use_regression: If True, use regression head; else use soft-argmax on segmentation\n",
    "    \n",
    "    Returns:\n",
    "        signal_pixel: [4, W] Y-coordinates in pixels\n",
    "    \"\"\"\n",
    "    # Crop and resize\n",
    "    image_cropped = image[Y0:Y1, X0:X1]\n",
    "    image_resized = cv2.resize(image_cropped, (TARGET_WIDTH, TARGET_HEIGHT), interpolation=cv2.INTER_LINEAR)\n",
    "    \n",
    "    # Convert to tensor\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",
    "    # Forward pass\n",
    "    with torch.amp.autocast('cuda', dtype=torch.bfloat16):\n",
    "        seg_logits, reg_coords = model(image_tensor)\n",
    "    \n",
    "    if use_regression:\n",
    "        # Use regression output (normalized 0-1)\n",
    "        signal_norm = reg_coords.float().cpu().numpy()[0]  # [4, W]\n",
    "        signal_pixel = signal_norm * (TARGET_HEIGHT - 1)\n",
    "    else:\n",
    "        # Use soft-argmax on segmentation\n",
    "        signal_norm = soft_argmax(seg_logits.float(), temperature=100.0).cpu().numpy()[0]  # [4, W]\n",
    "        signal_pixel = signal_norm * (TARGET_HEIGHT - 1)\n",
    "    \n",
    "    # Interpolate to match output region T0:T1\n",
    "    signal_out = np.zeros((4, OUTPUT_WIDTH), dtype=np.float32)\n",
    "    for i in range(4):\n",
    "        x_old = np.linspace(0, 1, signal_pixel.shape[1])\n",
    "        x_new = np.linspace(0, 1, OUTPUT_WIDTH)\n",
    "        signal_out[i] = np.interp(x_new, x_old, signal_pixel[i])\n",
    "    \n",
    "    return signal_out\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",
    "    \"\"\"\n",
    "    Convert 4-row series (mV) to 12-lead dictionary.\n",
    "    \n",
    "    Args:\n",
    "        series_mv: [4, W] signal in mV\n",
    "    \n",
    "    Returns:\n",
    "        dict: Lead name -> signal array\n",
    "    \"\"\"\n",
    "    leads = {}\n",
    "    segment_width = series_mv.shape[1] // 4\n",
    "    \n",
    "    # Rows 0-2: 4 leads each\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: Full II rhythm strip\n",
    "    leads['II'] = series_mv[3]\n",
    "    \n",
    "    return leads\n",
    "\n",
    "\n",
    "def process_full_pipeline(image_path, use_regression=True):\n",
    "    \"\"\"Run full pipeline on an image.\"\"\"\n",
    "    # Load image\n",
    "    image = cv2.imread(str(image_path))\n",
    "    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "    \n",
    "    # Stage 0\n",
    "    try:\n",
    "        normalized = process_stage0(image)\n",
    "    except Exception as e:\n",
    "        print(f\"Stage 0 failed: {e}\")\n",
    "        normalized = image\n",
    "    \n",
    "    # Stage 1\n",
    "    try:\n",
    "        rectified = process_stage1(normalized)\n",
    "    except Exception as e:\n",
    "        print(f\"Stage 1 failed: {e}\")\n",
    "        rectified = normalized\n",
    "    \n",
    "    # Stage 2\n",
    "    signal_pixel = process_stage2(rectified, use_regression=use_regression)\n",
    "    signal_mv = pixel_to_mv(signal_pixel)\n",
    "    \n",
    "    return signal_mv, rectified\n",
    "\n",
    "print(\"Inference functions ready!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2a7e769",
   "metadata": {},
   "source": [
    "## Test on Sample Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8404343e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test data path\n",
    "TEST_DIR = Path('/data/ecg-digitization/kaggle/test')\n",
    "test_images = list(TEST_DIR.glob('*.png'))\n",
    "print(f\"Found {len(test_images)} test images\")\n",
    "\n",
    "# Process first image\n",
    "if test_images:\n",
    "    test_img_path = test_images[0]\n",
    "    print(f\"\\nProcessing: {test_img_path.name}\")\n",
    "    \n",
    "    signal_mv, rectified = process_full_pipeline(test_img_path, use_regression=True)\n",
    "    leads = series_to_leads(signal_mv)\n",
    "    \n",
    "    print(f\"Signal shape: {signal_mv.shape}\")\n",
    "    print(f\"Leads extracted: {list(leads.keys())}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3bb9fbc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize results\n",
    "if test_images:\n",
    "    fig, axes = plt.subplots(5, 1, figsize=(16, 12))\n",
    "    \n",
    "    # Show rectified image\n",
    "    axes[0].imshow(rectified)\n",
    "    axes[0].set_title('Rectified Image')\n",
    "    axes[0].axis('off')\n",
    "    \n",
    "    # Plot each row's signal\n",
    "    row_names = ['Row 0 (I, aVR, V1, V4)', 'Row 1 (II, aVL, V2, V5)', \n",
    "                 'Row 2 (III, aVF, V3, V6)', 'Row 3 (II rhythm strip)']\n",
    "    for i in range(4):\n",
    "        axes[i+1].plot(signal_mv[i], linewidth=0.5)\n",
    "        axes[i+1].set_title(row_names[i])\n",
    "        axes[i+1].set_ylabel('mV')\n",
    "        axes[i+1].grid(True, alpha=0.3)\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8be96502",
   "metadata": {},
   "source": [
    "## Generate Submission"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92c8ca41",
   "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",
    "def create_submission_row(image_id, lead_name, signal, fs, num_samples):\n",
    "    \"\"\"Create submission rows for a single lead.\"\"\"\n",
    "    # Resample to target length\n",
    "    signal_resampled = resample_signal(signal, num_samples)\n",
    "    \n",
    "    # Create rows\n",
    "    rows = []\n",
    "    for i, val in enumerate(signal_resampled):\n",
    "        row_id = f\"{image_id}_{i}_{lead_name}\"\n",
    "        rows.append({'id': row_id, 'value': float(val)})\n",
    "    \n",
    "    return rows\n",
    "\n",
    "\n",
    "def generate_submission(test_csv_path, test_dir, output_path, use_regression=True):\n",
    "    \"\"\"Generate full submission file.\"\"\"\n",
    "    test_df = pd.read_csv(test_csv_path)\n",
    "    \n",
    "    # Get unique image IDs\n",
    "    image_ids = test_df['id'].unique()\n",
    "    print(f\"Processing {len(image_ids)} images...\")\n",
    "    \n",
    "    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\"Warning: {img_path} not found\")\n",
    "            continue\n",
    "        \n",
    "        # Get expected leads and lengths for this image\n",
    "        img_df = test_df[test_df['id'] == img_id]\n",
    "        \n",
    "        # Process image\n",
    "        try:\n",
    "            signal_mv, _ = process_full_pipeline(img_path, use_regression=use_regression)\n",
    "            leads = series_to_leads(signal_mv)\n",
    "        except Exception as e:\n",
    "            print(f\"Error processing {img_id}: {e}\")\n",
    "            # Create zero signals\n",
    "            for _, row in img_df.iterrows():\n",
    "                lead = row['lead']\n",
    "                num_samples = row['number_of_rows']\n",
    "                for i in range(num_samples):\n",
    "                    all_rows.append({'id': f\"{img_id}_{i}_{lead}\", 'value': 0.0})\n",
    "            continue\n",
    "        \n",
    "        # Create submission rows for each lead\n",
    "        for _, row in img_df.iterrows():\n",
    "            lead = row['lead']\n",
    "            fs = row['fs']\n",
    "            num_samples = row['number_of_rows']\n",
    "            \n",
    "            if lead in leads:\n",
    "                signal = leads[lead]\n",
    "                rows = create_submission_row(img_id, lead, signal, fs, num_samples)\n",
    "                all_rows.extend(rows)\n",
    "            else:\n",
    "                print(f\"Warning: Lead {lead} not found for image {img_id}\")\n",
    "                for i in range(num_samples):\n",
    "                    all_rows.append({'id': f\"{img_id}_{i}_{lead}\", 'value': 0.0})\n",
    "    \n",
    "    # Create submission DataFrame\n",
    "    submission_df = pd.DataFrame(all_rows)\n",
    "    \n",
    "    # Save\n",
    "    submission_df.to_parquet(output_path, index=False)\n",
    "    print(f\"\\nSubmission saved to {output_path}\")\n",
    "    print(f\"Total rows: {len(submission_df)}\")\n",
    "    \n",
    "    return submission_df\n",
    "\n",
    "print(\"Submission functions ready!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42370190",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate submission\n",
    "TEST_CSV = '/data/ecg-digitization/kaggle/test.csv'\n",
    "TEST_DIR = Path('/data/ecg-digitization/kaggle/test')\n",
    "OUTPUT_PATH = '/data/ecg-digitization/outputs/submission_v9_partial.parquet'\n",
    "\n",
    "# Create output directory\n",
    "Path('/data/ecg-digitization/outputs').mkdir(exist_ok=True)\n",
    "\n",
    "# Generate\n",
    "submission_df = generate_submission(TEST_CSV, TEST_DIR, OUTPUT_PATH, use_regression=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08c48192",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Preview submission\n",
    "print(\"Submission preview:\")\n",
    "print(submission_df.head(20))\n",
    "print(f\"\\nValue statistics:\")\n",
    "print(submission_df['value'].describe())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37786f5e",
   "metadata": {},
   "source": [
    "## Validation on Stage1 Training Data (with GT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "850344ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_snr(pred, gt):\n",
    "    \"\"\"Calculate SNR in dB.\"\"\"\n",
    "    signal_power = np.mean(gt ** 2)\n",
    "    noise_power = np.mean((pred - gt) ** 2)\n",
    "    if noise_power < 1e-10:\n",
    "        return 100.0\n",
    "    return 10 * np.log10(signal_power / noise_power)\n",
    "\n",
    "# Test on stage1 data (if available)\n",
    "STAGE1_DIR = Path('/data/ecg-digitization/stage1_data/train')\n",
    "if STAGE1_DIR.exists():\n",
    "    # Get a few samples\n",
    "    stage1_images = list(STAGE1_DIR.glob('*.png'))[:5]\n",
    "    \n",
    "    for img_path in stage1_images:\n",
    "        csv_path = img_path.with_suffix('.csv')\n",
    "        if not csv_path.exists():\n",
    "            continue\n",
    "        \n",
    "        print(f\"\\n{img_path.name}:\")\n",
    "        \n",
    "        # Get prediction\n",
    "        signal_mv, rectified = process_full_pipeline(img_path, use_regression=True)\n",
    "        \n",
    "        # Load GT\n",
    "        gt_df = pd.read_csv(csv_path)\n",
    "        \n",
    "        # Compare II lead (row 3)\n",
    "        if 'II' in gt_df.columns:\n",
    "            gt_signal = gt_df['II'].dropna().values\n",
    "            pred_signal = resample_signal(signal_mv[3], len(gt_signal))\n",
    "            snr = calculate_snr(pred_signal, gt_signal)\n",
    "            print(f\"  II SNR: {snr:.2f} dB\")\n",
    "else:\n",
    "    print(\"Stage1 data not found for validation\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6edc5cda",
   "metadata": {},
   "source": [
    "## Compare Regression vs Soft-Argmax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cebf1079",
   "metadata": {},
   "outputs": [],
   "source": [
    "if test_images:\n",
    "    test_img_path = test_images[0]\n",
    "    \n",
    "    # Get both outputs\n",
    "    signal_reg, rectified = process_full_pipeline(test_img_path, use_regression=True)\n",
    "    signal_soft, _ = process_full_pipeline(test_img_path, use_regression=False)\n",
    "    \n",
    "    # Plot comparison\n",
    "    fig, axes = plt.subplots(4, 1, figsize=(16, 10))\n",
    "    \n",
    "    for i in range(4):\n",
    "        axes[i].plot(signal_reg[i], label='Regression', linewidth=0.8, alpha=0.8)\n",
    "        axes[i].plot(signal_soft[i], label='Soft-Argmax', linewidth=0.8, alpha=0.8)\n",
    "        axes[i].set_title(f'Row {i}')\n",
    "        axes[i].set_ylabel('mV')\n",
    "        axes[i].legend()\n",
    "        axes[i].grid(True, alpha=0.3)\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb62ddda",
   "metadata": {},
   "source": [
    "## Upload Submission to Kaggle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9d403b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Upload submission (uncomment to run)\n",
    "import subprocess\n",
    "\n",
    "COMPETITION_NAME = 'physionet-ecg-image-digitization'\n",
    "SUBMISSION_MESSAGE = 'v9 model - partial training (epoch 15, 15.66 dB regression SNR)'\n",
    "\n",
    "# cmd = f'kaggle competitions submit -c {COMPETITION_NAME} -f {OUTPUT_PATH} -m \"{SUBMISSION_MESSAGE}\"'\n",
    "# result = subprocess.run(cmd, shell=True, capture_output=True, text=True)\n",
    "# print(result.stdout)\n",
    "# if result.stderr:\n",
    "#     print(f\"Error: {result.stderr}\")\n",
    "\n",
    "print(\"To submit, uncomment the above code or run:\")\n",
    "print(f'kaggle competitions submit -c {COMPETITION_NAME} -f {OUTPUT_PATH} -m \"{SUBMISSION_MESSAGE}\"')"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}