{"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":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":"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\n\nimport os, sys, gc, cv2, numpy as np, pandas as pd, torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision.transforms as T\nimport timm\nfrom scipy.signal import savgol_filter\n\nsys.path.append('/kaggle/input/hengck23-submit-physionet/hengck23-submit-physionet')\nimport stage0_common as s0c\nimport stage1_common as s1c\nimport stage2_common as s2c\nfrom stage0_model import Net as Stage0Net\nfrom stage1_model import Net as Stage1Net\nfrom stage2_model import MyCoordUnetDecoder, encode_with_resnet\n\nLEADS_ORDER = [\"I\",\"II\",\"III\",\"aVR\",\"aVL\",\"aVF\",\"V1\",\"V2\",\"V3\",\"V4\",\"V5\",\"V6\"]\n\ndef change_color(image_rgb):\n    hsv = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2HSV)\n    h, s, v = cv2.split(hsv)\n    v_denoised = cv2.fastNlMeansDenoising(v, h=5.46)\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 series_dict(series_4row):\n    series_4row = np.asarray(series_4row)\n    if series_4row.ndim == 3: series_4row = series_4row[0]\n    if series_4row.shape[0] != 4 and series_4row.shape[1] == 4: series_4row = series_4row.T\n\n    d = {}\n    names = [\n        ['I','aVR','V1','V4'],\n        ['II_short','aVL','V2','V5'],   \n        ['III','aVF','V3','V6'],\n    ]\n    for r in range(3):\n        for lead, arr in zip(names[r], np.array_split(series_4row[r], 4)):\n            d[lead] = np.asarray(arr, dtype=np.float32)\n\n    d['II'] = np.asarray(series_4row[3], dtype=np.float32)  # full 10s II\n    return d\n\n# ✅ FIX: Einthoven correction on SHORT leads only\ndef dw(d, alpha=0.33):\n    if all(k in d for k in ['I','II_short','III']):\n        L1, L2s, L3 = d['I'], d['II_short'], d['III']\n        e = L2s - (L1 + L3)\n        d['I']        = L1 + alpha*e\n        d['III']      = L3 + alpha*e\n        d['II_short'] = L2s - alpha*e\n    return d\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    return cv2.cvtColor(cv2.merge([l2, a, b]), 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)); g *= (m / (mg + 1e-6)); r *= (m / (mr + 1e-6))\n    return np.clip(cv2.merge([b, g, r]), 0, 255).astype(np.uint8)\n\ndef denoise_median(img_bgr, k=3):\n    k = int(k); k = k if k % 2 == 1 else k + 1\n    return cv2.medianBlur(img_bgr, k)\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 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); k = k if k % 2 == 1 else 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    return cv2.cvtColor(cv2.merge([l_corr, a, b]), cv2.COLOR_LAB2BGR)\n\ndef preprocess_by_source(img_bgr, source):\n    s = str(source)\n    if s == \"0001\": return img_bgr\n    if s == \"0003\": return clahe_luminance_bgr(grayworld_white_balance(img_bgr), clip=1.2, tile=8)\n    if s == \"0004\": return img_bgr\n    if s == \"0006\":\n        x = denoise_bilateral(img_bgr, d=5, sigmaColor=25, sigmaSpace=25)\n        return clahe_luminance_bgr(x, clip=1.2, tile=8)\n    if s == \"0005\":\n        x = img_bgr\n        if illumination_strength(x, sigma=35) > 0.14: x = bg_correct_lab_l(x, k=81)\n        if cv2.cvtColor(x, cv2.COLOR_BGR2GRAY).std() < 30: x = clahe_luminance_bgr(x, clip=1.1, tile=8)\n        return x\n    if s == \"0009\":\n        x = img_bgr\n        if illumination_strength(x, sigma=35) > 0.14: x = bg_correct_lab_l(x, k=101)\n        return denoise_median(x, k=3)\n    if s == \"0010\":\n        x = img_bgr\n        if illumination_strength(x, sigma=35) > 0.14: x = bg_correct_lab_l(x, k=81)\n        if cv2.cvtColor(x, cv2.COLOR_BGR2GRAY).std() < 30: x = clahe_luminance_bgr(x, clip=1.15, tile=8)\n        return x\n    if s == \"0011\": return clahe_luminance_bgr(grayworld_white_balance(img_bgr), clip=1.2, tile=8)\n    if s == \"0012\": return img_bgr\n    return img_bgr\n\ndef stage1_quality(s1_rgb):\n    g = cv2.cvtColor(s1_rgb.astype(np.uint8), cv2.COLOR_RGB2GRAY)\n    e = cv2.Canny(g, 50, 150)\n    density = e.mean() / 255.0\n    gx = cv2.Sobel(g, cv2.CV_32F, 1, 0, ksize=3)\n    gy = cv2.Sobel(g, cv2.CV_32F, 0, 1, ksize=3)\n    ax = float(np.mean(np.abs(gx))); ay = float(np.mean(np.abs(gy)))\n    anis = max(ax, ay) / (min(ax, ay) + 1e-6)\n    return float(density * 0.7 + np.tanh(anis - 1.0) * 0.3)\n\nclass Net3(nn.Module):\n    def __init__(self, pretrained=True):\n        super().__init__()\n        encoder_dim = [64, 128, 256, 512]\n        decoder_dim = [128, 64, 32, 16]\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        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        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        return self.pixel(last)\n\nclass PhysioPipeline:\n    def __init__(self, device=\"cuda:0\"):\n        self.device = device\n        self.stage0_net = self.stage1_net = self.stage2_net = None\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        self.resize = T.Resize((1696, 4352), interpolation=T.InterpolationMode.BILINEAR)\n\n    def load_models(self, stage0_w, stage1_w, stage2_w):\n        self.stage0_net = s0c.load_net(Stage0Net(pretrained=False), stage0_w).to(self.device).eval()\n        self.stage1_net = s1c.load_net(Stage1Net(pretrained=False), stage1_w).to(self.device).eval()\n        self.stage2_net = Net3(pretrained=False).to(self.device).eval()\n        st = torch.load(stage2_w, map_location=\"cpu\")\n        if isinstance(st, dict) and \"state_dict\" in st: st = st[\"state_dict\"]\n        self.stage2_net.load_state_dict(st, strict=True)\n\n    def run_stage0(self, img_bgr):\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        with torch.no_grad(), torch.amp.autocast(self.device.split(\":\")[0], dtype=torch.float32):\n            output = self.stage0_net(batch)\n        rotated, keypoint = s0c.output_to_predict(img_rgb, batch, output)\n        normalised, _, _ = s0c.normalise_by_homography(rotated, keypoint)\n        return normalised\n\n    def run_stage1(self, stage0_img_rgb):\n        image = stage0_img_rgb\n        batch = {'image': torch.from_numpy(np.ascontiguousarray(image.transpose(2, 0, 1))).unsqueeze(0)}\n        with torch.no_grad(), torch.amp.autocast(self.device.split(\":\")[0], dtype=torch.float32):\n            output = self.stage1_net(batch)\n        gridpoint_xy, _ = s1c.output_to_predict(image, batch, output)\n        return s1c.rectify_image(image, gridpoint_xy)\n\n    def run_stage2(self, stage1_img_rgb, length):\n        img = stage1_img_rgb[self.y0:self.y1, self.x0:self.x1] / 255.0\n        batch = self.resize(torch.from_numpy(np.ascontiguousarray(img.transpose(2, 0, 1))).unsqueeze(0)).float().to(self.device)\n        with torch.no_grad(), torch.amp.autocast(self.device.split(\":\")[0], dtype=torch.float32):\n            output = self.stage2_net(batch)\n        pixel = torch.sigmoid(output).float().cpu().numpy()[0]\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        for i in range(4):\n            series[i] = savgol_filter(series[i], window_length=7, polyorder=2)\n        return series\n\nCLS_MODEL_NAME=\"efficientnet_b2\"\nCLS_NUM_CLASSES=12\nCLS_RESOLUTION=256\nCLS_CKPT_PATH=\"/kaggle/input/physionet-image-multi-class-train/efficientnet_b2_full_train.pth\"\n\ndef build_classifier(device=\"cuda\"):\n    m = timm.create_model(CLS_MODEL_NAME, pretrained=False, num_classes=CLS_NUM_CLASSES)\n    st = torch.load(CLS_CKPT_PATH, map_location=\"cpu\")\n    if isinstance(st, dict) and \"state_dict\" in st: st = st[\"state_dict\"]\n    st2 = {k.replace(\"module.\",\"\"): v for k, v in st.items()}\n    m.load_state_dict(st2, strict=False)\n    return m.to(device).eval()\n\ndef cls_preprocess_bgr(img_bgr):\n    img = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)\n    img = cv2.resize(img, (CLS_RESOLUTION, CLS_RESOLUTION), interpolation=cv2.INTER_AREA).astype(np.float32)/255.0\n    mean = np.array([0.485,0.456,0.406], np.float32); std = np.array([0.229,0.224,0.225], np.float32)\n    img = (img - mean) / std\n    return torch.from_numpy(img).permute(2,0,1).unsqueeze(0)\n\n@torch.no_grad()\ndef predict_source_suffix(model, img_bgr, device=\"cuda\"):\n    x = cls_preprocess_bgr(img_bgr).to(device)\n    p = F.softmax(model(x), dim=1)[0]\n    cls = int(torch.argmax(p).item())\n    return f\"{cls+1:04d}\"\n\ndef select_stage1_with_source(pipeline, img_raw_bgr, pred_source_suffix, selector_margin=1.02):\n    img_pp = preprocess_by_source(img_raw_bgr.copy(), pred_source_suffix)\n    s1_raw = pipeline.run_stage1(pipeline.run_stage0(img_raw_bgr))\n    q_raw  = stage1_quality(s1_raw)\n    s1_pp  = pipeline.run_stage1(pipeline.run_stage0(img_pp))\n    q_pp   = stage1_quality(s1_pp)\n    return s1_pp if q_pp > q_raw * selector_margin else s1_raw\n\ndef make_submission_from_pred(base_id: str, fs: int, sig_len: int, d_series: dict):\n    base_id = str(base_id); fs = int(fs); sig_len = int(sig_len)\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) >= n: return y[:n]\n        if len(y) == 0: return np.zeros(n, np.float64)\n        return np.concatenate([y, np.full(n - len(y), y[-1], np.float64)])\n\n    rows = []\n    for lead in LEADS_ORDER:\n        y = np.asarray(d_series[lead], dtype=np.float64)\n        seg = take_segment(y, sig_len if lead==\"II\" else n_short)\n        rows.append(pd.DataFrame({\"id\":[f\"{base_id}_{i}_{lead}\" for i in range(len(seg))],\n                                  \"value\":seg.astype(np.float32)}))\n    return pd.concat(rows, ignore_index=True)\n\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)\nsample_submission = pd.read_parquet(f\"{WORK_DIR}/sample_submission.parquet\")[[\"id\"]]\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\npipeline = PhysioPipeline(device=\"cuda:0\" if device==\"cuda\" else \"cpu\")\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)\ncls_model = build_classifier(device=device)\n\nres = []\nfor sample_id, g in df_test.groupby(\"id\", sort=True):\n    img_path = f\"{WORK_DIR}/test/{sample_id}.png\"\n    img_raw = cv2.imread(img_path, cv2.IMREAD_COLOR)\n    if img_raw is None:\n        raise FileNotFoundError(img_path)\n\n    fs = int(g.fs.iloc[0])\n    sig_len = int(g.loc[g.lead==\"II\",\"number_of_rows\"].iloc[0])\n\n    pred_src = predict_source_suffix(cls_model, img_raw, device=device)\n    s1 = select_stage1_with_source(pipeline, img_raw, pred_src, selector_margin=1.02)\n    series_4row = pipeline.run_stage2(s1, length=sig_len)\n\n    d = dw(series_dict(series_4row))         # ✅ now safe (uses II_short)\n    res.append(make_submission_from_pred(sample_id, fs, sig_len, d))\n    gc.collect()\n\ndf_submission = pd.concat(res, ignore_index=True)\ndf_submission = df_submission.set_index(\"id\").reindex(sample_submission[\"id\"]).reset_index()\n\nassert df_submission[\"value\"].notna().all()\nassert (df_submission[\"id\"].values == sample_submission[\"id\"].values).all()\n\ndf_submission.to_csv(\"submission.csv\", index=False)\nprint(\"OK  submission.csv\", df_submission.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-15T12:32:40.854974Z","iopub.execute_input":"2026-01-15T12:32:40.855173Z","iopub.status.idle":"2026-01-15T12:33:51.202338Z","shell.execute_reply.started":"2026-01-15T12:32:40.855155Z","shell.execute_reply":"2026-01-15T12:33:51.201594Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}