{"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":[{"sourceType":"competition","sourceId":97984,"databundleVersionId":14096757},{"sourceType":"datasetVersion","sourceId":13746387,"datasetId":8747012,"databundleVersionId":14496454},{"sourceType":"datasetVersion","sourceId":15884760,"datasetId":10185101,"databundleVersionId":16838585},{"sourceType":"kernelVersion","sourceId":312109075}],"dockerImageVersionId":31153,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nimport cv2\nimport os\nimport torch\nfrom tabulate import tabulate\nfrom tqdm import tqdm\nimport warnings\nwarnings.filterwarnings('ignore')\n\n# Install cc3d silently\ntry:\n    import cc3d\nexcept:\n    import subprocess\n    subprocess.run([\n        'pip', 'install', 'connected-components-3d', '--no-index',\n        '--find-links=file:///kaggle/input/hengck23-submit-physionet/hengck23-submit-physionet/setup',\n        '-q'\n    ], capture_output=True)\n    import cc3d\nfrom scipy.signal import savgol_filter\n\nDEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\nFLOAT_TYPE = torch.float16\n\nBASE_PATH = '/kaggle/input/hengck23-submit-physionet/hengck23-submit-physionet'\nKAGGLE_DIR = '/kaggle/input/physionet-ecg-image-digitization'\nWEIGHT_DIR = f'{BASE_PATH}/weight'\nOUT_DIR = '/kaggle/working/output-combined'\n\nos.makedirs(OUT_DIR, exist_ok=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-23T04:41:38.099423Z","iopub.execute_input":"2026-04-23T04:41:38.100039Z","iopub.status.idle":"2026-04-23T04:41:38.106071Z","shell.execute_reply.started":"2026-04-23T04:41:38.100016Z","shell.execute_reply":"2026-04-23T04:41:38.105098Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Load model","metadata":{}},{"cell_type":"markdown","source":"## Hengck Model","metadata":{}},{"cell_type":"code","source":"import sys\nsys.path.insert(0, BASE_PATH)\n\nfrom stage0_model import Net as Stage0Net\nfrom stage0_common import load_net as s0_load, image_to_batch, output_to_predict as s0_output, normalise_by_homography\nfrom stage1_model import Net as Stage1Net\nfrom stage1_common import load_net as s1_load, output_to_predict as s1_output, rectify_image\n\nDEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\n\n# 1. Load Heng's Frozen Weights\nstage0_net = s0_load(Stage0Net(pretrained=False), f'{BASE_PATH}/weight/stage0-last.checkpoint.pth').to(DEVICE).eval()\nstage1_net = s1_load(Stage1Net(pretrained=False), f'{BASE_PATH}/weight/stage1-last.checkpoint.pth').to(DEVICE).eval()\n\n# 2. Setup Directories\nTRAIN_IMG_DIR = '/kaggle/input/physionet-ecg-image-digitization/train'\nOUT_DIR_RECTIFIED = '/kaggle/working/train_rectified_crops'\nos.makedirs(OUT_DIR_RECTIFIED, exist_ok=True)\n\ntrain_df = pd.read_csv('/kaggle/input/physionet-ecg-image-digitization/train.csv')\nunique_train_ids = train_df['id'].unique()\n\nprint(f\"Flattening {len(unique_train_ids)} training images...\")\n\n# 3. Process and Save\nfor sample_id in tqdm(unique_train_ids):\n    try:\n        # Load Raw Image\n        img_path = os.path.join(TRAIN_IMG_DIR, str(sample_id), f\"{sample_id}-0001.png\")\n        image = cv2.imread(img_path, cv2.IMREAD_COLOR_RGB)\n        \n        # --- RUN STAGE 0 (De-rotate) ---\n        batch0 = image_to_batch(image)\n        with torch.no_grad():\n            with torch.amp.autocast('cuda'):\n                out0 = stage0_net(batch0)\n            rotated, keypoint = s0_output(image, batch0, out0)\n            normalised, _, _ = normalise_by_homography(rotated, keypoint)\n        \n        # --- RUN STAGE 1 (Flatten Grid) ---\n        batch1 = {'image': torch.from_numpy(np.ascontiguousarray(normalised.transpose(2, 0, 1))).unsqueeze(0)}\n        with torch.no_grad():\n            with torch.amp.autocast('cuda'):\n                out1 = stage1_net(batch1)\n            gridpoint_xy, _ = s1_output(normalised, batch1, out1)\n            rectified = rectify_image(normalised, gridpoint_xy)\n            \n        # --- CROP THE 4 ROWS (Based on Heng's coordinates) ---\n        # Heng uses x: 0 to 2176, y: 0 to 1696 for the bounding box of the grid\n        crop = rectified[0:1696, 0:2176]\n        \n        # Save the perfectly flat, cropped image\n        save_path = os.path.join(OUT_DIR_RECTIFIED, f\"{sample_id}_flat.png\")\n        cv2.imwrite(save_path, cv2.cvtColor(crop, cv2.COLOR_RGB2BGR))\n        \n    except Exception as e:\n        print(f\"Failed on {sample_id}: {e}\")\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-23T04:42:16.032258Z","iopub.execute_input":"2026-04-23T04:42:16.033029Z","iopub.status.idle":"2026-04-23T05:32:40.061173Z","shell.execute_reply.started":"2026-04-23T04:42:16.033001Z","shell.execute_reply":"2026-04-23T05:32:40.060422Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# !pip install segmentation-models-pytorch\n\n!pip install --no-index --find-links /kaggle/input/notebooks/ravnoorsingh101/smp-offline-wheels/smp-wheels segmentation-models-pytorch --no-deps","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-23T03:30:03.520518Z","iopub.execute_input":"2026-04-23T03:30:03.52097Z","iopub.status.idle":"2026-04-23T03:30:05.499945Z","shell.execute_reply.started":"2026-04-23T03:30:03.520935Z","shell.execute_reply":"2026-04-23T03:30:05.498905Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport segmentation_models_pytorch as smp\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\nmodel = smp.Unet(\n    encoder_name=\"efficientnet-b0\",\n    encoder_weights=None,\n    in_channels=3,\n    classes=1,\n    activation=None,\n    decoder_attention_type=\"scse\"\n)\n\n# Load your trained weights\nmodel.load_state_dict(torch.load('/kaggle/input/datasets/ravnoor000/physionet-unet-weights-2/best_attention_unet (4).pth', map_location=device))\nmodel = model.to(device)\nmodel.eval()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-23T03:31:35.372745Z","iopub.execute_input":"2026-04-23T03:31:35.373095Z","iopub.status.idle":"2026-04-23T03:31:35.587615Z","shell.execute_reply.started":"2026-04-23T03:31:35.373073Z","shell.execute_reply":"2026-04-23T03:31:35.586656Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Main pipeline ","metadata":{}}]}