{ "cells": [ { "cell_type": "markdown", "id": "3fc80d50", "metadata": {}, "source": [ "# V20 ECG Digitization Inference (Integral Regression + TCN)\n", "\n", "This notebook performs inference using the V20 model architecture with:\n", "- **Encoder**: ConvNeXt-Base (pretrained on ImageNet-22k)\n", "- **Decoder**: U-Net style decoder with skip connections\n", "- **Temporal**: TCN (Temporal Convolutional Network) refiner\n", "- **Regression**: Integral regression with soft-argmax for smooth Y-coordinate prediction\n", "\n", "**Preprocessing**: Stage0 (orientation) + Stage1 (grid rectification) from hengck23 baseline\n", "\n", "**Dual-GPU Support**: Automatically detects and uses multiple GPUs for parallel inference" ] }, { "cell_type": "code", "execution_count": null, "id": "0fb88501", "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 threading\n", "import numpy as np\n", "import pandas as pd\n", "from pathlib import Path\n", "from tqdm import tqdm\n", "from scipy.signal import savgol_filter\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-v20-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", "# Detect available GPUs for dual-GPU inference\n", "if torch.cuda.is_available():\n", " num_gpus = torch.cuda.device_count()\n", " devices = [torch.device(f'cuda:{i}') for i in range(num_gpus)]\n", " print(f\"Found {num_gpus} GPU(s): {[torch.cuda.get_device_name(i) for i in range(num_gpus)]}\")\n", "else:\n", " devices = [torch.device('cpu')]\n", " print(\"No GPU available, using CPU\")\n", "\n", "print(f\"Devices: {devices}\")" ] }, { "cell_type": "code", "execution_count": null, "id": "3b1f0eab", "metadata": {}, "outputs": [], "source": [ "# =============================================================================\n", "# V20 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", "# Integral Regression parameters\n", "NUM_BINS = 128\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", "]\n", "\n", "# V20 baseline offsets (computed from epoch 14)\n", "BASELINE_OFFSETS = {\n", " 'I': 0.0052, 'II': 0.0071, 'III': 0.0049,\n", " 'aVR': 0.0067, 'aVL': 0.0059, 'aVF': 0.0067,\n", " 'V1': 0.0066, 'V2': 0.0061, 'V3': 0.0067,\n", " 'V4': 0.0068, 'V5': 0.0062, 'V6': 0.0065,\n", "}" ] }, { "cell_type": "markdown", "id": "f0ceb93a", "metadata": {}, "source": [ "## V20 Model Architecture (ConvNeXt-Base + U-Net + TCN + Integral Regression)" ] }, { "cell_type": "code", "execution_count": null, "id": "74b937bd", "metadata": {}, "outputs": [], "source": [ "class CoordConv2d(nn.Module):\n", " \"\"\"Conv2d with coordinate channels.\"\"\"\n", " def __init__(self, in_channels, out_channels, kernel_size, **kwargs):\n", " super().__init__()\n", " self.conv = nn.Conv2d(in_channels + 2, out_channels, kernel_size, **kwargs)\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 = torch.cat([x, yy, xx], dim=1)\n", " return self.conv(x)\n", "\n", "\n", "class TCNBlock(nn.Module):\n", " \"\"\"Temporal Convolutional Network block with dilated convolutions.\"\"\"\n", " def __init__(self, in_channels, out_channels, kernel_size=3, dilation=1, dropout=0.1):\n", " super().__init__()\n", " padding = (kernel_size - 1) * dilation // 2\n", " \n", " self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size, padding=padding, dilation=dilation)\n", " self.bn1 = nn.BatchNorm1d(out_channels)\n", " self.conv2 = nn.Conv1d(out_channels, out_channels, kernel_size, padding=padding, dilation=dilation)\n", " self.bn2 = nn.BatchNorm1d(out_channels)\n", " self.dropout = nn.Dropout(dropout)\n", " self.activation = nn.GELU()\n", " self.residual = nn.Conv1d(in_channels, out_channels, 1) if in_channels != out_channels else nn.Identity()\n", " \n", " def forward(self, x):\n", " residual = self.residual(x)\n", " out = self.dropout(self.activation(self.bn1(self.conv1(x))))\n", " out = self.activation(self.bn2(self.conv2(out)) + residual)\n", " return out\n", "\n", "\n", "class TCNRefiner(nn.Module):\n", " \"\"\"Multi-scale TCN for temporal refinement.\"\"\"\n", " def __init__(self, in_channels, hidden_channels=64, num_layers=4, kernel_size=5, dropout=0.1):\n", " super().__init__()\n", " dilations = [2**i for i in range(num_layers)]\n", " layers = [TCNBlock(in_channels, hidden_channels, kernel_size, dilations[0], dropout)]\n", " for dilation in dilations[1:]:\n", " layers.append(TCNBlock(hidden_channels, hidden_channels, kernel_size, dilation, dropout))\n", " self.tcn = nn.Sequential(*layers)\n", " self.output_proj = nn.Conv1d(hidden_channels, in_channels, 1)\n", " \n", " def forward(self, x):\n", " residual = x\n", " out = self.output_proj(self.tcn(x))\n", " return out + residual\n", "\n", "\n", "class UNetDecoderBlock(nn.Module):\n", " \"\"\"U-Net style decoder block with skip connections.\"\"\"\n", " def __init__(self, in_ch, skip_ch, out_ch):\n", " super().__init__()\n", " self.conv = nn.Sequential(\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", " self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)\n", " \n", " def forward(self, x, skip=None):\n", " x = self.upsample(x)\n", " if skip is not None:\n", " if x.shape[2:] != skip.shape[2:]:\n", " x = F.interpolate(x, size=skip.shape[2:], mode='bilinear', align_corners=True)\n", " x = torch.cat([x, skip], dim=1)\n", " return self.conv(x)\n", "\n", "\n", "class IntegralRegressionHead(nn.Module):\n", " \"\"\"Integral regression head using soft-argmax.\"\"\"\n", " def __init__(self, in_channels, num_bins=NUM_BINS, temperature=1.0):\n", " super().__init__()\n", " self.num_bins = num_bins\n", " self.temperature = temperature\n", " bin_centers = torch.linspace(0, 1, num_bins)\n", " self.register_buffer('bin_centers', bin_centers)\n", " \n", " self.heatmap_conv = nn.Sequential(\n", " CoordConv2d(in_channels, 128, 3, padding=1),\n", " nn.BatchNorm2d(128),\n", " nn.GELU(),\n", " nn.Conv2d(128, 64, 3, padding=1),\n", " nn.BatchNorm2d(64),\n", " nn.GELU(),\n", " nn.Conv2d(64, num_bins, 1),\n", " )\n", " self.height_attention = nn.Sequential(\n", " nn.Conv2d(in_channels, 64, 3, padding=1),\n", " nn.BatchNorm2d(64),\n", " nn.GELU(),\n", " nn.Conv2d(64, 1, 1),\n", " )\n", " \n", " def forward(self, x, temperature=None):\n", " B, C, H, W = x.shape\n", " temp = temperature if temperature is not None else self.temperature\n", " \n", " heatmap_2d = self.heatmap_conv(x)\n", " attn = F.softmax(self.height_attention(x), dim=2)\n", " heatmap = (heatmap_2d * attn).sum(dim=2)\n", " heatmap_prob = F.softmax(heatmap / temp, dim=1)\n", " coords = (heatmap_prob * self.bin_centers.view(1, -1, 1)).sum(dim=1)\n", " \n", " return coords, heatmap_prob\n", "\n", "\n", "class IntegralRegressionNet(nn.Module):\n", " \"\"\"V20: Integral Regression Network with TCN for ECG digitization.\"\"\"\n", " \n", " def __init__(self, encoder_name='convnext_base.fb_in22k_ft_in1k', pretrained=True, \n", " num_bins=NUM_BINS, temperature=1.0, use_tcn=True, tcn_layers=4):\n", " super().__init__()\n", " self.num_bins = num_bins\n", " self.use_tcn = use_tcn\n", " \n", " self.encoder = timm.create_model(encoder_name, pretrained=pretrained, \n", " features_only=True, out_indices=(0, 1, 2, 3))\n", " enc_channels = self.encoder.feature_info.channels()\n", " \n", " decoder_dims = [256, 128, 64, 32]\n", " self.dec_blocks = nn.ModuleList()\n", " in_ch = enc_channels[-1]\n", " skip_channels = enc_channels[:-1][::-1] + [0]\n", " \n", " for skip_ch, out_ch in zip(skip_channels, decoder_dims):\n", " self.dec_blocks.append(UNetDecoderBlock(in_ch, skip_ch, out_ch))\n", " in_ch = out_ch\n", " \n", " self.final_up = nn.Sequential(\n", " nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),\n", " nn.Conv2d(decoder_dims[-1], decoder_dims[-1], 3, padding=1, bias=False),\n", " nn.BatchNorm2d(decoder_dims[-1]),\n", " nn.GELU(),\n", " )\n", " \n", " self.height_pool = nn.Sequential(\n", " nn.Conv2d(decoder_dims[-1], 64, 3, padding=1),\n", " nn.BatchNorm2d(64),\n", " nn.GELU(),\n", " nn.Conv2d(64, 1, 1),\n", " )\n", " \n", " if use_tcn:\n", " self.tcn_refiner = TCNRefiner(decoder_dims[-1], 64, tcn_layers, 5, 0.1)\n", " \n", " self.integral_head = IntegralRegressionHead(decoder_dims[-1], num_bins, temperature)\n", " \n", " def forward(self, x, temperature=None):\n", " B, C, H, W = x.shape\n", " \n", " features = self.encoder(x)\n", " d = features[-1]\n", " skips = features[:-1][::-1] + [None]\n", " \n", " for block, skip in zip(self.dec_blocks, skips):\n", " d = block(d, skip)\n", " \n", " d = self.final_up(d)\n", " \n", " if d.shape[3] != W:\n", " d = F.interpolate(d, size=(d.shape[2], W), mode='bilinear', align_corners=True)\n", " \n", " if self.use_tcn:\n", " attn = F.softmax(self.height_pool(d), dim=2)\n", " d_1d = (d * attn).sum(dim=2)\n", " d_1d = self.tcn_refiner(d_1d)\n", " d = d_1d.unsqueeze(2).expand(-1, -1, d.shape[2], -1)\n", " \n", " coords, heatmap = self.integral_head(d, temperature)\n", " return coords, heatmap" ] }, { "cell_type": "markdown", "id": "20726808", "metadata": {}, "source": [ "## Load Models (Dual-GPU Support)" ] }, { "cell_type": "code", "execution_count": null, "id": "d119b858", "metadata": {}, "outputs": [], "source": [ "# Load Stage 0 (orientation correction) - on GPU 0\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(devices[0]).eval()\n", "\n", "# Load Stage 1 (grid rectification) - on GPU 0\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(devices[0]).eval()\n", "\n", "# Load V20 model on each GPU\n", "print(f\"Loading V20 model on {len(devices)} device(s)...\")\n", "v20_models = []\n", "checkpoint = torch.load(f'{WEIGHTS_PATH}/v20_integral_latest.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", "for device in devices:\n", " model = IntegralRegressionNet(encoder_name='convnext_base.fb_in22k_ft_in1k', pretrained=False,\n", " num_bins=NUM_BINS, temperature=1.0, use_tcn=True, tcn_layers=4)\n", " model.load_state_dict(state_dict)\n", " model.to(device).eval()\n", " v20_models.append(model)\n", " print(f\" V20 loaded on {device}\")\n", "\n", "epoch = checkpoint.get('epoch', '?')\n", "snr = checkpoint.get('snr', checkpoint.get('best_snr', 0))\n", "print(f\"Loaded V20 epoch {epoch}, SNR: {snr:.2f} dB\")\n", "print(f\"All models loaded! Using {len(devices)} GPU(s) for inference.\")" ] }, { "cell_type": "markdown", "id": "a99be5f0", "metadata": {}, "source": [ "## Inference Functions" ] }, { "cell_type": "code", "execution_count": null, "id": "6e52c9f4", "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(devices[0].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(devices[0].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, model, device, temperature=1.0):\n", " \"\"\"Run V20 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", " coords, _ = model(image_tensor, temperature=temperature)\n", " \n", " pred_y_crop = coords[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": "703054f9", "metadata": {}, "outputs": [], "source": [ "# =============================================================================\n", "# Post-Processing Functions\n", "# =============================================================================\n", "\n", "def apply_savgol_smoothing(signal_mv, window=7, polyorder=2):\n", " \"\"\"Apply Savitzky-Golay smoothing.\"\"\"\n", " if len(signal_mv) >= window:\n", " return savgol_filter(signal_mv, window_length=window, polyorder=polyorder)\n", " return signal_mv\n", "\n", "\n", "def apply_einthoven_correction(pred_mv_rows, alpha=0.33):\n", " \"\"\"Apply Einthoven's law correction: II = I + III.\"\"\"\n", " segment_width = len(pred_mv_rows[0]) // 4\n", " \n", " lead_I = pred_mv_rows[0][:segment_width].copy()\n", " lead_II_short = pred_mv_rows[1][:segment_width].copy()\n", " lead_III = pred_mv_rows[2][:segment_width].copy()\n", " \n", " derived_II = lead_I + lead_III\n", " error = lead_II_short - derived_II\n", " \n", " pred_mv_rows[0][:segment_width] = lead_I + alpha * error\n", " pred_mv_rows[2][:segment_width] = lead_III + alpha * error\n", " \n", " return pred_mv_rows\n", "\n", "\n", "def apply_baseline_correction(pred_mv_rows):\n", " \"\"\"Apply V20-specific baseline correction.\"\"\"\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", " offset = BASELINE_OFFSETS.get(lead_name, 0.0)\n", " seg_start = seg_idx * segment_width\n", " seg_end = (seg_idx + 1) * segment_width\n", " pred_mv_rows[row_idx][seg_start:seg_end] -= offset\n", " \n", " # Rhythm strip uses Lead II offset\n", " pred_mv_rows[3] -= BASELINE_OFFSETS.get('II', 0.0)\n", " \n", " return pred_mv_rows\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": "ab7272a2", "metadata": {}, "outputs": [], "source": [ "def process_image(image_path, model, device):\n", " \"\"\"Full V20 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. V20: per-row inference\n", " 6. Post-processing: smoothing, Einthoven, baseline correction\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", " # V20: 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, model, device, temperature=1.0)\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", " # Smooth\n", " pred_mv = apply_savgol_smoothing(pred_mv, window=7, polyorder=2)\n", " \n", " # 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", " # Einthoven correction\n", " pred_mv_rows = apply_einthoven_correction(pred_mv_rows, alpha=0.33)\n", " \n", " # Baseline correction\n", " pred_mv_rows = apply_baseline_correction(pred_mv_rows)\n", " \n", " return pred_mv_rows" ] }, { "cell_type": "markdown", "id": "5ef2bec0", "metadata": {}, "source": [ "## Generate Submission (Dual-GPU Parallel Processing)" ] }, { "cell_type": "code", "execution_count": null, "id": "fe540719", "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 on {len(devices)} GPU(s)...\")\n", "print(f\"Signal output width: {OUTPUT_WIDTH} (T0={T0} to T1={T1})\")" ] }, { "cell_type": "code", "execution_count": null, "id": "50c0b6d6", "metadata": {}, "outputs": [], "source": [ "# Process images with dual-GPU parallelism\n", "# Split work between GPUs\n", "\n", "results = {} # {img_id: leads dict or None}\n", "results_lock = threading.Lock()\n", "\n", "\n", "def process_batch(img_ids, model, device, gpu_idx):\n", " \"\"\"Process a batch of images on a specific GPU.\"\"\"\n", " local_results = {}\n", " desc = f\"GPU {gpu_idx}\" if len(devices) > 1 else \"Processing\"\n", " \n", " for img_id in tqdm(img_ids, desc=desc, position=gpu_idx):\n", " img_path = test_dir / f\"{img_id}.png\"\n", " \n", " if not img_path.exists():\n", " local_results[img_id] = None\n", " continue\n", " \n", " try:\n", " pred_mv_rows = process_image(img_path, model, device)\n", " leads = series_to_leads(pred_mv_rows)\n", " local_results[img_id] = leads\n", " except Exception as e:\n", " print(f\"Error {img_id}: {e}\")\n", " local_results[img_id] = None\n", " \n", " with results_lock:\n", " results.update(local_results)\n", "\n", "\n", "# Split image IDs across GPUs\n", "image_ids_list = list(image_ids)\n", "splits = np.array_split(image_ids_list, len(devices))\n", "\n", "if len(devices) > 1:\n", " # Multi-GPU: use threads\n", " threads = []\n", " for gpu_idx, (split, model, device) in enumerate(zip(splits, v20_models, devices)):\n", " t = threading.Thread(target=process_batch, args=(list(split), model, device, gpu_idx))\n", " threads.append(t)\n", " t.start()\n", " \n", " for t in threads:\n", " t.join()\n", "else:\n", " # Single GPU: process directly\n", " process_batch(image_ids_list, v20_models[0], devices[0], 0)\n", "\n", "print(f\"\\nProcessed {len(results)} images\")" ] }, { "cell_type": "code", "execution_count": null, "id": "5332c86c", "metadata": {}, "outputs": [], "source": [ "# Build submission rows\n", "all_rows = []\n", "\n", "for img_id in image_ids:\n", " img_df = test_df[test_df['id'] == img_id]\n", " leads = results.get(img_id)\n", " \n", " for _, row in img_df.iterrows():\n", " lead = row['lead']\n", " num_samples = row['number_of_rows']\n", " \n", " if leads is not None and lead in leads:\n", " signal = resample_signal(leads[lead], num_samples)\n", " for i in range(num_samples):\n", " val = float(signal[i])\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": "2e1bbc0b", "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 }