{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.12.12","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"}],"dockerImageVersionId":31260,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## Imports\n","metadata":{}},{"cell_type":"code","source":"import os\nimport cv2\nimport gc\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.amp import autocast, GradScaler\nfrom sklearn.model_selection import train_test_split\nimport timm\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nfrom tqdm import tqdm","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Configuration\n","metadata":{}},{"cell_type":"code","source":"CONFIG = {\n    \"seed\": 42,\n    \"img_size\": (512, 1024),\n    \"batch_size\": 8,\n    \"epochs\": 50,\n    \"lr\": 3e-4,\n    \"backbone\": \"tf_efficientnet_b3\",\n    \"num_leads\": 12,\n    \"output_len\": 5000,\n    \"device_type\": \"cuda\" if torch.cuda.is_available() else \"cpu\",\n    \"device\": torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\"),\n    \"base_path\": \"/kaggle/input/physionet-ecg-image-digitization\",\n}\n\n\ndef seed_everything(seed):\n    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n\n\nseed_everything(CONFIG[\"seed\"])","metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Preprocessing\n","metadata":{}},{"cell_type":"code","source":"def remove_grid_and_clean(image):\n    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)\n    thresh_val, binary = cv2.threshold(\n        gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU\n    )\n\n    kernel = np.ones((2, 2), np.uint8)\n    binary = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel)\n    edges = cv2.Canny(gray, 50, 150)\n\n    combined = np.stack([gray, binary, edges], axis=-1)\n    return combined","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Dataset Class\n","metadata":{}},{"cell_type":"code","source":"class ECGDataset(Dataset):\n    def __init__(self, df, img_dir, transform=None, mode=\"train\"):\n        self.df = df\n        self.img_dir = img_dir\n        self.transform = transform\n        self.mode = mode\n        self.leads = [\n            \"I\",\n            \"II\",\n            \"III\",\n            \"aVR\",\n            \"aVL\",\n            \"aVF\",\n            \"V1\",\n            \"V2\",\n            \"V3\",\n            \"V4\",\n            \"V5\",\n            \"V6\",\n        ]\n        self.exts = [\n            \"-0001\",\n            \"-0003\",\n            \"-0004\",\n            \"-0005\",\n            \"-0009\",\n            \"-0010\",\n            \"-0011\",\n            \"-0012\",\n        ]\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        base_id = str(row[\"id\"])\n\n        # Load Image\n        if self.mode == \"train\":\n            ext = np.random.choice(self.exts)\n            img_path = os.path.join(self.img_dir, base_id, f\"{base_id}{ext}.png\")\n            if not os.path.exists(img_path):\n                img_path = os.path.join(self.img_dir, base_id, f\"{base_id}-0001.png\")\n        else:\n            if self.mode == \"test\":\n                img_path = os.path.join(self.img_dir, f\"{base_id}.png\")\n            else:\n                img_path = os.path.join(self.img_dir, base_id, f\"{base_id}-0001.png\")\n\n        image = cv2.imread(img_path)\n        if image is None:\n            # Fallback black image\n            image = np.zeros(\n                (CONFIG[\"img_size\"][0], CONFIG[\"img_size\"][1], 3), dtype=np.uint8\n            )\n\n        image = cv2.resize(image, (CONFIG[\"img_size\"][1], CONFIG[\"img_size\"][0]))\n        image = remove_grid_and_clean(image)\n\n        if self.transform:\n            image = self.transform(image=image)[\"image\"]\n\n        # Return test data early\n        if self.mode == \"test\":\n            return {\"image\": image, \"id\": base_id}\n\n        # Load Signal\n        csv_path = os.path.join(self.img_dir, base_id, f\"{base_id}.csv\")\n        try:\n            sig_df = pd.read_csv(csv_path)\n            raw_signal = sig_df[self.leads].values\n\n            # Interpolate to fixed 5000 length\n            x_old = np.linspace(0, 1, len(raw_signal))\n            x_new = np.linspace(0, 1, CONFIG[\"output_len\"])\n            resampled_signal = np.zeros((CONFIG[\"output_len\"], 12), dtype=np.float32)\n\n            for i in range(12):\n                series = raw_signal[:, i]\n                if np.all(np.isnan(series)):\n                    resampled_signal[:, i] = 0.0\n                else:\n                    resampled_signal[:, i] = np.interp(\n                        x_new, x_old, np.nan_to_num(series)\n                    )\n\n            signal_tensor = torch.tensor(resampled_signal.T).float()\n            return {\"image\": image, \"signal\": signal_tensor}\n\n        except Exception:\n            dummy_sig = torch.zeros((12, CONFIG[\"output_len\"])).float()\n            return {\"image\": image, \"signal\": dummy_sig}","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Model Architecture\n","metadata":{}},{"cell_type":"code","source":"class ECGModel(nn.Module):\n    def __init__(self, backbone_name, num_leads=12, output_len=5000, pretrained=True):\n        super().__init__()\n        self.output_len = output_len\n\n        self.backbone = timm.create_model(\n            backbone_name, pretrained=pretrained, num_classes=0, global_pool=\"\"\n        )\n\n        # Get feature channel count\n        dummy = torch.randn(1, 3, CONFIG[\"img_size\"][0], CONFIG[\"img_size\"][1])\n        feat = self.backbone(dummy)\n        self.in_features = feat.shape[1]\n\n        # 1D Conversion Head\n        self.pool = nn.AdaptiveAvgPool2d((1, None))\n        self.conv1 = nn.Conv1d(self.in_features, 256, kernel_size=3, padding=1)\n        self.bn1 = nn.BatchNorm1d(256)\n        self.act = nn.ReLU()\n        self.conv2 = nn.Conv1d(256, num_leads, kernel_size=1)\n\n    def forward(self, x):\n        x = self.backbone(x)\n        x = self.pool(x).squeeze(2)\n        x = torch.nn.functional.interpolate(x, size=self.output_len, mode=\"linear\")\n        x = self.act(self.bn1(self.conv1(x)))\n        x = self.conv2(x)\n        return x","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Training Loop\n","metadata":{}},{"cell_type":"code","source":"def train_model():\n    train_df = pd.read_csv(f\"{CONFIG['base_path']}/train.csv\")\n    train_split, val_split = train_test_split(train_df, test_size=0.1, random_state=42)\n\n    aug = A.Compose([A.Normalize(), ToTensorV2()])\n\n    train_ds = ECGDataset(\n        train_split, f\"{CONFIG['base_path']}/train\", transform=aug, mode=\"train\"\n    )\n    val_ds = ECGDataset(\n        val_split, f\"{CONFIG['base_path']}/train\", transform=aug, mode=\"val\"\n    )\n\n    train_loader = DataLoader(\n        train_ds,\n        batch_size=CONFIG[\"batch_size\"],\n        shuffle=True,\n        num_workers=4,\n        pin_memory=True,\n    )\n    val_loader = DataLoader(\n        val_ds, batch_size=CONFIG[\"batch_size\"], shuffle=False, num_workers=4\n    )\n\n    model = ECGModel(CONFIG[\"backbone\"], output_len=CONFIG[\"output_len\"])\n    if torch.cuda.device_count() > 1:\n        model = nn.DataParallel(model)\n    model = model.to(CONFIG[\"device\"])\n\n    optimizer = optim.AdamW(model.parameters(), lr=CONFIG[\"lr\"])\n    scheduler = optim.lr_scheduler.OneCycleLR(\n        optimizer,\n        max_lr=CONFIG[\"lr\"],\n        steps_per_epoch=len(train_loader),\n        epochs=CONFIG[\"epochs\"],\n    )\n    criterion = nn.L1Loss()\n\n    scaler = GradScaler(CONFIG[\"device_type\"])\n\n    best_loss = float(\"inf\")\n\n    print(\"--- Training Start (with AMP) ---\")\n    for epoch in range(CONFIG[\"epochs\"]):\n        model.train()\n        train_loss = 0\n\n        for batch in tqdm(train_loader, desc=f\"Epoch {epoch+1}\"):\n            if batch is None:\n                continue\n            imgs = batch[\"image\"].to(CONFIG[\"device\"])\n            sigs = batch[\"signal\"].to(CONFIG[\"device\"])\n\n            optimizer.zero_grad()\n\n            with autocast(device_type=CONFIG[\"device_type\"]):\n                preds = model(imgs)\n                loss = criterion(preds, sigs)\n\n            scaler.scale(loss).backward()\n            scaler.step(optimizer)\n            scaler.update()\n\n            scheduler.step()\n            train_loss += loss.item()\n\n        # Validation\n        model.eval()\n        val_loss = 0\n        with torch.no_grad():\n            for batch in val_loader:\n                imgs = batch[\"image\"].to(CONFIG[\"device\"])\n                sigs = batch[\"signal\"].to(CONFIG[\"device\"])\n\n                with autocast(device_type=CONFIG[\"device_type\"]):\n                    preds = model(imgs)\n                    val_loss += criterion(preds, sigs).item()\n\n        avg_val = val_loss / len(val_loader)\n        print(f\"Epoch {epoch+1} | Val Loss: {avg_val:.5f}\")\n\n        if avg_val < best_loss:\n            best_loss = avg_val\n            torch.save(model.state_dict(), \"best_model.pth\")\n\n        torch.cuda.empty_cache()\n\n    return model","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Submission\n","metadata":{}},{"cell_type":"code","source":"def generate_submission(model_path=\"best_model.pth\"):\n    # Load Model\n    model = ECGModel(\n        CONFIG[\"backbone\"], output_len=CONFIG[\"output_len\"], pretrained=False\n    )\n    if torch.cuda.device_count() > 1:\n        model = nn.DataParallel(model)\n    model.load_state_dict(torch.load(model_path, map_location=CONFIG[\"device\"]))\n    model.to(CONFIG[\"device\"])\n    model.eval()\n\n    test_df = pd.read_csv(f\"{CONFIG['base_path']}/test.csv\")\n\n    meta_map = {}\n    for _, row in test_df.iterrows():\n        meta_map[(str(row[\"id\"]), row[\"lead\"])] = int(row[\"number_of_rows\"])\n\n    transform = A.Compose(\n        [\n            A.Resize(height=CONFIG[\"img_size\"][0], width=CONFIG[\"img_size\"][1]),\n            A.Normalize(),\n            ToTensorV2(),\n        ]\n    )\n\n    unique_ids_df = pd.DataFrame({\"id\": test_df[\"id\"].unique()})\n    test_ds = ECGDataset(\n        unique_ids_df, f\"{CONFIG['base_path']}/test\", transform=transform, mode=\"test\"\n    )\n    test_loader = DataLoader(\n        test_ds, batch_size=CONFIG[\"batch_size\"], shuffle=False, num_workers=4\n    )\n\n    leads = [\"I\", \"II\", \"III\", \"aVR\", \"aVL\", \"aVF\", \"V1\", \"V2\", \"V3\", \"V4\", \"V5\", \"V6\"]\n    submission_ids = []\n    submission_values = []\n\n    print(\"--- Generating Predictions ---\")\n    with torch.no_grad():\n        for batch in tqdm(test_loader):\n            imgs = batch[\"image\"].to(CONFIG[\"device\"])\n            batch_ids = batch[\"id\"]\n\n            preds = model(imgs).cpu().numpy()  # (B, 12, 5000)\n\n            for i, base_id in enumerate(batch_ids):\n                base_id = str(base_id)\n\n                for lead_idx, lead_name in enumerate(leads):\n                    raw_signal = preds[i, lead_idx, :]\n                    target_rows = meta_map.get((base_id, lead_name), 0)\n                    if target_rows == 0:\n                        continue\n\n                    # Interpolate to target row count\n                    x_curr = np.linspace(0, 1, CONFIG[\"output_len\"])\n                    x_targ = np.linspace(0, 1, target_rows)\n                    final_signal = np.interp(x_targ, x_curr, raw_signal)\n\n                    # Build ID strings efficiently\n                    row_ids = [f\"{base_id}_{r}_{lead_name}\" for r in range(target_rows)]\n                    submission_ids.extend(row_ids)\n                    submission_values.extend(final_signal)\n\n    print(\"... Saving CSV\")\n    sub_df = pd.DataFrame({\"id\": submission_ids, \"value\": submission_values})\n    sub_df.to_csv(\"submission.csv\", index=False)\n    print(\"Done.\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if __name__ == \"__main__\":\n    # 1. Train\n    trained_model = train_model()\n\n    # 2. Cleanup GPU memory\n    del trained_model\n    gc.collect()\n    torch.cuda.empty_cache()\n\n    # 3. Predict & Submit\n    generate_submission(\"best_model.pth\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}