{"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":271051632,"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'","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-01-03T14:30:48.689704Z","iopub.execute_input":"2026-01-03T14:30:48.689859Z","iopub.status.idle":"2026-01-03T14:31:10.84955Z","shell.execute_reply.started":"2026-01-03T14:30:48.689843Z","shell.execute_reply":"2026-01-03T14:31:10.848558Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport sys\nimport types\nimport cv2\nimport torch\nimport numpy as np\nimport pandas as pd\nfrom tqdm.auto import tqdm\nfrom scipy import ndimage\nfrom scipy.signal import resample, savgol_filter, butter, filtfilt\nimport torchvision.transforms as T\nfrom torch.cuda.amp import autocast\nfrom pathlib import Path\n\n# ====================================================\n# 1. SETUP: AUTO-FIND PATHS & BYPASS INSTALLS\n# ====================================================\nprint(\"⚙️ System Setup...\")\n\n# A. Bypass 'connected-components-3d'\ndef fake_connected_components(image, connectivity=26):\n    ndim = image.ndim\n    structure = ndimage.generate_binary_structure(ndim, ndim)\n    labeled_array, num_features = ndimage.label(image, structure=structure)\n    return labeled_array\n\ndummy_module = types.ModuleType(\"connected_components_3d\")\ndummy_module.connected_components = fake_connected_components\nsys.modules[\"connected_components_3d\"] = dummy_module\n\n# B. Auto-Find Paths\nCODE_PATH = None\nWEIGHTS_DIR = None\nNET3_PATH = None\n\nfor root, dirs, files in os.walk(\"/kaggle/input\"):\n    if \"stage0_model.py\" in files: CODE_PATH = root\n    if \"stage0-last.checkpoint.pth\" in files: WEIGHTS_DIR = root\n    if \"iter_0004200.pt\" in files: NET3_PATH = os.path.join(root, \"iter_0004200.pt\")\n\nif not (CODE_PATH and WEIGHTS_DIR):\n    raise FileNotFoundError(\"❌ Dataset missing! Add 'hengck23-submit-physionet' from sidebar.\")\n\nsys.path.append(CODE_PATH)\n\n# ====================================================\n# 2. IMPORT MODELS (FIXED: Added load_net)\n# ====================================================\ntry:\n    from stage0_model import Net as Stage0Net\n    # 🚨 FIX IS HERE: Added 'load_net' to imports\n    from stage0_common import image_to_batch, output_to_predict as out_s0, normalise_by_homography, load_net \n    from stage1_model import Net as Stage1Net\n    from stage1_common import output_to_predict as out_s1, rectify_image\n    from stage2_model import *\n    print(\"✅ Libraries Loaded Successfully.\")\nexcept ImportError as e:\n    print(f\"❌ Import Error: {e}\")\n    sys.exit(1)\n\n# ====================================================\n# 3. PIPELINE CLASS\n# ====================================================\nDEVICE = \"cuda:0\" if torch.cuda.is_available() else \"cpu\"\n\nclass EcgPipeline:\n    def __init__(self):\n        print(\"🚀 Loading Models...\")\n        self.net0 = Stage0Net(pretrained=False).to(DEVICE).eval()\n        self.net0 = load_net(self.net0, f'{WEIGHTS_DIR}/stage0-last.checkpoint.pth')\n        \n        self.net1 = Stage1Net(pretrained=False).to(DEVICE).eval()\n        self.net1 = load_net(self.net1, f'{WEIGHTS_DIR}/stage1-last.checkpoint.pth')\n        \n        self.net2 = Net3(pretrained=False).to(DEVICE).eval()\n        if NET3_PATH: self.net2.load_state_dict(torch.load(NET3_PATH, map_location=DEVICE))\n        \n        self.resize = T.Resize((1696, 4352), interpolation=T.InterpolationMode.BILINEAR)\n\n    def preprocess(self, img_rgb):\n        # CLAHE (Fix Lighting)\n        lab = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2LAB)\n        l, a, b = cv2.split(lab)\n        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))\n        cl = clahe.apply(l)\n        return cv2.cvtColor(cv2.merge((cl,a,b)), cv2.COLOR_LAB2RGB)\n\n    @torch.no_grad()\n    def process(self, image_path):\n        img_orig = cv2.imread(str(image_path))\n        if img_orig is None: return None\n        img_rgb = cv2.cvtColor(img_orig, cv2.COLOR_BGR2RGB)\n        \n        # Stage 0\n        img_enh = self.preprocess(img_rgb)\n        b0 = image_to_batch(img_enh)\n        with autocast(): out0 = self.net0(b0)\n        rot, kp = out_s0(img_rgb, b0, out0)\n        norm, _, _ = normalise_by_homography(rot, kp)\n        \n        # Stage 1\n        t1 = torch.from_numpy(np.ascontiguousarray(norm.transpose(2, 0, 1))).unsqueeze(0).to(DEVICE)\n        b1 = {'image': t1}\n        with autocast(): out1 = self.net1(b1)\n        grid_xy, _ = out_s1(norm, b1, out1)\n        rectified = rectify_image(norm, grid_xy)\n        \n        # Stage 2\n        x0, x1, y0, y1 = 0, 2176, 0, 1696\n        if rectified.shape[0] < y1 or rectified.shape[1] < x1:\n            rectified = cv2.resize(rectified, (x1, y1))\n        \n        crop = rectified[y0:y1, x0:x1] / 255.0\n        t2 = torch.from_numpy(np.ascontiguousarray(crop.transpose(2, 0, 1))).unsqueeze(0).float().to(DEVICE)\n        b2 = self.resize(t2)\n        \n        with autocast(): out2 = self.net2(b2)\n        return torch.sigmoid(out2).cpu().numpy()[0]\n\n# ====================================================\n# 4. ADVANCED SIGNAL PROCESSING (The Score Booster)\n# ====================================================\ndef robust_extraction(mask, zero_mv_indices, gain=78.66):\n    \"\"\"\n    1. Windowed Center-of-Mass (Ignores grid noise)\n    2. Interpolation (Fixes gaps)\n    \"\"\"\n    signals = []\n    window = 5 # +/- 5 pixels around peak\n    \n    for i in range(4):\n        channel = mask[i] # (H, W)\n        H, W = channel.shape\n        \n        # 1. Find the Peak (Argmax)\n        peaks = np.argmax(channel, axis=0)\n        \n        # 2. Windowed Center of Mass\n        y_centers = []\n        for x in range(W):\n            peak_y = peaks[x]\n            \n            # Define window bounds\n            start_y = max(0, peak_y - window)\n            end_y = min(H, peak_y + window + 1)\n            \n            # Extract slice\n            col_slice = channel[start_y:end_y, x]\n            mass = col_slice.sum()\n            \n            if mass < 0.1: # Confidence too low? Mark NaN\n                y_centers.append(np.nan)\n            else:\n                # Local CoM relative to start_y\n                local_y = (col_slice * np.arange(len(col_slice))).sum() / mass\n                y_centers.append(start_y + local_y)\n        \n        # 3. Interpolate Missing Data (Crucial for SNR)\n        y_series = pd.Series(y_centers).interpolate(method='linear', limit_direction='both')\n        y_final = y_series.to_numpy()\n        \n        # Convert to Voltage\n        voltage = (zero_mv_indices[i] - y_final) / gain\n        signals.append(np.nan_to_num(voltage)) # Fill remaining NaNs with 0\n        \n    return signals\n\ndef clinical_filter(sig, fs=500):\n    \"\"\"\n    High-pass filter to remove Baseline Wander (Breathing artifact).\n    Competition Metric penalizes drift heavily.\n    \"\"\"\n    try:\n        # 0.5Hz High-pass Butterworth\n        sos = butter(2, 0.5, 'hp', fs=fs, output='sos')\n        filtered = filtfilt(sos, sig)\n        return filtered\n    except:\n        return sig # Fallback\n\ndef split_leads(extracted_rows):\n    leads = {}\n    layout = [['I', 'aVR', 'V1', 'V4'], ['II', 'aVL', 'V2', 'V5'], ['III', 'aVF', 'V3', 'V6']]\n    for r_idx, names in enumerate(layout):\n        row_signal = extracted_rows[r_idx]\n        chunks = np.array_split(row_signal, 4)\n        for name, chunk in zip(names, chunks):\n            leads[name] = chunk\n    leads['II_long'] = extracted_rows[3]\n    return leads\n\n# ====================================================\n# 5. EXECUTION LOOP\n# ====================================================\npipeline = EcgPipeline()\nTEST_CSV = \"/kaggle/input/physionet-ecg-image-digitization/test.csv\"\nTEST_DIR = Path(\"/kaggle/input/physionet-ecg-image-digitization/test\")\ntest_df = pd.read_csv(TEST_CSV)\ntest_df['id'] = test_df['id'].astype(str)\n\nZERO_MV = [703.5, 987.5, 1271.5, 1531.5]\nGAIN = 78.66\n\nsubmission_rows = []\nprint(f\"🚀 Processing {len(test_df['id'].unique())} ECGs with Advanced Filtering...\")\n\nfor sample_id in tqdm(test_df['id'].unique()):\n    try:\n        path = TEST_DIR / f\"{sample_id}.png\"\n        mask = pipeline.process(path)\n        if mask is None: raise ValueError(\"Img Fail\")\n        \n        # Crop & Extract\n        active_mask = mask[..., 235:4161] \n        raw_signals = robust_extraction(active_mask, ZERO_MV, GAIN)\n        \n        # Split Leads\n        leads_dict = split_leads(raw_signals)\n        \n        requests = test_df[test_df['id'] == sample_id]\n        \n        for _, row in requests.iterrows():\n            lead_name = row['lead']\n            target_len = row['number_of_rows']\n            \n            # Select & Resample\n            if lead_name == 'II' and target_len > 2000:\n                sig = leads_dict['II_long']\n            else:\n                sig = leads_dict.get(lead_name, np.zeros(10))\n                \n            if len(sig) != target_len:\n                sig = resample(sig, target_len)\n            \n            # --- POST PROCESSING ---\n            # 1. High Pass Filter (Remove Drift)\n            sig = clinical_filter(sig, fs=row['fs'])\n            \n            # 2. Savitzky-Golay (Remove High-Freq Noise)\n            sig = savgol_filter(sig, window_length=7, polyorder=2)\n            \n            # 3. Median Center (Final Alignment)\n            sig = sig - np.median(sig)\n            \n            row_ids = [f\"{sample_id}_{i}_{lead_name}\" for i in range(target_len)]\n            submission_rows.extend(zip(row_ids, sig))\n            \n    except Exception as e:\n        reqs = test_df[test_df['id'] == sample_id]\n        for _, row in reqs.iterrows():\n            t_len = row['number_of_rows']\n            row_ids = [f\"{sample_id}_{i}_{row['lead']}\" for i in range(t_len)]\n            submission_rows.extend(zip(row_ids, np.zeros(t_len)))\n\nprint(\"💾 Saving Submission...\")\npd.DataFrame(submission_rows, columns=['id', 'value']).to_csv('submission.csv', index=False)\nprint(\"✅ Done!\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-03T14:39:01.965147Z","iopub.execute_input":"2026-01-03T14:39:01.966113Z","iopub.status.idle":"2026-01-03T14:39:12.196451Z","shell.execute_reply.started":"2026-01-03T14:39:01.966077Z","shell.execute_reply":"2026-01-03T14:39:12.195518Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}