{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":97984,"databundleVersionId":14096757,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":13746387,"sourceType":"datasetVersion","datasetId":8747012},{"sourceId":13816899,"sourceType":"datasetVersion","datasetId":8620533},{"sourceId":271051632,"sourceType":"kernelVersion"},{"sourceId":289965670,"sourceType":"kernelVersion"},{"sourceId":677607,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":513841,"modelId":528480}],"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip uninstall -y tensorflow\n!uv pip install --no-deps --system --no-index --find-links='/kaggle/input/hengck23-submit-physionet/hengck23-submit-physionet/setup' 'connected-components-3d'","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-01-04T13:51:56.162791Z","iopub.execute_input":"2026-01-04T13:51:56.163476Z","iopub.status.idle":"2026-01-04T13:52:39.840396Z","shell.execute_reply.started":"2026-01-04T13:51:56.163451Z","shell.execute_reply":"2026-01-04T13:52:39.839688Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import sys\nsys.path.append('/kaggle/input/hengck23-submit-physionet/hengck23-submit-physionet')\n\nimport os\nimport gc\nimport cv2\nimport torch\nimport traceback\nimport numpy as np\nimport pandas as pd\n\nfrom tqdm import tqdm\nfrom pathlib import Path\nfrom tqdm.auto import tqdm\nfrom shutil import copyfile\nfrom scipy.signal import resample\nimport torchvision.transforms as T\nfrom scipy.signal import savgol_filter, medfilt\ntest_dir = Path(\"/kaggle/input/physionet-ecg-image-digitization/test\")\n\ntest = pd.read_csv(\"/kaggle/input/physionet-ecg-image-digitization/test.csv\")\ntest['id'] = test['id'].astype(str) \ntest_id = test['id'].unique().tolist()\n\n\nglobal_dict = {\n    \"stage0_dir\": \"/kaggle/working/stage0\",\n    \"stage1_dir\": \"/kaggle/working/stage1\",\n    \"stage2_dir\": \"/kaggle/working/stage2\",\n}\ndef change_color(image_rgb):\n    hsv = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2HSV)\n    h, s, v = cv2.split(hsv)\n    \n    v_denoised = cv2.fastNlMeansDenoising(v, h=5.5)\n    \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    \n    hsv_enhanced = cv2.merge([h, s, v_enhanced])\n    return cv2.cvtColor(hsv_enhanced, cv2.COLOR_HSV2RGB)\n\ndef dw(series_dict, alpha=0.33):\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\n        error = L2 - (L1 + L3)\n\n        series_dict['I'] = L1 + (alpha * error)\n        series_dict['III'] = L3 + (alpha * error)\n        series_dict['II'] = L2 - (alpha * error)\n    \n    return series_dict\n\ndef series_to_dict_local(series_4row):\n    d = {}\n    names = [['I', 'aVR', 'V1', 'V4'], \n             ['II', 'aVL', 'V2', 'V5'], \n             ['III', 'aVF', 'V3', 'V6']]\n    \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    \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    \n    return new_series\n\ndef series_dict(series):\n    d = {}\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            d[k] = s\n    \n    d['II'] = series[3]\n    \n    return d\n    \nfrom stage0_model import Net as Stage0Net\nfrom stage0_common import *\n\nstage0_dir = Path(global_dict[\"stage0_dir\"])\nstage0_dir.mkdir(exist_ok=True)\n\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(\"cuda:0\")\nstage0_net.eval()\n\nfor n, sample_id in enumerate(tqdm(test_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    image_original = cv2.cvtColor(image_original, cv2.COLOR_BGR2RGB)\n    image_for_model = change_color(image_original)\n    \n    batch = image_to_batch(image_for_model)\n\n    try:\n        with torch.no_grad(), torch.amp.autocast('cuda', dtype=torch.float32):\n            output = stage0_net(batch)\n        \n        rotated, keypoint = output_to_predict(image_original, batch, output)\n        normalised, _, _ = normalise_by_homography(rotated, keypoint)\n        \n        cv2.imwrite(str(output_path), cv2.cvtColor(normalised, cv2.COLOR_RGB2BGR))\n    except Exception as e:\n        traceback.print_exc()\n        copyfile(path, output_path)\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)\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(\"cuda:0\")\n\nfor n, sample_id in enumerate(tqdm(test_id)):\n    path = stage0_dir / f'{sample_id}.png'\n    output_path = stage1_dir / f'{sample_id}.png'\n    image = cv2.imread(path, cv2.IMREAD_COLOR_RGB)\n    batch = {'image': torch.from_numpy(np.ascontiguousarray(image.transpose(2, 0, 1))).unsqueeze(0)}\n\n    try:\n        with torch.no_grad(), torch.amp.autocast('cuda', dtype=torch.float32):\n            output = stage1_net(batch)\n        gridpoint_xy, _ = output_to_predict(image, batch, output)\n        rectified = rectify_image(image, gridpoint_xy)\n        cv2.imwrite(output_path, cv2.cvtColor(rectified, cv2.COLOR_RGB2BGR))\n    except:\n        traceback.print_exc()\n        copyfile(path, output_path)\nfrom stage2_model import *\nfrom stage2_common import *\n\nclass Net3(nn.Module):\n    \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\n        self.encoder = timm.create_model(\n            model_name='resnet34.a3_in1k', pretrained=pretrained, in_chans=3, num_classes=0, global_pool=''\n        )\n\n        self.decoder = MyCoordUnetDecoder(\n            in_channel=encoder_dim[-1],\n            skip_channel=encoder_dim[:-1][::-1] + [0],\n            out_channel=decoder_dim,\n            scale=[2, 2, 2, 2]\n        )\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        pixel = self.pixel(last)\n        return pixel\n\nstage1_dir = Path(global_dict[\"stage1_dir\"])\nstage2_dir = Path(global_dict[\"stage2_dir\"])\nstage2_dir.mkdir(exist_ok=True)\n\nstage2_net = Net3(pretrained=False).to(\"cuda:0\")\nmodel_path = \"/kaggle/input/physio-seg-public/pytorch/net3_009_4200/1/iter_0004200.pt\"\nstage2_net.load_state_dict(torch.load(model_path))\nstage2_net.eval()\n\nx0, x1 = 0, 2176\ny0, y1 = 0, 1696\nzero_mv = [703.5, 987.5, 1271.5, 1531.5]\nmv_to_pixel = 78.66\nt0, t1 = 235, 4161\n\nresize = T.Resize((1696, 4352), interpolation=T.InterpolationMode.BILINEAR)\n\nfor n, sample_id in enumerate(tqdm(test_id)):\n    path = stage1_dir / f'{sample_id}.png'\n    output_path = stage2_dir / f'{sample_id}.npy'\n    image = cv2.imread(path, cv2.IMREAD_COLOR_RGB)\n    \n    length = test[(test['id']==sample_id) & (test['lead']=='II')].iloc[0].number_of_rows\n    \n    image = image[y0:y1, x0:x1] / 255\n    batch = resize(torch.from_numpy(np.ascontiguousarray(image.transpose(2, 0, 1))).unsqueeze(0)).float().to(\"cuda:0\")\n    \n    try:\n        with torch.no_grad(), torch.amp.autocast('cuda', dtype=torch.float32):\n            output = stage2_net(batch)\n        \n        pixel = torch.sigmoid(output).float().data.cpu().numpy()[0]\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 = dw(s_dict)\n        series = dict_to_series_local(s_dict, series.shape)\n\n        np.save(output_path, series)\n    except:\n        traceback.print_exc()\n        series = np.zeros((4, length)) \n        np.save(output_path, series)\nstage2_dir = Path(global_dict[\"stage2_dir\"])\n\nres = []\ngb = test.groupby('id')\n\nfor i, (sample_id, df) in enumerate(tqdm(gb)):\n    series = np.load(stage2_dir / f'{sample_id}.npy')\n    d_series = series_dict(series)\n\n    for _, d in df.iterrows():\n        s = d_series.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}_{x}_{d.lead}' for x in range(d.number_of_rows)]\n        res.append(pd.DataFrame({'id': row_id, 'value': s}))\n\n    if i % 100 == 0:\n        gc.collect()\n\nsubmission = pd.concat(res, axis=0, ignore_index=True)\nsubmission.to_csv('submission.csv', index=False)\nsubmission.head(30)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-04T13:53:11.373599Z","iopub.execute_input":"2026-01-04T13:53:11.373957Z","iopub.status.idle":"2026-01-04T13:53:54.926640Z","shell.execute_reply.started":"2026-01-04T13:53:11.373925Z","shell.execute_reply":"2026-01-04T13:53:54.926000Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}