{"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":"gpu","dataSources":[{"sourceId":97984,"databundleVersionId":14096757,"sourceType":"competition"},{"sourceId":13746387,"sourceType":"datasetVersion","datasetId":8747012},{"sourceId":13816899,"sourceType":"datasetVersion","datasetId":8620533},{"sourceId":272137252,"sourceType":"kernelVersion"},{"sourceId":677607,"sourceType":"modelInstanceVersion","modelInstanceId":513841,"modelId":528480}],"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Imports","metadata":{}},{"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":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T22:44:27.981357Z","iopub.execute_input":"2026-01-21T22:44:27.981677Z","iopub.status.idle":"2026-01-21T22:44:29.963467Z","shell.execute_reply.started":"2026-01-21T22:44:27.981651Z","shell.execute_reply":"2026-01-21T22:44:29.962690Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile cc3d.py\nimport numpy as np\nfrom scipy import ndimage as ndi\n\ndef _structure(ndim: int, connectivity: int):\n    if ndim == 2:\n        # 4 or 8\n        return ndi.generate_binary_structure(2, 1 if connectivity in (4,) else 2)\n    if ndim == 3:\n        # 6 / 18 / 26\n        if connectivity in (6,):\n            return ndi.generate_binary_structure(3, 1)\n        if connectivity in (18,):\n            return ndi.generate_binary_structure(3, 2)\n        if connectivity in (26,):\n            return np.ones((3,3,3), dtype=bool)\n        return np.ones((3,3,3), dtype=bool)\n    return ndi.generate_binary_structure(ndim, 1)\n\ndef connected_components(img, connectivity=26, return_N=False, out_dtype=np.int32):\n    img = np.asarray(img)\n    fg = (img != 0)\n    st = _structure(img.ndim, connectivity)\n    lab, N = ndi.label(fg, structure=st)\n    lab = lab.astype(out_dtype, copy=False)\n    return (lab, int(N)) if return_N else lab\n\ndef largest_k(img, k=1, connectivity=26, return_N=False, out_dtype=np.int32):\n    lab, N = connected_components(img, connectivity=connectivity, return_N=True, out_dtype=out_dtype)\n    if N == 0:\n        out = np.zeros_like(lab, dtype=out_dtype)\n        return (out, 0) if return_N else out\n\n    counts = np.bincount(lab.ravel())\n    counts[0] = 0\n    keep = np.argsort(counts)[::-1][: int(k)]\n    keep = keep[keep != 0]\n\n    out = np.zeros_like(lab, dtype=out_dtype)\n    for new_id, old_id in enumerate(keep, start=1):\n        out[lab == old_id] = new_id\n\n    return (out, int(len(keep))) if return_N else out","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T22:44:29.965280Z","iopub.execute_input":"2026-01-21T22:44:29.965520Z","iopub.status.idle":"2026-01-21T22:44:29.971992Z","shell.execute_reply.started":"2026-01-21T22:44:29.965493Z","shell.execute_reply":"2026-01-21T22:44:29.971309Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os, sys\nsys.path.insert(0, os.getcwd())  # so ./cc3d.py is found first\n\nimport numpy as np\nimport cc3d\n\n# Patch for Kaggle envs where cc3d.statistics is missing\nif not hasattr(cc3d, \"statistics\"):\n    def _cc3d_statistics(labels):\n        labels = np.asarray(labels)\n        max_lab = int(labels.max()) if labels.size else 0\n\n        # voxel counts, include background index 0\n        counts = np.bincount(labels.reshape(-1), minlength=max_lab + 1)\n\n        # centroids, include background index 0\n        centroids = np.full((max_lab + 1, labels.ndim), np.nan, dtype=np.float32)\n        for lab in range(1, max_lab + 1):\n            if counts[lab] == 0:\n                continue\n            coords = np.argwhere(labels == lab)\n            centroids[lab] = coords.mean(axis=0)\n\n        return {\"voxel_counts\": counts, \"centroids\": centroids}\n\n    cc3d.statistics = _cc3d_statistics\n    print(\"✅ Patched cc3d.statistics\")\nelse:\n    print(\"✅ cc3d.statistics exists, no patch needed\")\n\nprint(\"cc3d shim OK:\", cc3d)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T22:44:29.972918Z","iopub.execute_input":"2026-01-21T22:44:29.973182Z","iopub.status.idle":"2026-01-21T22:44:29.995232Z","shell.execute_reply.started":"2026-01-21T22:44:29.973160Z","shell.execute_reply":"2026-01-21T22:44:29.994527Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =========================\n# Imports & Paths\n# =========================\nimport sys\nsys.path.append('/kaggle/input/hengck23-submit-physionet/hengck23-submit-physionet')\n\nimport cv2\nimport gc\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd\nimport random\nimport scipy.signal\nimport scipy.optimize\nimport time\nimport timm\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport traceback\nfrom pathlib import Path\nfrom scipy.signal import medfilt\nfrom shutil import copyfile\nfrom sklearn.linear_model import Ridge\nfrom sklearn.preprocessing import StandardScaler\nfrom sklearn.pipeline import make_pipeline\nfrom tqdm.auto import tqdm\nfrom typing import Tuple","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T22:44:29.996623Z","iopub.execute_input":"2026-01-21T22:44:29.996947Z","iopub.status.idle":"2026-01-21T22:44:30.018624Z","shell.execute_reply.started":"2026-01-21T22:44:29.996925Z","shell.execute_reply":"2026-01-21T22:44:30.018089Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =========================\n# Import pipeline modules\n# =========================\nimport stage0_common as s0c\nimport stage1_common as s1c\nimport stage2_common as s2c\nimport torchvision.transforms as T\nfrom scipy.signal import savgol_filter\nfrom stage0_common import *\nfrom stage0_model import Net as Stage0Net\nfrom stage1_common import *\nfrom stage1_model import Net as Stage1Net\nfrom stage2_model import *\nfrom stage2_common import *","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T22:44:30.019495Z","iopub.execute_input":"2026-01-21T22:44:30.019740Z","iopub.status.idle":"2026-01-21T22:44:30.035163Z","shell.execute_reply.started":"2026-01-21T22:44:30.019695Z","shell.execute_reply":"2026-01-21T22:44:30.034488Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"os.environ[\"PYTORCH_CUDA_ALLOC_CONF\"] = \"expandable_segments:True\"\n\ndef cuda_cleanup():\n    # delete big objects you created\n    for name in [\"model\", \"net\", \"stage0_net\", \"stage1_net\", \"stage2_net\"]:\n        if name in globals():\n            del globals()[name]    \n    gc.collect()\n    torch.cuda.empty_cache()\n    torch.cuda.ipc_collect()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T22:44:30.035933Z","iopub.execute_input":"2026-01-21T22:44:30.036171Z","iopub.status.idle":"2026-01-21T22:44:30.055601Z","shell.execute_reply.started":"2026-01-21T22:44:30.036151Z","shell.execute_reply":"2026-01-21T22:44:30.054898Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Helper Functions","metadata":{}},{"cell_type":"code","source":"# =========================\n# Helper functions\n# =========================\nLEADS_ORDER = [\"I\", \"II\", \"III\", \"aVR\", \"aVL\", \"aVF\", \"V1\", \"V2\", \"V3\", \"V4\", \"V5\", \"V6\"]\n\ndef change_color(image_rgb):\n    \"\"\" Enhances contrast and removes noise on ECG scan for model robustness.\n    \"\"\"\n    hsv = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2HSV)\n    h, s, v = cv2.split(hsv)\n\n    v_denoised = cv2.fastNlMeansDenoising(v, h=5.46)\n\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\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.33):\n    \"\"\" Einthoven correction on SHORT leads only: II_short ≈ I + III\n    \"\"\"\n    if all(k in series_dict for k in ['I', 'II_short', 'III']):\n\n        L1 = series_dict['I']\n        L2 = series_dict['II_short']\n        L3 = series_dict['III']\n\n        # ✅ all same length now\n        error = L2 - (L1 + L3)\n\n        series_dict['I']        = L1 + (alpha * error)\n        series_dict['III']      = L3 + (alpha * error)\n        series_dict['II_short'] = L2 - (alpha * error)\n\n    return series_dict\n\ndef series_dict(series_4row):\n    \"\"\" Converts model output shape (4, length) into dict of 12 leads.\n        Each of first 3 rows is split into 4 equal chunks.\n    \"\"\"\n    series_4row = np.asarray(series_4row)\n\n    # ---- enforce correct shape ----\n    if series_4row.ndim == 3:\n        series_4row = series_4row[0]  # remove batch dim\n    if series_4row.shape[0] != 4 and series_4row.shape[1] == 4:\n        series_4row = series_4row.T  # auto-fix transpose\n\n    if series_4row.shape[0] != 4:\n        raise ValueError(f\"[series_dict] Expected (4, L), got {series_4row.shape}\")\n\n    L = series_4row.shape[1]\n    if L < 1000:\n        print(f\"[WARN] suspicious length: {L}\")\n\n    d = {}\n\n    names = [\n        ['I',   'aVR', 'V1', 'V4'],\n        ['II',  'aVL', 'V2', 'V5'],\n        ['III', 'aVF', 'V3', 'V6'],\n    ]\n\n    for row_idx in range(3):\n        splits = np.array_split(series_4row[row_idx], 4)\n        for lead, arr in zip(names[row_idx], splits):\n            d[lead] = np.asarray(arr, dtype=np.float32)\n\n    d['II'] = np.asarray(series_4row[3], dtype=np.float32)\n    return d\n\ndef dict_to_df_pred(d_series, df_signal, fs, sig_len):\n    pred = {lead: np.full(sig_len, np.nan, dtype=np.float32) for lead in LEADS_ORDER}\n\n    expected_short = int(round(fs * 2.5))\n\n    # --- Lead II is full length ---\n    y_ii = np.asarray(d_series.get(\"II\", np.zeros(sig_len)), dtype=np.float32)\n    if len(y_ii) != sig_len:\n        y_ii = np.interp(np.linspace(0, 1, sig_len), np.linspace(0, 1, len(y_ii)), y_ii)\n    pred[\"II\"] = y_ii\n\n    # --- all short leads: PLACE INTO GT SPAN ---\n    for lead in LEADS_ORDER:\n        if lead == \"II\":\n            continue\n\n        y = d_series.get(lead, None)\n        if y is None:\n            continue\n        y = np.asarray(y, dtype=np.float32)\n\n        # enforce 2.5s length\n        if len(y) != expected_short:\n            y = np.interp(np.linspace(0, 1, expected_short), np.linspace(0, 1, len(y)), y)\n\n        # locate GT span (where label exists)\n        mask = df_signal[lead].notna().values\n        idx = np.flatnonzero(mask)\n        if idx.size == 0:\n            continue\n\n        start, end = idx[0], idx[-1] + 1\n        n_seg = end - start\n\n        y_seg = y[:n_seg]\n        if len(y_seg) < n_seg:\n            y_seg = np.pad(y_seg, (0, n_seg - len(y_seg)), mode=\"edge\")\n\n        pred[lead][start:end] = y_seg\n\n    return pd.DataFrame(pred, index=df_signal.index, columns=LEADS_ORDER)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T22:44:30.056569Z","iopub.execute_input":"2026-01-21T22:44:30.056870Z","iopub.status.idle":"2026-01-21T22:44:30.075696Z","shell.execute_reply.started":"2026-01-21T22:44:30.056829Z","shell.execute_reply":"2026-01-21T22:44:30.075143Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Model Pipeline","metadata":{}},{"cell_type":"code","source":"# =========================\n# Stage2 Net definition (Net3)\n# =========================\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\n        self.encoder = timm.create_model(\n            model_name='resnet34.a3_in1k',\n            pretrained=pretrained,\n            in_chans=3,\n            num_classes=0,\n            global_pool=''\n        )\n\n        self.decoder = MyCoordUnetDecoder(\n            in_channel=encoder_dim[-1],\n            skip_channel=encoder_dim[:-1][::-1] + [0],\n            out_channel=decoder_dim,\n            scale=[2, 2, 2, 2]\n        )\n\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# =========================\n# Pipeline Wrapper\n# =========================\nclass PhysioPipeline:\n    def __init__(self, device=\"cuda:0\"):\n        self.device = device\n\n        self.stage0_net = None\n        self.stage1_net = None\n        self.stage2_net = None\n\n        # stage2 constants\n        self.x0, self.x1 = 0, 2176\n        self.y0, self.y1 = 0, 1696\n        self.zero_mv = [703.5, 987.5, 1271.5, 1531.5]\n        self.mv_to_pixel = 78.8\n        self.t0, self.t1 = 235, 4161\n\n        self.resize = T.Resize((1696, 4352), interpolation=T.InterpolationMode.BILINEAR)\n\n    def load_models(self, stage0_w, stage1_w, stage2_w):\n        # ---- stage0 ----\n        self.stage0_net = Stage0Net(pretrained=False)\n        self.stage0_net = s0c.load_net(self.stage0_net, stage0_w)\n        self.stage0_net.to(self.device)\n        self.stage0_net.eval()\n\n        # ---- stage1 ----\n        self.stage1_net = Stage1Net(pretrained=False)\n        self.stage1_net = s1c.load_net(self.stage1_net, stage1_w)\n        self.stage1_net.to(self.device)\n        self.stage1_net.eval()\n\n        # ---- stage2 ----\n        self.stage2_net = Net3(pretrained=False).to(self.device)\n        self.stage2_net.load_state_dict(torch.load(stage2_w))\n        self.stage2_net.eval()\n\n    def run_stage0(self, img_bgr):\n        \"\"\" Stage0: normalize rotation + homography\"\"\"\n        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)\n        img_for_model = change_color(img_rgb)\n        batch = s0c.image_to_batch(img_for_model)\n\n        with torch.no_grad(), torch.amp.autocast(self.device.split(\":\")[0], dtype=torch.float32):\n            output = self.stage0_net(batch)\n\n        # ✅ must use stage0_common.output_to_predict\n        rotated, keypoint = s0c.output_to_predict(img_rgb, batch, output)\n        normalised, _, _ = s0c.normalise_by_homography(rotated, keypoint)\n\n        return normalised\n\n    def run_stage1(self, stage0_img_rgb):\n        \"\"\" Stage1: rectify grid (perspective alignment)\"\"\"\n        image = stage0_img_rgb\n        batch = {'image': torch.from_numpy(np.ascontiguousarray(image.transpose(2, 0, 1))).unsqueeze(0)}\n\n        with torch.no_grad(), torch.amp.autocast(self.device.split(\":\")[0], dtype=torch.float32):\n            output = self.stage1_net(batch)\n\n        # ✅ must use stage1_common.output_to_predict\n        gridpoint_xy, _ = s1c.output_to_predict(image, batch, output)\n        rectified = s1c.rectify_image(image, gridpoint_xy)\n\n        return rectified\n\n    def run_stage2(self, stage1_img_rgb, length):\n        \"\"\" Stage2: segmentation → pixel_to_series\"\"\"\n        img = stage1_img_rgb\n        img = img[self.y0:self.y1, self.x0:self.x1] / 255.0\n\n        batch = self.resize(torch.from_numpy(np.ascontiguousarray(img.transpose(2, 0, 1))).unsqueeze(0)).float().to(self.device)\n\n        with torch.no_grad(), torch.amp.autocast(self.device.split(\":\")[0], dtype=torch.float32):\n            output = self.stage2_net(batch)\n\n        pixel = torch.sigmoid(output).float().data.cpu().numpy()[0]  # (4,H,W)\n\n        series_in_pixel = s2c.pixel_to_series(pixel[..., self.t0:self.t1], self.zero_mv, length)\n        series = (np.array(self.zero_mv).reshape(4, 1) - series_in_pixel) / self.mv_to_pixel\n\n        # smooth\n        for i in range(4):\n            series[i] = savgol_filter(series[i], window_length=7, polyorder=2)\n\n        return pixel, series\n    \n    def run_full(self, img_bgr, df_signal, fs, sig_len):\n        \"\"\" Full pipeline:\n                stage0 -> stage1 -> stage2 -> series_dict -> df_pred\n        \n            Returns:\n                df_pred: (sig_len, 12) prediction dataframe\n                d_series: dict of 12 leads\n                pixel: stage2 segmentation map (for debug)\n        \"\"\"\n    \n        # ---- stage0 ----\n        stage0_img_rgb = self.run_stage0(img_bgr)\n    \n        # ---- stage1 ----\n        stage1_img_rgb = self.run_stage1(stage0_img_rgb)\n    \n        # ---- stage2 ----\n        length = sig_len   # Lead II length is sig_len = fs*10\n        out = self.run_stage2(stage1_img_rgb, length)\n    \n        # ✅ support both old and new stage2 behavior\n        if isinstance(out, tuple) and len(out) == 2:\n            pixel, series_4row = out\n        else:\n            pixel = None\n            series_4row = out\n    \n        # ---- convert 4-row -> 12-lead dict ----\n        d_series = series_dict(series_4row)   # your function\n    \n        # ---- convert dict -> df_pred (sig_len x 12) ----\n        df_pred = dict_to_df_pred(d_series, df_signal, fs, sig_len)\n        return df_pred, d_series, pixel","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T22:44:30.076501Z","iopub.execute_input":"2026-01-21T22:44:30.076679Z","iopub.status.idle":"2026-01-21T22:44:30.097306Z","shell.execute_reply.started":"2026-01-21T22:44:30.076662Z","shell.execute_reply":"2026-01-21T22:44:30.096755Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Evaluate","metadata":{}},{"cell_type":"code","source":"# ============================================================\n# Core metric functions (same logic as competition notebook)\n# ============================================================\n\n# ---- Competition constants ----\nMAX_TIME_SHIFT = 0.2\nPERFECT_SCORE = 384\nLEADS_ORDER = [\"I\", \"II\", \"III\", \"aVR\", \"aVL\", \"aVF\", \"V1\", \"V2\", \"V3\", \"V4\", \"V5\", \"V6\"]\n\n# ---- Competition-style exception ----\nclass ParticipantVisibleError(Exception):\n    pass\n\ndef compute_power(label: np.ndarray, prediction: np.ndarray) -> Tuple[float, float]:\n    if label.ndim != 1 or prediction.ndim != 1:\n        raise ParticipantVisibleError(\"Inputs must be 1-dimensional arrays.\")\n\n    finite_mask = np.isfinite(prediction)\n    if not np.any(finite_mask):\n        raise ParticipantVisibleError(\"The 'prediction' array contains no finite values (all NaN or inf).\")\n\n    prediction = prediction.copy()\n    prediction[~np.isfinite(prediction)] = 0\n\n    noise = label - prediction\n    p_signal = np.sum(label**2)\n    p_noise = np.sum(noise**2)\n    return p_signal, p_noise\n\ndef compute_snr(signal: float, noise: float) -> float:\n    if noise == 0:\n        return PERFECT_SCORE\n    elif signal == 0:\n        return 0\n    else:\n        return min(signal / noise, PERFECT_SCORE)\n\ndef align_signals(label: np.ndarray, pred: np.ndarray, max_shift: int) -> np.ndarray:\n    \"\"\" Align pred to label via cross-correlation + vertical shift correction.\n        Matches official evaluation.\n    \"\"\"\n    if np.any(~np.isfinite(label)):\n        raise ParticipantVisibleError(\"values in label should all be finite\")\n    if np.sum(np.isfinite(pred)) == 0:\n        raise ParticipantVisibleError(\"prediction cannot all be infinite\")\n\n    label_arr = np.asarray(label, dtype=np.float64)\n    pred_arr = np.asarray(pred, dtype=np.float64)\n\n    label_mean = np.mean(label_arr)\n    pred_mean = np.mean(pred_arr)\n\n    label_centered = label_arr - label_mean\n    pred_centered = pred_arr - pred_mean\n\n    correlation = scipy.signal.correlate(label_centered, pred_centered, mode=\"full\")\n\n    n_label = len(label_arr)\n    n_pred = len(pred_arr)\n    lags = scipy.signal.correlation_lags(n_label, n_pred, mode=\"full\")\n\n    valid = (lags >= -max_shift) & (lags <= max_shift)\n    max_corr = np.nanmax(correlation[valid])\n\n    all_max_idx = np.flatnonzero(correlation == max_corr)\n    best_idx = min(all_max_idx, key=lambda i: abs(lags[i]))\n\n    time_shift = lags[best_idx]\n\n    start_padding_len = max(time_shift, 0)\n    pred_slice_start = max(-time_shift, 0)\n    pred_slice_end = min(n_label - time_shift, n_pred)\n    end_padding_len = max(n_label - n_pred - time_shift, 0)\n\n    aligned_pred = np.concatenate((\n        np.full(start_padding_len, np.nan),\n        pred_arr[pred_slice_start:pred_slice_end],\n        np.full(end_padding_len, np.nan)\n    ))\n\n    # vertical alignment\n    def objective(v_shift):\n        return np.nansum((label_arr - (aligned_pred - v_shift)) ** 2)\n\n    if np.any(np.isfinite(label_arr) & np.isfinite(aligned_pred)):\n        res = scipy.optimize.minimize_scalar(objective, method=\"Brent\")\n        aligned_pred -= res.x\n\n    return aligned_pred\n\ndef _calculate_image_score(group: pd.DataFrame) -> float:\n    unique_fs = group[\"fs\"].unique()\n    if len(unique_fs) != 1:\n        raise ParticipantVisibleError(\"Sampling frequency should be consistent across each ECG.\")\n    fs = int(unique_fs[0])\n\n    ii = group[group[\"lead\"] == \"II\"]\n    if ii.empty:\n        raise ParticipantVisibleError(\"Missing lead II in this image group.\")\n\n    expected_len = fs * 10\n    if len(ii) != expected_len:\n        raise ParticipantVisibleError(f\"Lead II length {len(ii)} != expected {expected_len} (fs={fs})\")\n\n    sum_signal = 0.0\n    sum_noise = 0.0\n\n    for lead in LEADS_ORDER:\n        sub = group[group[\"lead\"] == lead]\n        if sub.empty:\n            raise ParticipantVisibleError(f\"Missing lead {lead} in this image group.\")\n\n        label = sub[\"value_true\"].to_numpy()\n        pred  = sub[\"value_pred\"].to_numpy()\n\n        aligned_pred = align_signals(label, pred, int(fs * MAX_TIME_SHIFT))\n        p_signal, p_noise = compute_power(label, aligned_pred)\n\n        sum_signal += float(p_signal)\n        sum_noise  += float(p_noise)\n\n    return float(compute_snr(sum_signal, sum_noise))\n\ndef snr_score(solution: pd.DataFrame, submission: pd.DataFrame, row_id_column_name: str = \"id\") -> float:\n    \"\"\" Returns the public metric score.\n        Exactly follows the official evaluation.\n    \"\"\"\n    for df in [solution, submission]:\n        if row_id_column_name not in df.columns:\n            raise ParticipantVisibleError(f\"'{row_id_column_name}' column not found.\")\n        if df[\"value\"].isna().any():\n            raise ParticipantVisibleError(\"NaN exists in solution/submission\")\n        if not np.isfinite(df[\"value\"]).all():\n            raise ParticipantVisibleError(\"Infinity exists in solution/submission\")\n\n    submission = submission[[row_id_column_name, \"value\"]]\n    merged = pd.merge(solution, submission, on=row_id_column_name, suffixes=(\"_true\", \"_pred\"))\n\n    merged[\"image_id\"] = merged[row_id_column_name].str.split(\"_\").str[0]\n    merged[\"row_id\"] = merged[row_id_column_name].str.split(\"_\").str[1].astype(\"int64\")\n    merged[\"lead\"] = merged[row_id_column_name].str.split(\"_\").str[2]\n\n    merged.sort_values(by=[\"image_id\", \"row_id\", \"lead\"], inplace=True)\n\n    image_scores = merged.groupby(\"image_id\").apply(_calculate_image_score, include_groups=False)\n\n    return max(float(10 * np.log10(image_scores.mean())), -PERFECT_SCORE)\n\n# ============================================================\n# Build solution/submission DF from train pred/gt\n# ============================================================\ndef make_eval_dfs_from_container(base_id: str, fs: int, df_signal: pd.DataFrame, df_pred: pd.DataFrame):\n    \"\"\" Converts df_signal + df_pred into the exact (solution, submission) format expected by snr_score.\n    \"\"\"\n    sol_rows = []\n    sub_rows = []\n\n    def extract_short_segment(arr_true, arr_pred):\n        mask = np.isfinite(arr_true)\n        idx = np.flatnonzero(mask)\n        if idx.size == 0:\n            return None, None\n\n        start, end = idx[0], idx[-1] + 1\n        seg_true = arr_true[start:end]\n        seg_pred = arr_pred[start:end]\n\n        n_expected = int(np.floor(fs * 2.5))\n        if len(seg_true) > n_expected:\n            seg_true = seg_true[:n_expected]\n            seg_pred = seg_pred[:n_expected]\n        elif len(seg_true) < n_expected:\n            pad = n_expected - len(seg_true)\n            seg_true = np.concatenate([seg_true, np.full(pad, seg_true[-1])])\n            seg_pred = np.concatenate([seg_pred, np.full(pad, seg_pred[-1])])\n\n        return seg_true, seg_pred\n\n    for lead in LEADS_ORDER:\n        y_true = df_signal[lead].values.astype(np.float64)\n        y_pred = df_pred[lead].values.astype(np.float64)\n\n        if lead == \"II\":\n            seg_true, seg_pred = y_true, y_pred\n        else:\n            seg_true, seg_pred = extract_short_segment(y_true, y_pred)\n            if seg_true is None:\n                continue\n\n        for row_idx, (vt, vp) in enumerate(zip(seg_true, seg_pred)):\n            row_id = f\"{base_id}_{row_idx}_{lead}\"\n            sol_rows.append({\"id\": row_id, \"fs\": fs, \"value\": float(vt)})\n            sub_rows.append({\"id\": row_id, \"value\": float(vp)})\n\n    solution = pd.DataFrame(sol_rows)\n    submission = pd.DataFrame(sub_rows)\n    return solution, submission","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T22:44:30.098271Z","iopub.execute_input":"2026-01-21T22:44:30.098531Z","iopub.status.idle":"2026-01-21T22:44:30.125416Z","shell.execute_reply.started":"2026-01-21T22:44:30.098510Z","shell.execute_reply":"2026-01-21T22:44:30.124847Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Preprocessing Functions","metadata":{}},{"cell_type":"code","source":"# =========================\n# Preprocessing functions\n# =========================\ndef moire_soften(img_bgr, sigma=0.7):\n    # tiny blur in gray, then re-inject into LAB L\n    lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2LAB)\n    l, a, b = cv2.split(lab)\n    l2 = cv2.GaussianBlur(l, (0, 0), sigmaX=sigma, sigmaY=sigma)\n    out = cv2.merge([l2, a, b])\n    return cv2.cvtColor(out, cv2.COLOR_LAB2BGR)\n\ndef illumination_strength(img_bgr, sigma=35):\n    gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY).astype(np.float32) / 255.0\n    blur = cv2.GaussianBlur(gray, (0, 0), sigma)\n    return float(np.std(blur))\n\ndef bg_correct_lab_l(img_bgr, k=81):\n    lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2LAB)\n    l, a, b = cv2.split(lab)\n    k = int(k)\n    if k % 2 == 0:\n        k += 1\n    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k))\n    bg = cv2.morphologyEx(l, cv2.MORPH_OPEN, kernel)\n    l_corr = cv2.subtract(l, bg)\n    l_corr = cv2.normalize(l_corr, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)\n    lab2 = cv2.merge([l_corr, a, b])\n    return cv2.cvtColor(lab2, cv2.COLOR_LAB2BGR)\n    \ndef clahe_L_only(img_bgr, clip=0.6, tile=8):\n    lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2LAB)\n    l,a,b = cv2.split(lab)\n    clahe = cv2.createCLAHE(clipLimit=float(clip), tileGridSize=(int(tile), int(tile)))\n    l2 = clahe.apply(l)\n    return cv2.cvtColor(cv2.merge([l2,a,b]), cv2.COLOR_LAB2BGR)\n\ndef downscale_antialias_then_restore(img_bgr, scale=0.70):\n    \"\"\"\n    Instead of 0.5 down/up, try softer antialias: 0.7 with LANCZOS/AREA.\n    \"\"\"\n    h, w = img_bgr.shape[:2]\n    nh, nw = int(h * scale), int(w * scale)\n    small = cv2.resize(img_bgr, (nw, nh), interpolation=cv2.INTER_AREA)\n    back = cv2.resize(small, (w, h), interpolation=cv2.INTER_LANCZOS4)\n    return back\n\ndef _to_lab(img_bgr): return cv2.cvtColor(img_bgr, cv2.COLOR_BGR2LAB)\ndef _from_lab(lab):   return cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)\n    \ndef adjust_gamma_bgr(img_bgr, gamma=1.1):\n    inv = 1.0 / float(gamma)\n    table = (np.arange(256) / 255.0) ** inv\n    table = np.clip(table * 255.0, 0, 255).astype(np.uint8)\n    return cv2.LUT(img_bgr, table)\n\ndef unsharp_mask_bgr(img_bgr, sigma=1.0, amount=0.7):\n    blur = cv2.GaussianBlur(img_bgr, (0, 0), sigmaX=sigma, sigmaY=sigma)\n    out = cv2.addWeighted(img_bgr, 1.0 + amount, blur, -amount, 0)\n    return np.clip(out, 0, 255).astype(np.uint8)\n\ndef bilateral_light(img_bgr, d=5, sigmaColor=12, sigmaSpace=12):\n    return cv2.bilateralFilter(img_bgr, int(d), float(sigmaColor), float(sigmaSpace))\n\ndef clahe_luminance_bgr(img_bgr, clip=2.0, tile=8):\n    lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2LAB)\n    l, a, b = cv2.split(lab)\n    clahe = cv2.createCLAHE(clipLimit=float(clip), tileGridSize=(int(tile), int(tile)))\n    l2 = clahe.apply(l)\n    lab2 = cv2.merge([l2, a, b])\n    return cv2.cvtColor(lab2, cv2.COLOR_LAB2BGR)\n\ndef grayworld_white_balance(img_bgr):\n    img = img_bgr.astype(np.float32)\n    b, g, r = cv2.split(img)\n    mb, mg, mr = b.mean(), g.mean(), r.mean()\n    m = (mb + mg + mr) / 3.0\n    b *= (m / (mb + 1e-6))\n    g *= (m / (mg + 1e-6))\n    r *= (m / (mr + 1e-6))\n    out = cv2.merge([b, g, r])\n    return np.clip(out, 0, 255).astype(np.uint8)\n\ndef denoise_bilateral(img_bgr, d=7, sigmaColor=50, sigmaSpace=50):\n    return cv2.bilateralFilter(img_bgr, d=int(d), sigmaColor=float(sigmaColor), sigmaSpace=float(sigmaSpace))\n\ndef pp_clahe_if_flat(img):\n    gstd = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY).std()\n    if gstd < 35:\n        return clahe_luminance_bgr(img, clip=1.0, tile=8)\n    return img\n\ndef _clahe_L(img_bgr, clip=1.0, tile=8):\n    lab = _to_lab(img_bgr)\n    l, a, b = cv2.split(lab)\n    clahe = cv2.createCLAHE(clipLimit=float(clip), tileGridSize=(int(tile), int(tile)))\n    l2 = clahe.apply(l)\n    return _from_lab(cv2.merge([l2, a, b]))\n\ndef _gamma(img_bgr, gamma=1.0):\n    inv = 1.0 / float(gamma)\n    table = (np.arange(256) / 255.0) ** inv\n    table = np.clip(table * 255.0, 0, 255).astype(np.uint8)\n    return cv2.LUT(img_bgr, table)\n\ndef _unsharp(img_bgr, sigma=1.0, amount=0.2):\n    blur = cv2.GaussianBlur(img_bgr, (0, 0), sigmaX=float(sigma), sigmaY=float(sigma))\n    out = cv2.addWeighted(img_bgr, 1.0 + float(amount), blur, -float(amount), 0)\n    return np.clip(out, 0, 255).astype(np.uint8)\n\ndef _bilat(img, d=5, sc=14, ss=14):\n    return cv2.bilateralFilter(img, d=int(d), sigmaColor=float(sc), sigmaSpace=float(ss))\n\ndef _gray_std(img):\n    return float(cv2.cvtColor(img, cv2.COLOR_BGR2GRAY).std())\n\ndef _gray_mean(img):\n    return float(cv2.cvtColor(img, cv2.COLOR_BGR2GRAY).mean())\n\ndef top_hat_L(img_bgr, k=71, w=0.35):\n    lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2LAB)\n    l, a, b = cv2.split(lab)\n    k = int(k); k = k + 1 if k % 2 == 0 else k\n    ker = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k))\n    th = cv2.morphologyEx(l, cv2.MORPH_TOPHAT, ker)\n    l2 = cv2.addWeighted(l, 1.0, th, float(w), 0)\n    l2 = np.clip(l2, 0, 255).astype(np.uint8)\n    return cv2.cvtColor(cv2.merge([l2, a, b]), cv2.COLOR_LAB2BGR)\n\ndef _red_suppress(img_bgr, alpha=0.18):\n    b, g, r = cv2.split(img_bgr.astype(np.float32))\n    r2 = r * (1.0 - float(alpha))\n    out = cv2.merge([b, g, np.clip(r2, 0, 255)])\n    return out.astype(np.uint8)\n\ndef antialias_downup(img_bgr, scale=0.70):\n    h, w = img_bgr.shape[:2]\n    nh, nw = max(2, int(h * scale)), max(2, int(w * scale))\n    small = cv2.resize(img_bgr, (nw, nh), interpolation=cv2.INTER_AREA)\n    back  = cv2.resize(small, (w, h), interpolation=cv2.INTER_LANCZOS4)\n    return back\n\ndef soften_L_only(img_bgr, sigma=0.9):\n    lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2LAB)\n    l, a, b = cv2.split(lab)\n    l2 = cv2.GaussianBlur(l, (0, 0), sigmaX=float(sigma), sigmaY=float(sigma))\n    return cv2.cvtColor(cv2.merge([l2, a, b]), cv2.COLOR_LAB2BGR)\n\ndef clahe_gray_only(img_bgr, clip=1.0, tile=8):\n    g = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)\n    clahe = cv2.createCLAHE(clipLimit=float(clip), tileGridSize=(int(tile), int(tile)))\n    g2 = clahe.apply(g)\n    return cv2.cvtColor(g2, cv2.COLOR_GRAY2BGR)\n\ndef laplacian_sharpen(img_bgr, k=0.12):\n    # mild Laplacian boost\n    g = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY).astype(np.float32)\n    lap = cv2.Laplacian(g, cv2.CV_32F, ksize=3)\n    g2 = g + float(k) * lap\n    g2 = np.clip(g2, 0, 255).astype(np.uint8)\n    return cv2.cvtColor(g2, cv2.COLOR_GRAY2BGR)\n\ndef adaptive_gamma_by_mean(img_bgr, target=135.0, lo=0.93, hi=1.10):\n    gmean = float(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY).mean())\n    # if dark -> brighten (gamma < 1), if bright -> darken (gamma > 1)\n    if gmean < target:\n        gamma = lo\n    else:\n        gamma = hi\n    inv = 1.0 / float(gamma)\n    table = (np.arange(256) / 255.0) ** inv\n    table = np.clip(table * 255.0, 0, 255).astype(np.uint8)\n    return cv2.LUT(img_bgr, table)\n\ndef _u8(x): \n    return np.clip(x, 0, 255).astype(np.uint8)\n    \ndef clip_highlights(img_bgr, p=99.6):\n    # reduce blown highlights (phone glare) without touching mids much\n    x = img_bgr.astype(np.float32)\n    hi = np.percentile(x, float(p))\n    if hi <= 1: \n        return img_bgr\n    scale = 255.0 / hi\n    x = np.minimum(x * scale, 255.0)\n    return _u8(x)\n\ndef _moire_soften_L(img_bgr, sigma=0.7):\n    lab = _to_lab(img_bgr)\n    l, a, b = cv2.split(lab)\n    l2 = cv2.GaussianBlur(l, (0, 0), sigmaX=float(sigma), sigmaY=float(sigma))\n    return _from_lab(cv2.merge([l2, a, b]))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T22:44:30.127410Z","iopub.execute_input":"2026-01-21T22:44:30.127625Z","iopub.status.idle":"2026-01-21T22:44:30.158024Z","shell.execute_reply.started":"2026-01-21T22:44:30.127605Z","shell.execute_reply":"2026-01-21T22:44:30.157450Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Preprocess by Source","metadata":{}},{"cell_type":"code","source":"# =========================\n# Preprocessing by source\n# =========================\ndef preprocess_by_source_v2(img_bgr, source):\n    s = str(source)\n\n    if s == \"0001\":\n        return adjust_gamma_bgr(img_bgr, gamma=1.03)\n\n    # 0003: color scan\n    if s == \"0003\":\n        return pp_clahe_if_flat(grayworld_white_balance(img_bgr))\n\n    # 0004: bw scan\n    if s == \"0004\":\n        return img_bgr\n\n    # 0005: mobile photo of printed page\n    if s == \"0005\": # pp_bilat_light\n        return bilateral_light(img_bgr, d=5, sigmaColor=10, sigmaSpace=10)\n\n    # 0006: laptop screen photos\n    if s == \"0006\": # pp_bilat_light_moireL_claheL\n        x = downscale_antialias_then_restore(img_bgr, scale=0.70)\n        x = moire_soften(x, sigma=0.7)\n        x = clahe_luminance_bgr(x, clip=1.1, tile=8)\n        return x\n\n    # 0009: stained/soaked\n    if s == \"0009\":\n        return img_bgr\n    \n    # 0010: damage\n    if s == \"0010\": # pp_basic_then_claheLonly\n        x = img_bgr\n        if illumination_strength(x, sigma=35) > 0.14:\n            x = bg_correct_lab_l(x, k=81)\n        if cv2.cvtColor(x, cv2.COLOR_BGR2GRAY).std() < 30:\n            x = clahe_luminance_bgr(x, clip=1.08, tile=8)\n        x = clahe_L_only(x, clip=0.60, tile=8)\n        return x\n\n    # 0011: mold color scan\n    if s == \"0011\": # pp_basic\n        x = grayworld_white_balance(img_bgr)\n        x = clahe_luminance_bgr(x, clip=1.2, tile=8)\n        return x\n\n    # 0012: mold bw scan \n    if s == \"0012\":\n        return adjust_gamma_bgr(img_bgr, gamma=1.05)\n\n    return img_bgr","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T22:44:30.158742Z","iopub.execute_input":"2026-01-21T22:44:30.158974Z","iopub.status.idle":"2026-01-21T22:44:30.184194Z","shell.execute_reply.started":"2026-01-21T22:44:30.158953Z","shell.execute_reply":"2026-01-21T22:44:30.183620Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Do Submission","metadata":{}},{"cell_type":"code","source":"# load training dataset & set pipeline\nTRAIN_DIR = Path(\"/kaggle/input/physionet-ecg-image-digitization/train\")\ndf_train = pd.read_csv(\"/kaggle/input/physionet-ecg-image-digitization/train.csv\")\n\npipeline = PhysioPipeline(device=\"cuda:0\")\npipeline.load_models(\n    stage0_w=\"/kaggle/input/hengck23-submit-physionet/hengck23-submit-physionet/weight/stage0-last.checkpoint.pth\",\n    stage1_w=\"/kaggle/input/hengck23-submit-physionet/hengck23-submit-physionet/weight/stage1-last.checkpoint.pth\",\n    stage2_w=\"/kaggle/input/physio-seg-public/pytorch/net3_009_4200/1/iter_0004200.pt\",\n)\n\n# set benchmarking sample data\ndf_bench = pd.DataFrame()\ndf_bench[\"id\"] = pd.Series(['1404687891', '3646789790', '4066709567', '1395556771', '3817987033', '474136527', '3166474998', '3926202474', '88853795', '1536586741', '3110772881', '4099267973', '1736488789', '2450104795', '3042142152', '2589114748', '3591987991', '4229829773', '1420515841', '1919286802', '2800037729', '1278116713', '4246336200', '474136527', '1665441750', '3828947205', '4019229265'])\ndf_bench[\"source\"] = pd.Series(['0001', '0001', '0001', '0003', '0003', '0003', '0004', '0004', '0004', '0005', '0005', '0005', '0006', '0006', '0006', '0009', '0009', '0009', '0010', '0010', '0010', '0011', '0011', '0011', '0012', '0012', '0012'])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T22:44:30.259304Z","iopub.execute_input":"2026-01-21T22:44:30.259846Z","iopub.status.idle":"2026-01-21T22:44:31.360297Z","shell.execute_reply.started":"2026-01-21T22:44:30.259822Z","shell.execute_reply":"2026-01-21T22:44:31.359481Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# load test dataset\nWORK_DIR = \"/kaggle/input/physionet-ecg-image-digitization\"\ndf_test = pd.read_csv(f\"{WORK_DIR}/test.csv\")\ndf_test['id'] = df_test['id'].astype(str) \n\n# load submission dataset\nsample_submission = pd.read_parquet(f\"{WORK_DIR}/sample_submission.parquet\")\nsample_submission[['id_root', 'row', 'lead']] = sample_submission['id'].str.split('_', expand=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T22:44:31.361168Z","iopub.execute_input":"2026-01-21T22:44:31.361477Z","iopub.status.idle":"2026-01-21T22:44:31.499143Z","shell.execute_reply.started":"2026-01-21T22:44:31.361434Z","shell.execute_reply":"2026-01-21T22:44:31.498499Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# Submission Script\n# ============================================================\n# ---- set these to match your classifier training config ----\nCLS_MODEL_NAME = \"efficientnet_b2\"\nCLS_NUM_CLASSES = 12\nCLS_RESOLUTION = 256\n\n# path to trained weights\nCLS_CKPT_PATH = \"/kaggle/input/physionet-image-multi-class-train/efficientnet_b2_full_train.pth\"\n\n\n# -------------------------\n# Classifier\n# -------------------------\ndef build_classifier(device=\"cuda\"):\n    model = timm.create_model(CLS_MODEL_NAME, pretrained=False, num_classes=CLS_NUM_CLASSES)\n    ckpt = torch.load(CLS_CKPT_PATH, map_location=\"cpu\")\n    sd = ckpt[\"state_dict\"] if isinstance(ckpt, dict) and \"state_dict\" in ckpt else ckpt\n    sd = {k.replace(\"module.\", \"\"): v for k, v in sd.items()}\n    missing, unexpected = model.load_state_dict(sd, strict=False)\n    print(\"---- LOAD DIAGNOSTICS ----\")\n    print(\"missing:\", len(missing))\n    print(\"unexpected:\", len(unexpected))\n    print(\"missing head-ish:\", [k for k in missing if \"classifier\" in k or k.endswith(\"fc.weight\") or k.endswith(\"fc.bias\")][:20])\n    print(\"unexpected head-ish:\", [k for k in unexpected if \"classifier\" in k or k.endswith(\"fc.weight\") or k.endswith(\"fc.bias\")][:20])\n    model.to(device).eval()\n    return model\n\ndef cls_preprocess_bgr(img_bgr):\n    # IMPORTANT: training used cv2.imread (BGR) and did NOT convert to RGB.\n    img = cv2.resize(img_bgr, (CLS_RESOLUTION, CLS_RESOLUTION), interpolation=cv2.INTER_LINEAR)\n    img = img.astype(np.float32) / 255.0\n    # Same mean/std as training notebook (Albumentations Normalize defaults)\n    mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)  # applied to B,G,R in training\n    std  = np.array([0.229, 0.224, 0.225], dtype=np.float32)\n    img = (img - mean) / std\n    # to CHW\n    x = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0)\n    return x\n\n@torch.no_grad()\ndef predict_source_suffix(model, img_bgr, device=\"cuda\"):\n    x = cls_preprocess_bgr(img_bgr).to(device, non_blocking=True)\n    logits = model(x)\n    probs = F.softmax(logits, dim=1)[0]\n    cls = int(torch.argmax(probs).item())          # 0..11\n    conf = float(probs[cls].item())\n    suffix = f\"{cls+1:04d}\"                        # 0001..0012\n    return suffix, conf\n\ndef eval_cls_on_train_samples(df_bench, n=25):\n    ok = 0\n    rows = []\n    for r in df_bench[:n].itertuples(index=False):\n        train_id = str(r.id)\n        true_src = str(r.source)\n        img_path = TRAIN_DIR / train_id / f\"{train_id}-{true_src}.png\"\n        img = cv2.imread(str(img_path), cv2.IMREAD_COLOR)\n        if img is None:\n            continue\n        pred, conf = predict_source_suffix(cls_model, img, device=\"cuda\")\n        rows.append((true_src, pred, conf))\n        ok += int(pred == true_src)\n    acc = ok / len(rows)\n    print(\"Train-like accuracy:\", acc, \"n=\", len(rows))\n    return rows\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\ncls_model = build_classifier(device=device)\nprint(\"Classifier ready:\", type(cls_model))\nrows = eval_cls_on_train_samples(df_bench, n=25)\nrows\n\n\n# -------------------------\n# Submission DF builder\n# -------------------------\ndef make_submission_from_pred(base_id: str, fs: int, sig_len: int, d_series: dict):\n    sub_rows = []\n    base_id = str(base_id)\n    fs = int(fs)\n    sig_len = int(sig_len)\n\n    n_short = int(np.floor(fs * 2.5))\n\n    def take_segment(y, n):\n        y = np.asarray(y, dtype=np.float64)\n        if len(y) == 0:\n            return np.zeros(n, dtype=np.float64)\n        if len(y) >= n:\n            return y[:n]\n        pad = n - len(y)\n        return np.concatenate([y, np.full(pad, y[-1], dtype=np.float64)])\n\n    for lead in LEADS_ORDER:\n        y_pred = np.asarray(d_series[lead], dtype=np.float64)\n\n        if lead == \"II\":\n            seg = take_segment(y_pred, sig_len)\n        else:\n            seg = take_segment(y_pred, n_short)\n\n        for row_idx, vp in enumerate(seg):\n            row_id = f\"{base_id}_{row_idx}_{lead}\"\n            sub_rows.append({\"id\": row_id, \"value\": float(vp)})\n\n    return pd.DataFrame(sub_rows)\n\n\n# -------------------------\n# Single prediction (UPDATED: uses preprocess_by_source + no raw-vs-pp selector)\n# -------------------------\n@torch.no_grad()\ndef predict_one_test(\n    sample_id,\n    image_path,\n    df_meta,          # df_test\n    pipeline,\n    cls_model,\n    device=\"cuda\",\n):\n    img_raw = cv2.imread(str(image_path), cv2.IMREAD_COLOR)\n    if img_raw is None:\n        raise FileNotFoundError(image_path)\n\n    # infer length from metadata\n    length = int(df_meta[(df_meta[\"id\"] == sample_id) & (df_meta[\"lead\"] == \"II\")].iloc[0].number_of_rows)\n\n    # classify source -> choose preprocessing\n    pred_src, conf = predict_source_suffix(cls_model, img_raw, device=device)\n\n    # preprocess once\n    print(sample_id, \"\\t\", pred_src, \"\\t\", conf)\n    img_pp = preprocess_by_source_v2(img_raw.copy(), pred_src)        \n\n    # stage0/1/2\n    s0 = pipeline.run_stage0(img_pp)\n    s1 = pipeline.run_stage1(s0)\n    pixel, series_4row = pipeline.run_stage2(s1, length=length)\n   \n    dbg = {\n        \"pred_source\": pred_src,\n        \"conf\": float(conf),\n        \"length\": int(length),\n    }\n    cuda_cleanup()\n    return series_4row, dbg\n\n\n# ============================================================\n# RUN SUBMISSION\n# ============================================================\nres = []\nfor sample_id, df_meta in tqdm(df_test.groupby(\"id\"), total=df_test[\"id\"].nunique()):\n    img_path = f\"{WORK_DIR}/test/{sample_id}.png\"\n    fs = int(df_meta.iloc[0].fs)\n    sig_len = int(df_meta[df_meta[\"lead\"] == \"II\"].iloc[0].number_of_rows)\n    series_4row, dbg = predict_one_test(\n        sample_id=sample_id,\n        image_path=img_path,\n        df_meta=df_test,\n        pipeline=pipeline,\n        cls_model=cls_model,\n        device=device,\n    )\n    # submission rows\n    d_series = dw(series_dict(series_4row))\n    sub_df = make_submission_from_pred(sample_id, fs, sig_len, d_series)\n    res.append(sub_df)\n\n# do submission\ndf_submission = pd.concat(res, ignore_index=True)\nprint(all(df_submission.id == sample_submission.id))\ndf_submission.to_csv(\"submission.csv\", index=False)\nprint(\"submission shape:\", df_submission.shape)\ndf_submission.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-21T22:44:31.500009Z","iopub.execute_input":"2026-01-21T22:44:31.500211Z","iopub.status.idle":"2026-01-21T22:45:15.951921Z","shell.execute_reply.started":"2026-01-21T22:44:31.500191Z","shell.execute_reply":"2026-01-21T22:45:15.951332Z"}},"outputs":[],"execution_count":null}]}