{"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,"isSourceIdPinned":false},{"sourceType":"datasetVersion","sourceId":13731160,"datasetId":8733970,"databundleVersionId":14479231},{"sourceType":"datasetVersion","sourceId":15281722,"datasetId":9775234,"databundleVersionId":16183481},{"sourceType":"datasetVersion","sourceId":14695416,"datasetId":9387663,"databundleVersionId":15539655}],"dockerImageVersionId":31328,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# import os\n\n# path1 = '/kaggle/input/datasets/hengck23/hengck23-demo-submit-physionet'\n# path2 = '/kaggle/input/datasets/takashisomeya/physionet-final-submission-models'\n\n# print(\"=== hengck23-demo-submit-physionet ===\")\n# for f in sorted(os.listdir(path1)):\n#     full = os.path.join(path1, f)\n#     if os.path.isdir(full):\n#         print(f\"  [{f}/]\")\n#         for ff in sorted(os.listdir(full)):\n#             print(f\"    {ff}\")\n#     else:\n#         print(f\"  {f}\")\n\n# print(f\"\\n=== physionet-final-submission-models ===\")\n# for f in sorted(os.listdir(path2)):\n#     print(f\"  {f}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-22T01:34:11.911885Z","iopub.execute_input":"2026-03-22T01:34:11.912698Z","iopub.status.idle":"2026-03-22T01:34:11.939756Z","shell.execute_reply.started":"2026-03-22T01:34:11.912662Z","shell.execute_reply":"2026-03-22T01:34:11.939031Z"},"jupyter":{"outputs_hidden":true,"source_hidden":true},"collapsed":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 1: Setup & Install\n!pip install connected-components-3d segmentation-models-pytorch --no-index --find-links=file:///kaggle/input/datasets/tylerde/my-pip-packages/","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-22T01:14:41.275719Z","iopub.execute_input":"2026-03-22T01:14:41.276466Z","iopub.status.idle":"2026-03-22T01:14:44.826389Z","shell.execute_reply.started":"2026-03-22T01:14:41.276431Z","shell.execute_reply":"2026-03-22T01:14:44.825642Z"},"collapsed":true,"jupyter":{"outputs_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 2: Imports & Config\nimport cc3d\nimport cv2\nimport pandas as pd\nimport numpy as np\nfrom scipy import signal\nimport torch\nimport torch.nn as nn\nimport matplotlib.pyplot as plt\nimport shutil\nimport copy\nimport multiprocessing as mp\nimport pickle\nimport os\nimport sys\n\nprint('import ok!!!')\n\nFLOAT_TYPE = torch.float16\n\n# === PATHS ===\nKAGGLE_DIR   = '/kaggle/input/competitions/physionet-ecg-image-digitization'\nHENGCK23_DIR = '/kaggle/input/datasets/hengck23/hengck23-demo-submit-physionet'\nSOMEYA_DIR   = '/kaggle/input/datasets/takashisomeya/physionet-final-submission-models'\nWEIGHT_DIR   = f'{HENGCK23_DIR}/weight'\nOUT_DIR      = '/kaggle/working/outputs'\n\n# Add to sys.path\nsys.path.append(HENGCK23_DIR)\nsys.path.append(SOMEYA_DIR)\n\n# === Read test metadata ===\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()\n\ndef read_image(sample_id):\n    return cv2.imread(f'{KAGGLE_DIR}/test/{sample_id}.png', cv2.IMREAD_COLOR_RGB)\n\ndef read_sampling_length(sample_id):\n    d = valid_df[\n        (valid_df['id'] == sample_id) & (valid_df['lead'] == 'II')\n    ].iloc[0]\n    return d.number_of_rows\n\nprint(f'valid_id: {len(valid_id)}')\nprint(f'\\t{valid_id[:3]} ...')\nprint('setting ok!!!\\n')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-22T01:14:44.827963Z","iopub.execute_input":"2026-03-22T01:14:44.828211Z","iopub.status.idle":"2026-03-22T01:14:54.453471Z","shell.execute_reply.started":"2026-03-22T01:14:44.828181Z","shell.execute_reply":"2026-03-22T01:14:54.452822Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 3: Stage 0 — Keypoint Detection & Normalization\nprint('*** STARTING STAGE0 ***')\n\nfrom stage0_model import Net as Stage0Net\nfrom stage0_common import *\n\nos.makedirs(f'{OUT_DIR}/normalised', exist_ok=True)\n\ndef run_stage0(gpu_id=0, assigned_ids=None, fail_id_file=None):\n    device = f'cuda:{gpu_id}'\n    if assigned_ids is None:\n        assigned_ids = valid_id\n\n    local_fail_id = []\n\n    stage0_net = Stage0Net(pretrained=False)\n    stage0_net = load_net(stage0_net, f'{WEIGHT_DIR}/stage0-last.checkpoint.pth')\n    stage0_net.to(device)\n\n    start_timer = timer()\n    for n, sample_id in enumerate(assigned_ids):\n        timestamp = time_to_str(timer() - start_timer, 'sec')\n        print(f'\\r\\t [GPU{gpu_id}] {n:4d}/{len(assigned_ids)} {sample_id}', timestamp, end='', flush=True)\n\n        image = read_image(sample_id)\n        batch = image_to_batch(image)\n\n        with torch.amp.autocast('cuda', dtype=FLOAT_TYPE):\n            with torch.no_grad():\n                output = stage0_net(batch)\n                try:\n                    rotated, keypoint = output_to_predict(image, batch, output)\n                    normalised, keypoint, homo = normalise_by_homography(rotated, keypoint)\n                    cv2.imwrite(f'{OUT_DIR}/normalised/{sample_id}.norm.png', cv2.cvtColor(normalised, cv2.COLOR_RGB2BGR))\n                    np.save(f'{OUT_DIR}/normalised/{sample_id}.homo.npy', homo)\n                except:\n                    local_fail_id.append(sample_id)\n\n        torch.cuda.empty_cache()\n\n        if n < 3 and gpu_id == 0:\n            overlay = draw_results_stage0(rotated, keypoint)\n            print(f'\\ndemo results for stage0: {sample_id}')\n            plt.imshow(image); plt.show()\n            plt.imshow(overlay); plt.show()\n            plt.imshow(normalised); plt.show()\n\n    print(f'\\n[GPU{gpu_id}] Stage0 completed. Failed: {len(local_fail_id)}')\n\n    if fail_id_file:\n        with open(fail_id_file, 'wb') as f:\n            pickle.dump(local_fail_id, f)\n\n    return local_fail_id\n\n\ndef run_stage0_parallel():\n    print('*** STARTING STAGE0 (2GPU PARALLEL) ***')\n    n_gpus = torch.cuda.device_count()\n\n    if n_gpus < 2:\n        print(f'Only {n_gpus} GPU(s) available, running single GPU')\n        return run_stage0(gpu_id=0)\n\n    mid_idx = len(valid_id) // 2\n    ids_gpu0, ids_gpu1 = valid_id[:mid_idx], valid_id[mid_idx:]\n    print(f'GPU0: {len(ids_gpu0)} | GPU1: {len(ids_gpu1)}')\n\n    fail_file_0 = f'{OUT_DIR}/fail_stage0_gpu0.pkl'\n    fail_file_1 = f'{OUT_DIR}/fail_stage0_gpu1.pkl'\n\n    p0 = mp.Process(target=run_stage0, args=(0, ids_gpu0, fail_file_0))\n    p1 = mp.Process(target=run_stage0, args=(1, ids_gpu1, fail_file_1))\n    p0.start(); p1.start()\n    p0.join();  p1.join()\n\n    fail_id = []\n    for ff in [fail_file_0, fail_file_1]:\n        if os.path.exists(ff):\n            with open(ff, 'rb') as f:\n                fail_id.extend(pickle.load(f))\n\n    print(f'FAIL_ID (Stage0): {fail_id}')\n    return fail_id\n\n\nFAIL_ID_STAGE0 = run_stage0_parallel()\nprint('Stage0 done!\\n')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-22T01:14:54.454244Z","iopub.execute_input":"2026-03-22T01:14:54.454513Z","iopub.status.idle":"2026-03-22T01:15:10.326043Z","shell.execute_reply.started":"2026-03-22T01:14:54.454487Z","shell.execute_reply":"2026-03-22T01:15:10.325383Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 4: Stage 1 — Grid Detection & Rectification\nprint('*** STARTING STAGE1 ***')\n\nfrom stage1_model import Net as Stage1Net\nfrom stage1_common import *\n\nos.makedirs(f'{OUT_DIR}/rectified', exist_ok=True)\n\ndef run_stage1(gpu_id=0, assigned_ids=None, prev_fail_ids=None, fail_id_file=None):\n    device = f'cuda:{gpu_id}'\n    if assigned_ids is None:\n        assigned_ids = valid_id\n    if prev_fail_ids is None:\n        prev_fail_ids = []\n\n    local_fail_id = []\n\n    stage1_net = Stage1Net(pretrained=False)\n    stage1_net = load_net(stage1_net, f'{WEIGHT_DIR}/stage1-last.checkpoint.pth')\n    stage1_net.to(device)\n\n    start_timer = timer()\n    for n, sample_id in enumerate(assigned_ids):\n        timestamp = time_to_str(timer() - start_timer, 'sec')\n        print(f'\\r\\t [GPU{gpu_id}] {n:4d}/{len(assigned_ids)} {sample_id}', timestamp, end='', flush=True)\n\n        if sample_id in prev_fail_ids:\n            continue\n\n        image = cv2.imread(f'{OUT_DIR}/normalised/{sample_id}.norm.png', cv2.IMREAD_COLOR_RGB)\n        batch = {\n            'image': torch.from_numpy(np.ascontiguousarray(image.transpose(2, 0, 1))).unsqueeze(0),\n        }\n\n        with torch.amp.autocast('cuda', dtype=FLOAT_TYPE):\n            with torch.no_grad():\n                output = stage1_net(batch)\n                try:\n                    gridpoint_xy, more = output_to_predict(image, batch, output)\n                    rectified = rectify_image(image, gridpoint_xy)\n                    cv2.imwrite(f'{OUT_DIR}/rectified/{sample_id}.rect.png', cv2.cvtColor(rectified, cv2.COLOR_RGB2BGR))\n                    np.save(f'{OUT_DIR}/rectified/{sample_id}.gridpoint_xy.npy', gridpoint_xy)\n                except:\n                    local_fail_id.append(sample_id)\n\n        torch.cuda.empty_cache()\n\n        if n < 3 and gpu_id == 0:\n            overlay = draw_mapping(image, gridpoint_xy)\n            ghfiltered, gvfiltered = draw_results_stage1(more)\n            print(f'\\ndemo results for stage1: {sample_id}')\n            plt.imshow(overlay); plt.show()\n            plt.imshow(gvfiltered); plt.show()\n            plt.imshow(ghfiltered); plt.show()\n            plt.imshow(rectified); plt.show()\n\n    print(f'\\n[GPU{gpu_id}] Stage1 completed. Failed: {len(local_fail_id)}')\n\n    if fail_id_file:\n        with open(fail_id_file, 'wb') as f:\n            pickle.dump(local_fail_id, f)\n\n    return local_fail_id\n\n\ndef run_stage1_parallel(prev_fail_ids=None):\n    print('*** STARTING STAGE1 (2GPU PARALLEL) ***')\n    n_gpus = torch.cuda.device_count()\n\n    if n_gpus < 2:\n        print(f'Only {n_gpus} GPU(s) available, running single GPU')\n        return run_stage1(gpu_id=0, prev_fail_ids=prev_fail_ids)\n\n    mid_idx = len(valid_id) // 2\n    ids_gpu0, ids_gpu1 = valid_id[:mid_idx], valid_id[mid_idx:]\n    print(f'GPU0: {len(ids_gpu0)} | GPU1: {len(ids_gpu1)}')\n\n    fail_file_0 = f'{OUT_DIR}/fail_stage1_gpu0.pkl'\n    fail_file_1 = f'{OUT_DIR}/fail_stage1_gpu1.pkl'\n\n    p0 = mp.Process(target=run_stage1, args=(0, ids_gpu0, prev_fail_ids, fail_file_0))\n    p1 = mp.Process(target=run_stage1, args=(1, ids_gpu1, prev_fail_ids, fail_file_1))\n    p0.start(); p1.start()\n    p0.join();  p1.join()\n\n    fail_id = []\n    for ff in [fail_file_0, fail_file_1]:\n        if os.path.exists(ff):\n            with open(ff, 'rb') as f:\n                fail_id.extend(pickle.load(f))\n    if prev_fail_ids:\n        fail_id.extend(prev_fail_ids)\n\n    print(f'FAIL_ID (Stage1): {fail_id}')\n    return fail_id\n\n\nFAIL_ID_STAGE1 = run_stage1_parallel(prev_fail_ids=FAIL_ID_STAGE0)\nprint('Stage1 done!\\n')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-22T01:15:10.327534Z","iopub.execute_input":"2026-03-22T01:15:10.328326Z","iopub.status.idle":"2026-03-22T01:15:16.849162Z","shell.execute_reply.started":"2026-03-22T01:15:10.328293Z","shell.execute_reply":"2026-03-22T01:15:16.848085Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 5: Stage 2 — Segmentation & Signal Extraction\nprint('*** STARTING STAGE2 ***')\n\nfrom stage2_smp_model import Net as WholeModel\nfrom stage2_lead_model import Net as LeadModel\nfrom stage2_model import prob_to_series_by_max\nfrom stage2_common import *\n\nos.makedirs(f'{OUT_DIR}/digitalised', exist_ok=True)\n\n# === Stage 2 Constants ===\nWINDOW_SIZE = 240\nOFFSET = 416\nIGNORE_EDGE = 8\nx_scale = 5000 / (2080 - 118)\nadd_x = 1\ny_scale = 1\nIMG_H, IMG_W = int(1700 * y_scale), int(2200 * x_scale) + add_x\n\ntta = [0, 2]  # no flip, horizontal flip\n\nx0, x1 = 0, 5600\ny0, y1 = 0, 1696\nzero_mv = [703.5, 987.5, 1271.5, 1531.5]\nzero_mv_trimed = [pos - OFFSET for pos in zero_mv]\nzero_mv_croped = [WINDOW_SIZE + 0.5 for _ in range(4)]\nmv_to_pixel = 79.0\nt0 = int(118 * x_scale) + add_x\nt1 = int(2080 * x_scale) + add_x\n\n# Pre-compute ensemble regions\nheight_after_trimed = y1 - OFFSET\nens_regions = []\nfor zmv in zero_mv_trimed:\n    trim_upper = int(zmv) - WINDOW_SIZE\n    trim_lower = int(zmv) + WINDOW_SIZE\n    lead_upper = IGNORE_EDGE\n    lead_lower = -IGNORE_EDGE\n    if trim_lower > height_after_trimed:\n        lead_lower = (trim_lower - height_after_trimed + IGNORE_EDGE) * -1\n        trim_lower = height_after_trimed\n    trim_upper += IGNORE_EDGE\n    trim_lower -= IGNORE_EDGE\n    ens_regions.append([trim_upper, trim_lower, lead_upper, lead_lower])\n\nprint(f'IMG: ({IMG_H}, {IMG_W}), timespan: [{t0}:{t1}]')\nprint(f'ens_regions: {ens_regions}')\n\n\ndef pixel_to_series_exp(pixel, zero_mv, length):\n    _, H, W = pixel.shape\n    eps = 1e-8\n    y_idx = np.arange(H, dtype=np.float32)[:, None]\n\n    series = []\n    for j in [0, 1, 2, 3]:\n        p = pixel[j]\n        denom = p.sum(axis=0)\n        y_exp = (p * y_idx).sum(axis=0) / (denom + eps)\n        series.append(y_exp)\n    series = np.stack(series).astype(np.float32)\n\n    if length is not None and length != W:\n        resampled_series = []\n        for s in series:\n            rs = signal.resample(s, length).astype(np.float32)\n            resampled_series.append(rs)\n        series = np.stack(resampled_series)\n\n    return series\n\n\ndef get_whole_model(encoder_name, weight_path, device):\n    model = WholeModel(\n        encoder_name=encoder_name,\n        encoder_weights=None,\n        decoder_name=\"unet\",\n        use_coord_conv=True,\n        pretrained=False\n    )\n    state_dict = torch.load(weight_path, map_location='cpu')\n    print(model.load_state_dict(state_dict, strict=False))\n    model.to(device)\n    model.eval()\n    model.output_type = ['infer']\n    return model\n\n\ndef get_lead_model(encoder_name, weight_path, fusion_type, device):\n    model = LeadModel(\n        encoder_name=encoder_name,\n        encoder_weights=None,\n        fusion_type=fusion_type,\n    )\n    state_dict = torch.load(weight_path, map_location='cpu')\n    print(model.load_state_dict(state_dict, strict=False))\n    model.to(device)\n    model.eval()\n    model.output_type = ['infer']\n    return model\n\n\ndef read_images(path):\n    image = cv2.imread(path, cv2.IMREAD_COLOR)\n    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n    image = cv2.resize(image, (IMG_W, IMG_H), interpolation=cv2.INTER_LINEAR)\n    trim_image = image.copy()[OFFSET:y1, x0:x1]\n\n    # Make lead images\n    image = image[y0:y1, x0:x1]\n    H, W, _ = image.shape\n    lead_images = []\n    for i, zmv in enumerate(zero_mv):\n        h0, h1 = int(zmv) - WINDOW_SIZE, int(zmv) + WINDOW_SIZE\n        src_h0 = max(0, h0)\n        src_h1 = min(H, h1)\n        dst_h0 = src_h0 - h0\n        dst_h1 = dst_h0 + (src_h1 - src_h0)\n\n        lead_img = np.zeros((WINDOW_SIZE * 2, W, 3))\n        lead_img[dst_h0:dst_h1, :, :] = image[src_h0:src_h1, :, :]\n        lead_images.append(lead_img)\n\n    lead_images = np.stack(lead_images)  # (4, H, W, 3)\n    return trim_image, lead_images\n\n\ndef run_stage2(gpu_id=0, assigned_ids=None, prev_fail_ids=None, fail_id_file=None):\n    device = f'cuda:{gpu_id}'\n    if assigned_ids is None:\n        assigned_ids = valid_id\n    if prev_fail_ids is None:\n        prev_fail_ids = []\n\n    local_fail_id = []\n\n    # Load whole models\n    whole_models = [\n        get_whole_model(\"tu-timm/tf_efficientnet_b7.ns_jft_in1k\",\n                        f\"{SOMEYA_DIR}/whole_b7_lb22.93.pth\", device),\n        get_whole_model(\"tu-timm/tf_efficientnetv2_l.in21k\",\n                        f\"{SOMEYA_DIR}/whole_v2_l_lb22.60.pth\", device),\n    ]\n\n    # Load lead/series models\n    lead_models = [\n        get_lead_model(\"tu-timm/tf_efficientnet_b6.ns_jft_in1k\",\n                       f\"{SOMEYA_DIR}/series_b6_shared_conv2d_lb23.10.pth\", \"shared_conv2d\", device),\n        get_lead_model(\"tu-timm/tf_efficientnet_b6.ns_jft_in1k\",\n                       f\"{SOMEYA_DIR}/series_b6_shared_conv2d_lb23.00.pth\", \"shared_conv2d\", device),\n        get_lead_model(\"tu-timm/tf_efficientnetv2_l.in21k\",\n                       f\"{SOMEYA_DIR}/series_v2_l_conv3d_lb22.92.pth\", \"conv3d\", device),\n        get_lead_model(\"tu-timm/tf_efficientnetv2_l.in21k\",\n                       f\"{SOMEYA_DIR}/series_v2_l_conv2d_lb22.85.pth\", \"conv2d\", device),\n    ]\n\n    start_timer = timer()\n    for n, sample_id in enumerate(assigned_ids):\n        timestamp = time_to_str(timer() - start_timer, 'sec')\n        print(f'\\r\\t [GPU{gpu_id}] {n:4d}/{len(assigned_ids)} {sample_id}', timestamp, end='', flush=True)\n\n        if sample_id in prev_fail_ids:\n            continue\n\n        length = read_sampling_length(sample_id)\n        trim_image, lead_images = read_images(f'{OUT_DIR}/rectified/{sample_id}.rect.png')\n\n        pixel_ens = np.zeros((4, trim_image.shape[0], trim_image.shape[1])) * 1.0\n\n        # --- Whole models ---\n        batch = {\n            'image': torch.from_numpy(\n                np.ascontiguousarray(trim_image.transpose(2, 0, 1))\n            ).unsqueeze(0),\n        }\n        batch_tta = {\n            'image': torch.from_numpy(\n                np.ascontiguousarray(np.fliplr(trim_image).copy().transpose(2, 0, 1))\n            ).unsqueeze(0),\n        }\n\n        with torch.amp.autocast('cuda', dtype=FLOAT_TYPE):\n            with torch.no_grad():\n                for model in whole_models:\n                    for flip in tta:\n                        if flip:\n                            output = model(batch_tta)\n                            pixel = output['pixel'].float().data.cpu().numpy()[0]\n                            pixel = np.flip(pixel, axis=flip)\n                        else:\n                            output = model(batch)\n                            pixel = output['pixel'].float().data.cpu().numpy()[0]\n                        pixel_ens += pixel\n\n        # --- Lead/series models ---\n        lead_tensor = torch.from_numpy(\n            lead_images.transpose(0, 3, 1, 2)\n        ).contiguous()  # (4, 3, H, W)\n        batch = {'image': lead_tensor.unsqueeze(0)}\n        batch_tta = {'image': torch.flip(lead_tensor, dims=[3]).unsqueeze(0)}\n\n        with torch.amp.autocast('cuda', dtype=FLOAT_TYPE):\n            with torch.no_grad():\n                for model in lead_models:\n                    for flip in tta:\n                        if flip:\n                            output = model(batch_tta)\n                            pixel = output['pixel'].float().data.cpu().numpy()[0].squeeze(1)\n                            pixel = np.flip(pixel, axis=flip)\n                        else:\n                            output = model(batch)\n                            pixel = output['pixel'].float().data.cpu().numpy()[0].squeeze(1)\n\n                        for i in range(4):\n                            trim_upper, trim_lower, lead_upper, lead_lower = ens_regions[i]\n                            pixel_ens[i][trim_upper:trim_lower] += pixel[i][lead_upper:lead_lower]\n\n        # --- Weighted average ---\n        ens_weight = np.ones((trim_image.shape[0], trim_image.shape[1])) * len(whole_models) * len(tta)\n        for i in range(4):\n            trim_upper, trim_lower, _, _ = ens_regions[i]\n            ens_weight[trim_upper:trim_lower] += len(lead_models) * len(tta)\n        pixel_ens /= ens_weight\n\n        # --- Convert pixel → time series ---\n        try:\n            series_in_pixel = pixel_to_series_exp(pixel_ens[..., t0:t1], zero_mv_trimed, length)\n            series = (np.array(zero_mv_trimed).reshape(4, 1) - series_in_pixel) / mv_to_pixel\n            np.save(f'{OUT_DIR}/digitalised/{sample_id}.series.npy', series)\n        except:\n            local_fail_id.append(sample_id)\n\n        if n < 3 and gpu_id == 0:\n            print()\n            print(f\"check max intensity: {np.max(pixel_ens)}\")\n            overlay = draw_lead_pixel(trim_image, pixel_ens)\n            plt.imshow(overlay); plt.show()\n            t = np.arange(len(series[0]))\n            fig, axes = plt.subplots(4, 1, figsize=(12, 10))\n            for j in range(4):\n                axes[j].plot(t, series[j], alpha=1.0, color='blue', linewidth=1, label='predict')\n                axes[j].legend()\n            plt.tight_layout(); plt.show()\n\n    print(f'\\n[GPU{gpu_id}] Stage2 completed. Failed: {len(local_fail_id)}')\n\n    if fail_id_file:\n        with open(fail_id_file, 'wb') as f:\n            pickle.dump(local_fail_id, f)\n\n    return local_fail_id\n\n\ndef run_stage2_parallel(prev_fail_ids=None):\n    print('*** STARTING STAGE2 (2GPU PARALLEL) ***')\n    n_gpus = torch.cuda.device_count()\n\n    if n_gpus < 2:\n        print(f'Only {n_gpus} GPU(s) available, running single GPU')\n        return run_stage2(gpu_id=0, prev_fail_ids=prev_fail_ids)\n\n    mid_idx = len(valid_id) // 2\n    ids_gpu0, ids_gpu1 = valid_id[:mid_idx], valid_id[mid_idx:]\n    print(f'GPU0: {len(ids_gpu0)} | GPU1: {len(ids_gpu1)}')\n\n    fail_file_0 = f'{OUT_DIR}/fail_stage2_gpu0.pkl'\n    fail_file_1 = f'{OUT_DIR}/fail_stage2_gpu1.pkl'\n\n    p0 = mp.Process(target=run_stage2, args=(0, ids_gpu0, prev_fail_ids, fail_file_0))\n    p1 = mp.Process(target=run_stage2, args=(1, ids_gpu1, prev_fail_ids, fail_file_1))\n    p0.start(); p1.start()\n    p0.join();  p1.join()\n\n    fail_id = []\n    for ff in [fail_file_0, fail_file_1]:\n        if os.path.exists(ff):\n            with open(ff, 'rb') as f:\n                fail_id.extend(pickle.load(f))\n    if prev_fail_ids:\n        fail_id.extend(prev_fail_ids)\n\n    print(f'FAIL_ID (Stage2): {fail_id}')\n    return fail_id\n\n\nFAIL_ID_STAGE2 = run_stage2_parallel(prev_fail_ids=FAIL_ID_STAGE1)\nprint('Stage2 done!\\n')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-22T01:15:16.852484Z","iopub.execute_input":"2026-03-22T01:15:16.853216Z","iopub.status.idle":"2026-03-22T01:16:26.483544Z","shell.execute_reply.started":"2026-03-22T01:15:16.853181Z","shell.execute_reply":"2026-03-22T01:16:26.482725Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 6: Build Submission\ndef make_submission():\n    print('===========================================')\n    print('Making submission CSV ...')\n\n    submit_df = []\n    gb = valid_df.groupby('id')\n\n    for i, (sample_id, df) in enumerate(gb):\n        try:\n            series = np.load(f'{OUT_DIR}/digitalised/{sample_id}.series.npy')\n\n            series_by_lead = {}\n            for l in range(3):\n                lead = [\n                    ['I',   'aVR', 'V1', 'V4'],\n                    ['II',  'aVL', 'V2', 'V5'],\n                    ['III', 'aVF', 'V3', 'V6'],\n                ][l]\n\n                length = [\n                    df[df['lead'] == lead[j]].iloc[0].number_of_rows\n                    for j in range(4)\n                ]\n                if lead[0] == 'II':\n                    length[0] = length[0] - sum(length[1:])\n\n                index = np.cumsum(length)[:-1]\n                split = np.split(series[l], index)\n                for (k, s) in zip(lead, split):\n                    series_by_lead[k] = s\n\n            series_by_lead['II'] = series[3]\n\n        except:\n            series_by_lead = {}\n            for j, d in df.iterrows():\n                series_by_lead[d.lead] = np.zeros(d.number_of_rows)\n\n        for j, d in df.iterrows():\n            # Safety pad: ensure correct length\n            series_by_lead[d.lead] = np.concatenate([\n                series_by_lead[d.lead], np.zeros_like(series_by_lead[d.lead])\n            ])[:d.number_of_rows]\n            assert len(series_by_lead[d.lead]) == d.number_of_rows\n\n            print(f'\\r\\t {i} {sample_id} : {d.lead}', end='', flush=True)\n\n            row_id = [f'{sample_id}_{i}_{d.lead}' for i in range(d.number_of_rows)]\n            this_df = pd.DataFrame({\n                'id': row_id,\n                'value': series_by_lead[d.lead].astype(np.float32),\n            })\n            submit_df.append(this_df)\n\n    print('')\n    submit_df = pd.concat(submit_df, axis=0, ignore_index=True, sort=False, copy=False)\n    print(submit_df)\n    print(f'\\nTotal rows: {len(submit_df)}')\n    submit_df.to_csv('submission.csv', index=False)\n    print('Saved submission.csv')\n\n\nmake_submission()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-22T01:16:26.484912Z","iopub.execute_input":"2026-03-22T01:16:26.485183Z","iopub.status.idle":"2026-03-22T01:16:26.954418Z","shell.execute_reply.started":"2026-03-22T01:16:26.485153Z","shell.execute_reply":"2026-03-22T01:16:26.953671Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ## Cell 7: Cleanup\n\n# %%\nshutil.rmtree(OUT_DIR)\nprint('Cleanup done!')\nprint('Ready to submit submission.csv')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-22T01:16:26.955369Z","iopub.execute_input":"2026-03-22T01:16:26.955677Z","iopub.status.idle":"2026-03-22T01:16:26.964636Z","shell.execute_reply.started":"2026-03-22T01:16:26.955650Z","shell.execute_reply":"2026-03-22T01:16:26.963983Z"}},"outputs":[],"execution_count":null}]}