{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":97984,"databundleVersionId":14096757,"sourceType":"competition"},{"sourceId":623160,"sourceType":"modelInstanceVersion","modelInstanceId":468838,"modelId":484689},{"sourceId":623170,"sourceType":"modelInstanceVersion","modelInstanceId":468846,"modelId":484698}],"dockerImageVersionId":31153,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Please leave a ⭐ and fork the repository if you want.\n\nhttps://github.com/Ahus-AIM/Open-ECG-Digitizer\n\nhttps://arxiv.org/abs/2510.19590\n\nNo segmentation model parameters have been adapted to the kaggle training set.\n\nNo hyperparameters have been rigorously tuned on the kaggle training set (we just made sure it looked OK on folder no. 100642785).","metadata":{}},{"cell_type":"markdown","source":"# Imports","metadata":{}},{"cell_type":"code","source":"%cd /kaggle/input/open-ecg-digitizer/pytorch/default/1\nimport torch\nimport glob\nimport os.path\nimport yaml\nimport torch.nn.functional as F\nimport pandas as pd\nimport numpy as np\nfrom torchvision.io import read_image\nfrom torch import Tensor\n# Imports from https://github.com/Ahus-AIM/Open-ECG-Digitizer\nfrom src.model.unet import UNet\nfrom src.model.perspective_detector import PerspectiveDetector\nfrom src.model.cropper import Cropper\nfrom src.model.pixel_size_finder import PixelSizeFinder\nfrom src.model.signal_extractor import SignalExtractor\nfrom src.model.lead_identifier import LeadIdentifier\nfrom src.model.lead_identifier import LeadIdentifier","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-02T17:13:18.079755Z","iopub.execute_input":"2025-11-02T17:13:18.080329Z","iopub.status.idle":"2025-11-02T17:13:18.088165Z","shell.execute_reply.started":"2025-11-02T17:13:18.080304Z","shell.execute_reply":"2025-11-02T17:13:18.087544Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Hyperparameters - please optimize these using the training set!\n\nSome descriptions of the modules: https://arxiv.org/abs/2510.19590","metadata":{}},{"cell_type":"code","source":"# Resample size: trade-off between accuracy and speed/vram\nresample_size = 3000 # maximum no. pixels along the longest dimension\n\n# PixelSizeFinder\nPixelSizeFinderKwargs = dict(\n    mm_between_grid_lines=5,\n    samples=1000,\n    min_number_of_grid_lines=20,\n    max_number_of_grid_lines=100,\n    max_zoom=10,\n    zoom_factor=10.0,\n    lower_grid_line_factor=0.1,\n)\n\n# SignalExtractor\nSignalExtractorKwargs = dict(\n    threshold_sum=10.0,\n    threshold_line_in_mask=0.8,\n    label_thresh=0.05,\n    max_iterations=4,\n    split_num_stripes=4,\n    candidate_span=10,\n    lam=0.5,\n    min_line_width=30,\n)\n\n# PerspectiveDetector\nPerspectiveDetectorKwargs = dict(\n    num_thetas=200, # Higher -> more accurate but slower and more VRAM\n    max_num_nonzero=10_000,\n)\n\n# LeadIdentifier\nLeadIdentifierKwargs = dict(\n    layouts=None,  # should be a dict[str, Any]\n    unet=None,     # should be a torch.nn.Module\n    device=None,   # should be a torch.device\n    possibly_flipped=True,\n    target_num_samples=None,\n    required_valid_samples=2,\n    debug=False,\n)\n\n# Cropper\nCropperKwargs = dict(\n    granularity=50,\n    percentiles=(0.03, 0.97),\n    alpha=0.95,\n)\n\n# Layout UNet (you CAN change this but then you need to retrain the net)\n# https://github.com/Ahus-AIM/Open-ECG-Digitizer/blob/main/src/train.py\n# https://github.com/Ahus-AIM/Open-ECG-Digitizer/blob/main/src/config/lead_name_unet.yml\nLayoutUNetKwargs = dict(\n    weights_path=\"/kaggle/input/open-ecg-digitizer-weights/pytorch/default/1/lead_name_unet_weights_07072025.pt\",\n    num_in_channels=1,\n    num_out_channels=13,\n    dims=[32, 64, 128, 256, 256],\n    depth=2,\n)\n\n# Segmentation UNet (you CAN change this but then you need to retrain the net)\n# https://github.com/Ahus-AIM/Open-ECG-Digitizer/blob/main/src/train.py\n# https://github.com/Ahus-AIM/Open-ECG-Digitizer/blob/main/src/config/unet.yml\nLoadModelKwargs = dict(\n    weights_path=\"/kaggle/input/open-ecg-digitizer-weights/pytorch/default/1/unet_weights_07072025.pt\",\n    num_in_channels=3,\n    num_out_channels=4,\n    dims=[32, 64, 128, 256, 320, 320, 320, 320],\n    depth=2,\n)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-02T17:13:18.089160Z","iopub.execute_input":"2025-11-02T17:13:18.089396Z","iopub.status.idle":"2025-11-02T17:13:18.101731Z","shell.execute_reply.started":"2025-11-02T17:13:18.089382Z","shell.execute_reply":"2025-11-02T17:13:18.101039Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Utility code for plotting and orchestrating the modules\n\nYou can also use https://github.com/Ahus-AIM/Open-ECG-Digitizer/blob/main/src/model/inference_wrapper.py for orchestrating the modules if you want.","metadata":{}},{"cell_type":"code","source":"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\ndef add_noise_to_image(input_img, sigma=1.0, opacity=0.85):\n    noise = torch.sigmoid(torch.randn_like(input_img) * sigma)\n    input_img = (opacity)*input_img + (1-opacity) * noise\n    return input_img\n\ndef load_model(**kwargs):\n    weights_path = kwargs.pop(\"weights_path\", None)  # safely extract\n    model = UNet(**kwargs)\n    state_dict = torch.load(weights_path, map_location=device)\n    state_dict = {k.replace(\"_orig_mod.\", \"\"): v for k, v in state_dict.items()}\n    model.load_state_dict(state_dict)\n    model.eval().to(device)\n    return model\n\n\ndef load_png_file(path):\n    img = read_image(path)\n    img = img.float() / 255.0\n    img = img.unsqueeze(0)\n    if img.shape[1] > 3:\n        img = img[:, :3, :, :]\n    return img\n\n\ndef _crop_y(\n    image: Tensor,\n    signal_prob: Tensor,\n    grid_prob: Tensor,\n    text_prob: Tensor,\n) -> tuple[Tensor, Tensor, Tensor, Tensor]:\n    def get_bounds(tensor: Tensor) -> tuple[int, int]:\n        prob = torch.clamp(\n            tensor.squeeze().sum(dim=tensor.dim() - 3)\n            - tensor.squeeze().sum(dim=tensor.dim() - 3).mean(),\n            min=0,\n        )\n        non_zero = (prob > 0).nonzero(as_tuple=True)[0]\n        if non_zero.numel() == 0:\n            return 0, tensor.shape[2] - 1\n        return int(non_zero[0].item()), int(non_zero[-1].item())\n\n    y1, y2 = get_bounds(signal_prob + grid_prob)\n\n    slices = (slice(None), slice(None), slice(y1, y2 + 1), slice(None))\n    return (\n        image[slices],\n        signal_prob[slices],\n        grid_prob[slices],\n        text_prob[slices],\n    )\n\n\ndef _align_feature_maps(\n    cropper: Cropper,\n    image: Tensor,\n    signal_prob: Tensor,\n    grid_prob: Tensor,\n    text_prob: Tensor,\n    source_points: Tensor,\n) -> tuple[Tensor, Tensor, Tensor, Tensor]:\n    aligned_signal_prob = cropper.apply_perspective(\n        signal_prob,\n        source_points,\n        fill_value=0,\n    )\n    aligned_image = cropper.apply_perspective(\n        image,\n        source_points,\n        fill_value=0,\n    )\n    aligned_grid_prob = cropper.apply_perspective(\n        grid_prob,\n        source_points,\n        fill_value=0,\n    )\n    aligned_text_prob = cropper.apply_perspective(\n        text_prob,\n        source_points,\n        fill_value=0,\n    )\n    (\n        aligned_image,\n        aligned_signal_prob,\n        aligned_grid_prob,\n        aligned_text_prob,\n    ) = _crop_y(\n        aligned_image,\n        aligned_signal_prob,\n        aligned_grid_prob,\n        aligned_text_prob,\n    )\n\n    return (\n        aligned_image,\n        aligned_signal_prob,\n        aligned_grid_prob,\n        aligned_text_prob,\n    )\n\n\ndef plot_segmentation_and_image(\n    image,\n    segmentation,\n    aligned_signal,\n    aligned_grid,\n    lines,\n):\n    import matplotlib.pyplot as plt\n    import numpy as np\n\n    image_np = image.squeeze(0).permute(1, 2, 0).cpu().numpy()\n    probs = segmentation.squeeze(0).cpu()\n\n    show_featuremap = torch.ones(probs.shape[1], probs.shape[2], 3)\n    probs[2] /= probs[2].max()\n    show_featuremap[:, :, [0, 1, 2]] -= 2 * probs[2].unsqueeze(-1)\n    show_featuremap[:, :, [1, 2]] -= probs[0].unsqueeze(-1)\n    show_featuremap = torch.clamp(show_featuremap, 0, 1).numpy()\n\n    straightened_featuremap = torch.ones(\n        aligned_signal.shape[2],\n        aligned_signal.shape[3],\n        3,\n        device=aligned_signal.device,\n    )\n    aligned_signal /= aligned_signal.max()\n    straightened_featuremap[:, :, [0, 1, 2]] -= 2 * aligned_signal[0, 0].unsqueeze(-1)\n    aligned_grid /= aligned_grid.max()\n    straightened_featuremap[:, :, [1, 2]] -= aligned_grid[0, 0].unsqueeze(-1)\n    straightened_featuremap = torch.clamp(straightened_featuremap, 0, 1)\n\n    fig, ax = plt.subplots(2, 2, figsize=(16, 12))\n    ax[0, 0].imshow(image_np)\n    ax[0, 0].axis(\"off\")\n\n    ax[0, 1].imshow(show_featuremap)\n    ax[0, 1].axis(\"off\")\n\n    ax[1, 0].imshow(straightened_featuremap.cpu())\n    ax[1, 0].axis(\"off\")\n\n    offsets = [-0, -10.5, -7, -0, -3.5, -7, -0, -3.5, -7, -0, -3.5, -7]\n    if lines.numel() > 0:\n        ax[1, 1].plot(lines.T.cpu().numpy() + offsets[: lines.shape[0]])\n    ax[1, 1].axis(\"off\")\n    plt.tight_layout()\n    plt.show()\n\n\ndef crop_image(image, probs):\n    perspective_detector = PerspectiveDetector(**PerspectiveDetectorKwargs)\n\n    cropper = Cropper(**CropperKwargs)\n\n    alignment_params = perspective_detector(probs[0, 0])\n\n    source_points = cropper(probs[0, 1], alignment_params)\n\n    signal_prob, grid_prob, text_prob = (\n        probs[:, [2]],\n        probs[:, [0]],\n        probs[:, [1]],\n    )\n\n    (\n        aligned_image,\n        aligned_signal_prob,\n        aligned_grid_prob,\n        aligned_text_prob,\n    ) = _align_feature_maps(\n        cropper,\n        image,\n        signal_prob,\n        grid_prob,\n        text_prob,\n        source_points,\n    )\n\n    return (\n        aligned_image,\n        aligned_signal_prob,\n        aligned_grid_prob,\n        aligned_text_prob,\n    )\n\n\ndef extract_signals(\n    aligned_signal_prob: Tensor,\n    aligned_grid_prob: Tensor,\n    aligned_text_prob: Tensor,\n    target_num_samples: int,\n) -> Tensor:\n    pixel_size_finder = PixelSizeFinder(**PixelSizeFinderKwargs)\n    signal_extractor = SignalExtractor(**SignalExtractorKwargs)\n\n    layout_unet = load_model(\n        **LayoutUNetKwargs,\n    )\n\n    layouts = yaml.safe_load(\n        open(\"src/config/lead_layouts_george-moody-2024.yml\", \"r\"),\n    )\n\n    identifier_kwargs = LeadIdentifierKwargs.copy()\n    identifier_kwargs.update(\n        dict(\n            layouts=layouts,\n            unet=layout_unet,\n            device=torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\"),\n            target_num_samples=target_num_samples,\n        )\n    )\n    identifier = LeadIdentifier(**identifier_kwargs)\n    mm_per_pixel_x, mm_per_pixel_y = pixel_size_finder(aligned_grid_prob)\n\n    avg_pixel_per_mm = (1 / mm_per_pixel_x + 1 / mm_per_pixel_y) / 2 # Is there a better way?\n    signals = signal_extractor(aligned_signal_prob.squeeze())\n    \n    signals = identifier(\n        signals,\n        aligned_text_prob,\n        avg_pixel_per_mm=avg_pixel_per_mm,\n    )\n\n    return signals\n\ndef resample_image(image: Tensor, resample_size: int) -> Tensor:\n    height, width = image.shape[2], image.shape[3]\n    min_dim = min(height, width)\n    max_dim = max(height, width)\n\n    if isinstance(resample_size, int):\n        if max_dim > resample_size:\n            scale = resample_size / max_dim\n            new_size = (int(height * scale), int(width * scale))\n            return F.interpolate(image, size=new_size, mode=\"bilinear\", align_corners=False, antialias=True)\n        return image\n\n    if isinstance(resample_size, tuple):\n        interpolated = F.interpolate(\n            image, size=resample_size, mode=\"bilinear\", align_corners=False, antialias=True\n        )\n        return interpolated\n\n    raise ValueError(f\"Invalid resample_size: {resample_size}. Expected int or tuple of (height, width).\")\n\nleads_names = ['I','II','III','aVR','aVL','aVF','V1','V2','V3','V4','V5','V6']\ndef get_slice(lead_name: str, number_of_rows: int) -> slice:\n    assert lead_name in leads_names\n    if lead_name in (\"II\",):\n        return slice(0, number_of_rows)\n    if lead_name in ((\"I\", \"III\")):\n        return slice(0, number_of_rows)\n    if lead_name in ((\"aVR\", \"aVF\", \"aVL\")):\n        return slice(1*number_of_rows, 2*number_of_rows)\n    if lead_name in ((\"V1\", \"V2\", \"V3\")):\n        return slice(2*number_of_rows, 3*number_of_rows)\n    if lead_name in ((\"V4\", \"V5\", \"V6\")):\n        return slice(3*number_of_rows, 4*number_of_rows)\n\ndef digitize_image(input_img: Tensor, resample_size: int, target_num_samples: int) -> Tensor:\n    input_img = add_noise_to_image(input_img) # The UNet is trained for \"real\" images. Sometimes it performs better with added noise for generated images.\n    input_img = resample_image(image=input_img, resample_size=resample_size) # higher resample size is (probably) better but watch out for VRAM and time consraints\n    \n    with torch.no_grad():\n        logits = model(input_img.to(device))\n        output_probs = torch.softmax(logits, dim=1)\n        aligned_image, aligned_signal, aligned_grid, aligned_text = crop_image(input_img, output_probs)\n        lines = extract_signals(aligned_signal, aligned_grid, aligned_text, target_num_samples=target_num_samples)\n        lines = lines[\"canonical_lines\"] * 1e-3  # microvolt to millivolt\n\n    return output_probs, aligned_signal, aligned_grid, lines.float()\n    \n\nmodel = load_model(**LoadModelKwargs)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-02T17:13:18.293245Z","iopub.execute_input":"2025-11-02T17:13:18.293563Z","iopub.status.idle":"2025-11-02T17:13:18.639429Z","shell.execute_reply.started":"2025-11-02T17:13:18.293540Z","shell.execute_reply":"2025-11-02T17:13:18.638633Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Visualize a few example images.","metadata":{}},{"cell_type":"code","source":"paths = [\n    \"/kaggle/input/physionet-ecg-image-digitization/train/1006427285/1006427285-0004.png\",\n    \"/kaggle/input/physionet-ecg-image-digitization/train/1006427285/1006427285-0005.png\",\n    \"/kaggle/input/physionet-ecg-image-digitization/train/1006427285/1006427285-0006.png\",\n]\n\nfor path in paths:\n    input_img = load_png_file(path)\n    if not os.path.exists(path):\n        continue\n\n    ### SIGNAL EXTRACTION ###\n    print(f\"Segmenting {path}...\")\n    output_probs, aligned_signal, aligned_grid, lines = digitize_image(input_img, 2500, 10000)\n    plot_segmentation_and_image(input_img, output_probs, aligned_signal, aligned_grid, lines)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-02T17:13:18.640937Z","iopub.execute_input":"2025-11-02T17:13:18.641196Z","iopub.status.idle":"2025-11-02T17:14:19.770497Z","shell.execute_reply.started":"2025-11-02T17:13:18.641179Z","shell.execute_reply":"2025-11-02T17:14:19.769825Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Run the digitization","metadata":{}},{"cell_type":"code","source":"test = pd.read_csv('/kaggle/input/physionet-ecg-image-digitization/test.csv')\noutput_path = '/kaggle/working/submission.csv'\n\n# Prepare the file (write header once)\nif os.path.exists(output_path):\n    os.remove(output_path)\npd.DataFrame(columns=[\"id\", \"value\"]).to_csv(output_path, index=False)\n\nold_id = None\n\nfor index, row in test.iterrows():\n    if row.id != old_id:\n        old_id = row.id\n    \n        path = f\"/kaggle/input/physionet-ecg-image-digitization/test/{row.id}.png\"\n        target_num_samples = row.fs * 10 # We assume 10 second signals, not optimal.\n        input_img = load_png_file(path)\n\n        ### SIGNAL EXTRACTION ###\n        print(f\"Segmenting {path}...\")\n        output_probs, aligned_signal, aligned_grid, lines = digitize_image(input_img, resample_size, target_num_samples)\n\n        ### (optional) VISUALIZATION ###\n        if False:\n            plot_segmentation_and_image(input_img, output_probs, aligned_signal, aligned_grid, lines)\n\n    ### SAVING ###\n    file_id = row.id\n    lead_name = row.lead\n    number_of_rows_in_lead = row.number_of_rows\n\n    index = leads_names.index(lead_name)\n\n    lead_data = lines[index]\n    lead_data = lead_data[get_slice(lead_name, number_of_rows_in_lead)]\n    \n    mean_val = np.nanmean(lead_data)\n    if np.isnan(mean_val):\n        mean_val = 0.0\n    lead_data = np.nan_to_num(lead_data, nan=mean_val)\n\n    assert len(lead_data) == number_of_rows_in_lead\n    \n    chunk = []\n    for t in range(number_of_rows_in_lead):\n        chunk.append({\"id\": f\"{file_id}_{t}_{lead_name}\", \"value\": float(lead_data[t])})\n\n    if chunk:\n        pd.DataFrame(chunk).to_csv(output_path, mode='a', index=False, header=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-02T17:14:19.771329Z","iopub.execute_input":"2025-11-02T17:14:19.771560Z","iopub.status.idle":"2025-11-02T17:14:49.427414Z","shell.execute_reply.started":"2025-11-02T17:14:19.771543Z","shell.execute_reply":"2025-11-02T17:14:49.426844Z"}},"outputs":[],"execution_count":null}]}