{"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":13731160,"sourceType":"datasetVersion","datasetId":8733970},{"sourceId":14695416,"sourceType":"datasetVersion","datasetId":9387663}],"dockerImageVersionId":31193,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Reference\n- https://www.kaggle.com/code/hengck23/demo-submission","metadata":{}},{"cell_type":"code","source":"!ls /kaggle/input/hengck23-demo-submit-physionet/setup\n!pip install connected-components-3d --no-index --find-links=file:///kaggle/input/hengck23-demo-submit-physionet/setup/","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-01T08:17:29.834279Z","iopub.execute_input":"2026-02-01T08:17:29.834608Z","iopub.status.idle":"2026-02-01T08:17:33.085946Z","shell.execute_reply.started":"2026-02-01T08:17:29.834586Z","shell.execute_reply":"2026-02-01T08:17:33.085222Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import cc3d\nimport cv2\nimport pandas as pd\nimport numpy as np\nfrom scipy import signal\nimport torch\nimport matplotlib.pyplot as plt\nimport matplotlib\n#matplotlib.use('TkAgg')\nimport shutil\nimport copy\nimport multiprocessing as mp\nimport pickle\nimport os\n\nimport sys\nsys.path.append('/kaggle/input/hengck23-demo-submit-physionet')\n\nprint('import ok!!!')","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-02-01T08:17:33.087533Z","iopub.execute_input":"2026-02-01T08:17:33.088158Z","iopub.status.idle":"2026-02-01T08:17:33.093359Z","shell.execute_reply.started":"2026-02-01T08:17:33.088131Z","shell.execute_reply":"2026-02-01T08:17:33.092708Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"FLOAT_TYPE = torch.float16 #torch.bfloat16\n\nKAGGLE_DIR = \\\n\t'/kaggle/input/physionet-ecg-image-digitization'\nWEIGHT_DIR = \\\n\t'/kaggle/input/hengck23-demo-submit-physionet/weight'\nOUT_DIR = \\\n    f'/kaggle/working/outputs'\n    \n#--------------------------------------\n\ndef read_image(sample_id):\n    image_id = sample_id\n    image = cv2.imread(f'{KAGGLE_DIR}/test/{image_id}.png', cv2.IMREAD_COLOR_RGB)\n    return image\n\ndef read_sampling_length(sample_id):\n    image_id = sample_id\n    d = valid_df[\n        (valid_df['id']==image_id) & (valid_df['lead']=='II')\n    ].iloc[0]\n    length = d.number_of_rows\n    return length\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()\n\n#valid_id = valid_id[:300]\nprint('valid_id:', len(valid_id))\nprint('\\t', valid_id[:3], '...')\nprint('setting ok!!!\\n')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-01T08:17:33.094119Z","iopub.execute_input":"2026-02-01T08:17:33.094350Z","iopub.status.idle":"2026-02-01T08:17:33.119952Z","shell.execute_reply.started":"2026-02-01T08:17:33.094335Z","shell.execute_reply":"2026-02-01T08:17:33.119221Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Stage0","metadata":{}},{"cell_type":"code","source":"# stage0\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\tdevice = f'cuda:{gpu_id}'\n\t\n\tif assigned_ids is None:\n\t\tassigned_ids = valid_id\n\t\n\tlocal_fail_id = []\n\t\n\tstage0_net = Stage0Net(pretrained=False)\n\tstage0_net = load_net(stage0_net, f'{WEIGHT_DIR}/stage0-last.checkpoint.pth')\n\tstage0_net.to(device)\n\n\tstart_timer = timer()\n\tfor n, sample_id in enumerate(assigned_ids):\n\t\ttimestamp = time_to_str(timer() - start_timer, 'sec')\n\t\tprint(f'\\r\\t [GPU{gpu_id}] {n:4d}/{len(assigned_ids)} {sample_id}', timestamp, end='', flush=True)\n\n\t\timage = read_image(sample_id)\n\t\tbatch = image_to_batch(image)\n\n\t\twith torch.amp.autocast('cuda', dtype=FLOAT_TYPE):\n\t\t\twith torch.no_grad():\n\t\t\t\toutput = stage0_net(batch)\n\n\t\t\t\ttry:\n\t\t\t\t\trotated, keypoint = output_to_predict(image, batch, output)\n\t\t\t\t\tnormalised, keypoint, homo = normalise_by_homography(rotated, keypoint)\n\t\t\t\t\t# ---\n\t\t\t\t\tcv2.imwrite(f'{OUT_DIR}/normalised/{sample_id}.norm.png', cv2.cvtColor(normalised, cv2.COLOR_RGB2BGR))\n\t\t\t\t\tnp.save(f'{OUT_DIR}/normalised/{sample_id}.homo.npy', homo)\n\t\t\t\texcept:\n\t\t\t\t\tlocal_fail_id.append(sample_id)\n\n\t\ttorch.cuda.empty_cache()\n\t\t\n\t\t# Visualization only on GPU0\n\t\tif n<10 and gpu_id==0:\n\t\t\toverlay = draw_results_stage0(rotated, keypoint)\n\t\t\tprint('')\n\t\t\tprint('demo results for stage0--------------')\n\t\t\tprint(sample_id)\n\t\t\tplt.imshow(image);plt.show()\n\t\t\tplt.imshow(overlay);plt.show()\n\t\t\tplt.imshow(normalised);plt.show()\n\t\t\t\n\tprint(f'\\n[GPU{gpu_id}] Stage0 completed. Failed: {len(local_fail_id)}')\n\t\n\t# Save FAIL_ID\n\tif fail_id_file:\n\t\twith open(fail_id_file, 'wb') as f:\n\t\t\tpickle.dump(local_fail_id, f)\n\t\n\treturn local_fail_id\n\n\ndef run_stage0_parallel():\n\t\"\"\"Run Stage0 in parallel on 2 GPUs\"\"\"\n\tprint('*** STARTING STAGE0 (2GPU PARALLEL) ***')\n\t\n\t# Split valid_id into 2 parts\n\tmid_idx = len(valid_id) // 2\n\tids_gpu0 = valid_id[:mid_idx]\n\tids_gpu1 = valid_id[mid_idx:]\n\t\n\tprint(f'GPU0: {len(ids_gpu0)} samples')\n\tprint(f'GPU1: {len(ids_gpu1)} samples')\n\t\n\t# FAIL_ID file paths\n\tfail_file_0 = f'{OUT_DIR}/fail_stage0_gpu0.pkl'\n\tfail_file_1 = f'{OUT_DIR}/fail_stage0_gpu1.pkl'\n\t\n\t# Run in parallel with 2 processes\n\tp0 = mp.Process(target=run_stage0, args=(0, ids_gpu0, fail_file_0))\n\tp1 = mp.Process(target=run_stage0, args=(1, ids_gpu1, fail_file_1))\n\t\n\tp0.start()\n\tp1.start()\n\t\n\tp0.join()\n\tp1.join()\n\t\n\t# Merge FAIL_IDs\n\tfail_id = []\n\tif os.path.exists(fail_file_0):\n\t\twith open(fail_file_0, 'rb') as f:\n\t\t\tfail_id.extend(pickle.load(f))\n\tif os.path.exists(fail_file_1):\n\t\twith open(fail_file_1, 'rb') as f:\n\t\t\tfail_id.extend(pickle.load(f))\n\t\n\tprint(f'FAIL_ID (Stage0): {fail_id}')\n\tprint('run_stage0_parallel() ok!!!\\n')\n\t\n\treturn fail_id\n\n\n# Run (using parallel version)\nFAIL_ID_STAGE0 = run_stage0_parallel()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-01T08:17:33.121688Z","iopub.execute_input":"2026-02-01T08:17:33.121889Z","iopub.status.idle":"2026-02-01T08:17:37.114073Z","shell.execute_reply.started":"2026-02-01T08:17:33.121873Z","shell.execute_reply":"2026-02-01T08:17:37.113361Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Stage1","metadata":{}},{"cell_type":"code","source":"# stage1\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\t\"\"\"\n\tStage1 processing (GPU parallelization supported)\n\t\n\tArgs:\n\t\tgpu_id: GPU number to use (0 or 1)\n\t\tassigned_ids: List of sample_ids to process (if None, use all valid_id)\n\t\tprev_fail_ids: IDs that failed in Stage0 (to be skipped)\n\t\tfail_id_file: File path to save FAIL_ID\n\t\"\"\"\n\t# GPU configuration\n\tdevice = f'cuda:{gpu_id}'\n\t\n\t# Determine processing targets\n\tif assigned_ids is None:\n\t\tassigned_ids = valid_id\n\tif prev_fail_ids is None:\n\t\tprev_fail_ids = []\n\t\n\t# Local FAIL_ID\n\tlocal_fail_id = []\n\t\n\tstage1_net = Stage1Net(pretrained=False)\n\tstage1_net = load_net(stage1_net, f'{WEIGHT_DIR}/stage1-last.checkpoint.pth')\n\tstage1_net.to(device)\n\n\tstart_timer = timer()\n\tfor n, sample_id in enumerate(assigned_ids):\n\t\ttimestamp = time_to_str(timer() - start_timer, 'sec')\n\t\tprint(f'\\r\\t [GPU{gpu_id}] {n:4d}/{len(assigned_ids)} {sample_id}', timestamp, end='', flush=True)\n\t\t\n\t\t# Skip IDs that failed in Stage0\n\t\tif sample_id in prev_fail_ids:\n\t\t\tcontinue\n\n\t\timage = cv2.imread(f'{OUT_DIR}/normalised/{sample_id}.norm.png', cv2.IMREAD_COLOR_RGB)\n\t\tbatch = {\n\t\t\t'image': torch.from_numpy(np.ascontiguousarray(image.transpose(2, 0, 1))).unsqueeze(0),\n\t\t}\n\t\tnum_tta = 1\n\n\t\twith torch.amp.autocast('cuda', dtype=FLOAT_TYPE):\n\t\t\twith torch.no_grad():\n\t\t\t\toutput = stage1_net(batch)\n\n\t\t\t\ttry:\n\t\t\t\t\tgridpoint_xy, more = output_to_predict(image, batch, output)\n\t\t\t\t\trectified = rectify_image(image, gridpoint_xy)\n\t\t\t\t\t# ---\n\t\t\t\t\tcv2.imwrite(f'{OUT_DIR}/rectified/{sample_id}.rect.png', cv2.cvtColor(rectified, cv2.COLOR_RGB2BGR))\n\t\t\t\t\tnp.save(f'{OUT_DIR}/rectified/{sample_id}.gridpoint_xy.npy',gridpoint_xy)\n\t\t\t\texcept:\n\t\t\t\t\tlocal_fail_id.append(sample_id)\n\n\t\ttorch.cuda.empty_cache()\n\t\t\n\t\t# Visualization only on GPU0\n\t\tif n<10 and gpu_id==0:\n\t\t\toverlay = draw_mapping(image, gridpoint_xy)\n\t\t\tghfiltered, gvfiltered = draw_results_stage1(more)\n\t\t\t\n\t\t\tprint('')\n\t\t\tprint('demo results for stage1--------------')\n\t\t\tprint(sample_id)\n\t\t\tplt.imshow(overlay);plt.show()\n\t\t\tplt.imshow(gvfiltered);plt.show()\n\t\t\tplt.imshow(ghfiltered);plt.show()\n\t\t\tplt.imshow(rectified);plt.show()\n             \n\tprint(f'\\n[GPU{gpu_id}] Stage1 completed. Failed: {len(local_fail_id)}')\n\t\n\t# Save FAIL_ID\n\tif fail_id_file:\n\t\twith open(fail_id_file, 'wb') as f:\n\t\t\tpickle.dump(local_fail_id, f)\n\t\n\treturn local_fail_id\n\n\ndef run_stage1_parallel(prev_fail_ids=None):\n\t\"\"\"Run Stage1 in parallel on 2 GPUs\"\"\"\n\tprint('*** STARTING STAGE1 (2GPU PARALLEL) ***')\n\t\n\t# Split valid_id into 2 parts\n\tmid_idx = len(valid_id) // 2\n\tids_gpu0 = valid_id[:mid_idx]\n\tids_gpu1 = valid_id[mid_idx:]\n\t\n\tprint(f'GPU0: {len(ids_gpu0)} samples')\n\tprint(f'GPU1: {len(ids_gpu1)} samples')\n\t\n\t# FAIL_ID file paths\n\tfail_file_0 = f'{OUT_DIR}/fail_stage1_gpu0.pkl'\n\tfail_file_1 = f'{OUT_DIR}/fail_stage1_gpu1.pkl'\n\t\n\t# Run in parallel with 2 processes\n\tp0 = mp.Process(target=run_stage1, args=(0, ids_gpu0, prev_fail_ids, fail_file_0))\n\tp1 = mp.Process(target=run_stage1, args=(1, ids_gpu1, prev_fail_ids, fail_file_1))\n\t\n\tp0.start()\n\tp1.start()\n\t\n\tp0.join()\n\tp1.join()\n\t\n\t# Merge FAIL_IDs\n\tfail_id = []\n\tif os.path.exists(fail_file_0):\n\t\twith open(fail_file_0, 'rb') as f:\n\t\t\tfail_id.extend(pickle.load(f))\n\tif os.path.exists(fail_file_1):\n\t\twith open(fail_file_1, 'rb') as f:\n\t\t\tfail_id.extend(pickle.load(f))\n\t\n\t# Include Stage0 failed IDs as well\n\tif prev_fail_ids:\n\t\tfail_id.extend(prev_fail_ids)\n\t\n\tprint(f'FAIL_ID (Stage1): {fail_id}')\n\tprint('run_stage1_parallel() ok!!!\\n')\n\t\n\treturn fail_id\n\n\n# Run (using parallel version)\nFAIL_ID_STAGE1 = run_stage1_parallel(prev_fail_ids=FAIL_ID_STAGE0)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-01T08:17:37.115122Z","iopub.execute_input":"2026-02-01T08:17:37.115352Z","iopub.status.idle":"2026-02-01T08:17:42.734891Z","shell.execute_reply.started":"2026-02-01T08:17:37.115330Z","shell.execute_reply":"2026-02-01T08:17:42.733938Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Stage2","metadata":{}},{"cell_type":"code","source":"def 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]  # (H, 1) for broadcasting\n    \n    series = []\n    for j in [0, 1, 2, 3]:\n        p = pixel[j]\n        denom = p.sum(axis=0)  # (W,)\n        y_exp = (p * y_idx).sum(axis=0) / (denom + eps)  # (W,)\n        \n        series.append(y_exp)\n    series = np.stack(series).astype(np.float32)\n\n    if length is not None:\n        if 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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-01T08:17:42.742067Z","iopub.execute_input":"2026-02-01T08:17:42.742824Z","iopub.status.idle":"2026-02-01T08:17:42.760360Z","shell.execute_reply.started":"2026-02-01T08:17:42.742796Z","shell.execute_reply":"2026-02-01T08:17:42.759732Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_whole_model(encoder_name, weight_path, device):\n    stage2_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=lambda storage, loc: storage)\n    print(stage2_model.load_state_dict(state_dict, strict=False))\n    stage2_model.to(device)\n    stage2_model.eval()\n    stage2_model.output_type = ['infer']\n    return stage2_model","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-01T08:17:42.761233Z","iopub.execute_input":"2026-02-01T08:17:42.761449Z","iopub.status.idle":"2026-02-01T08:17:42.780689Z","shell.execute_reply.started":"2026-02-01T08:17:42.761435Z","shell.execute_reply":"2026-02-01T08:17:42.779971Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_lead_model(encoder_name, weight_path, fusion_type, device):\n    stage2_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=lambda storage, loc: storage)\n    print(stage2_model.load_state_dict(state_dict, strict=False))\n    stage2_model.to(device)\n    stage2_model.eval()\n    stage2_model.output_type = ['infer']\n    return stage2_model","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-01T08:17:42.781403Z","iopub.execute_input":"2026-02-01T08:17:42.781569Z","iopub.status.idle":"2026-02-01T08:17:42.797254Z","shell.execute_reply.started":"2026-02-01T08:17:42.781557Z","shell.execute_reply":"2026-02-01T08:17:42.796548Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def 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\n    for i, zmv in enumerate(zero_mv):\n        # Crop within zero_mv +- WINDOW_SIZE range (pad with black for out-of-bounds areas)\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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-01T08:17:42.799153Z","iopub.execute_input":"2026-02-01T08:17:42.799600Z","iopub.status.idle":"2026-02-01T08:17:42.816123Z","shell.execute_reply.started":"2026-02-01T08:17:42.799582Z","shell.execute_reply":"2026-02-01T08:17:42.815608Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"WINDOW_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]\n\nx0, x1 = 0, 5600\ny0, y1 = 0, 1696\nprint(x0, x1)\nprint(y0, y1)\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,t1 = timespan = int(118*x_scale)+add_x, int(2080*x_scale)+add_x","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-01T08:17:42.816782Z","iopub.execute_input":"2026-02-01T08:17:42.816962Z","iopub.status.idle":"2026-02-01T08:17:42.835128Z","shell.execute_reply.started":"2026-02-01T08:17:42.816949Z","shell.execute_reply":"2026-02-01T08:17:42.834603Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"height_after_trimed = y1 - OFFSET\nprint(\"height_after_trimed :\", height_after_trimed)\nens_regions = []\nfor zmv in zero_mv_trimed:\n    trim_upper = int(zmv) - WINDOW_SIZE\n    trim_lower = int(zmv) + WINDOW_SIZE\n\n    # Do not use the top and bottom IGNORE_EDGE pixels of lead prediction for ensemble\n    lead_upper = IGNORE_EDGE\n    lead_lower = -IGNORE_EDGE\n    if trim_lower > height_after_trimed:\n        print(trim_lower)\n        # (black padding area + edge) * -1\n        lead_lower = (trim_lower - height_after_trimed + IGNORE_EDGE) * -1\n        trim_lower = height_after_trimed\n\n    trim_upper += IGNORE_EDGE\n    trim_lower -= IGNORE_EDGE\n    ens_regions.append([trim_upper, trim_lower, lead_upper, lead_lower])\nens_regions","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-01T08:17:42.835783Z","iopub.execute_input":"2026-02-01T08:17:42.835984Z","iopub.status.idle":"2026-02-01T08:17:42.851972Z","shell.execute_reply.started":"2026-02-01T08:17:42.835961Z","shell.execute_reply":"2026-02-01T08:17:42.851287Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# stage2\nprint('*** STARTING STAGE2 ***')\n\nsys.path.append(\"/kaggle/input/physionet-final-submission-models\")\nimport torch.nn as nn\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#os.makedirs(f'{OUT_DIR}/debug', exist_ok=True)\n\ndef run_stage2(gpu_id=0, assigned_ids=None, prev_fail_ids=None, fail_id_file=None):\n    \"\"\"\n    Stage2 processing (GPU parallelization supported)\n    \n    Args:\n        gpu_id: GPU number to use (0 or 1)\n        assigned_ids: List of sample_ids to process (if None, use all valid_id)\n        prev_fail_ids: IDs that failed in Stage0/1 (to be skipped)\n        fail_id_file: File path to save FAIL_ID\n    \"\"\"\n    # GPU configuration\n    device = f'cuda:{gpu_id}'\n    \n    # Determine processing targets\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    local_fail_id = []\n\n    whole_models = [\n        # b7, 22.93\n        get_whole_model(\"tu-timm/tf_efficientnet_b7.ns_jft_in1k\", \"/kaggle/input/physionet-final-submission-models/whole_b7_lb22.93.pth\", device),\n        # v2-l, 22.60\n        get_whole_model(\"tu-timm/tf_efficientnetv2_l.in21k\", \"/kaggle/input/physionet-final-submission-models/whole_v2_l_lb22.60.pth\", device),\n        # v2-m, 22.58\n        #get_whole_model(\"tu-timm/tf_efficientnetv2_m.in21k_ft_in1k\", \"/kaggle/input/physionet-final-submission-models/whole_v2_m_lb22.58.pth\", device),\n        # v2-m, 22.52\n        #get_whole_model(\"tu-timm/tf_efficientnetv2_m.in21k\", \"/kaggle/input/physionet-final-submission-models/whole_v2_m_lb22.52.pth\", device),\n    ]\n    \n    lead_models = [\n        # b6-shared-conv2d, 23.10\n        get_lead_model(\"tu-timm/tf_efficientnet_b6.ns_jft_in1k\", \"/kaggle/input/physionet-final-submission-models/series_b6_shared_conv2d_lb23.10.pth\", \"shared_conv2d\", device),\n        # b6-shared-conv2d, 23.00\n        get_lead_model(\"tu-timm/tf_efficientnet_b6.ns_jft_in1k\", \"/kaggle/input/physionet-final-submission-models/series_b6_shared_conv2d_lb23.00.pth\", \"shared_conv2d\", device),\n        # b4-shared-conv2d, 22.93\n        #get_lead_model(\"tu-timm/tf_efficientnet_b4.ns_jft_in1k\", \"/kaggle/input/physionet-final-submission-models/series_b4_shared_conv2d_lb22.93.pth\", \"shared_conv2d\", device),\n        # v2-l-3d, 22.92\n        get_lead_model(\"tu-timm/tf_efficientnetv2_l.in21k\", \"/kaggle/input/physionet-final-submission-models/series_v2_l_conv3d_lb22.92.pth\", \"conv3d\", device),\n        # v2-m-3d, 22.73\n        #get_lead_model(\"tu-timm/tf_efficientnetv2_m.in21k\", \"/kaggle/input/physionet-final-submission-models/series_v2_m_conv3d_lb22.73.pth\", \"conv3d\", device),\n        # v2-l-2d, 22.85\n        get_lead_model(\"tu-timm/tf_efficientnetv2_l.in21k\", \"/kaggle/input/physionet-final-submission-models/series_v2_l_conv2d_lb22.85.pth\", \"conv2d\", device),\n        # v2-m-2d, 22.76\n        #get_lead_model(\"tu-timm/tf_efficientnetv2_m.in21k\", \"/kaggle/input/physionet-final-submission-models/series_v2_m_conv2d_lb22.76.pth\", \"conv2d\", device),\n\n    ]\n        \n    start_timer = timer()\n    for n, sample_id in enumerate(valid_id):    \n        timestamp = time_to_str(timer() - start_timer, 'sec')\n        print(f'\\r\\t {n:4d} {sample_id}', timestamp, end='', flush=True)\n        if sample_id in prev_fail_ids: continue\n\n        length = read_sampling_length(sample_id)\n    \n        trim_image, lead_images = read_images(f'{OUT_DIR}/rectified/{sample_id}.rect.png')\n\n        \n        pixel_ens = np.zeros((4, trim_image.shape[0], trim_image.shape[1])) * 1.0\n        # Infer whole models\n        \n        batch = {\n            'image': torch.from_numpy(np.ascontiguousarray(trim_image.transpose(2, 0, 1))).unsqueeze(0),\n        }\n        batch_tta = {\n            'image': torch.from_numpy(np.ascontiguousarray(np.fliplr(trim_image).copy().transpose(2, 0, 1))).unsqueeze(0),\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        # Infer lead models\n        lead_images = torch.from_numpy(lead_images.transpose(0, 3, 1, 2)).contiguous() # (4, 3, H, W)\n        batch = {\n            'image': lead_images.unsqueeze(0),\n        }\n        batch_tta = {\n            'image': torch.flip(lead_images, 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                        for i in range(4):\n                            # Lead i\n                            trim_upper, trim_lower, lead_upper, lead_lower = ens_regions[i]\n                            \n                            pixel_ens[i][trim_upper:trim_lower] += pixel[i][lead_upper:lead_lower]\n\n        # Average per pixel\n        # Weight for whole model (entire trimmed image)\n        ens_weight = np.ones((trim_image.shape[0], trim_image.shape[1])) * len(whole_models) * len(tta)\n        # Weight for lead model (lead regions only)\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\n        pixel_ens /= ens_weight\n        \n    \n        #---\n        try:\n        #if 1:\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            #series = filter_series_by_limits_modified(series)\n    \n            # ---\n            #cv2.imwrite(f'{OUT_DIR}/digitalised/{sample_id}.lead.png', cv2.cvtColor(overlay, cv2.COLOR_RGB2BGR))\n            np.save(f'{OUT_DIR}/digitalised/{sample_id}.series.npy', series)\n    \n        except:\n            local_fail_id.append(sample_id)\n    \n        if n<10 and gpu_id==0: # optional: show results\n            print()\n            print(\"check max intensity : \", np.max(pixel_ens))\n            print()\n            \n            overlay = draw_lead_pixel(trim_image, pixel_ens)\n            plt.imshow(overlay); plt.show()\n    \n            t = np.arange(len(series[0]))\n            fig, axes = plt.subplots(4, 1, figsize=(12, 10))\n            for j in range(4):\n                snr=0\n                axes[j].plot(t, series[j], alpha=1.0, color='blue', linewidth=1, label='predict')\n                axes[j].set_title(f'snr {snr:8.3f}')\n                axes[j].legend()\n            plt.show()\n    print('')\n    print(f'\\n[GPU{gpu_id}] Stage2 completed. Failed: {len(local_fail_id)}')\n    \n    # Save FAIL_ID\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    \ndef run_stage2_parallel(prev_fail_ids=None):\n\t\"\"\"Run Stage2 in parallel on 2 GPUs\"\"\"\n\tprint('*** STARTING STAGE2 (2GPU PARALLEL) ***')\n\t\n\t# Split valid_id into 2 parts\n\tmid_idx = len(valid_id) // 2\n\tids_gpu0 = valid_id[:mid_idx]\n\tids_gpu1 = valid_id[mid_idx:]\n\t\n\tprint(f'GPU0: {len(ids_gpu0)} samples')\n\tprint(f'GPU1: {len(ids_gpu1)} samples')\n\t\n\t# FAIL_ID file paths\n\tfail_file_0 = f'{OUT_DIR}/fail_stage2_gpu0.pkl'\n\tfail_file_1 = f'{OUT_DIR}/fail_stage2_gpu1.pkl'\n\t\n\t# Run in parallel with 2 processes\n\tp0 = mp.Process(target=run_stage2, args=(0, ids_gpu0, prev_fail_ids, fail_file_0))\n\tp1 = mp.Process(target=run_stage2, args=(1, ids_gpu1, prev_fail_ids, fail_file_1))\n\t\n\tp0.start()\n\tp1.start()\n\t\n\tp0.join()\n\tp1.join()\n\t\n\t# Merge FAIL_IDs\n\tfail_id = []\n\tif os.path.exists(fail_file_0):\n\t\twith open(fail_file_0, 'rb') as f:\n\t\t\tfail_id.extend(pickle.load(f))\n\tif os.path.exists(fail_file_1):\n\t\twith open(fail_file_1, 'rb') as f:\n\t\t\tfail_id.extend(pickle.load(f))\n\t\n\t# Include Stage0/1 failed IDs as well\n\tif prev_fail_ids:\n\t\tfail_id.extend(prev_fail_ids)\n\t\n\tprint(f'FAIL_ID (Stage2): {fail_id}')\n\tprint('run_stage2_parallel() ok!!!\\n')\n\t\n\treturn fail_id\n\n\n# Run (using parallel version)\nFAIL_ID_STAGE2 = run_stage2_parallel(prev_fail_ids=FAIL_ID_STAGE1)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-01T08:17:42.852780Z","iopub.execute_input":"2026-02-01T08:17:42.853045Z","iopub.status.idle":"2026-02-01T08:18:27.333876Z","shell.execute_reply.started":"2026-02-01T08:17:42.853030Z","shell.execute_reply":"2026-02-01T08:18:27.332996Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Submission","metadata":{}},{"cell_type":"code","source":"#make sbmission csv\ndef make_submission():\n\tprint('===========================================')\n\tprint('making submission csv ...')\n\n\tsubmit_df=[]\n\tgb = valid_df.groupby('id')\n\tfor i,(sample_id, df) in enumerate(gb): \n\t\ttry:\n\t\t\tseries = np.load(f'{OUT_DIR}/digitalised/{sample_id}.series.npy')\n\n\t\t\tseries_by_lead={}\n\t\t\tfor l in range(3):\n\t\t\t\tlead = [\n\t\t\t\t\t['I',   'aVR', 'V1', 'V4'],\n\t\t\t\t\t['II',  'aVL', 'V2', 'V5'],\n\t\t\t\t\t['III', 'aVF', 'V3', 'V6'],\n\t\t\t\t][l]\n\n\t\t\t\tlength=[\n\t\t\t\t\tdf[df['lead']==lead[j]].iloc[0].number_of_rows\n\t\t\t\t\tfor j in range(4)\n\t\t\t\t]\n\t\t\t\tif lead[0]=='II':\n\t\t\t\t\tlength[0] = length[0]-sum(length[1:])\n\n\t\t\t\tindex = np.cumsum(length)[:-1]\n\t\t\t\tsplit = np.split(series[l], index)\n\t\t\t\tfor (k, s) in zip(lead, split):\n\t\t\t\t\tseries_by_lead[k] = s\n\t\t\tseries_by_lead['II'] = series[3]\n    \n\t\texcept: \n\t\t\tseries_by_lead = {}\n\t\t\tfor j,d in df.iterrows():\n\t\t\t\tseries_by_lead[d.lead] = np.zeros(d.number_of_rows)\n\n\t\tfor j,d in df.iterrows():\n\n\t\t\tseries_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\t\t\tassert(len(series_by_lead[d.lead])==d.number_of_rows) \n\t\t\tprint(f'\\r\\t {i} {sample_id} : {d.lead}', end='', flush=True)\n\n\t\t\trow_id = [\n\t\t\t\tf'{sample_id}_{i}_{d.lead}' for i in range(d.number_of_rows)\n\t\t\t]\n\t\t\tthis_df = pd.DataFrame({\n\t\t\t\t'id':row_id,\n\t\t\t\t'value': series_by_lead[d.lead].astype(np.float32),\n\t\t\t})\n\t\t\tsubmit_df.append(this_df)\n\n\tprint('')\n\tsubmit_df = pd.concat(submit_df, axis=0, ignore_index=True, sort=False, copy=False)\n\tprint(submit_df)\n\tsubmit_df.to_csv('submission.csv',index=False)\n\nmake_submission()\nprint('make_submission() ok!!!\\n')\nshutil.rmtree(OUT_DIR)\n!ls\n#!rm -rf {OUT_DIR}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-01T08:18:27.335085Z","iopub.execute_input":"2026-02-01T08:18:27.335381Z","iopub.status.idle":"2026-02-01T08:18:27.948514Z","shell.execute_reply.started":"2026-02-01T08:18:27.335358Z","shell.execute_reply":"2026-02-01T08:18:27.947569Z"}},"outputs":[],"execution_count":null}]}