{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":97984,"databundleVersionId":14096757,"isSourceIdPinned":false},{"sourceType":"datasetVersion","sourceId":13731160,"datasetId":8733970,"databundleVersionId":14479231}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install connected-components-3d -q\n\nimport os, sys, shutil\nimport cv2\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom scipy import signal as scipy_signal\nfrom sklearn.model_selection import KFold\n\nsys.path.append('/kaggle/input/datasets/hengck23/hengck23-demo-submit-physionet')\nimport stage0_common\nfrom stage0_model import Net as Stage0Net\nimport stage1_common\nfrom stage1_model import Net as Stage1Net\n\nKAGGLE_DIR = '/kaggle/input/competitions/physionet-ecg-image-digitization'\nWEIGHT_DIR = '/kaggle/input/datasets/hengck23/hengck23-demo-submit-physionet/weight'\nOUT_DIR = '/kaggle/working/data'\nRECT_DIR = f'{OUT_DIR}/rectified'\nMASK_DIR = f'{OUT_DIR}/masks'\nECG_DIR = f'{OUT_DIR}/ecg_csv'\nos.makedirs(RECT_DIR, exist_ok=True)\nos.makedirs(MASK_DIR, exist_ok=True)\nos.makedirs(ECG_DIR, exist_ok=True)\n\n# SEGMENTS = ['0001']\n# SEGMENTS = ['0004', '0009', '0010']\n# SEGMENTS = ['0004', '0009']\n# SEGMENTS = ['0010', '0011']\nSEGMENTS = ['0012']\n\ndevice = 'cuda:0'\ntrain_df = pd.read_csv(f'{KAGGLE_DIR}/train.csv')\ntrain_df['id'] = train_df['id'].astype(str)\ntrain_ids = train_df['id'].tolist()\nprint(f'Records: {len(train_ids)}')\nprint(f'Segments: {SEGMENTS}')\nprint(f'Expected images: {len(train_ids) * len(SEGMENTS)}')\n\n# === Step 1: Rectify images ===\nprint('\\n=== Step 1: Rectifying images ===')\ns0 = Stage0Net(pretrained=False)\ns0 = stage0_common.load_net(s0, f'{WEIGHT_DIR}/stage0-last.checkpoint.pth')\ns0.to(device)\ns1 = Stage1Net(pretrained=False)\ns1 = stage1_common.load_net(s1, f'{WEIGHT_DIR}/stage1-last.checkpoint.pth')\ns1.to(device)\n\nfor seg in SEGMENTS:\n    print(f'\\n  --- Segment {seg} ---')\n    fail_ids = []\n    for i, sid in enumerate(train_ids):\n        out_path = f'{RECT_DIR}/{sid}-{seg}.rect.png'\n        if os.path.exists(out_path):\n            continue\n        src_path = f'{KAGGLE_DIR}/train/{sid}/{sid}-{seg}.png'\n        if not os.path.exists(src_path):\n            fail_ids.append(sid)\n            continue\n        try:\n            image = cv2.imread(src_path, cv2.IMREAD_COLOR_RGB)\n            batch = stage0_common.image_to_batch(image)\n            with torch.amp.autocast('cuda', dtype=torch.float16):\n                with torch.no_grad():\n                    out = s0(batch)\n                    rot, kp = stage0_common.output_to_predict(image, batch, out)\n                    norm, kp, homo = stage0_common.normalise_by_homography(rot, kp)\n            batch1 = {'image': torch.from_numpy(\n                np.ascontiguousarray(norm.transpose(2, 0, 1))\n            ).unsqueeze(0)}\n            with torch.amp.autocast('cuda', dtype=torch.float16):\n                with torch.no_grad():\n                    out1 = s1(batch1)\n                    gp, more = stage1_common.output_to_predict(norm, batch1, out1)\n                    rect = stage1_common.rectify_image(norm, gp)\n            cv2.imwrite(out_path, cv2.cvtColor(rect, cv2.COLOR_RGB2BGR))\n        except:\n            fail_ids.append(sid)\n        torch.cuda.empty_cache()\n        if (i + 1) % 100 == 0:\n            print(f'    {i+1}/{len(train_ids)} done, fails: {len(fail_ids)}')\n    print(f'  Segment {seg} done, failed: {len(fail_ids)}')\n\ndel s0, s1\ntorch.cuda.empty_cache()\nprint(f'  Total rectified images: {len(os.listdir(RECT_DIR))}')\n\n# === Step 2: Generate masks ===\nprint('\\n=== Step 2: Generating masks ===')\nMASK_H, MASK_W = 1700, 5600\nT0, T1, SIGNAL_W = 301, 5301, 5000\nZERO_MV = [703.5, 987.5, 1271.5, 1531.5]\nMV_TO_PIXEL = 79.0\nROW_TO_LEADS = [['I','aVR','V1','V4'], ['II','aVL','V2','V5'], ['III','aVF','V3','V6']]\n\ndef create_mask(ecg_csv_path):\n    ecg_df = pd.read_csv(ecg_csv_path)\n    mask = np.zeros((4, MASK_H, MASK_W), dtype=np.float32)\n    for row_idx in range(3):\n        leads = ROW_TO_LEADS[row_idx]\n        zero_y = ZERO_MV[row_idx]\n        row_values = []\n        for lead_name in leads:\n            row_values.append(ecg_df[lead_name].dropna().values.astype(np.float32))\n        if leads[0] == 'II':\n            ii_short_len = len(ecg_df['I'].dropna())\n            row_values[0] = row_values[0][:ii_short_len]\n        all_values = np.concatenate(row_values)\n        if len(all_values) != SIGNAL_W:\n            all_values = scipy_signal.resample(all_values, SIGNAL_W).astype(np.float32)\n        y_positions = zero_y - all_values * MV_TO_PIXEL\n        for col_idx in range(SIGNAL_W):\n            y = y_positions[col_idx]\n            if y < 0 or y >= MASK_H - 1: continue\n            yf = int(np.floor(y)); frac = y - yf; x = T0 + col_idx\n            mask[row_idx, yf, x] = 1.0 - frac\n            mask[row_idx, yf + 1, x] = frac\n    zero_y = ZERO_MV[3]\n    lead_ii = ecg_df['II'].dropna().values.astype(np.float32)\n    if len(lead_ii) != SIGNAL_W:\n        lead_ii = scipy_signal.resample(lead_ii, SIGNAL_W).astype(np.float32)\n    y_positions = zero_y - lead_ii * MV_TO_PIXEL\n    for col_idx in range(SIGNAL_W):\n        y = y_positions[col_idx]\n        if y < 0 or y >= MASK_H - 1: continue\n        yf = int(np.floor(y)); frac = y - yf; x = T0 + col_idx\n        mask[3, yf, x] = 1.0 - frac\n        mask[3, yf + 1, x] = frac\n    return mask\n\ndef save_mask_coo(mask, output_path):\n    C, H, W = mask.shape\n    d = {'shape': np.array([C, H, W])}\n    for i in range(C):\n        ys, xs = np.nonzero(mask[i])\n        d[f'ch{i}_y'] = ys.astype(np.int32)\n        d[f'ch{i}_x'] = xs.astype(np.int32)\n        d[f'ch{i}_v'] = mask[i, ys, xs].astype(np.float32)\n    np.savez_compressed(output_path, **d)\n\nfor i, sid in enumerate(train_ids):\n    out_path = f'{MASK_DIR}/{sid}.mask-coo.npz'\n    if os.path.exists(out_path): continue\n    try:\n        create_mask_result = create_mask(f'{KAGGLE_DIR}/train/{sid}/{sid}.csv')\n        save_mask_coo(create_mask_result, out_path)\n    except Exception as e:\n        print(f'  FAIL {sid}: {e}')\n    if (i + 1) % 200 == 0:\n        print(f'  {i+1}/{len(train_ids)}')\nprint(f'  Done: {len(os.listdir(MASK_DIR))} masks')\n\n# === Step 3: Copy ECG CSVs ===\nprint('\\n=== Step 3: Copying ECG CSVs ===')\nfor sid in train_ids:\n    src = f'{KAGGLE_DIR}/train/{sid}/{sid}.csv'\n    dst = f'{ECG_DIR}/{sid}.csv'\n    if not os.path.exists(dst): shutil.copy2(src, dst)\nprint(f'  Done: {len(os.listdir(ECG_DIR))} CSVs')\n\n# === Step 4: Generate fold CSV ===\nprint('\\n=== Step 4: Generating fold CSV ===')\nkf = KFold(n_splits=5, shuffle=True, random_state=42)\nfold_map = {}\nfor fold_idx, (_, val_idx) in enumerate(kf.split(train_ids)):\n    for idx in val_idx:\n        fold_map[train_ids[idx]] = fold_idx\n\nrows = []\nfor sid in train_ids:\n    sig_len = train_df[train_df['id'] == sid].iloc[0]['sig_len']\n    fold = fold_map[sid]\n    for seg in SEGMENTS:\n        if os.path.exists(f'{RECT_DIR}/{sid}-{seg}.rect.png'):\n            rows.append({\n                'id': sid, 'type_id': seg, 'fold': fold,\n                'sig_len': sig_len, 'is_synthesis': False,\n                'image_path': f'data/rectified/{sid}-{seg}.rect.png',\n            })\n\nfold_csv = pd.DataFrame(rows)\nfold_csv.to_csv(f'{OUT_DIR}/train_fold.csv', index=False)\nprint(f'  Done: {len(fold_csv)} rows')\nfor seg in SEGMENTS:\n    n = len(fold_csv[fold_csv['type_id'] == seg])\n    print(f'    Segment {seg}: {n} rows')\n\n# === Summary ===\nprint(f'\\n=== SUMMARY ===')\nprint(f'Rectified images: {len(os.listdir(RECT_DIR))}')\nprint(f'Masks:            {len(os.listdir(MASK_DIR))}')\nprint(f'ECG CSVs:         {len(os.listdir(ECG_DIR))}')\nrect_size = sum(os.path.getsize(os.path.join(RECT_DIR, f)) for f in os.listdir(RECT_DIR))\nprint(f'Total size:       ~{rect_size/1e9:.1f} GB')\nprint('DONE! Save this output as dataset: ecg-training-data-batch5')","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null}]}