{"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,"sourceType":"competition"},{"sourceId":13746387,"sourceType":"datasetVersion","datasetId":8747012},{"sourceId":13816899,"sourceType":"datasetVersion","datasetId":8620533},{"sourceId":677607,"sourceType":"modelInstanceVersion","modelInstanceId":513841,"modelId":528480}],"dockerImageVersionId":31154,"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":"2025-12-22T13:58:49.900984Z","iopub.execute_input":"2025-12-22T13:58:49.901749Z","iopub.status.idle":"2025-12-22T13:59:11.299522Z","shell.execute_reply.started":"2025-12-22T13:58:49.901713Z","shell.execute_reply":"2025-12-22T13:59:11.298763Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import sys\nsys.path.append('/kaggle/input/hengck23-submit-physionet/hengck23-submit-physionet')\n\nimport os, gc, cv2, torch, psutil, timm\nimport numpy as np\nimport pandas as pd\nimport torchvision.transforms as T\nfrom tqdm.auto import tqdm\nfrom scipy.signal import savgol_filter\nfrom pathlib import Path\n\n# --- CUSTOM MODULES ---\nfrom stage0_model import Net as Stage0Net\nfrom stage0_common import *\nfrom stage1_model import Net as Stage1Net\nfrom stage1_common import *\nfrom stage2_model import *\nfrom stage2_common import *\n\n# --- HYPERPARAMETERS & PRECISION CONSTANTS ---\ndevice = torch.device(\"cuda:0\")\nmv_to_pixel = 78.74   # Standard 200 DPI scale\nt0, t1 = 235, 4161    # Verified Golden Window\nx0, x1, y0, y1 = 0, 2176, 0, 1696\nref_zeros = [703.5, 987.5, 1271.5, 1531.5]\n\n# --- PRECISION UTILITIES ---\n\ndef extract_soft_argmax(heatmap, centers, window=25):\n    \"\"\"\n    SUB-PIXEL EXTRACTION: Calculates vertical center of mass.\n    This eliminates 1-pixel rounding errors that cap SNR at 18dB.\n    \"\"\"\n    C, H, W = heatmap.shape\n    series = np.zeros((C, W), dtype=np.float32)\n    y_coords = np.arange(H, dtype=np.float32)\n    for c in range(C):\n        cy = int(centers[c])\n        y_s, y_e = max(0, cy - window), min(H, cy + window)\n        # Vertical probability slice around the expected lead\n        slice_p = heatmap[c, y_s:y_e, :] \n        weights = y_coords[y_s:y_e].reshape(-1, 1)\n        # Soft-Argmax: Sum(Y * Prob) / Sum(Prob)\n        denom = np.sum(slice_p, axis=0) + 1e-9\n        series[c] = np.sum(slice_p * weights, axis=0) / denom\n    return series\n\n# --- MODEL INITIALIZATION ---\n\nclass Net3(nn.Module):\n    def __init__(self):\n        super(Net3, self).__init__()\n        self.encoder = timm.create_model('resnet34.a3_in1k', pretrained=False, in_chans=3, num_classes=0, global_pool='')\n        self.decoder = MyCoordUnetDecoder(in_channel=512, skip_channel=[256, 128, 64, 0], out_channel=[128, 64, 32, 16], scale=[2, 2, 2, 2])\n        self.pixel = nn.Conv2d(16, 4, 1)\n    def forward(self, x):\n        e = encode_with_resnet(self.encoder, x)\n        last, _ = self.decoder(feature=e[-1], skip=e[:-1][::-1] + [None])\n        return self.pixel(last)\n\ns0_net = Stage0Net(pretrained=False).to(device)\ns0_net = load_net(s0_net, '/kaggle/input/hengck23-submit-physionet/hengck23-submit-physionet/weight/stage0-last.checkpoint.pth').eval()\ns1_net = Stage1Net(pretrained=False).to(device)\ns1_net = load_net(s1_net, '/kaggle/input/hengck23-submit-physionet/hengck23-submit-physionet/weight/stage1-last.checkpoint.pth').eval()\ns2_net = Net3().to(device)\ns2_net.load_state_dict(torch.load(\"/kaggle/input/physio-seg-public/pytorch/net3_009_4200/1/iter_0004200.pt\"))\ns2_net.eval()\n\nresize_s2 = T.Resize((1696, 4352), interpolation=T.InterpolationMode.BILINEAR)\n\n# --- INCREMENTAL PROCESSING ---\ntest_dir = Path(\"/kaggle/input/physionet-ecg-image-digitization/test\")\ntest = pd.read_csv(\"/kaggle/input/physionet-ecg-image-digitization/test.csv\")\ntest['id'] = test['id'].astype(str)\ntest_id = test['id'].unique().tolist()\ntest_gb = test.groupby('id')\n\nSUBMISSION_FILE = 'submission.csv'\nheader = True\n\nfor i, sample_id in enumerate(tqdm(test_id, desc=\"Processing\")):\n    current_batch = []\n    try:\n        # 1. In-Memory 3-Stage Pipeline\n        img = cv2.cvtColor(cv2.imread(str(test_dir/f'{sample_id}.png')), cv2.COLOR_BGR2RGB)\n        with torch.no_grad(), torch.amp.autocast('cuda'):\n            # Stage 0 & 1: Alignment & Rectification\n            b0 = image_to_batch(change_color(img)).to(device)\n            out0 = s0_net(b0)\n            rot, kp = output_to_predict(img, b0, out0)\n            img_norm, _, _ = normalise_by_homography(rot, kp)\n            \n            b1 = {'image': torch.from_numpy(np.ascontiguousarray(img_norm.transpose(2,0,1))).unsqueeze(0).to(device)}\n            out1 = s1_net(b1)\n            grid, _ = output_to_predict(img_norm, b1, out1)\n            img_rect = rectify_image(img_norm, grid)\n            \n            # Stage 2: Segmentation Heatmap\n            df_sub = test_gb.get_group(sample_id)\n            target_len = df_sub[df_sub['lead']=='II'].iloc[0].number_of_rows\n            crop = (img_rect[y0:y1, x0:x1] / 255.0).transpose(2, 0, 1)\n            batch2 = resize_s2(torch.from_numpy(np.ascontiguousarray(crop)).unsqueeze(0)).float().to(device)\n            out2 = s2_net(batch2)\n            # SIGMOID is required for soft-argmax probability distribution\n            probs = torch.sigmoid(out2).cpu().numpy()[0]\n\n        # 2. Precision Sub-Pixel Extraction\n        s_raw = extract_soft_argmax(probs[:, :, t0:t1], ref_zeros)\n        \n        final_series = np.zeros((4, target_len), dtype=np.float32)\n        for row_i in range(4):\n            # Safe-Edge Dynamic Baseline (Match the paper DC offset)\n            quiet_zone = np.concatenate([s_raw[row_i, :50], s_raw[row_i, -50:]])\n            actual_zero = np.median(quiet_zone)\n            \n            lead_mv = (actual_zero - s_raw[row_i]) / mv_to_pixel\n            # Window 3: Minimum cleaning to preserve R-peak fidelity\n            lead_mv = savgol_filter(lead_mv, window_length=3, polyorder=2).astype(np.float32)\n            \n            if len(lead_mv) != target_len:\n                final_series[row_i] = np.interp(np.linspace(0, 1, target_len), \n                                                np.linspace(0, 1, len(lead_mv)), \n                                                lead_mv).astype(np.float32)\n            else:\n                final_series[row_i] = lead_mv\n\n        # 3. Lead Mapping & CSV Writing\n        d_series = {'II': final_series[3]}\n        names = [['I','aVR','V1','V4'], ['II','aVL','V2','V5'], ['III','aVF','V3','V6']]\n        for r in range(3):\n            splits = np.array_split(final_series[r], 4)\n            for k, s in zip(names[r], splits): d_series[k] = s\n\n        for _, row in df_sub.iterrows():\n            val = d_series.get(row.lead, np.zeros(row.number_of_rows, dtype=np.float32))\n            if len(val) != row.number_of_rows:\n                val = np.interp(np.linspace(0, 1, row.number_of_rows), \n                                np.linspace(0, 1, len(val)), val).astype(np.float32)\n            current_batch.append(pd.DataFrame({'id': [f'{sample_id}_{x}_{row.lead}' for x in range(row.number_of_rows)], 'value': val}))\n\n    except Exception:\n        df_sub = test_gb.get_group(sample_id)\n        for _, row in df_sub.iterrows():\n            current_batch.append(pd.DataFrame({'id': [f'{sample_id}_{x}_{row.lead}' for x in range(row.number_of_rows)], 'value': np.zeros(row.number_of_rows, dtype=np.float32)}))\n\n    # --- MEMORY SAFETY: Append to Disk & Flush RAM ---\n    if current_batch:\n        batch_df = pd.concat(current_batch, axis=0, ignore_index=True)\n        batch_df.to_csv(SUBMISSION_FILE, mode='a', index=False, header=header)\n        header = False \n    \n    del current_batch\n    if (i + 1) % 50 == 0:\n        gc.collect()\n        torch.cuda.empty_cache()\n\nprint(\"Fidelity-Elite Submission Complete.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-22T13:59:11.300953Z","iopub.execute_input":"2025-12-22T13:59:11.301201Z","iopub.status.idle":"2025-12-22T14:01:00.648241Z","shell.execute_reply.started":"2025-12-22T13:59:11.301177Z","shell.execute_reply":"2025-12-22T14:01:00.647614Z"}},"outputs":[],"execution_count":null}]}