{"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":[{"sourceType":"competition","sourceId":97984,"databundleVersionId":14096757},{"sourceType":"datasetVersion","sourceId":13746387,"datasetId":8747012,"databundleVersionId":14496454},{"sourceType":"datasetVersion","sourceId":13816899,"datasetId":8620533,"databundleVersionId":14574490},{"sourceType":"modelInstanceVersion","sourceId":677607,"databundleVersionId":14853021,"modelInstanceId":513841},{"sourceType":"kernelVersion","sourceId":271051632},{"sourceType":"kernelVersion","sourceId":290179953}],"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# =========================================================================================\n# 1. INSTALLATION & SETUP\n# =========================================================================================\nimport os\nos.system(\"pip uninstall -y tensorflow\")\nos.system(\"uv pip install --no-deps --system --no-index --find-links='/kaggle/input/hengck23-submit-physionet/hengck23-submit-physionet/setup' 'connected-components-3d'\")\n\nimport sys\nsys.path.append('/kaggle/input/hengck23-submit-physionet/hengck23-submit-physionet')\n\nimport gc\nimport cv2\nimport torch\nimport torch.nn as nn\nimport traceback\nimport numpy as np\nimport pandas as pd\nimport timm \nimport torchvision.transforms as T\n\nfrom tqdm.auto import tqdm\nfrom pathlib import Path\nfrom shutil import copyfile\nfrom scipy.signal import butter, filtfilt, medfilt, resample\n\n# =========================================================================================\n# 2. CONFIGURATION & GLOBALS\n# =========================================================================================\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()\n\nglobal_dict = {\n    \"stage0_dir\": \"/kaggle/working/stage0\",\n    \"stage1_dir\": \"/kaggle/working/stage1\",\n    \"stage2_dir\": \"/kaggle/working/stage2\",\n}\n\nfor d in global_dict.values():\n    Path(d).mkdir(exist_ok=True, parents=True)\n\n# Constants for Signal Extraction (Tuned for the 1696x4352 resolution)\nIMAGE_H, IMAGE_W = 1696, 4352\nZERO_MV = [703.5, 987.5, 1271.5, 1531.5] # Y-coordinates of the 0mV line for the 4 rows\nMV_TO_PIXEL = 78.66                  # Scale factor\nT0, T1 = 235, 4161                   # Crop margins for time\nFS_EFFECTIVE = IMAGE_W / 10.0        # Approx sampling rate in pixel domain\n\n# =========================================================================================\n# 3. HELPER FUNCTIONS\n# =========================================================================================\n\ndef change_color(image_rgb):\n    \"\"\"Enhance image contrast/color for Stage 0\"\"\"\n    hsv = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2HSV)\n    h, s, v = cv2.split(hsv)\n    v_denoised = cv2.fastNlMeansDenoising(v, h=5.5)\n    std = np.std(v_denoised)\n    clip_limit = max(1.0, min(3.5, 2.0 + std / 25))\n    clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=(8, 8))\n    v_enhanced = clahe.apply(v_denoised)\n    hsv_enhanced = cv2.merge([h, s, v_enhanced])\n    return cv2.cvtColor(hsv_enhanced, cv2.COLOR_HSV2RGB)\n\ndef dw(series_dict, alpha=0.4):\n    \"\"\"Einthoven's Law Correction: II = I + III\"\"\"\n    if all(k in series_dict for k in ['I', 'II', 'III']):\n        L1 = series_dict['I']\n        L2 = series_dict['II']\n        L3 = series_dict['III']\n        # Calculate consistency error\n        error = L2 - (L1 + L3)\n        # Distribute error correction\n        series_dict['I']   = L1 + (alpha * error)\n        series_dict['III'] = L3 + (alpha * error)\n        series_dict['II']  = L2 - (alpha * error)\n    return series_dict\n\ndef butter_lowpass_filter(data, cutoff, fs, order=2):\n    \"\"\"Apply Butterworth Lowpass filter to smooth pixel noise\"\"\"\n    nyq = 0.5 * fs\n    normal_cutoff = cutoff / nyq\n    b, a = butter(order, normal_cutoff, btype='low', analog=False)\n    # filtfilt applies filter forward and backward to ensure zero phase shift\n    y = filtfilt(b, a, data)\n    return y\n\ndef extract_signal_subpixel(prob_map, zero_mv_y, scale, tracking_window=40):\n    \"\"\"\n    Extracts signal using Weighted Center of Mass (Sub-pixel) \n    and Viterbi-like tracking to prevent jumping to grid lines.\n    \"\"\"\n    h, w = prob_map.shape\n    signal = np.zeros(w)\n    \n    # 1. Find Anchor: Column with highest confidence peak\n    col_maxs = prob_map.max(axis=0)\n    start_col = np.argmax(col_maxs)\n    current_y = np.argmax(prob_map[:, start_col]) # Integer start\n    \n    # 2. Forward Pass\n    for x in range(start_col, w):\n        y_min = max(0, int(current_y) - tracking_window)\n        y_max = min(h, int(current_y) + tracking_window)\n        \n        slice_col = prob_map[y_min:y_max, x]\n        mass = np.sum(slice_col)\n        \n        if mass < 1e-5:\n            signal[x] = current_y # Hold value if uncertain\n        else:\n            # Sub-pixel calculation: Sum(val * idx) / Sum(val)\n            local_idxs = np.arange(y_min, y_max)\n            center = np.sum(slice_col * local_idxs) / mass\n            signal[x] = center\n            current_y = center\n\n    # 3. Backward Pass\n    current_y = signal[start_col]\n    for x in range(start_col - 1, -1, -1):\n        y_min = max(0, int(current_y) - tracking_window)\n        y_max = min(h, int(current_y) + tracking_window)\n        \n        slice_col = prob_map[y_min:y_max, x]\n        mass = np.sum(slice_col)\n        \n        if mass < 1e-5:\n            signal[x] = current_y\n        else:\n            local_idxs = np.arange(y_min, y_max)\n            center = np.sum(slice_col * local_idxs) / mass\n            signal[x] = center\n            current_y = center\n            \n    # Convert pixels to Millivolts\n    # (Zero - Signal) because image Y is inverted (0 is top)\n    return (zero_mv_y - signal) / scale\n\ndef series_to_dict_local(series_4row):\n    d = {}\n    names = [['I', 'aVR', 'V1', 'V4'], \n             ['II', 'aVL', 'V2', 'V5'], \n             ['III', 'aVF', 'V3', 'V6']]\n    \n    # Split the first 3 rows into 4 chunks each\n    for i in range(3):\n        splits = np.array_split(series_4row[i], 4)\n        for name, data in zip(names[i], splits):\n            d[name] = data\n    # The 4th row is the long Lead II\n    d['II_Long'] = series_4row[3]\n    return d\n\ndef dict_to_series_local(d, original_shape):\n    new_series = np.zeros(original_shape)\n    new_series[0] = np.concatenate([d['I'], d['aVR'], d['V1'], d['V4']])\n    new_series[1] = np.concatenate([d['II'], d['aVL'], d['V2'], d['V5']])\n    new_series[2] = np.concatenate([d['III'], d['aVF'], d['V3'], d['V6']])\n    new_series[3] = d.get('II_Long', d.get('II', np.zeros(original_shape[1])))\n    return new_series\n\ndef series_dict_final(series):\n    d = {}\n    lead_names = [\n        ['I',   'aVR', 'V1', 'V4'],\n        ['II',  'aVL', 'V2', 'V5'],\n        ['III', 'aVF', 'V3', 'V6'],\n    ]\n    for l in range(3):\n        split = np.array_split(series[l], 4)\n        for (k, s) in zip(lead_names[l], split):\n            d[k] = s\n    d['II'] = series[3] # Use the long lead II for the 'II' prediction\n    return d\n\n# =========================================================================================\n# 4. MODEL LOADING\n# =========================================================================================\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# --- Define Stage 2 Net Explicitly (often customized) ---\nclass Net3(nn.Module):\n    def __init__(self, pretrained=True):\n        super(Net3, self).__init__()\n        encoder_dim = [64, 128, 256, 512]\n        decoder_dim = [128, 64, 32, 16]\n        self.encoder = timm.create_model('resnet34.a3_in1k', pretrained=pretrained, in_chans=3, num_classes=0, global_pool='')\n        self.decoder = MyCoordUnetDecoder(in_channel=encoder_dim[-1], skip_channel=encoder_dim[:-1][::-1] + [0], out_channel=decoder_dim, scale=[2, 2, 2, 2])\n        self.pixel = nn.Conv2d(decoder_dim[-1], 4, 1)\n\n    def forward(self, image):\n        encode = encode_with_resnet(self.encoder, image)\n        last, _ = self.decoder(feature=encode[-1], skip=encode[:-1][::-1] + [None])\n        pixel = self.pixel(last)\n        return pixel\n\n# Load Stage 0\nstage0_net = Stage0Net(pretrained=False)\nstage0_net = load_net(stage0_net, '/kaggle/input/hengck23-submit-physionet/hengck23-submit-physionet/weight/stage0-last.checkpoint.pth')\nstage0_net.to(\"cuda:0\").eval()\n\n# Load Stage 1\nstage1_net = Stage1Net(pretrained=False)\nstage1_net = load_net(stage1_net, '/kaggle/input/hengck23-submit-physionet/hengck23-submit-physionet/weight/stage1-last.checkpoint.pth')\nstage1_net.to(\"cuda:0\").eval()\n\n# Load Stage 2\nstage2_net = Net3(pretrained=False).to(\"cuda:0\")\nstage2_net.load_state_dict(torch.load(\"/kaggle/input/physio-seg-public/pytorch/net3_009_4200/1/iter_0004200.pt\"))\nstage2_net.eval()\n\nprint(\"Models loaded successfully.\")\n\n# =========================================================================================\n# 5. INFERENCE LOOP\n# =========================================================================================\n\nresize = T.Resize((IMAGE_H, IMAGE_W), interpolation=T.InterpolationMode.BILINEAR)\nstage2_dir = Path(global_dict[\"stage2_dir\"])\n\nfor n, sample_id in enumerate(tqdm(test_id)):\n    \n    # ---------------- STAGE 0: Rotation & Homography ----------------\n    try:\n        path = test_dir / f'{sample_id}.png'\n        output_path_s0 = Path(global_dict[\"stage0_dir\"]) / f'{sample_id}.png'\n        \n        image_original = cv2.imread(str(path), cv2.IMREAD_COLOR)\n        image_original = cv2.cvtColor(image_original, cv2.COLOR_BGR2RGB)\n        image_for_model = change_color(image_original)\n        \n        batch = image_to_batch(image_for_model)\n        with torch.no_grad(), torch.amp.autocast('cuda', dtype=torch.float32):\n            output = stage0_net(batch)\n            \n        rotated, keypoint = output_to_predict(image_original, batch, output)\n        normalised, _, _ = normalise_by_homography(rotated, keypoint)\n        \n        # Save for Stage 1\n        cv2.imwrite(str(output_path_s0), cv2.cvtColor(normalised, cv2.COLOR_RGB2BGR))\n        \n    except Exception:\n        # Fallback: copy original\n        copyfile(path, output_path_s0)\n\n    # ---------------- STAGE 1: Rectification ----------------\n    try:\n        path_s1 = output_path_s0\n        output_path_s1 = Path(global_dict[\"stage1_dir\"]) / f'{sample_id}.png'\n        \n        image = cv2.imread(str(path_s1), cv2.IMREAD_COLOR) # BGR\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        batch = {'image': torch.from_numpy(np.ascontiguousarray(image.transpose(2, 0, 1))).unsqueeze(0)}\n        \n        with torch.no_grad(), torch.amp.autocast('cuda', dtype=torch.float32):\n            output = stage1_net(batch)\n        \n        gridpoint_xy, _ = output_to_predict(image, batch, output)\n        rectified = rectify_image(image, gridpoint_xy)\n        cv2.imwrite(str(output_path_s1), cv2.cvtColor(rectified, cv2.COLOR_RGB2BGR))\n        \n    except Exception:\n        copyfile(path_s1, output_path_s1)\n\n    # ---------------- STAGE 2: Signal Extraction (Optimized) ----------------\n    try:\n        path_s2 = output_path_s1\n        output_path_npy = stage2_dir / f'{sample_id}.npy'\n        \n        image = cv2.imread(str(path_s2), cv2.IMREAD_COLOR)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        \n        # Crop based on training config\n        x0, x1 = 0, 2176\n        y0, y1 = 0, 1696\n        image = image[y0:y1, x0:x1] \n        \n        # Normalize & Resize\n        img_tensor = torch.from_numpy(np.ascontiguousarray(image.transpose(2, 0, 1))).unsqueeze(0).float() / 255.0\n        batch = resize(img_tensor).to(\"cuda:0\")\n        \n        with torch.no_grad(), torch.amp.autocast('cuda', dtype=torch.float16): # fp16 for speed\n            output = stage2_net(batch)\n            probs = torch.sigmoid(output).float().cpu().numpy()[0]\n        \n        # Time crop\n        probs = probs[..., T0:T1]\n        \n        series_extracted = []\n        for i in range(4): # 4 rows\n            prob_map = probs[i]\n            \n            # SUB-PIXEL EXTRACTION + TRACKING\n            raw_signal = extract_signal_subpixel(prob_map, ZERO_MV[i], MV_TO_PIXEL)\n            \n            # FILTERING: Median (spikes) + Butterworth (noise)\n            sig_med = medfilt(raw_signal, kernel_size=3)\n            sig_final = butter_lowpass_filter(sig_med, cutoff=45, fs=FS_EFFECTIVE, order=2)\n            \n            series_extracted.append(sig_final)\n            \n        series = np.array(series_extracted)\n        \n        # Apply Einthoven's Law (II = I + III)\n        s_dict = series_to_dict_local(series)\n        s_dict = dw(s_dict)\n        series = dict_to_series_local(s_dict, series.shape)\n        \n        np.save(output_path_npy, series)\n        \n    except Exception:\n        traceback.print_exc()\n        # Fallback zeros\n        length = test[(test['id']==sample_id) & (test['lead']=='II')].iloc[0].number_of_rows\n        np.save(output_path_npy, np.zeros((4, length)))\n        \n    # GC to prevent OOM\n    if n % 50 == 0:\n        gc.collect()\n\n# =========================================================================================\n# 6. SUBMISSION GENERATION\n# =========================================================================================\nprint(\"Generating submission file...\")\nres = []\ngb = test.groupby('id')\n\nfor sample_id, df in tqdm(gb):\n    try:\n        npy_path = stage2_dir / f'{sample_id}.npy'\n        if not npy_path.exists():\n            series = np.zeros((4, 4000)) # Fallback shape\n        else:\n            series = np.load(npy_path)\n            \n        d_series = series_dict_final(series)\n        \n        for _, row in df.iterrows():\n            # Get extracted signal for specific lead\n            s = d_series.get(row.lead, np.zeros(row.number_of_rows))\n            \n            # Resample to Match Target Frequency/Length\n            target_len = row.number_of_rows\n            if len(s) != target_len:\n                x_old = np.linspace(0, 1, len(s))\n                x_new = np.linspace(0, 1, target_len)\n                s = np.interp(x_new, x_old, s)\n            \n            # Create Row IDs\n            row_ids = [f'{sample_id}_{x}_{row.lead}' for x in range(target_len)]\n            \n            # Append to list (more memory efficient than concat inside loop)\n            res.append(pd.DataFrame({'id': row_ids, 'value': s}))\n            \n    except Exception as e:\n        print(f\"Error packing {sample_id}: {e}\")\n\nsubmission = pd.concat(res, axis=0, ignore_index=True)\nsubmission.to_csv('submission.csv', index=False)\nsubmission.to_parquet('submission.parquet', index=False) # Parquet is safer/faster\n\nprint(\"Submission saved. Head:\")\nprint(submission.head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T18:59:19.580382Z","iopub.execute_input":"2026-01-13T18:59:19.580625Z","iopub.status.idle":"2026-01-13T19:00:33.260515Z","shell.execute_reply.started":"2026-01-13T18:59:19.580605Z","shell.execute_reply":"2026-01-13T19:00:33.259824Z"}},"outputs":[],"execution_count":null}]}