{"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":[{"sourceId":97984,"databundleVersionId":14096757,"sourceType":"competition"}],"dockerImageVersionId":31193,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport glob\nimport cv2\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\n\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms\nfrom PIL import Image\nimport matplotlib.pyplot as plt","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-09T18:42:26.237749Z","iopub.execute_input":"2025-12-09T18:42:26.237907Z","iopub.status.idle":"2025-12-09T18:42:36.411513Z","shell.execute_reply.started":"2025-12-09T18:42:26.237892Z","shell.execute_reply":"2025-12-09T18:42:36.410702Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"IMG_SIZE = 256\nBATCH_SIZE = 8\nEPOCHS = 10\nLR = 1e-4\n\nDATA_DIR = \"/kaggle/input/physionet-ecg-image-digitization\"\nTRAIN_DIR = os.path.join(DATA_DIR, \"train\")\nTRAIN_META_CSV = os.path.join(DATA_DIR, \"train.csv\")\n\ndevice=torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"Device: {device}\")\nif torch.cuda.is_available():\n    print(f\"GPU Name: {torch.cuda.get_device_name(0)}\")\n    print(f\"GPU Available: True\")\nelse:\n    print(\"GPU Available: False - Using CPU\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-09T18:42:38.998391Z","iopub.execute_input":"2025-12-09T18:42:38.998674Z","iopub.status.idle":"2025-12-09T18:42:39.114952Z","shell.execute_reply.started":"2025-12-09T18:42:38.998652Z","shell.execute_reply":"2025-12-09T18:42:39.114176Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"LEAD_ORDER = [\n    'I','II','III','aVR',\n    'aVL','aVF','V1','V2',\n    'V3','V4','V5','V6'\n]\n\ndef split_into_12_leads(img):\n    \"\"\"\n    img: np.array H×W×3 (np.uint8)\n    zwraca listę 12 patchy: [lead0_img, lead1_img, ..., lead11_img]\n    \"\"\"\n    H, W, _ = img.shape\n    h = H // 3\n    w = W // 4\n    leads = []\n    for r in range(3):\n        for c in range(4):\n            y1, y2 = r*h, (r+1)*h\n            x1, x2 = c*w, (c+1)*w\n            crop = img[y1:y2, x1:x2]\n            leads.append(crop)\n    return leads","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-09T18:42:47.925496Z","iopub.execute_input":"2025-12-09T18:42:47.926238Z","iopub.status.idle":"2025-12-09T18:42:47.931096Z","shell.execute_reply.started":"2025-12-09T18:42:47.926212Z","shell.execute_reply":"2025-12-09T18:42:47.930435Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def generate_mask_from_signal(signal, img_size=IMG_SIZE, thickness=3):\n    signal = np.array(signal, dtype=float)\n    not_nan = ~np.isnan(signal)\n    if not np.any(not_nan):\n        return np.zeros((img_size, img_size), dtype=np.uint8)\n    \n    sig = signal[not_nan]\n    sig_norm = (sig - np.nanmin(sig)) / (np.nanmax(sig) - np.nanmin(sig) + 1e-8)\n    \n    y = (1 - sig_norm) * (img_size - 1)\n    y = y.astype(int)\n    xs = np.linspace(0, img_size - 1, len(y)).astype(int)\n\n    mask = np.zeros((img_size, img_size), dtype=np.uint8)\n    \n    points = np.stack([xs, y], axis=1)\n    \n    for i in range(1, len(points)):\n        pt1 = tuple(points[i-1])\n        pt2 = tuple(points[i])\n        cv2.line(mask, pt1, pt2, 255, thickness=thickness) \n        \n    return mask","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-09T18:42:49.398801Z","iopub.execute_input":"2025-12-09T18:42:49.399400Z","iopub.status.idle":"2025-12-09T18:42:49.405216Z","shell.execute_reply.started":"2025-12-09T18:42:49.399371Z","shell.execute_reply":"2025-12-09T18:42:49.404539Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_meta = pd.read_csv(TRAIN_META_CSV)\nprint(train_meta.head())\n\ntrain_samples = []  \nfor row in tqdm(train_meta.itertuples(), total=len(train_meta)):\n    rec_id = str(row.id)\n    folder = os.path.join(TRAIN_DIR, rec_id)\n    if not os.path.isdir(folder):\n        continue\n    \n    png_files = sorted(glob.glob(os.path.join(folder, \"*.png\")))\n    if len(png_files) == 0:\n        continue\n    img_path = png_files[0]\n    \n    csv_path = os.path.join(folder, f\"{rec_id}.csv\")\n    if not os.path.isfile(csv_path):\n        continue\n    \n    df_sig = pd.read_csv(csv_path)\n    for lead_idx, lead_name in enumerate(LEAD_ORDER):\n        if lead_name not in df_sig.columns:\n            continue\n        col = df_sig[lead_name].values\n        if np.all(np.isnan(col)):\n            continue \n        train_samples.append({\n            \"img_path\": img_path,\n            \"csv_path\": csv_path,\n            \"lead_idx\": lead_idx,\n            \"lead_name\": lead_name\n        })\n\nprint(\"Liczba przykładów treningowych:\", len(train_samples))\nprint(train_samples[0])\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-09T18:42:52.053255Z","iopub.execute_input":"2025-12-09T18:42:52.054055Z","iopub.status.idle":"2025-12-09T18:43:08.853832Z","shell.execute_reply.started":"2025-12-09T18:42:52.054030Z","shell.execute_reply":"2025-12-09T18:43:08.853108Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ECGLeadDataset(Dataset):\n    def __init__(self, samples, transform_img=None):\n        self.samples = samples\n        self.transform_img = transform_img\n\n    def __len__(self):\n        return len(self.samples)\n\n    def __getitem__(self, idx):\n        s = self.samples[idx]\n        img_path = s[\"img_path\"]\n        csv_path = s[\"csv_path\"]\n        lead_idx = s[\"lead_idx\"]\n        lead_name = s[\"lead_name\"]\n        \n        img = cv2.imread(img_path)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        lead_imgs = split_into_12_leads(img)\n        lead_img = lead_imgs[lead_idx]  # patch dla danego leada\n        \n        lead_img = cv2.resize(lead_img, (IMG_SIZE, IMG_SIZE))\n        lead_pil = Image.fromarray(lead_img)\n        \n        if self.transform_img:\n            img_tensor = self.transform_img(lead_pil)\n        else:\n            img_tensor = transforms.ToTensor()(lead_pil)\n        \n        df_sig = pd.read_csv(csv_path)\n        signal = df_sig[lead_name].values\n        mask_np = generate_mask_from_signal(signal, IMG_SIZE)\n        mask_pil = Image.fromarray(mask_np)\n        mask_tensor = transforms.ToTensor()(mask_pil)  \n        \n        return img_tensor, mask_tensor\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-09T18:43:11.708072Z","iopub.execute_input":"2025-12-09T18:43:11.708876Z","iopub.status.idle":"2025-12-09T18:43:11.715351Z","shell.execute_reply.started":"2025-12-09T18:43:11.708844Z","shell.execute_reply":"2025-12-09T18:43:11.714622Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"transform_img = transforms.Compose([\n    transforms.ToTensor(),\n])\n\ndataset = ECGLeadDataset(train_samples, transform_img)\nprint(\"Dataset size:\", len(dataset))\n\n# szybki test\nimg_t, mask_t = dataset[0]\nprint(\"img_t:\", img_t.shape, \"mask_t:\", mask_t.shape)\ntrain_loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-09T18:43:12.814434Z","iopub.execute_input":"2025-12-09T18:43:12.815104Z","iopub.status.idle":"2025-12-09T18:43:12.951744Z","shell.execute_reply.started":"2025-12-09T18:43:12.815078Z","shell.execute_reply":"2025-12-09T18:43:12.950995Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class DoubleConv(nn.Module):\n    def __init__(self, in_ch, out_ch):\n        super().__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(in_ch, out_ch, 3, padding=1),\n            nn.BatchNorm2d(out_ch),\n            nn.ReLU(inplace=True),\n\n            nn.Conv2d(out_ch, out_ch, 3, padding=1),\n            nn.BatchNorm2d(out_ch),\n            nn.ReLU(inplace=True),\n        )\n\n    def forward(self, x):\n        return self.conv(x)\n\n\nclass UNet(nn.Module):\n    def __init__(self, in_ch=3, out_ch=1):\n        super().__init__()\n        self.down1 = DoubleConv(in_ch, 64)\n        self.pool1 = nn.MaxPool2d(2)\n\n        self.down2 = DoubleConv(64, 128)\n        self.pool2 = nn.MaxPool2d(2)\n\n        self.down3 = DoubleConv(128, 256)\n        self.pool3 = nn.MaxPool2d(2)\n\n        self.down4 = DoubleConv(256, 512)\n        self.pool4 = nn.MaxPool2d(2)\n\n        self.bottleneck = DoubleConv(512, 1024)\n\n        self.up4 = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2)\n        self.conv4 = DoubleConv(1024, 512)\n\n        self.up3 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)\n        self.conv3 = DoubleConv(512, 256)\n\n        self.up2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)\n        self.conv2 = DoubleConv(256, 128)\n\n        self.up1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)\n        self.conv1 = DoubleConv(128, 64)\n\n        self.out_conv = nn.Conv2d(64, out_ch, kernel_size=1)\n\n    def forward(self, x):\n        c1 = self.down1(x)\n        c2 = self.down2(self.pool1(c1))\n        c3 = self.down3(self.pool2(c2))\n        c4 = self.down4(self.pool3(c3))\n\n        b = self.bottleneck(self.pool4(c4))\n\n        x = self.up4(b)\n        x = self.conv4(torch.cat([x, c4], dim=1))\n\n        x = self.up3(x)\n        x = self.conv3(torch.cat([x, c3], dim=1))\n\n        x = self.up2(x)\n        x = self.conv2(torch.cat([x, c2], dim=1))\n\n        x = self.up1(x)\n        x = self.conv1(torch.cat([x, c1], dim=1))\n\n        return self.out_conv(x)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-09T18:43:13.655823Z","iopub.execute_input":"2025-12-09T18:43:13.656474Z","iopub.status.idle":"2025-12-09T18:43:13.666631Z","shell.execute_reply.started":"2025-12-09T18:43:13.656444Z","shell.execute_reply":"2025-12-09T18:43:13.665985Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model = UNet(in_ch=3, out_ch=1).to(device)\ncriterion = nn.BCEWithLogitsLoss()\noptimizer = optim.Adam(model.parameters(), lr=LR)\n\nprint(\"Liczba parametrów:\", sum(p.numel() for p in model.parameters() if p.requires_grad))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-09T18:43:16.308843Z","iopub.execute_input":"2025-12-09T18:43:16.309388Z","iopub.status.idle":"2025-12-09T18:43:16.764607Z","shell.execute_reply.started":"2025-12-09T18:43:16.309364Z","shell.execute_reply":"2025-12-09T18:43:16.763845Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"for epoch in range(EPOCHS):\n    model.train()\n    running_loss = 0.0\n\n    for imgs, masks in tqdm(train_loader, desc=f\"Epoch {epoch+1}/{EPOCHS}\"):\n        imgs = imgs.to(device)\n        masks = masks.to(device)\n\n        optimizer.zero_grad()\n        logits = model(imgs)\n        loss = criterion(logits, masks)\n        loss.backward()\n        optimizer.step()\n\n        running_loss += loss.item() * imgs.size(0)\n\n    epoch_loss = running_loss / len(dataset)\n    print(f\"Epoch {epoch+1}/{EPOCHS} - loss: {epoch_loss:.4f}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-09T18:43:17.619098Z","iopub.execute_input":"2025-12-09T18:43:17.619678Z","iopub.status.idle":"2025-12-09T20:41:44.166361Z","shell.execute_reply.started":"2025-12-09T18:43:17.619654Z","shell.execute_reply":"2025-12-09T20:41:44.165495Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"torch.save(model.state_dict(), \"unet_ecg_256.pth\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-09T20:41:44.167841Z","iopub.execute_input":"2025-12-09T20:41:44.168050Z","iopub.status.idle":"2025-12-09T20:41:44.357827Z","shell.execute_reply.started":"2025-12-09T20:41:44.168031Z","shell.execute_reply":"2025-12-09T20:41:44.357240Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model.eval()\nwith torch.no_grad():\n    imgs, masks = next(iter(train_loader))\n    imgs = imgs.to(device)\n    masks = masks.to(device)\n\n    logits = model(imgs)\n    probs = torch.sigmoid(logits)\n    preds = (probs > 0.5).float()\n\nN = min(5, imgs.size(0))\n\nplt.figure(figsize=(15, 3*N))\n\nfor i in range(N):\n    img_np = imgs[i].cpu().permute(1,2,0).numpy()\n    mask_true = masks[i,0].cpu().numpy()\n    mask_pred = preds[i,0].cpu().numpy()\n\n    plt.subplot(N, 3, i*3 + 1)\n    plt.title(f\"Image {i+1}\")\n    plt.imshow(img_np)\n    plt.axis(\"off\")\n\n    plt.subplot(N, 3, i*3 + 2)\n    plt.title(\"GT mask\")\n    plt.imshow(mask_true, cmap=\"gray\")\n    plt.axis(\"off\")\n\n    plt.subplot(N, 3, i*3 + 3)\n    plt.title(\"Pred mask\")\n    plt.imshow(mask_pred, cmap=\"gray\")\n    plt.axis(\"off\")\n\nplt.tight_layout()\nplt.show()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-09T20:41:44.358540Z","iopub.execute_input":"2025-12-09T20:41:44.358737Z","iopub.status.idle":"2025-12-09T20:41:47.208681Z","shell.execute_reply.started":"2025-12-09T20:41:44.358721Z","shell.execute_reply":"2025-12-09T20:41:47.207833Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}