{"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":"none","dataSources":[{"sourceId":97984,"databundleVersionId":14096757,"sourceType":"competition"},{"sourceId":623160,"sourceType":"modelInstanceVersion","modelInstanceId":468838,"modelId":484689},{"sourceId":623170,"sourceType":"modelInstanceVersion","modelInstanceId":468846,"modelId":484698}],"dockerImageVersionId":31154,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"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-01T07:09:01.488859Z","iopub.execute_input":"2025-11-01T07:09:01.489317Z","iopub.status.idle":"2025-11-01T07:09:01.497365Z","shell.execute_reply.started":"2025-11-01T07:09:01.489296Z","shell.execute_reply":"2025-11-01T07:09:01.49668Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def add_noise_to_image(input_img, sigma=2, opacity=0.2):\n    noise = torch.sigmoid(torch.randn_like(input_img) * sigma)\n    input_img = (1-opacity)*input_img + opacity * noise\n    return input_img\n\ndef load_model(weights_path, **kwargs):\n    model = UNet(\n        num_in_channels=kwargs.get(\"num_in_channels\", 3),\n        num_out_channels=kwargs.get(\"num_out_channels\", 4),\n        dims=kwargs.get(\n            \"dims\",\n            [32, 64, 128, 256, 320, 320, 320, 320],\n        ),\n        depth=kwargs.get(\"depth\", 2),\n    )\n    state_dict = torch.load(weights_path, map_location=device)\n    # replace _orig_model. with nothing in all keys (model was trained with torch.compile)\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  # Normalize to [0, 1]\n    img = img.unsqueeze(0)  # Add batch dimension\n    # ensure only 3 channels\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    # Make the black pixels white\n    show_featuremap = torch.ones(probs.shape[1], probs.shape[2], 3)\n    probs[2] /= probs[2].max()  # Normalize the third channel to [0, 1]\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, -9, -6, -0, -3, -6, -0, -3, -6, -0, -3, -6]\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(num_thetas=200)\n\n    cropper = Cropper(percentiles=(0.03, 0.97), alpha=0.99)\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(\n        min_number_of_grid_lines=30,\n        max_number_of_grid_lines=100,\n        lower_grid_line_factor=0.1,\n    )\n    signal_extractor = SignalExtractor()\n\n    num_in_channels = 1\n    num_out_channels = 13\n    dims = [32, 64, 128, 256, 256]\n    depth = 2\n    layout_unet = load_model(\n        weights_path=\"/kaggle/input/open-ecg-digitizer-weights/pytorch/default/1/lead_name_unet_weights_07072025.pt\",\n        num_in_channels=num_in_channels,\n        num_out_channels=num_out_channels,\n        dims=dims,\n        depth=depth,\n        device=device,\n    )\n\n    layouts = yaml.safe_load(\n        open(\"src/config/lead_layouts_george-moody-2024.yml\", \"r\"),\n    )\n\n    identifier = LeadIdentifier(\n        layouts=layouts,\n        unet=layout_unet,\n        device=torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\"),\n        possibly_flipped=False,\n        target_num_samples=target_num_samples,\n        required_valid_samples=2,\n    )\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\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\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):\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        \n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nnum_in_channels = 3\nnum_out_channels = 4\ndims = [32, 64, 128, 256, 320, 320, 320, 320]\ndepth = 2\nmodel = load_model(weights_path=\"/kaggle/input/open-ecg-digitizer-weights/pytorch/default/1/unet_weights_07072025.pt\", num_in_channels=num_in_channels, num_out_channels=num_out_channels, dims=dims, depth=depth, device=device)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-01T07:11:08.400632Z","iopub.execute_input":"2025-11-01T07:11:08.40133Z","iopub.status.idle":"2025-11-01T07:11:08.688324Z","shell.execute_reply.started":"2025-11-01T07:11:08.401306Z","shell.execute_reply":"2025-11-01T07:11:08.687741Z"}},"outputs":[],"execution_count":null},{"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        #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=3200) # 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            print(f\"Extracting signal from segmented {path}...\")\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        ### (optional) VISUALIZATION ###\n        if index == 0:\n            if index == 0:\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-01T07:12:13.020328Z","iopub.execute_input":"2025-11-01T07:12:13.020814Z","iopub.status.idle":"2025-11-01T07:12:46.680254Z","shell.execute_reply.started":"2025-11-01T07:12:13.020792Z","shell.execute_reply":"2025-11-01T07:12:46.679473Z"}},"outputs":[],"execution_count":null}]}