{"cells": [{"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "!pip uninstall -y tensorflow -q\n!pip install connected-components-3d -q"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "%%writefile constant.py\n\nimport warnings\nwarnings.filterwarnings(\"ignore\", category=UserWarning, module=\"pydantic\")\n\nimport os\nimport sys\nimport random\nimport numpy as np\nimport torch\n\n# Set seeds for reproducibility\nseed = 42\nrandom.seed(seed)\nnp.random.seed(seed)\ntorch.manual_seed(seed)\nif torch.cuda.is_available():\n    torch.cuda.manual_seed_all(seed)\n\nsys.path.append('/kaggle/input/hengck23-submit-physionet/hengck23-submit-physionet')\n\nimport traceback\nfrom pathlib import Path\nfrom shutil import copyfile\nimport cv2\nimport pandas as pd\nfrom tqdm.auto import tqdm\n\nCUDA0 = 'cuda:0' if torch.cuda.is_available() else 'cpu'\nFLOAT_TYPE = torch.float32\n\nif_submit = os.getenv('KAGGLE_IS_COMPETITION_RERUN')\n\nif if_submit:\n    test_meta = Path(\"/kaggle/input/physionet-ecg-image-digitization/test.csv\")\n    test_dir = Path(\"/kaggle/input/physionet-ecg-image-digitization/test\")\nelse:\n    test_meta = Path(\"/kaggle/input/physionet-ecg-image-digitization/test.csv\")\n    test_dir = Path(\"/kaggle/input/physionet-ecg-image-digitization/test\")\n\nvalid_df = pd.read_csv(test_meta)\nvalid_df['id'] = valid_df['id'].astype(str)\nvalid_id = valid_df['id'].unique().tolist()\n\nglobal_dict = {\n    \"stage0_dir\": \"/kaggle/working/stage0\",\n    \"stage1_dir\": \"/kaggle/working/stage1\",\n    \"stage2_dir\": \"/kaggle/working/stage2\",\n}\n\nprint(f'Device: {CUDA0}')\nprint(f'Test images: {len(valid_id)}')"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "%%writefile stage0.py\nfrom constant import *\nfrom stage0_model import Net as Stage0Net\nfrom stage0_common import *\nimport torch.nn.functional as F\n\ndef apply_grayscale_guidance(image_rgb):\n    gray = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2GRAY)\n    denoised = cv2.fastNlMeansDenoising(gray, h=10)\n    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8))\n    contrast_enhanced = clahe.apply(denoised)\n    guidance_img = cv2.cvtColor(contrast_enhanced, cv2.COLOR_GRAY2RGB)\n    return guidance_img\n\ndef to_device(batch, device):\n    if isinstance(batch, dict):\n        return {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}\n    return batch.to(device)\n\nstage0_dir = Path(global_dict[\"stage0_dir\"])\nstage0_dir.mkdir(exist_ok=True, parents=True)\n\nprint(\"Loading Stage 0 Model...\")\nstage0_net = Stage0Net(pretrained=False)\nstage0_net = load_net(stage0_net, '/kaggle/input/hengck23-submit-physionet/hengck23-submit-physionet/weight/stage0-last.checkpoint.pth')\nstage0_net.to(CUDA0)\nstage0_net.eval()\n\nprint(\"Starting Stage 0 Processing...\")\nfor n, sample_id in enumerate(tqdm(valid_id)):\n    path = test_dir / f'{sample_id}.png'\n    output_path = stage0_dir / f'{sample_id}.png'\n    \n    image_original = cv2.imread(str(path), cv2.IMREAD_COLOR)\n    if image_original is None:\n        continue\n    image_original = cv2.cvtColor(image_original, cv2.COLOR_BGR2RGB)\n    \n    image_for_model = apply_grayscale_guidance(image_original)\n    batch = image_to_batch(image_for_model)\n    batch = to_device(batch, CUDA0)\n\n    try:\n        with torch.no_grad(), torch.amp.autocast('cuda', dtype=FLOAT_TYPE):\n            output = stage0_net(batch)\n        rotated, keypoint = output_to_predict(image_original, batch, output)\n        normalised, _, _ = normalise_by_homography(rotated, keypoint)\n        cv2.imwrite(str(output_path), cv2.cvtColor(normalised, cv2.COLOR_RGB2BGR))\n    except Exception as e:\n        print(f\"Stage0 error {sample_id}: {e}\")\n        copyfile(path, output_path)\n\nprint(f\"Stage 0 finished.\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "%%writefile stage1.py\nfrom constant import *\nfrom stage1_model import Net as Stage1Net\nfrom stage1_common import *\n\nstage0_dir = Path(global_dict[\"stage0_dir\"])\nstage1_dir = Path(global_dict[\"stage1_dir\"])\nstage1_dir.mkdir(exist_ok=True, parents=True)\n\nstage1_net = Stage1Net(pretrained=False)\nstage1_net = load_net(stage1_net, '/kaggle/input/hengck23-submit-physionet/hengck23-submit-physionet/weight/stage1-last.checkpoint.pth')\nstage1_net.to(CUDA0)\nstage1_net.eval()\n\nprint(\"Starting Stage 1 Processing...\")\nfor n, sample_id in enumerate(tqdm(valid_id)):\n    path = stage0_dir / f'{sample_id}.png'\n    output_path = stage1_dir / f'{sample_id}.png'\n    \n    image = cv2.imread(str(path), cv2.IMREAD_COLOR)\n    if image is None:\n        image = cv2.imread(str(test_dir / f'{sample_id}.png'), cv2.IMREAD_COLOR)\n    \n    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n    batch = {'image': torch.from_numpy(np.ascontiguousarray(image.transpose(2, 0, 1))).unsqueeze(0).to(CUDA0)}\n\n    try:\n        with torch.no_grad(), torch.amp.autocast('cuda', dtype=FLOAT_TYPE):\n            output = stage1_net(batch)\n        gridpoint_xy, _ = output_to_predict(image, batch, output)\n        rectified = rectify_image(image, gridpoint_xy)\n        cv2.imwrite(str(output_path), cv2.cvtColor(rectified, cv2.COLOR_RGB2BGR))\n    except Exception as e:\n        print(f\"Stage1 error {sample_id}: {e}\")\n        copyfile(path, output_path)\n\nprint(f\"Stage 1 finished.\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "%%writefile stage2.py\nfrom constant import *\nimport timm\nimport torch.nn as nn\nimport torchvision.transforms as T\nfrom stage2_model import *\nfrom stage2_common import *\nfrom scipy.signal import savgol_filter\n\nclass MedicalConstraintRefiner:\n    def __init__(self, alpha=0.33):\n        self.alpha = alpha\n\n    def apply_einthoven_law(self, series_dict):\n        if all(k in series_dict for k in ['I', 'II', 'III']):\n            L1 = series_dict['I']\n            L2 = series_dict['II']\n            L3 = series_dict['III']\n            error = L2 - (L1 + L3)\n            series_dict['I'] = L1 + (self.alpha * error)\n            series_dict['III'] = L3 + (self.alpha * error)\n            series_dict['II'] = L2 - (self.alpha * error)\n        return series_dict\n\nclass Net3(nn.Module):\n    def __init__(self, pretrained=True):\n        super(Net3, self).__init__()\n        encoder_dim = [64, 128, 256, 512]\n        decoder_dim = [128, 64, 32, 16]\n        self.encoder = timm.create_model('resnet34.a3_in1k', pretrained=pretrained, in_chans=3, num_classes=0, global_pool='')\n        self.decoder = MyCoordUnetDecoder(in_channel=encoder_dim[-1], skip_channel=encoder_dim[:-1][::-1] + [0], out_channel=decoder_dim, scale=[2, 2, 2, 2])\n        self.pixel = nn.Conv2d(decoder_dim[-1], 4, 1)\n\n    def forward(self, image):\n        encode = encode_with_resnet(self.encoder, image)\n        last, _ = self.decoder(feature=encode[-1], skip=encode[:-1][::-1] + [None])\n        return self.pixel(last)\n\nstage1_dir = Path(global_dict[\"stage1_dir\"])\nstage2_dir = Path(global_dict[\"stage2_dir\"])\nstage2_dir.mkdir(exist_ok=True, parents=True)\n\nstage2_net = Net3(pretrained=False).to(CUDA0)\nmodel_path = \"/kaggle/input/physio-seg-public/pytorch/net3_009_4200/1/iter_0004200.pt\"\nstage2_net.load_state_dict(torch.load(model_path, map_location=CUDA0))\nstage2_net.eval()\n\nresize = T.Resize((1696, 4352), interpolation=T.InterpolationMode.BILINEAR)\nrefiner = MedicalConstraintRefiner()\n\nmv_to_pixel = 78.5\nzero_mv = [703.5, 987.5, 1271.5, 1531.5]\nt0, t1 = 235, 4161\n\ndef series_to_dict_local(series_4row):\n    d = {}\n    names = [['I', 'aVR', 'V1', 'V4'], ['II', 'aVL', 'V2', 'V5'], ['III', 'aVF', 'V3', 'V6']]\n    for i in range(3):\n        splits = np.array_split(series_4row[i], 4)\n        for name, data in zip(names[i], splits):\n            d[name] = data\n    d['II_Long'] = series_4row[3]\n    return d\n\ndef dict_to_series_local(d, original_shape):\n    new_series = np.zeros(original_shape)\n    new_series[0] = np.concatenate([d['I'], d['aVR'], d['V1'], d['V4']])\n    new_series[1] = np.concatenate([d['II'], d['aVL'], d['V2'], d['V5']])\n    new_series[2] = np.concatenate([d['III'], d['aVF'], d['V3'], d['V6']])\n    new_series[3] = d['II_Long']\n    return new_series\n\nprint(\"Starting Stage 2...\")\nfor n, sample_id in enumerate(tqdm(valid_id)):\n    path = stage1_dir / f'{sample_id}.png'\n    output_path = stage2_dir / f'{sample_id}.npy'\n    \n    image = cv2.imread(str(path), cv2.IMREAD_COLOR)\n    if image is None:\n        length = valid_df[(valid_df['id']==sample_id) & (valid_df['lead']=='II')].iloc[0].number_of_rows\n        np.save(output_path, np.zeros((4, length)))\n        continue\n\n    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n    img_input = (image[:1696, :2176] / 255.0)\n    batch = resize(torch.from_numpy(np.ascontiguousarray(img_input.transpose(2, 0, 1))).unsqueeze(0)).float().to(CUDA0)\n    \n    try:\n        with torch.no_grad(), torch.amp.autocast('cuda', dtype=FLOAT_TYPE):\n            output = stage2_net(batch)\n        \n        pixel = torch.sigmoid(output).float().data.cpu().numpy()[0]\n        length = valid_df[(valid_df['id']==sample_id) & (valid_df['lead']=='II')].iloc[0].number_of_rows\n        series_in_pixel = pixel_to_series(pixel[..., t0:t1], zero_mv, length)\n        series = (np.array(zero_mv).reshape(4, 1) - series_in_pixel) / mv_to_pixel\n\n        for i in range(series.shape[0]):\n            series[i] = savgol_filter(series[i], window_length=7, polyorder=2)\n        \n        s_dict = series_to_dict_local(series)\n        s_dict = refiner.apply_einthoven_law(s_dict)\n        series_corrected = dict_to_series_local(s_dict, series.shape)\n        np.save(output_path, series_corrected)\n        \n    except Exception as e:\n        print(f\"Stage2 error {sample_id}: {e}\")\n        length = valid_df[(valid_df['id']==sample_id) & (valid_df['lead']=='II')].iloc[0].number_of_rows\n        np.save(output_path, np.zeros((4, length)))\n\nprint(\"Stage 2 finished!\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "!python stage0.py\n!python stage1.py\n!python stage2.py"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "import gc\nimport pandas as pd\nimport numpy as np\nfrom tqdm import tqdm\nfrom pathlib import Path\n\ntest_meta = Path(\"/kaggle/input/physionet-ecg-image-digitization/test.csv\")\nvalid_df = pd.read_csv(test_meta)\nvalid_df['id'] = valid_df['id'].astype(str)\n\nstage2_dir = Path(\"/kaggle/working/stage2\")\n\ndef series_dict(series):\n    series_by_lead = dict()\n    for l in range(3):\n        lead_names = [\n            ['I', 'aVR', 'V1', 'V4'],\n            ['II', 'aVL', 'V2', 'V5'],\n            ['III', 'aVF', 'V3', 'V6'],\n        ][l]\n        split = np.array_split(series[l], 4)\n        for (k, s) in zip(lead_names, split):\n            series_by_lead[k] = s\n    series_by_lead['II'] = series[3]\n    return series_by_lead\n\nsubmit_df = list()\ngb = valid_df.groupby('id')\n\nprint(\"Generating submission file...\")\nfor rec_idx, (sample_id, df) in enumerate(tqdm(gb)):\n    try:\n        series = np.load(stage2_dir / f'{sample_id}.npy')\n        series_by_lead = series_dict(series)\n\n        for _, d in df.iterrows():\n            s = series_by_lead.get(d.lead, np.zeros(d.number_of_rows))\n            \n            if len(s) != d.number_of_rows:\n                x_old = np.linspace(0, 1, len(s))\n                x_new = np.linspace(0, 1, d.number_of_rows)\n                s = np.interp(x_new, x_old, s)\n            \n            row_id = [f'{sample_id}_{t}_{d.lead}' for t in range(d.number_of_rows)]\n            submit_df.append(pd.DataFrame({'id': row_id, 'value': s}))\n            \n    except Exception as e:\n        print(f\"Error {sample_id}: {e}\")\n\n    if rec_idx % 100 == 0:\n        gc.collect()\n\nif submit_df:\n    final_df = pd.concat(submit_df, axis=0, ignore_index=True)\n    final_df.to_csv('submission.csv', index=False)\n    print(f\"Done! Shape: {final_df.shape}\")\n    print(final_df.head())\n    print(final_df['value'].describe())\nelse:\n    print(\"Error: No predictions!\")"}], "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, "language_info": {"name": "python", "version": "3.10.0"}}, "nbformat": 4, "nbformat_minor": 4}