{"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":"none","dataSources":[{"sourceId":97984,"databundleVersionId":14096757,"sourceType":"competition"},{"sourceId":14433542,"sourceType":"datasetVersion","datasetId":9219202}],"dockerImageVersionId":31239,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"\"\"\"\nGenerate 2x1 resolution pixel mask labels for Stage D training.\n\nResolution: 4400 x 1696 (2x width, 1x height, cropped to multiple of 32)\n\"\"\"\n\nimport numpy as np\nimport pandas as pd\nimport cv2\nimport matplotlib.pyplot as plt\nfrom scipy.interpolate import PchipInterpolator\n\n# Output dimensions\nOUTPUT_WIDTH = 4400\nFULL_HEIGHT = 1700\nOUTPUT_HEIGHT = 1696  # (1700 // 32) * 32\nY_OFFSET = 2  # Pixels cropped from top of 1700\n\n# Signal extraction parameters\nT0 = 236   # 118 * 2\nT1 = 4161  # 2080 * 2 + 1 (corrected for stretch)\nSIGNAL_WIDTH = 3925  # T1 - T0\n\n# Baseline y positions (from 2x2 values / 2, includes Y_OFFSET and drift correction)\n# 2x2: [1403, 1972, 2540, 3060] / 2 = [701.5, 986.0, 1270.0, 1530.0]\nZERO_MV = [701.5, 986.0, 1270.0, 1530.0]\nMV_TO_PIXEL = 79.0\n\nLEAD_NAMES = ['I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6']\n\n# Lead layout: (row_index, start_fraction, end_fraction)\nLEAD_LAYOUT = {\n    'I':    (0, 0.00, 0.25),\n    'aVR':  (0, 0.25, 0.50),\n    'V1':   (0, 0.50, 0.75),\n    'V4':   (0, 0.75, 1.00),\n    'II':   (1, 0.00, 0.25),\n    'aVL':  (1, 0.25, 0.50),\n    'V2':   (1, 0.50, 0.75),\n    'V5':   (1, 0.75, 1.00),\n    'III':  (2, 0.00, 0.25),\n    'aVF':  (2, 0.25, 0.50),\n    'V3':   (2, 0.50, 0.75),\n    'V6':   (2, 0.75, 1.00),\n}\n\n\ndef load_signals_from_csv(csv_path):\n    \"\"\"Load ECG signals from competition ground truth CSV file.\"\"\"\n    df = pd.read_csv(csv_path)\n\n    total_length = len(df)\n    quarter_length = total_length // 4\n\n    signals = {}\n\n    for col in df.columns:\n        if col not in LEAD_NAMES:\n            continue\n\n        values = df[col].values\n\n        if col == 'II':\n            signals['II'] = np.array(values[:quarter_length])\n            signals['fullII'] = np.array(values)\n        else:\n            valid_values = values[~np.isnan(values)]\n            if len(valid_values) > 0:\n                signals[col] = np.array(valid_values)\n\n    return signals\n\n\ndef interpolate_signal_to_pixels(signal, n_pixels):\n    \"\"\"Interpolate signal to pixel resolution using PCHIP.\"\"\"\n    n_samples = len(signal)\n    x_samples = np.arange(n_samples)\n    x_pixels = np.linspace(0, n_samples - 1, n_pixels)\n    interpolator = PchipInterpolator(x_samples, signal)\n    return interpolator(x_pixels)\n\n\ndef set_subpixel_y(mask, row, x, y_float):\n    \"\"\"Set two-pixel subpixel y position in mask.\"\"\"\n    if x < 0 or x >= mask.shape[2]:\n        return\n\n    y_floor = int(np.floor(y_float))\n    y_ceil = y_floor + 1\n    frac = y_float - y_floor\n\n    if 0 <= y_floor < mask.shape[1]:\n        mask[row, y_floor, x] = 1.0 - frac\n    if 0 <= y_ceil < mask.shape[1]:\n        mask[row, y_ceil, x] = frac\n\n\ndef generate_pixel_mask(csv_path):\n    \"\"\"Generate pixel mask from ground truth CSV.\"\"\"\n    signals = load_signals_from_csv(csv_path)\n    pixel_mask = np.zeros((4, OUTPUT_HEIGHT, OUTPUT_WIDTH), dtype=np.float32)\n\n    for lead_name, (row, x_start_frac, x_end_frac) in LEAD_LAYOUT.items():\n        if lead_name not in signals:\n            continue\n\n        signal = signals[lead_name]\n        if len(signal) == 0:\n            continue\n\n        x_start = T0 + int(x_start_frac * SIGNAL_WIDTH)\n        x_end = T0 + int(x_end_frac * SIGNAL_WIDTH)\n        n_pixels = x_end - x_start\n\n        signal_interp = interpolate_signal_to_pixels(signal, n_pixels)\n\n        for i, val in enumerate(signal_interp):\n            x = x_start + i\n            y = ZERO_MV[row] - val * MV_TO_PIXEL\n            set_subpixel_y(pixel_mask, row, x, y)\n\n    if 'fullII' in signals:\n        signal = signals['fullII']\n        if len(signal) > 0:\n            row = 3\n            signal_interp = interpolate_signal_to_pixels(signal, SIGNAL_WIDTH)\n\n            for i, val in enumerate(signal_interp):\n                x = T0 + i\n                y = ZERO_MV[row] - val * MV_TO_PIXEL\n                set_subpixel_y(pixel_mask, row, x, y)\n\n    return pixel_mask\n\n\ndef verify_mask_alignment(pixel_mask, image_id, rectified_dir):\n    \"\"\"Overlay pixel mask on rectified image to verify alignment.\"\"\"\n    import os\n\n    colors = [\n        [255, 0, 0],\n        [0, 255, 0],\n        [0, 0, 255],\n        [255, 255, 0],\n    ]\n\n    rectified_path = os.path.join(rectified_dir, f\"{image_id}-0001.png\")\n    if not os.path.exists(rectified_path):\n        print(f\"Rectified image not found: {rectified_path}\")\n        return\n\n    rectified = cv2.imread(rectified_path)\n    rectified = cv2.cvtColor(rectified, cv2.COLOR_BGR2RGB)\n\n    # Resize to full dimensions, then crop to output height\n    rectified_scaled = cv2.resize(rectified, (OUTPUT_WIDTH, FULL_HEIGHT), interpolation=cv2.INTER_LINEAR)\n    rectified_cropped = rectified_scaled[Y_OFFSET:Y_OFFSET + OUTPUT_HEIGHT, :]\n\n    overlay = rectified_cropped.copy().astype(np.float32)\n\n    for row_idx in range(4):\n        mask = pixel_mask[row_idx] > 0.1\n        for c in range(3):\n            overlay[:, :, c] = np.where(mask, colors[row_idx][c], overlay[:, :, c])\n\n    plt.figure(figsize=(16, 6))\n    plt.imshow(overlay.astype(np.uint8))\n    plt.title(f'2x1 Mask Overlay: {image_id}')\n    plt.axis('off')\n    plt.show()\n\n\ndef generate_training_data(input_dir, output_dir, train_csv_path=None, rectified_dir=None, limit=None):\n    \"\"\"Generate training labels for all competition images.\"\"\"\n    import os\n    from glob import glob\n    from tqdm import tqdm\n\n    os.makedirs(output_dir, exist_ok=True)\n\n    if train_csv_path is None:\n        train_csv_path = os.path.join(os.path.dirname(input_dir), 'train.csv')\n\n    train_df = pd.read_csv(train_csv_path)\n    signal_info = train_df.set_index('id')[['fs', 'sig_len']].to_dict('index')\n\n    csv_files = sorted(glob(os.path.join(input_dir, '*', '*.csv')))\n\n    if limit:\n        csv_files = csv_files[:limit]\n\n    print(f\"Found {len(csv_files)} CSV files\")\n    print(f\"Output resolution: {OUTPUT_WIDTH} x {OUTPUT_HEIGHT} (2x1)\")\n    print(f\"Signal region: x=[{T0}, {T1}], width={SIGNAL_WIDTH}\")\n\n    verified_first = False\n    for csv_path in tqdm(csv_files, desc=\"Generating 2x1\"):\n        image_id = os.path.basename(csv_path).replace('.csv', '')\n        npz_path = os.path.join(output_dir, f\"{image_id}.npz\")\n\n        if os.path.exists(npz_path):\n            continue\n\n        try:\n            pixel_mask = generate_pixel_mask(csv_path)\n\n            if pixel_mask is None:\n                print(f\"Warning: Failed to generate for {image_id}\")\n                continue\n\n            np.savez_compressed(npz_path, pixel_mask=pixel_mask)\n\n            if not verified_first and rectified_dir:\n                verify_mask_alignment(pixel_mask, image_id, rectified_dir)\n                verified_first = True\n\n        except Exception as e:\n            print(f\"Error processing {csv_path}: {e}\")\n            import traceback\n            traceback.print_exc()\n            continue\n\n    print(f\"Done! Saved to {output_dir}\")\n\n\nif __name__ == '__main__':\n    INPUT_DIR = '/kaggle/input/physionet-ecg-image-digitization/train'\n    OUTPUT_DIR = '/kaggle/working'\n    RECTIFIED_DIR = '/kaggle/input/rectified-images-0001/rectified'\n\n    generate_training_data(INPUT_DIR, OUTPUT_DIR, rectified_dir=RECTIFIED_DIR)\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-01-09T13:17:54.065700Z","iopub.execute_input":"2026-01-09T13:17:54.066736Z"}},"outputs":[],"execution_count":null}]}