{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.7.9","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## Install dependencies","metadata":{}},{"cell_type":"code","source":"!pip install segmentation-models-pytorch","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-14T10:08:45.570274Z","iopub.execute_input":"2026-08-14T10:08:45.570669Z","iopub.status.idle":"2026-08-14T10:08:50.996278Z","shell.execute_reply.started":"2026-08-14T10:08:45.570626Z","shell.execute_reply":"2026-08-14T10:08:50.995444Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# Imports\n# ============================================================\nimport os\nimport zipfile\nimport random\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom tqdm import tqdm\nfrom skimage.io import imread\nfrom skimage.transform import resize\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import TensorDataset, DataLoader\n\nimport segmentation_models_pytorch as smp","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-14T10:08:57.163839Z","iopub.execute_input":"2026-08-14T10:08:57.164155Z","iopub.status.idle":"2026-08-14T10:08:57.169572Z","shell.execute_reply.started":"2026-08-14T10:08:57.164129Z","shell.execute_reply":"2026-08-14T10:08:57.168406Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# 1. Configuration\n# ============================================================\n\nSEED = 42\n\nrandom.seed(SEED)\nnp.random.seed(SEED)\ntorch.manual_seed(SEED)\n\nif torch.cuda.is_available():\n    torch.cuda.manual_seed_all(SEED)\n\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\nIMG_WIDTH = 128\nIMG_HEIGHT = 128\nIMG_CHANNELS = 3\n\nBATCH_SIZE = 16\nEPOCHS = 25\nLEARNING_RATE = 1e-3\n\nDATA_PATH = \"../input/data-science-bowl-2018/\"\nTRAIN_PATH = \"./train/\"\nTEST_PATH = \"./test/\"\n\nprint(\"Device:\", DEVICE)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-14T10:08:59.528782Z","iopub.execute_input":"2026-08-14T10:08:59.529092Z","iopub.status.idle":"2026-08-14T10:08:59.583834Z","shell.execute_reply.started":"2026-08-14T10:08:59.529066Z","shell.execute_reply":"2026-08-14T10:08:59.582893Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# 2. Unzip dataset\n# ============================================================\n\nos.makedirs(TRAIN_PATH, exist_ok=True)\nos.makedirs(TEST_PATH, exist_ok=True)\n\ntrain_zip = os.path.join(DATA_PATH, \"stage1_train.zip\")\ntest_zip = os.path.join(DATA_PATH, \"stage1_test.zip\")\n\nif os.path.exists(train_zip):\n    with zipfile.ZipFile(train_zip, \"r\") as zip_ref:\n        zip_ref.extractall(TRAIN_PATH)\n\nif os.path.exists(test_zip):\n    with zipfile.ZipFile(test_zip, \"r\") as zip_ref:\n        zip_ref.extractall(TEST_PATH)\n\ntrain_ids = next(os.walk(TRAIN_PATH))[1]\ntest_ids = next(os.walk(TEST_PATH))[1]\n\nprint(\"Training images:\", len(train_ids))\nprint(\"Test images:\", len(test_ids))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-14T10:09:03.004343Z","iopub.execute_input":"2026-08-14T10:09:03.004638Z","iopub.status.idle":"2026-08-14T10:09:07.206232Z","shell.execute_reply.started":"2026-08-14T10:09:03.004614Z","shell.execute_reply":"2026-08-14T10:09:07.205483Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# 3. Load training images + masks\n# ============================================================\n\nX_train = np.zeros(\n    (len(train_ids), IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS),\n    dtype=np.float32\n)\n\nY_train = np.zeros(\n    (len(train_ids), IMG_HEIGHT, IMG_WIDTH, 1),\n    dtype=np.float32\n)\n\n\nfor n, id_ in tqdm(\n    enumerate(train_ids),\n    total=len(train_ids),\n    desc=\"Loading training data\"\n):\n\n    path = os.path.join(TRAIN_PATH, id_)\n\n    # -------------------------\n    # Image\n    # -------------------------\n\n    image_path = os.path.join(\n        path,\n        \"images\",\n        id_ + \".png\"\n    )\n\n    img = imread(image_path)\n\n    # Keep RGB channels\n    img = img[:, :, :IMG_CHANNELS]\n\n    img = resize(\n        img,\n        (IMG_HEIGHT, IMG_WIDTH),\n        mode=\"constant\",\n        preserve_range=True\n    )\n\n    X_train[n] = img / 255.0\n\n\n    # -------------------------\n    # Mask\n    # -------------------------\n\n    mask = np.zeros(\n        (IMG_HEIGHT, IMG_WIDTH),\n        dtype=np.float32\n    )\n\n    mask_path = os.path.join(path, \"masks\")\n\n    for mask_file in os.listdir(mask_path):\n\n        mask_file_path = os.path.join(\n            mask_path,\n            mask_file\n        )\n\n        mask_ = imread(mask_file_path)\n\n        mask_ = resize(\n            mask_,\n            (IMG_HEIGHT, IMG_WIDTH),\n            mode=\"constant\",\n            preserve_range=True\n        )\n\n        # Combine individual nucleus masks\n        mask = np.maximum(mask, mask_)\n\n    # Convert to binary\n    mask = (mask > 0).astype(np.float32)\n\n    Y_train[n, :, :, 0] = mask\n\n\nprint(\"X_train:\", X_train.shape)\nprint(\"Y_train:\", Y_train.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-14T10:09:10.046376Z","iopub.execute_input":"2026-08-14T10:09:10.046703Z","iopub.status.idle":"2026-08-14T10:13:21.424657Z","shell.execute_reply.started":"2026-08-14T10:09:10.046676Z","shell.execute_reply":"2026-08-14T10:13:21.423808Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# 4. Convert NHWC → NCHW\n# ============================================================\n\n# NumPy / TensorFlow:\n#\n#     N × H × W × C\n#\n# PyTorch:\n#\n#     N × C × H × W\n\nX_train = torch.tensor(X_train, dtype=torch.float32)\nY_train = torch.tensor(Y_train, dtype=torch.float32)\n\nX_train = X_train.permute(0, 3, 1, 2)\nY_train = Y_train.permute(0, 3, 1, 2)\n\nprint(\"PyTorch X_train:\", X_train.shape)\nprint(\"PyTorch Y_train:\", Y_train.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-14T10:14:22.085532Z","iopub.execute_input":"2026-08-14T10:14:22.085833Z","iopub.status.idle":"2026-08-14T10:14:22.169824Z","shell.execute_reply.started":"2026-08-14T10:14:22.085808Z","shell.execute_reply":"2026-08-14T10:14:22.169052Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# 5. Train / validation split\n# ============================================================\n\nnum_samples = len(X_train)\n\nindices = torch.randperm(num_samples)\n\nsplit = int(0.9 * num_samples)\n\ntrain_indices = indices[:split]\nval_indices = indices[split:]\n\nX_tr = X_train[train_indices]\nY_tr = Y_train[train_indices]\n\nX_val = X_train[val_indices]\nY_val = Y_train[val_indices]\n\nprint(\"Training:\", X_tr.shape[0])\nprint(\"Validation:\", X_val.shape[0])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-14T10:14:25.755815Z","iopub.execute_input":"2026-08-14T10:14:25.756125Z","iopub.status.idle":"2026-08-14T10:14:25.840408Z","shell.execute_reply.started":"2026-08-14T10:14:25.756099Z","shell.execute_reply":"2026-08-14T10:14:25.839385Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# 6. Create DataLoaders\n# ============================================================\n\ntrain_dataset = TensorDataset(\n    X_tr,\n    Y_tr\n)\n\nval_dataset = TensorDataset(\n    X_val,\n    Y_val\n)\n\ntrain_loader = DataLoader(\n    train_dataset,\n    batch_size=BATCH_SIZE,\n    shuffle=True\n)\n\nval_loader = DataLoader(\n    val_dataset,\n    batch_size=BATCH_SIZE,\n    shuffle=False\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-14T10:14:28.209947Z","iopub.execute_input":"2026-08-14T10:14:28.210257Z","iopub.status.idle":"2026-08-14T10:14:28.214675Z","shell.execute_reply.started":"2026-08-14T10:14:28.210234Z","shell.execute_reply":"2026-08-14T10:14:28.213850Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# 7. Create U-Net or DeepLabv3+ using Segmentation Models PyTorch\n# ============================================================\n\n# model = smp.Unet(\n#     encoder_name=\"resnet34\",\n#     encoder_weights=\"imagenet\",\n#     in_channels=3,\n#     classes=1,\n#     activation=None\n# )\n\nmodel = smp.DeepLabV3Plus(\n    encoder_name=\"resnet34\",\n    encoder_weights=\"imagenet\",\n    in_channels=3,\n    classes=1,\n    activation=None\n)\n\nmodel = model.to(DEVICE)\n\n# print(\"\\nModel:\")\n# print(model)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-14T10:16:56.841255Z","iopub.execute_input":"2026-08-14T10:16:56.841555Z","iopub.status.idle":"2026-08-14T10:16:57.464711Z","shell.execute_reply.started":"2026-08-14T10:16:56.841530Z","shell.execute_reply":"2026-08-14T10:16:57.464054Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# 8. Loss + optimizer\n# ============================================================\n\n# BCEWithLogitsLoss already includes sigmoid internally\nbce_loss = nn.BCEWithLogitsLoss()\n\n# Dice loss from SMP\ndice_loss = smp.losses.DiceLoss(\n    mode=\"binary\",\n    from_logits=True\n)\n\n\ndef loss_function(predictions, masks):\n\n    return (\n        bce_loss(predictions, masks)\n        + dice_loss(predictions, masks)\n    )\n\n\noptimizer = torch.optim.Adam(\n    model.parameters(),\n    lr=LEARNING_RATE\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-14T10:16:58.381237Z","iopub.execute_input":"2026-08-14T10:16:58.381528Z","iopub.status.idle":"2026-08-14T10:16:58.388598Z","shell.execute_reply.started":"2026-08-14T10:16:58.381504Z","shell.execute_reply":"2026-08-14T10:16:58.387748Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# 9. Training\n# ============================================================\n\nbest_val_loss = float(\"inf\")\n\nfor epoch in range(EPOCHS):\n\n    # -------------------------\n    # TRAIN\n    # -------------------------\n\n    model.train()\n\n    train_loss = 0.0\n\n    for images, masks in train_loader:\n\n        images = images.to(DEVICE)\n        masks = masks.to(DEVICE)\n\n        # Clear previous gradients\n        optimizer.zero_grad()\n\n        # Forward pass\n        predictions = model(images)\n\n        # Calculate loss\n        loss = loss_function(\n            predictions,\n            masks\n        )\n\n        # Backpropagation\n        loss.backward()\n\n        # Update weights\n        optimizer.step()\n\n        train_loss += loss.item()\n\n\n    train_loss /= len(train_loader)\n\n\n    # -------------------------\n    # VALIDATION\n    # -------------------------\n\n    model.eval()\n\n    val_loss = 0.0\n\n    with torch.no_grad():\n\n        for images, masks in val_loader:\n\n            images = images.to(DEVICE)\n            masks = masks.to(DEVICE)\n\n            predictions = model(images)\n\n            loss = loss_function(\n                predictions,\n                masks\n            )\n\n            val_loss += loss.item()\n\n\n    val_loss /= len(val_loader)\n\n\n    print(\n        f\"Epoch [{epoch+1}/{EPOCHS}] \"\n        f\"Train Loss: {train_loss:.4f} \"\n        f\"Val Loss: {val_loss:.4f}\"\n    )\n\n\n    # -------------------------\n    # Save best model\n    # -------------------------\n\n    if val_loss < best_val_loss:\n\n        best_val_loss = val_loss\n\n        torch.save(\n            model.state_dict(),\n            \"best_unet.pth\"\n        )\n\n        print(\"  ✓ Saved best model\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-14T10:17:00.043353Z","iopub.execute_input":"2026-08-14T10:17:00.043685Z","iopub.status.idle":"2026-08-14T10:18:03.126348Z","shell.execute_reply.started":"2026-08-14T10:17:00.043651Z","shell.execute_reply":"2026-08-14T10:18:03.125425Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# 10. Load best model\n# ============================================================\n\nmodel.load_state_dict(\n    torch.load(\n        \"best_unet.pth\",\n        map_location=DEVICE\n    )\n)\n\nmodel.eval()\n\nprint(\"\\nBest model loaded.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-14T10:18:09.370970Z","iopub.execute_input":"2026-08-14T10:18:09.371293Z","iopub.status.idle":"2026-08-14T10:18:09.461890Z","shell.execute_reply.started":"2026-08-14T10:18:09.371265Z","shell.execute_reply":"2026-08-14T10:18:09.460975Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# 11. Prediction function\n# ============================================================\n\ndef predict(model, images):\n\n    model.eval()\n\n    with torch.no_grad():\n\n        images = images.to(DEVICE)\n\n        logits = model(images)\n\n        # Convert logits → probabilities\n        probabilities = torch.sigmoid(logits)\n\n        # Convert probabilities → binary mask\n        masks = (probabilities > 0.5).float()\n\n    return masks.cpu()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-14T10:18:11.515797Z","iopub.execute_input":"2026-08-14T10:18:11.516117Z","iopub.status.idle":"2026-08-14T10:18:11.520684Z","shell.execute_reply.started":"2026-08-14T10:18:11.516091Z","shell.execute_reply":"2026-08-14T10:18:11.519725Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# 12. Get predictions\n# ============================================================\n\ntrain_predictions = predict(\n    model,\n    X_tr\n)\n\nval_predictions = predict(\n    model,\n    X_val\n)\n\n\nprint(\"Prediction shape:\", val_predictions.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-14T10:18:13.581636Z","iopub.execute_input":"2026-08-14T10:18:13.582040Z","iopub.status.idle":"2026-08-14T10:18:23.054554Z","shell.execute_reply.started":"2026-08-14T10:18:13.582009Z","shell.execute_reply":"2026-08-14T10:18:23.053581Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# 13. Dice score\n# ============================================================\n\ndef dice_score(pred, target):\n\n    pred = pred.flatten()\n    target = target.flatten()\n\n    intersection = (pred * target).sum()\n\n    dice = (\n        2.0 * intersection + 1e-7\n    ) / (\n        pred.sum() + target.sum() + 1e-7\n    )\n\n    return dice.item()\n\n\ndice_scores = []\n\nfor i in range(len(Y_val)):\n\n    score = dice_score(\n        val_predictions[i],\n        Y_val[i]\n    )\n\n    dice_scores.append(score)\n\n\nprint(\n    f\"\\nMean Validation Dice: \"\n    f\"{np.mean(dice_scores):.4f}\"\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-14T10:18:23.057473Z","iopub.execute_input":"2026-08-14T10:18:23.057839Z","iopub.status.idle":"2026-08-14T10:18:23.074372Z","shell.execute_reply.started":"2026-08-14T10:18:23.057799Z","shell.execute_reply":"2026-08-14T10:18:23.073499Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================================================\n# 14. Visualize prediction\n# ============================================================\n\nidx = random.randint(\n    0,\n    len(X_val) - 1\n)\n\nimage = X_val[idx]\n\nground_truth = Y_val[idx]\n\nprediction = val_predictions[idx]\n\n\n# Convert CHW → HWC for matplotlib\nimage = image.permute(1, 2, 0)\n\nground_truth = ground_truth.squeeze()\n\nprediction = prediction.squeeze()\n\n\nplt.figure(figsize=(12, 4))\n\n\n# Input\nplt.subplot(1, 3, 1)\n\nplt.imshow(image)\n\nplt.title(\"Input Image\")\n\nplt.axis(\"off\")\n\n\n# Ground truth\nplt.subplot(1, 3, 2)\n\nplt.imshow(\n    ground_truth,\n    cmap=\"gray\"\n)\n\nplt.title(\"Ground Truth\")\n\nplt.axis(\"off\")\n\n\n# Prediction\nplt.subplot(1, 3, 3)\n\nplt.imshow(\n    prediction,\n    cmap=\"gray\"\n)\n\nplt.title(\n    f\"Prediction\\nDice={dice_scores[idx]:.3f}\"\n)\n\nplt.axis(\"off\")\n\n\nplt.tight_layout()\n\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-14T10:18:23.075778Z","iopub.execute_input":"2026-08-14T10:18:23.076103Z","iopub.status.idle":"2026-08-14T10:18:23.303980Z","shell.execute_reply.started":"2026-08-14T10:18:23.076077Z","shell.execute_reply":"2026-08-14T10:18:23.302903Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}