{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","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":"none","dataSources":[{"sourceType":"competition","sourceId":97984,"databundleVersionId":14096757}],"dockerImageVersionId":31328,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# =============================\n# 🚀 FULL ECG PIPELINE (SAMPLED DATA VERSION)\n# =============================\n\nimport os\nimport cv2\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nimport torch.optim as optim\nfrom scipy.signal import savgol_filter\nimport matplotlib.pyplot as plt\nfrom sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score\n\n# =============================\n# CONFIG\n# =============================\n\n# 🔥 USING YOUR 5GB SAMPLED DATASET\nTRAIN_DIR = \"/kaggle/working/sample_dataset/train\"\nTEST_DIR = \"/kaggle/input/competitions/physionet-ecg-image-digitization/test\"\n\nIMG_SIZE = 256\nSIGNAL_LEN = 1024\nBATCH_SIZE = 8\nEPOCHS = 10   # reduce for Kaggle speed\nLR = 5e-5\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n# =============================\n# DATASET\n# =============================\n\nclass ECGDataset(Dataset):\n    def __init__(self, root_dir):\n        self.samples = []\n\n        for folder in os.listdir(root_dir):\n            path = os.path.join(root_dir, folder)\n            if not os.path.isdir(path):\n                continue\n\n            csv_path = os.path.join(path, f\"{folder}.csv\")\n            if not os.path.exists(csv_path):\n                continue\n\n            gt = pd.read_csv(csv_path)\n\n            # 🔹 Selecting ONE lead (sampling across leads)\n            signal = gt.select_dtypes(include=['number']).iloc[:, 0].values.astype(np.float32)\n\n            signal = signal[~np.isnan(signal)]\n            if len(signal) < 10:\n                continue\n\n            # 🔹 Noise removal\n            signal = savgol_filter(signal, 11, 2)\n\n            # 🔹 Normalization\n            signal = (signal - np.mean(signal)) / (np.std(signal) + 1e-8)\n\n            # 🔥 SIGNAL SAMPLING (VERY IMPORTANT)\n            # Convert variable-length signal → fixed 1024 length\n            signal = np.interp(\n                np.linspace(0, len(signal) - 1, SIGNAL_LEN),\n                np.arange(len(signal)),\n                signal\n            )\n\n            for f in os.listdir(path):\n                if f.endswith(\".png\"):\n                    self.samples.append((os.path.join(path, f), signal))\n\n    def __len__(self):\n        return len(self.samples)\n\n    def __getitem__(self, idx):\n        img_path, signal = self.samples[idx]\n\n        img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)\n        h = img.shape[0]\n\n        # 🔥 IMAGE SAMPLING → bottom 25%\n        img = img[3*h//4:h, :]\n\n        # 🔥 IMAGE RESAMPLING → fixed size\n        img = cv2.resize(img, (IMG_SIZE, IMG_SIZE))\n\n        img = img / 255.0\n\n        # 🔹 Edge enhancement\n        edges = cv2.Canny((img * 255).astype(np.uint8), 50, 150)\n        img = img + edges / 255.0\n\n        img = (img - 0.5) / 0.5\n        img = np.expand_dims(img, axis=0)\n\n        return torch.tensor(img, dtype=torch.float32), \\\n               torch.tensor(signal, dtype=torch.float32)\n\n# =============================\n# MODEL\n# =============================\n\nclass ECGModel(nn.Module):\n    def __init__(self):\n        super().__init__()\n\n        self.cnn = nn.Sequential(\n            nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),\n            nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),\n            nn.Conv2d(64, 128, 3, padding=1), nn.ReLU(),\n            nn.AdaptiveAvgPool2d((16, 16))\n        )\n\n        self.flatten = nn.Flatten(2)\n\n        encoder_layer = nn.TransformerEncoderLayer(\n            d_model=128,\n            nhead=4,\n            dim_feedforward=256,\n            dropout=0.1,\n            batch_first=True\n        )\n\n        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=2)\n\n        self.fc = nn.Sequential(\n            nn.Linear(128 * 256, 1024),\n            nn.ReLU(),\n            nn.Linear(1024, SIGNAL_LEN)\n        )\n\n    def forward(self, x):\n        x = self.cnn(x)\n        x = self.flatten(x)\n        x = x.permute(0, 2, 1)\n        x = self.transformer(x)\n        x = x.flatten(1)\n        x = self.fc(x)\n        return x\n\n# =============================\n# TRAINING SETUP\n# =============================\n\ndataset = ECGDataset(TRAIN_DIR)\nloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)\n\nmodel = ECGModel().to(DEVICE)\noptimizer = optim.Adam(model.parameters(), lr=LR)\nscheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)\n\nloss_history = []\n\nprint(f\"Training on {len(dataset)} samples\")\n\n# =============================\n# TRAIN LOOP\n# =============================\n\nfor epoch in range(EPOCHS):\n    model.train()\n    total_loss = 0\n\n    for imgs, signals in loader:\n        imgs = imgs.to(DEVICE)\n        signals = signals.to(DEVICE)\n\n        preds = model(imgs)\n\n        # 🔥 Clamp output\n        preds = torch.clamp(preds, -5, 5)\n\n        # 🔥 Weighted MSE (important for peaks)\n        mse = (preds - signals) ** 2\n        peak_weight = 1 + torch.abs(signals)\n        loss = (mse * peak_weight).mean()\n\n        optimizer.zero_grad()\n        loss.backward()\n\n        # 🔹 Gradient clipping\n        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n\n        optimizer.step()\n        total_loss += loss.item()\n\n    scheduler.step()\n    epoch_loss = total_loss / len(loader)\n    loss_history.append(epoch_loss)\n\n    print(f\"Epoch {epoch+1}/{EPOCHS} | Loss: {epoch_loss:.4f}\")\n\n# =============================\n# SAVE MODEL\n# =============================\n\ntorch.save(model.state_dict(), \"ecg_model.pth\")\nprint(\"Model saved!\")\n\n# =============================\n# LOSS CURVE\n# =============================\n\nplt.plot(loss_history)\nplt.title(\"Loss Curve\")\nplt.show()\n\n# =============================\n# TEST + SUBMISSION\n# =============================\n\ntest_meta = pd.read_csv(\"/kaggle/input/competitions/physionet-ecg-image-digitization/test.csv\")\nrows = []\n\nfor img_id in test_meta[\"id\"].unique():\n    img_path = os.path.join(TEST_DIR, f\"{img_id}.png\")\n\n    img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)\n    h = img.shape[0]\n\n    img = img[3*h//4:h, :]\n    img = cv2.resize(img, (IMG_SIZE, IMG_SIZE))\n    img = img / 255.0\n    img = (img - 0.5) / 0.5\n    img = np.expand_dims(img, axis=0)\n\n    inp = torch.tensor(img, dtype=torch.float32).unsqueeze(0).to(DEVICE)\n\n    with torch.no_grad():\n        pred = model(inp).cpu().numpy().flatten()\n\n    subset = test_meta[test_meta[\"id\"] == img_id]\n\n    for _, row in subset.iterrows():\n        length = row[\"number_of_rows\"]\n\n        # 🔥 OUTPUT SAMPLING → match required length\n        sig = np.interp(\n            np.linspace(0, len(pred)-1, length),\n            np.arange(len(pred)),\n            pred\n        )\n\n        # 🔹 Post-processing\n        sig = savgol_filter(sig, 11, 2)\n        sig = (sig - np.mean(sig)) / (np.std(sig) + 1e-8)\n\n        for i in range(len(sig)):\n            rows.append({\n                \"id\": f\"{img_id}_{i}_{row['lead']}\",\n                \"value\": float(sig[i])\n            })\n\nsubmission = pd.DataFrame(rows)\nsubmission.to_parquet(\"submission.parquet\", index=False)\n\nprint(\"✅ submission.parquet generated!\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-01T04:11:07.280794Z","iopub.execute_input":"2026-05-01T04:11:07.281616Z","iopub.status.idle":"2026-05-01T04:14:15.630657Z","shell.execute_reply.started":"2026-05-01T04:11:07.281579Z","shell.execute_reply":"2026-05-01T04:14:15.629609Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# -----------------------------\n# LOSS CURVE\n# -----------------------------\nimport matplotlib.pyplot as plt\n\nplt.figure(figsize=(8,5))\nplt.plot(loss_history, marker='o')\nplt.title(\"Training Loss Curve\")\nplt.xlabel(\"Epoch\")\nplt.ylabel(\"Loss\")\nplt.grid()import matplotlib.pyplot as plt\nimport numpy as np\n\ndef plot_prediction_vs_gt(model, dataset, device):\n    model.eval()\n\n    # Take one sample\n    img, gt = dataset[0]\n\n    # Predict\n    with torch.no_grad():\n        pred = model(img.unsqueeze(0).to(device)).cpu().numpy().flatten()\n\n    gt = gt.numpy()\n\n    # Optional: make them same length (just in case)\n    if len(pred) != len(gt):\n        pred = np.interp(\n            np.linspace(0, len(pred)-1, len(gt)),\n            np.arange(len(pred)),\n            pred\n        )\n\n    # Plot\n    plt.figure(figsize=(12,5))\n\n    plt.plot(gt, label=\"GT\", linewidth=2)      # Ground Truth (blue)\n    plt.plot(pred, label=\"Pred\", linewidth=1.5)  # Prediction (orange)\n\n    plt.title(\"Prediction vs Ground Truth\")\n    plt.xlabel(\"Time Steps\")\n    plt.ylabel(\"Amplitude\")\n\n    plt.legend()\n    plt.grid(alpha=0.3)\n\n    plt.show()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-01T04:15:55.530732Z","iopub.execute_input":"2026-05-01T04:15:55.531165Z","iopub.status.idle":"2026-05-01T04:15:55.670087Z","shell.execute_reply.started":"2026-05-01T04:15:55.531125Z","shell.execute_reply":"2026-05-01T04:15:55.669196Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import matplotlib.pyplot as plt\nimport numpy as np\n\ndef plot_prediction_vs_gt(model, dataset, device):\n    model.eval()\n\n    # Take one sample\n    img, gt = dataset[0]\n\n    # Predict\n    with torch.no_grad():\n        pred = model(img.unsqueeze(0).to(device)).cpu().numpy().flatten()\n\n    gt = gt.numpy()\n\n    # Optional: make them same length (just in case)\n    if len(pred) != len(gt):\n        pred = np.interp(\n            np.linspace(0, len(pred)-1, len(gt)),\n            np.arange(len(pred)),\n            pred\n        )\n\n    # Plot\n    plt.figure(figsize=(12,5))\n\n    plt.plot(gt, label=\"GT\", linewidth=2)      # Ground Truth (blue)\n    plt.plot(pred, label=\"Pred\", linewidth=1.5)  # Prediction (orange)\n\n    plt.title(\"Prediction vs Ground Truth\")\n    plt.xlabel(\"Time Steps\")\n    plt.ylabel(\"Amplitude\")\n\n    plt.legend()\n    plt.grid(alpha=0.3)\n\n    plt.show()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =============================\n# 🚀 FINAL ECG PIPELINE (NO INTERNAL SAMPLING)\n# =============================\n\nimport os\nimport cv2\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nimport torch.optim as optim\nimport matplotlib.pyplot as plt\nfrom sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score\n\n# =============================\n# CONFIG\n# =============================\n\nTRAIN_DIR = \"/kaggle/working/sample_dataset/train\"\nTEST_DIR  = \"/kaggle/input/competitions/physionet-ecg-image-digitization/test\"\n\nIMG_SIZE = 256\nSIGNAL_LEN = 1024   # expected by model\nBATCH_SIZE = 8\nEPOCHS = 10\nLR = 5e-5\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n# =============================\n# DATASET\n# =============================\n\nclass ECGDataset(Dataset):\n    def __init__(self, root_dir):\n        self.samples = []\n\n        for folder in os.listdir(root_dir):\n            path = os.path.join(root_dir, folder)\n            if not os.path.isdir(path):\n                continue\n\n            csv_path = os.path.join(path, f\"{folder}.csv\")\n            if not os.path.exists(csv_path):\n                continue\n\n            gt = pd.read_csv(csv_path)\n\n            # 🔹 Take first lead (already sampled dataset)\n            signal = gt.select_dtypes(include=['number']).iloc[:, 0].values.astype(np.float32)\n\n            signal = signal[~np.isnan(signal)]\n            if len(signal) < 10:\n                continue\n\n            # 🔹 Normalize only (NO sampling here)\n            signal = (signal - np.mean(signal)) / (np.std(signal) + 1e-8)\n\n            # 🔥 Only resize if needed (model requires fixed size)\n            if len(signal) != SIGNAL_LEN:\n                signal = np.interp(\n                    np.linspace(0, len(signal)-1, SIGNAL_LEN),\n                    np.arange(len(signal)),\n                    signal\n                )\n\n            for f in os.listdir(path):\n                if f.endswith(\".png\"):\n                    self.samples.append((os.path.join(path, f), signal))\n\n    def __len__(self):\n        return len(self.samples)\n\n    def __getitem__(self, idx):\n        img_path, signal = self.samples[idx]\n\n        img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)\n        h = img.shape[0]\n\n        # 🔹 Crop bottom region\n        img = img[3*h//4:h, :]\n\n        # 🔹 Resize image\n        img = cv2.resize(img, (IMG_SIZE, IMG_SIZE))\n\n        img = img / 255.0\n\n        # 🔹 Edge enhancement\n        edges = cv2.Canny((img*255).astype(np.uint8), 50, 150)\n        img = img + edges/255.0\n\n        img = (img - 0.5)/0.5\n        img = np.expand_dims(img, axis=0)\n\n        return torch.tensor(img, dtype=torch.float32), torch.tensor(signal, dtype=torch.float32)\n\n# =============================\n# MODEL\n# =============================\n\nclass ECGModel(nn.Module):\n    def __init__(self):\n        super().__init__()\n\n        self.cnn = nn.Sequential(\n            nn.Conv2d(1,32,3,padding=1), nn.ReLU(), nn.MaxPool2d(2),\n            nn.Conv2d(32,64,3,padding=1), nn.ReLU(), nn.MaxPool2d(2),\n            nn.Conv2d(64,128,3,padding=1), nn.ReLU(),\n            nn.AdaptiveAvgPool2d((16,16))\n        )\n\n        self.flatten = nn.Flatten(2)\n\n        encoder_layer = nn.TransformerEncoderLayer(\n            d_model=128, nhead=4, dim_feedforward=256,\n            dropout=0.1, batch_first=True\n        )\n\n        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=2)\n\n        self.fc = nn.Sequential(\n            nn.Linear(128*256,1024),\n            nn.ReLU(),\n            nn.Linear(1024,SIGNAL_LEN)\n        )\n\n    def forward(self,x):\n        x = self.cnn(x)\n        x = self.flatten(x)\n        x = x.permute(0,2,1)\n        x = self.transformer(x)\n        x = x.flatten(1)\n        x = self.fc(x)\n        return x\n\n# =============================\n# EVALUATION\n# =============================\n\ndef evaluate_model(model, dataset):\n    model.eval()\n    total_loss = 0\n\n    with torch.no_grad():\n        for i in range(min(50, len(dataset))):\n            img, gt = dataset[i]\n            img = img.unsqueeze(0).to(DEVICE)\n            gt = gt.to(DEVICE)\n\n            pred = model(img).squeeze()\n            loss = torch.mean((pred - gt)**2)\n            total_loss += loss.item()\n\n    return total_loss / min(50, len(dataset))\n\n\ndef evaluate_full_metrics(model, dataset):\n    model.eval()\n    preds_all, gts_all = [], []\n\n    with torch.no_grad():\n        for i in range(len(dataset)):\n            img, gt = dataset[i]\n            img = img.unsqueeze(0).to(DEVICE)\n\n            pred = model(img).cpu().numpy().flatten()\n            gt = gt.numpy().flatten()\n\n            preds_all.extend(pred)\n            gts_all.extend(gt)\n\n    mse = mean_squared_error(gts_all, preds_all)\n    rmse = np.sqrt(mse)\n    mae = mean_absolute_error(gts_all, preds_all)\n    r2 = r2_score(gts_all, preds_all)\n\n    print(\"\\n===== METRICS =====\")\n    print(f\"MSE  : {mse:.4f}\")\n    print(f\"RMSE : {rmse:.4f}\")\n    print(f\"MAE  : {mae:.4f}\")\n    print(f\"R2   : {r2:.4f}\")\n\n    return mse\n\n\ndef plot_prediction(model, dataset):\n    model.eval()\n    img, gt = dataset[0]\n\n    with torch.no_grad():\n        pred = model(img.unsqueeze(0).to(DEVICE)).cpu().numpy().flatten()\n\n    plt.figure(figsize=(10,4))\n    plt.plot(gt.numpy(), label=\"GT\")\n    plt.plot(pred, label=\"Pred\")\n    plt.legend()\n    plt.title(\"Prediction vs Ground Truth\")\n    plt.show()\n\n# =============================\n# TRAINING\n# =============================\n\ndataset = ECGDataset(TRAIN_DIR)\nloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)\n\nmodel = ECGModel().to(DEVICE)\noptimizer = optim.Adam(model.parameters(), lr=LR)\n\nloss_history = []\n\nprint(f\"Training on {len(dataset)} samples\")\n\nfor epoch in range(EPOCHS):\n    model.train()\n    total_loss = 0\n\n    for imgs, signals in loader:\n        imgs, signals = imgs.to(DEVICE), signals.to(DEVICE)\n\n        preds = model(imgs)\n        preds = torch.clamp(preds, -5, 5)\n\n        loss = ((preds - signals)**2 * (1 + torch.abs(signals))).mean()\n\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n\n        total_loss += loss.item()\n\n    epoch_loss = total_loss / len(loader)\n    loss_history.append(epoch_loss)\n\n    print(f\"Epoch {epoch+1}/{EPOCHS} | Loss: {epoch_loss:.4f}\")\n\n# =============================\n# LOSS CURVE\n# =============================\n\nplt.plot(loss_history)\nplt.title(\"Loss Curve\")\nplt.show()\n\n# =============================\n# EVALUATION\n# =============================\n\nval_loss = evaluate_model(model, dataset)\nprint(f\"\\nValidation MSE: {val_loss:.4f}\")\n\nplot_prediction(model, dataset)\nevaluate_full_metrics(model, dataset)\n\n# =============================\n# TEST + SUBMISSION\n# =============================\n\ntest_meta = pd.read_csv(\"/kaggle/input/competitions/physionet-ecg-image-digitization/test.csv\")\n\nrows = []\n\nfor img_id in test_meta[\"id\"].unique():\n    img_path = os.path.join(TEST_DIR, f\"{img_id}.png\")\n\n    img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)\n    h = img.shape[0]\n\n    img = img[3*h//4:h, :]\n    img = cv2.resize(img, (IMG_SIZE, IMG_SIZE))\n    img = img/255.0\n    img = (img-0.5)/0.5\n    img = np.expand_dims(img, axis=0)\n\n    inp = torch.tensor(img, dtype=torch.float32).unsqueeze(0).to(DEVICE)\n\n    with torch.no_grad():\n        pred = model(inp).cpu().numpy().flatten()\n\n    subset = test_meta[test_meta[\"id\"] == img_id]\n\n    for _, row in subset.iterrows():\n        length = row[\"number_of_rows\"]\n\n        # 🔹 Only output resize (required for Kaggle)\n        sig = np.interp(\n            np.linspace(0,len(pred)-1,length),\n            np.arange(len(pred)),\n            pred\n        )\n\n        for i in range(len(sig)):\n            rows.append({\n                \"id\": f\"{img_id}_{i}_{row['lead']}\",\n                \"value\": float(sig[i])\n            })\n\nsubmission = pd.DataFrame(rows)\nsubmission.to_parquet(\"submission.parquet\", index=False)\n\nprint(\"✅ submission.parquet generated!\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}