{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":97984,"databundleVersionId":14096757},{"sourceType":"datasetVersion","sourceId":15613723,"datasetId":9992354,"databundleVersionId":16547648},{"sourceType":"datasetVersion","sourceId":15658841,"datasetId":10023761,"databundleVersionId":16595325},{"sourceType":"datasetVersion","sourceId":15705742,"datasetId":10061224,"databundleVersionId":16645301},{"sourceType":"datasetVersion","sourceId":15810903,"datasetId":10135127,"databundleVersionId":16758858},{"sourceType":"datasetVersion","sourceId":13731160,"datasetId":8733970,"databundleVersionId":14479231},{"sourceType":"datasetVersion","sourceId":15281722,"datasetId":9775234,"databundleVersionId":16183481}],"dockerImageVersionId":31329,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# === Setup ===\n\n!pip install connected-components-3d --no-index --find-links=file:///kaggle/input/datasets/tylerde/my-pip-packages/ -q\n\nimport os, sys, cv2\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom scipy import signal as scipy_signal\nimport matplotlib.pyplot as plt\n\nKAGGLE_DIR = '/kaggle/input/competitions/physionet-ecg-image-digitization'\nHENGCK23_DIR = '/kaggle/input/datasets/hengck23/hengck23-demo-submit-physionet'\nMODEL_DIR = '/kaggle/input/datasets/tylerde/ecg-model-all-samples'\nWEIGHT_DIR = f'{HENGCK23_DIR}/weight'\nOUT_DIR = '/kaggle/working/outputs'\n\nsys.path.append(HENGCK23_DIR)\nsys.path.append(MODEL_DIR)\n\ndevice = 'cuda:0'\n\nvalid_df = pd.read_csv(f'{KAGGLE_DIR}/test.csv')\nvalid_df['id'] = valid_df['id'].astype(str)\nvalid_id = valid_df['id'].unique().tolist()\nprint(f'Test images: {len(valid_id)}')","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-04-13T07:38:58.044904Z","iopub.execute_input":"2026-04-13T07:38:58.045179Z","iopub.status.idle":"2026-04-13T07:39:09.064830Z","shell.execute_reply.started":"2026-04-13T07:38:58.045141Z","shell.execute_reply":"2026-04-13T07:39:09.064085Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# === Stage 0 — Keypoint Detection ===\n\nfrom stage0_model import Net as Stage0Net\nfrom stage0_common import *\n\nos.makedirs(f'{OUT_DIR}/normalised', exist_ok=True)\n\ns0 = Stage0Net(pretrained=False)\ns0 = load_net(s0, f'{WEIGHT_DIR}/stage0-last.checkpoint.pth')\ns0.to(device)\n\nfail_ids_s0 = []\nfor n, sid in enumerate(valid_id):\n    image = cv2.imread(f'{KAGGLE_DIR}/test/{sid}.png', cv2.IMREAD_COLOR_RGB)\n    batch = image_to_batch(image)\n    with torch.amp.autocast('cuda', dtype=torch.float16):\n        with torch.no_grad():\n            try:\n                out = s0(batch)\n                rot, kp = output_to_predict(image, batch, out)\n                norm, kp, homo = normalise_by_homography(rot, kp)\n                cv2.imwrite(\n                    f'{OUT_DIR}/normalised/{sid}.norm.png',\n                    cv2.cvtColor(norm, cv2.COLOR_RGB2BGR)\n                )\n            except:\n                fail_ids_s0.append(sid)\n    torch.cuda.empty_cache()\n    if (n + 1) % 100 == 0:\n        print(f'  {n+1}/{len(valid_id)} done')\n\ndel s0; torch.cuda.empty_cache()\nprint(f'Stage 0 done. Failed: {len(fail_ids_s0)}')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-13T07:39:39.695142Z","iopub.execute_input":"2026-04-13T07:39:39.695751Z","iopub.status.idle":"2026-04-13T07:39:52.274807Z","shell.execute_reply.started":"2026-04-13T07:39:39.695713Z","shell.execute_reply":"2026-04-13T07:39:52.274049Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# === Stage 1 — Grid Detection + Rectification ===\n\nfrom stage1_model import Net as Stage1Net\nfrom stage1_common import *\n\nos.makedirs(f'{OUT_DIR}/rectified', exist_ok=True)\n\ns1 = Stage1Net(pretrained=False)\ns1 = load_net(s1, f'{WEIGHT_DIR}/stage1-last.checkpoint.pth')\ns1.to(device)\n\nfail_ids_s1 = list(fail_ids_s0)\nfor n, sid in enumerate(valid_id):\n    if sid in fail_ids_s0:\n        continue\n    image = cv2.imread(f'{OUT_DIR}/normalised/{sid}.norm.png', cv2.IMREAD_COLOR_RGB)\n    batch = {'image': torch.from_numpy(\n        np.ascontiguousarray(image.transpose(2, 0, 1))\n    ).unsqueeze(0)}\n    with torch.amp.autocast('cuda', dtype=torch.float16):\n        with torch.no_grad():\n            try:\n                out = s1(batch)\n                gp, more = output_to_predict(image, batch, out)\n                rect = rectify_image(image, gp)\n                cv2.imwrite(\n                    f'{OUT_DIR}/rectified/{sid}.rect.png',\n                    cv2.cvtColor(rect, cv2.COLOR_RGB2BGR)\n                )\n            except:\n                fail_ids_s1.append(sid)\n    torch.cuda.empty_cache()\n    if (n + 1) % 100 == 0:\n        print(f'  {n+1}/{len(valid_id)} done')\n\ndel s1; torch.cuda.empty_cache()\nprint(f'Stage 1 done. Failed: {len(fail_ids_s1)}')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-13T07:39:54.574494Z","iopub.execute_input":"2026-04-13T07:39:54.575377Z","iopub.status.idle":"2026-04-13T07:40:01.214164Z","shell.execute_reply.started":"2026-04-13T07:39:54.575347Z","shell.execute_reply":"2026-04-13T07:40:01.213507Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# === Stage 2 — Model ===\n\nfrom model import ECGRowNet\nfrom config import CFG\nfrom inference import predict_image\n\nmodel = ECGRowNet(CFG).to(device)\nckpt = torch.load(f'{MODEL_DIR}/best_fold0.pth', map_location=device, weights_only=False)\nmodel.load_state_dict(ckpt['state_dict'])\nmodel.eval()\nprint(f'Loaded: epoch {ckpt[\"epoch\"]}, SNR {ckpt[\"snr\"]:.2f} dB')\n\nos.makedirs(f'{OUT_DIR}/digitalised', exist_ok=True)\n\nfor n, sid in enumerate(valid_id):\n    if sid in fail_ids_s1:\n        continue\n    rect_path = f'{OUT_DIR}/rectified/{sid}.rect.png'\n    series = predict_image(model, rect_path, device, use_tta=True)\n    np.save(f'{OUT_DIR}/digitalised/{sid}.series.npy', series)\n    if (n + 1) % 100 == 0:\n        print(f'  {n+1}/{len(valid_id)} done')\n\ndel model; torch.cuda.empty_cache()\nprint(f'Stage 2 done.')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-13T07:40:01.215446Z","iopub.execute_input":"2026-04-13T07:40:01.215827Z","iopub.status.idle":"2026-04-13T07:40:07.326880Z","shell.execute_reply.started":"2026-04-13T07:40:01.215804Z","shell.execute_reply":"2026-04-13T07:40:07.326198Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# === Build Submission ===\n\nROW_TO_LEADS = [\n    ['I', 'aVR', 'V1', 'V4'],\n    ['II', 'aVL', 'V2', 'V5'],\n    ['III', 'aVF', 'V3', 'V6'],\n]\nsubmit_parts = []\ngb = valid_df.groupby('id')\nfor i, (sid, df) in enumerate(gb):\n    try:\n        series = np.load(f'{OUT_DIR}/digitalised/{sid}.series.npy')\n        series_by_lead = {}\n        for row_idx in range(3):\n            leads = ROW_TO_LEADS[row_idx]\n            lengths = [df[df['lead'] == lead].iloc[0].number_of_rows for lead in leads]\n            if leads[0] == 'II':\n                lengths[0] = lengths[0] - sum(lengths[1:])\n            total_len = sum(lengths)\n            row_resampled = scipy_signal.resample(series[row_idx], total_len).astype(np.float32)\n            idx = np.cumsum(lengths)[:-1]\n            splits = np.split(row_resampled, idx)\n            for lead, s in zip(leads, splits):\n                series_by_lead[lead] = s\n        ii_len = df[df['lead'] == 'II'].iloc[0].number_of_rows\n        series_by_lead['II'] = scipy_signal.resample(series[3], ii_len).astype(np.float32)\n    except:\n        series_by_lead = {}\n        for _, d in df.iterrows():\n            series_by_lead[d.lead] = np.zeros(d.number_of_rows, dtype=np.float32)\n    for _, d in df.iterrows():\n        lead_data = series_by_lead[d.lead]\n        lead_data = np.concatenate([lead_data, np.zeros(d.number_of_rows)])[:d.number_of_rows]\n        row_ids = [f'{sid}_{i}_{d.lead}' for i in range(d.number_of_rows)]\n        submit_parts.append(pd.DataFrame({'id': row_ids, 'value': lead_data.astype(np.float32)}))\nsubmit_df = pd.concat(submit_parts, ignore_index=True)\nsubmit_df.to_csv('submission.csv', index=False)\nprint(f'Submission saved: {len(submit_df)} rows')","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}