{"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":"gpu","dataSources":[{"sourceId":97984,"databundleVersionId":14096757,"sourceType":"competition"},{"sourceId":13746387,"sourceType":"datasetVersion","datasetId":8747012}],"dockerImageVersionId":31193,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\n\nbase_path = '/kaggle/input/hengck23-submit-physionet'\n\nprint(\"Contents of base path:\")\nfor item in os.listdir(base_path):\n    full_path = os.path.join(base_path, item)\n    print(f\"  {item}/ \" if os.path.isdir(full_path) else f\"  {item}\")\n    \n    if os.path.isdir(full_path):\n        for sub_item in os.listdir(full_path):\n            sub_path = os.path.join(full_path, sub_item)\n            print(f\"    {sub_item}/\" if os.path.isdir(sub_path) else f\"    {sub_item}\")\n            \n            if os.path.isdir(sub_path):\n                for sub_sub in os.listdir(sub_path)[:10]:  # Limit to first 10\n                    print(f\"      {sub_sub}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-25T13:21:09.182391Z","iopub.execute_input":"2025-11-25T13:21:09.182728Z","iopub.status.idle":"2025-11-25T13:21:09.205965Z","shell.execute_reply.started":"2025-11-25T13:21:09.182705Z","shell.execute_reply":"2025-11-25T13:21:09.205220Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Advanced PhysioNet ECG Image Digitization with ARCO/RCI Denoising\n# Based on @hengck23's 3-stage pipeline with validation framework\n# ARCO/RCI denoising applied post-Stage 2 to refine 12-lead signals\n\nimport subprocess\nimport sys\nimport os\nimport math\nimport cv2\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom fractions import Fraction\nfrom scipy import signal as scipy_signal\nimport inspect\nimport warnings\n\nwarnings.filterwarnings('ignore')\n\n# ===========================================\n# 1. Install dependencies from hengck23's setup\n# ===========================================\nHENGCK23_PATH = '/kaggle/input/hengck23-submit-physionet/hengck23-submit-physionet'\nSETUP_PATH = f'{HENGCK23_PATH}/setup'\n\nprint(\"Installing dependencies...\")\ntry:\n    subprocess.check_call([\n        sys.executable, '-m', 'pip', 'install', '--quiet',\n        f'{SETUP_PATH}/connected_components_3d-3.26.1-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl'\n    ])\n    print(\"Dependencies installed.\")\nexcept Exception as e:\n    print(f\"Dependency installation skipped or failed (might be already installed): {e}\")\n\n# Add hengck23's code to path\nsys.path.insert(0, HENGCK23_PATH)\n\n# ===========================================\n# 2. Dynamic Import & Discovery\n# ===========================================\nprint(\"Discovering module contents...\")\nimport stage0_model\nimport stage1_model\nimport stage2_model\nimport stage0_common\nimport stage1_common\nimport stage2_common\n\ndef find_net_class(module):\n    \"\"\"Find the main neural network class in a module.\"\"\"\n    module_name = module.__name__\n    candidates = []\n    for name, obj in inspect.getmembers(module):\n        if inspect.isclass(obj) and issubclass(obj, nn.Module) and obj is not nn.Module:\n            if hasattr(obj, '__module__') and obj.__module__ == module_name:\n                candidates.append((name, obj))\n    \n    if not candidates:\n        for name, obj in inspect.getmembers(module):\n            if inspect.isclass(obj) and issubclass(obj, nn.Module) and 'Net' in name:\n                candidates.append((name, obj))\n    \n    if candidates:\n        for name, obj in candidates:\n            if 'Net' in name: return obj\n        return candidates[0][1]\n    return None\n\ndef find_function(module, possible_names):\n    \"\"\"Find a function by trying multiple possible names.\"\"\"\n    for name in possible_names:\n        if hasattr(module, name): return getattr(module, name)\n    return None\n\n# Discover Net classes\nStage0Net = find_net_class(stage0_model)\nStage1Net = find_net_class(stage1_model)\nStage2Net = find_net_class(stage2_model)\n\n# Discover helper functions\nimage_to_batch = find_function(stage0_common, ['image_to_batch', 'to_batch', 'make_batch'])\noutput_to_predict = find_function(stage0_common, ['output_to_predict', 'to_predict', 'predict'])\nnormalise_by_homography = find_function(stage0_common, ['normalise_by_homography', 'normalize_by_homography'])\nstage1_output_to_predict = find_function(stage1_common, ['stage1_output_to_predict', 'output_to_predict'])\nrectify_image = find_function(stage1_common, ['rectify_image', 'rectify', 'warp_image'])\npixel_to_series = find_function(stage2_common, ['pixel_to_series', 'to_series'])\nfilter_series_by_limits = find_function(stage2_common, ['filter_series_by_limits', 'filter_series'])\n\n# Check critical components\nif None in [Stage0Net, Stage1Net, Stage2Net, image_to_batch, rectify_image, pixel_to_series]:\n    raise RuntimeError(\"Critical models or functions could not be dynamically discovered.\")\n\n# ===========================================\n# 3. Configuration & Model Loading\n# ===========================================\nKAGGLE_DIR = '/kaggle/input/physionet-ecg-image-digitization'\nWEIGHT_DIR = f'{HENGCK23_PATH}/weight'\nDEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\nFLOAT_TYPE = torch.float32\nLEADS = ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6']\nTYPE_IDS = ['0001', '0003', '0004', '0005', '0006', '0009', '0010', '0011', '0012']\nBASE_FS = 500\n\ndef load_net(model, path):\n    checkpoint = torch.load(path, map_location='cpu')\n    state_dict = checkpoint.get('state_dict', checkpoint.get('model_state_dict', checkpoint))\n    model.load_state_dict(state_dict, strict=False)\n    return model\n\nprint('Loading models...')\nstage0_net = load_net(Stage0Net(pretrained=False), f'{WEIGHT_DIR}/stage0-last.checkpoint.pth').to(DEVICE).eval()\nstage1_net = load_net(Stage1Net(pretrained=False), f'{WEIGHT_DIR}/stage1-last.checkpoint.pth').to(DEVICE).eval()\nstage2_net = load_net(Stage2Net(pretrained=False), f'{WEIGHT_DIR}/stage2-00005810.checkpoint.pth').to(DEVICE).eval()\nprint('Models loaded successfully!')\n\n# ===========================================\n# 4. ARCO/RCI Denoising Implementation\n# ===========================================\ndef make_rational_frequencies(max_q: int = 30) -> list:\n    arcs = set()\n    for q in range(2, max_q + 1):\n        for a in range(1, q):\n            if math.gcd(a, q) == 1:\n                f = Fraction(a, q)\n                if f > Fraction(1, 2): f = 1 - f\n                arcs.add(f)\n    return sorted(arcs, key=float)\n\ndef _bin_frequencies(n: int, device):\n    return torch.arange(0, n // 2 + 1, device=device, dtype=torch.float32) / float(n)\n\ndef gaussian_arc_weights(n: int, centers, q_vals, sigma_factor=0.5):\n    device = centers.device\n    freqs = _bin_frequencies(n, device)\n    sigma = sigma_factor / (q_vals ** 2 + 1e-6)\n    diff = freqs[None, :] - centers[:, None]\n    W = torch.exp(-0.5 * (diff / (sigma[:, None] + 1e-9)) ** 2)\n    if W.shape[1] > 0: W[:, 0] = 0.0\n    return W / (W.sum(dim=0, keepdim=True) + 1e-9)\n\ndef power_spectrum(x, use_hann=True):\n    x = x.float()\n    B, T = x.shape\n    if use_hann:\n        w = torch.hann_window(T, device=x.device, dtype=torch.float32).unsqueeze(0)\n        xw = x * w\n    else:\n        xw = x\n    X = torch.fft.rfft(xw, dim=1)\n    P = (X.real ** 2 + X.imag ** 2)\n    return P / (P.sum(dim=1, keepdim=True) + 1e-12)\n\nclass ArcRCI:\n    def __init__(self, max_q=30, sigma_factor=0.5, device='cpu'):\n        fracs = make_rational_frequencies(max_q)\n        self.centers = torch.tensor([float(f) for f in fracs], dtype=torch.float32, device=device)\n        self.q_vals = torch.tensor([f.denominator for f in fracs], dtype=torch.float32, device=device)\n        self.sigma_factor = sigma_factor\n\n    def rci(self, x, use_hann=True):\n        P = power_spectrum(x, use_hann)\n        W = gaussian_arc_weights(x.shape[1], self.centers, self.q_vals, self.sigma_factor)\n        arcogram = torch.matmul(P, W.t())\n        return torch.clamp(arcogram.sum(dim=1), 0.0, 1.0)\n\ndef denoised_multi_lead_arco(signals_dict, fs=500, window_size=200, step=50, max_q=30):\n    lead_names = list(signals_dict.keys())\n    signals_list = [signals_dict[name] for name in lead_names]\n    L_max = max(len(s) for s in signals_list)\n    sigs = np.stack([np.pad(s, (0, L_max - len(s)), mode='edge') if len(s) < L_max else s[:L_max] for s in signals_list], axis=0).astype(np.float32)\n    \n    n_leads, n_samples = sigs.shape\n    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n    arc = ArcRCI(max_q, device=device)\n    \n    kern_size = min(window_size // 4, 51)\n    kern = np.ones(kern_size) / kern_size\n    smooth_signals = np.array([np.convolve(lead, kern, mode='same') for lead in sigs])\n    refined = np.zeros_like(sigs)\n    b, a = scipy_signal.butter(4, 40 / (fs / 2), 'low')\n    \n    for i in range(n_leads):\n        lead = sigs[i]\n        lead_t = torch.from_numpy(lead).unsqueeze(0).to(device)\n        rci_profile = []\n        end = 0\n        for start in range(0, n_samples - window_size + 1, step):\n            end = start + window_size\n            rci = arc.rci(lead_t[:, start:end]).item()\n            rci_profile.append(rci)\n            denom = max(rci_profile) if rci_profile else 1.0\n            w = min(1.0, rci / denom) if denom > 0 else 0.5\n            refined[i, start:end] = w * lead[start:end] + (1 - w) * smooth_signals[i, start:end]\n        \n        if end < n_samples:\n            last_w = min(1.0, rci_profile[-1] / max(rci_profile)) if rci_profile else 0.5\n            refined[i, end:] = last_w * lead[end:] + (1 - last_w) * smooth_signals[i, end:]\n        \n        refined[i] = scipy_signal.filtfilt(b, a, refined[i])\n    \n    return {name: refined[i, :len(signals_dict[name])] for i, name in enumerate(lead_names)}\n\n# ===========================================\n# 5. Processing Functions (With Fixes)\n# ===========================================\ndef add_dummy_targets_stage0(batch, device):\n    B, C, H, W = batch['image'].shape\n    if 'marker' not in batch: batch['marker'] = torch.zeros((B, H, W), device=device, dtype=torch.long)\n    if 'orientation' not in batch: batch['orientation'] = torch.zeros((B,), device=device, dtype=torch.float32)\n    return batch\n\ndef add_dummy_targets_stage1(batch, device):\n    B, C, H, W = batch['image'].shape\n    if 'marker' not in batch: batch['marker'] = torch.zeros((B, H, W), device=device, dtype=torch.long)\n    if 'gridpoint' not in batch: batch['gridpoint'] = torch.zeros((B, 1, H, W), device=device, dtype=torch.float32)\n    for k in ['gridhline', 'gridvline']:\n        if k not in batch: batch[k] = torch.zeros((B, H, W), device=device, dtype=torch.long)\n    return batch\n\ndef add_dummy_targets_stage2(batch, device):\n    B, C, H, W = batch['image'].shape\n    if 'pixel' not in batch: batch['pixel'] = torch.zeros((B, 4, H, W), device=device, dtype=torch.float32)\n    return batch\n\ndef process_stage0(image, stage0_net, device='cuda', float_type=None):\n    batch = image_to_batch(image)\n    batch['image'] = batch['image'].to(device)\n    batch = add_dummy_targets_stage0(batch, device) \n    \n    with torch.amp.autocast('cuda', dtype=float_type):\n        with torch.no_grad():\n            output = stage0_net(batch)\n            rotated, keypoint = output_to_predict(image, batch, output)\n            normalised, keypoint, homography = normalise_by_homography(rotated, keypoint)\n    torch.cuda.empty_cache()\n    return normalised, keypoint, homography\n\ndef process_stage1(normalised, stage1_net, device='cuda', float_type=None, num_tta=4):\n    batch = {'image': torch.from_numpy(np.ascontiguousarray(normalised.transpose(2, 0, 1))).unsqueeze(0)}\n    batch = add_dummy_targets_stage1(batch, device)\n    \n    crop = batch['image'].clone().byte().to(device)\n    trial_batch = {'image': crop}\n    trial_batch = add_dummy_targets_stage1(trial_batch, device)\n    \n    with torch.amp.autocast('cuda', dtype=float_type):\n        with torch.no_grad():\n            output = stage1_net(trial_batch)\n            gridpoint_xy, more = stage1_output_to_predict(normalised, batch, output)\n            rectified = rectify_image(normalised, gridpoint_xy)\n    torch.cuda.empty_cache()\n    return rectified, gridpoint_xy\n\ndef process_stage2(rectified, signal_length, stage2_net, device='cuda', float_type=None):\n    x0, x1, y0, y1 = 0, 2176, 0, 1696\n    zero_mv = [703.5, 987.5, 1271.5, 1531.5]\n    mv_to_pixel = 79.0\n    t0, t1 = 118, 2080\n    \n    crop = rectified[y0:y1, x0:x1]\n    batch = {'image': torch.from_numpy(np.ascontiguousarray(crop.transpose(2, 0, 1))).unsqueeze(0)}\n    batch['image'] = batch['image'].to(device)\n    batch = add_dummy_targets_stage2(batch, device)\n    \n    with torch.amp.autocast('cuda', dtype=float_type):\n        with torch.no_grad():\n            output = stage2_net(batch)\n            pixel = output['pixel'].data.cpu().numpy()[0]\n            \n    series_in_pixel = pixel_to_series(pixel[..., t0:t1], zero_mv, signal_length)\n    series = (np.array(zero_mv).reshape(4, 1) - series_in_pixel) / mv_to_pixel\n    series = filter_series_by_limits(series)\n    torch.cuda.empty_cache()\n    return series\n\ndef extract_leads_from_series(series):\n    \"\"\"Standard extraction from 4-row series.\"\"\"\n    length = series.shape[1]\n    quarter = length // 4\n    return {\n        'I': series[0, 0:quarter], 'aVR': series[0, quarter:2*quarter], 'V1': series[0, 2*quarter:3*quarter], 'V4': series[0, 3*quarter:],\n        'II': series[3, :], 'aVL': series[1, quarter:2*quarter], 'V2': series[1, 2*quarter:3*quarter], 'V5': series[1, 3*quarter:],\n        'III': series[2, 0:quarter], 'aVF': series[2, quarter:2*quarter], 'V3': series[2, 2*quarter:3*quarter], 'V6': series[2, 3*quarter:]\n    }\n\n# ===========================================\n# 6. Main Pipeline (Integrated)\n# ===========================================\ndef process_image_pipeline(image, lead_records, stage0_net, stage1_net, stage2_net, use_arco=True):\n    # Data Integrity: Ensure lead_records works for both schema types\n    # If we have the \"compact\" schema (fs, sig_len), we don't filter by lead='II'\n    \n    is_detailed_schema = 'lead' in lead_records.columns\n    \n    # Stage 0\n    normalised, _, _ = process_stage0(image, stage0_net, DEVICE, FLOAT_TYPE)\n    # Stage 1\n    rectified, _ = process_stage1(normalised, stage1_net, DEVICE, FLOAT_TYPE)\n    \n    # Stage 2 Prep: Determine Signal Length\n    signal_length = 5000 # Default fallback\n    \n    if is_detailed_schema:\n        lead_i_record = lead_records[lead_records['lead'] == 'II']\n        if not lead_i_record.empty:\n            signal_length = int(lead_i_record.iloc[0]['number_of_rows'])\n    else:\n        # Compact schema handling: use 'sig_len' if available\n        if 'sig_len' in lead_records.columns:\n            signal_length = int(lead_records.iloc[0]['sig_len'])\n    \n    # Stage 2\n    series = process_stage2(rectified, signal_length, stage2_net, DEVICE, FLOAT_TYPE)\n    predicted_leads = extract_leads_from_series(series)\n    \n    # ARCO/RCI Denoising\n    if use_arco:\n        predicted_leads = denoised_multi_lead_arco(predicted_leads, fs=BASE_FS)\n        \n    return predicted_leads, normalised, rectified\n\n# ===========================================\n# 7. Helpers for Submission/Scoring\n# ===========================================\ndef build_submission_rows(predicted_leads, image_id, lead_records):\n    rows = []\n    \n    # Handle both schemas\n    if 'lead' in lead_records.columns:\n        # Detailed schema (one row per lead)\n        iterator = lead_records.iterrows()\n    else:\n        # Compact schema: Create synthetic iterator for 12 leads\n        # sig_len might be in the dataframe\n        sig_len = int(lead_records.iloc[0]['sig_len']) if 'sig_len' in lead_records.columns else 5000\n        # Create a generator of (index, row-like dict)\n        synthetic_data = []\n        for lead in LEADS:\n            synthetic_data.append({'lead': lead, 'number_of_rows': sig_len})\n        iterator = enumerate(synthetic_data) # enumerating list of dicts acts like iterrows somewhat\n\n    for _, row in iterator:\n        # Access keys safely (row can be Series or dict)\n        lead = row['lead']\n        num_rows = int(row['number_of_rows'])\n        \n        sig = predicted_leads.get(lead, np.zeros(num_rows))\n        \n        if len(sig) != num_rows:\n            sig = np.interp(np.linspace(0, 1, num_rows), np.linspace(0, 1, len(sig)), sig)\n            \n        for r in range(num_rows):\n            rows.append({'id': f\"{image_id}_{r}_{lead}\", 'value': float(sig[r])})\n    return rows\n\ndef load_ground_truth(image_id, kaggle_dir):\n    path = f'{kaggle_dir}/train/{image_id}/{image_id}.csv'\n    if not os.path.exists(path): return {}\n    df = pd.read_csv(path)\n    return {lead: df[lead].ffill().bfill().values for lead in LEADS if lead in df.columns}\n\ndef score(sol_df, sub_df):\n    merged = sol_df.merge(sub_df, on='id', suffixes=('_true', '_pred'))\n    t, p = merged['value_true'].values, merged['value_pred'].values\n    noise_var = np.var(t - p)\n    return 10 * np.log10(np.var(t) / noise_var) if noise_var > 0 else 0.0\n\n# ===========================================\n# 8. Inference & Validation Loops\n# ===========================================\nprint('\\n' + '='*40 + '\\nINFERENCE\\n' + '='*40)\ntest_path = f'{KAGGLE_DIR}/test.csv'\nif os.path.exists(test_path):\n    test_df = pd.read_csv(test_path)\n    test_df['id'] = test_df['id'].astype(str)\n    image_ids = test_df['id'].unique().tolist()\n    submission_rows = []\n    \n    for n, image_id in enumerate(image_ids):\n        print(f'\\rProcessing {n+1}/{len(image_ids)}: {image_id}', end='', flush=True)\n        filename = f'{KAGGLE_DIR}/test/{image_id}.png'\n        if not os.path.exists(filename): continue\n        \n        try:\n            image = cv2.imread(filename, cv2.IMREAD_COLOR_RGB)\n            lead_records = test_df[test_df['id'] == image_id].copy()\n            predicted_leads, _, _ = process_image_pipeline(image, lead_records, stage0_net, stage1_net, stage2_net)\n            if predicted_leads:\n                submission_rows.extend(build_submission_rows(predicted_leads, image_id, lead_records))\n        except Exception as e:\n            print(f\" Error: {e}\")\n\n    if submission_rows:\n        pd.DataFrame(submission_rows).to_csv('submission.csv', index=False)\n        print('\\nSubmission saved to submission.csv')\n\nprint('\\n' + '='*40 + '\\nVALIDATION\\n' + '='*40)\ntrain_path = f'{KAGGLE_DIR}/train.csv'\nif not os.getenv('KAGGLE_IS_COMPETITION_RERUN') and os.path.exists(train_path):\n    train_df = pd.read_csv(train_path)\n    \n    # --- ADAPTER: Handle Compact Schema ---\n    if 'lead' not in train_df.columns and 'sig_len' in train_df.columns:\n        print(\"Notice: Detected 'compact' train.csv schema (['id', 'fs', 'sig_len']). Using adapter.\")\n        # We don't need to explode the dataframe here, we just need to ensure\n        # process_image_pipeline and build_submission_rows handle the missing 'lead' column gracefully\n        # which we updated above.\n    elif 'lead' not in train_df.columns:\n        print(f\"FATAL ERROR: Unknown schema. Columns found: {list(train_df.columns)}\")\n        # Stop validation if schema is completely unknown\n        train_df = pd.DataFrame() \n\n    if not train_df.empty:\n        train_df['id'] = train_df['id'].astype(str)\n        val_ids = train_df['id'].unique()[:5].tolist() \n        scores = []\n        \n        for image_id in val_ids:\n            lead_records = train_df[train_df['id'] == image_id].copy()\n            if lead_records.empty: continue\n            \n            found_file = False\n            for type_id in TYPE_IDS:\n                filename = f'{KAGGLE_DIR}/train/{image_id}/{image_id}-{type_id}.png'\n                if os.path.exists(filename):\n                    print(f\"Validating {image_id} ({type_id})...\")\n                    image = cv2.imread(filename, cv2.IMREAD_COLOR_RGB)\n                    \n                    # Pipeline handles schema differences internally now\n                    preds, _, _ = process_image_pipeline(image, lead_records, stage0_net, stage1_net, stage2_net)\n                    \n                    if preds:\n                        gt = load_ground_truth(image_id, KAGGLE_DIR)\n                        # We need a 'lead_records' that works for building rows. \n                        # If we have the compact schema, we need to pass a compliant structure or let the function handle it.\n                        # The updated build_submission_rows handles the compact schema by checking columns.\n                        \n                        # Create solution rows (Ground Truth)\n                        # Ground truth loading might yield different keys than preds if file is partial\n                        sol = pd.DataFrame(build_submission_rows(gt, image_id, lead_records))\n                        \n                        # Create submission rows (Prediction)\n                        sub = pd.DataFrame(build_submission_rows(preds, image_id, lead_records))\n                        \n                        if not sol.empty and not sub.empty:\n                            s = score(sol, sub)\n                            scores.append(s)\n                            print(f\" -> SNR: {s:.2f} dB\")\n                    found_file = True\n                    break\n            if not found_file: print(f\"No image found for {image_id}\")\n\n        if scores: print(f\"\\nAverage Validation SNR: {np.mean(scores):.2f} dB\")\n\nprint('\\nDone!')","metadata":{"_uuid":"aa383469-aebc-4fc7-8faf-a657138ee73b","_cell_guid":"1de96977-61d7-4d2a-a09f-eaac90d2a7d9","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-11-25T13:21:09.295782Z","iopub.execute_input":"2025-11-25T13:21:09.296066Z","iopub.status.idle":"2025-11-25T13:21:47.351260Z","shell.execute_reply.started":"2025-11-25T13:21:09.296043Z","shell.execute_reply":"2025-11-25T13:21:47.350503Z"}},"outputs":[],"execution_count":null}]}