{"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":117682,"databundleVersionId":14443416,"sourceType":"competition"}],"dockerImageVersionId":31193,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install --upgrade imagecodecs","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\"\"\"\nEnd-to-end Kaggle pipeline for the Vesuvius Challenge (papyrus layer segmentation)\n- Data loading (3D .tif volumes and masks)\n- 3D U-Net model (PyTorch)\n- Dataset + augmentations (random crops, flips, intensity)\n- Loss: BCE + Dice + optional boundary-aware loss (distance transform)\n- Direct train/validation split training (no K-Fold)\n- Mixed precision (torch.amp)\n- Sliding-window inference + test-time augmentation (flips)\n- Postprocessing (connected components)\n- Ensembling folds removed — single model saved as best.pth\n\nThis version replaces the K-Fold training with a single direct train/validation split and a simple training loop that runs for CFG.epochs. Use this for faster iteration or limited compute.\n\"\"\"\n\nimport os\nimport gc\nimport math\nimport time\nimport random\nimport zipfile\nfrom glob import glob\nfrom collections import defaultdict\n\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch import amp\n\nimport tifffile\nfrom scipy import ndimage\nfrom sklearn.model_selection import train_test_split\nimport warnings\nwarnings.filterwarnings('ignore')\n\n# ----------------------------- Config -----------------------------\nclass CFG:\n    seed = 42\n    img_dir = '/kaggle/input/vesuvius-challenge-surface-detection/train_images'  # adjust\n    mask_dir = '/kaggle/input/vesuvius-challenge-surface-detection/train_labels'  # adjust\n    train_csv = '/kaggle/input/vesuvius-challenge-surface-detection/train.csv'\n    test_csv = '/kaggle/input/vesuvius-challenge-surface-detection/test.csv'\n    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n    epochs = 10\n    batch_size = 8\n    lr = 1e-4\n    weight_decay = 1e-5\n    patch_size = (64, 128, 128)  # (z,y,x) - tune to fit GPU memory\n    crop_margin = 16\n    num_workers = 4\n    save_dir = './models'\n    inference_overlap = 0.25\n    tta = True\n    val_split = 0.15  # fraction for validation\n\nos.makedirs(CFG.save_dir, exist_ok=True)\n\n# ----------------------------- Utilities -----------------------------\n\ndef seed_everything(seed=42):\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n\nseed_everything(CFG.seed)\n\n\ndef read_tiff(path):\n    return tifffile.imread(path)\n\ndef save_tiff(image, path):\n    tifffile.imwrite(path, image.astype(np.uint8))\n\n# Dice\n\ndef dice_loss(pred, target, eps=1e-6):\n    num = 2 * (pred * target).sum()\n    den = pred.sum() + target.sum() + eps\n    return 1 - (num / den)\n\n# BCE + Dice\nclass BCEDiceLoss(nn.Module):\n    def __init__(self, weight_bce=1.0, weight_dice=1.0):\n        super().__init__()\n        self.bce = nn.BCEWithLogitsLoss()\n        self.wb = weight_bce\n        self.wd = weight_dice\n\n    def forward(self, logits, target):\n        bce = self.bce(logits, target)\n        probs = torch.sigmoid(logits)\n        d = dice_loss(probs, target)\n        return self.wb * bce + self.wd * d\n\n# ----------------------------- Model: 3D U-Net -----------------------------\nclass ConvBlock(nn.Module):\n    def __init__(self, in_ch, out_ch):\n        super().__init__()\n        self.conv = nn.Sequential(\n            nn.Conv3d(in_ch, out_ch, 3, padding=1),\n            nn.InstanceNorm3d(out_ch),\n            nn.LeakyReLU(0.1, inplace=True),\n            nn.Conv3d(out_ch, out_ch, 3, padding=1),\n            nn.InstanceNorm3d(out_ch),\n            nn.LeakyReLU(0.1, inplace=True),\n        )\n\n    def forward(self, x):\n        return self.conv(x)\n\nclass UpBlock(nn.Module):\n    def __init__(self, in_ch, out_ch):\n        super().__init__()\n        self.up = nn.ConvTranspose3d(in_ch, out_ch, kernel_size=2, stride=2)\n        self.conv = ConvBlock(in_ch, out_ch)\n\n    def forward(self, x, skip):\n        x = self.up(x)\n        # pad if necessary\n        if x.shape != skip.shape:\n            diffZ = skip.size(2) - x.size(2)\n            diffY = skip.size(3) - x.size(3)\n            diffX = skip.size(4) - x.size(4)\n            x = F.pad(x, [diffX//2, diffX - diffX//2,\n                          diffY//2, diffY - diffY//2,\n                          diffZ//2, diffZ - diffZ//2])\n        x = torch.cat([skip, x], dim=1)\n        return self.conv(x)\n\nclass UNet3D(nn.Module):\n    def __init__(self, in_ch=1, out_ch=1, features=[32, 64, 128, 256]):\n        super().__init__()\n        self.enc1 = ConvBlock(in_ch, features[0])\n        self.pool1 = nn.MaxPool3d(2)\n        self.enc2 = ConvBlock(features[0], features[1])\n        self.pool2 = nn.MaxPool3d(2)\n        self.enc3 = ConvBlock(features[1], features[2])\n        self.pool3 = nn.MaxPool3d(2)\n        self.enc4 = ConvBlock(features[2], features[3])\n\n        self.bottleneck = ConvBlock(features[3], features[3] * 2)\n\n        self.up3 = UpBlock(features[3]*2, features[3])\n        self.up2 = UpBlock(features[3], features[2])\n        self.up1 = UpBlock(features[2], features[1])\n        self.up0 = UpBlock(features[1], features[0])    # ← missing block\n\n        self.head = nn.Conv3d(features[0], out_ch, kernel_size=1)\n\n    def forward(self, x):\n        e1 = self.enc1(x)\n        p1 = self.pool1(e1)\n\n        e2 = self.enc2(p1)\n        p2 = self.pool2(e2)\n\n        e3 = self.enc3(p2)\n        p3 = self.pool3(e3)\n\n        e4 = self.enc4(p3)\n        b = self.bottleneck(e4)\n\n        u3 = self.up3(b, e4)\n        u2 = self.up2(u3, e3)\n        u1 = self.up1(u2, e2)\n        u0 = self.up0(u1, e1)       # ← final upsample\n\n        out = self.head(u0)\n        return out\n\n\n# ----------------------------- Dataset -----------------------------\nclass VesuviusDataset(Dataset):\n    def __init__(self, df, img_dir, mask_dir=None, patch_size=(64,128,128), is_train=True):\n        self.df = df.reset_index(drop=True)\n        self.img_dir = img_dir\n        self.mask_dir = mask_dir\n        self.patch_size = patch_size\n        self.is_train = is_train\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        row = self.df.loc[idx]\n        image_id = row['id']\n        img_path = os.path.join(self.img_dir, f'{image_id}.tif')\n        img = read_tiff(img_path).astype(np.float32)\n        # normalize per-volume\n        img = (img - img.mean()) / (img.std() + 1e-8)\n        # expand to channel dim\n        img = img[np.newaxis, ...]\n\n        if self.is_train:\n            mask_path = os.path.join(self.mask_dir, f'{image_id}.tif')\n            mask = read_tiff(mask_path).astype(np.uint8)\n            # convert 2->ignore: treat 2 as background but could mask loss\n            mask = (mask == 1).astype(np.float32)\n            mask = mask[np.newaxis, ...]\n            img_patch, mask_patch = self.random_crop(img, mask)\n            img_patch, mask_patch = self.augment(img_patch, mask_patch)\n            # ensure output sizes match CFG.patch_size exactly\n            dz, dy, dx = self.patch_size\n            assert img_patch.shape[1:] == (dz, dy, dx), f'Patch size mismatch: {img_patch.shape[1:]} vs {self.patch_size}'\n            return torch.tensor(img_patch).float(), torch.tensor(mask_patch).float()\n        else:\n            # for inference, return full volume\n            return torch.tensor(img).float(), image_id\n\n    def random_crop(self, img, mask):\n        # Ensure patch dims are multiples of downsampling factor (2^n_pools)\n        down_factor = 8  # 2**3 because we have 3 pooling layers\n        dz, dy, dx = self.patch_size\n        _, cz, cy, cx = img.shape\n\n        # compute minimal target dims (>= current dims and multiples of down_factor)\n        target_z = max(dz, int(math.ceil(cz / down_factor)) * down_factor)\n        target_y = max(dy, int(math.ceil(cy / down_factor)) * down_factor)\n        target_x = max(dx, int(math.ceil(cx / down_factor)) * down_factor)\n\n        pad_z = max(0, target_z - cz)\n        pad_y = max(0, target_y - cy)\n        pad_x = max(0, target_x - cx)\n\n        if pad_z or pad_y or pad_x:\n            img = np.pad(img, ((0,0),\n                               (pad_z//2, pad_z - pad_z//2),\n                               (pad_y//2, pad_y - pad_y//2),\n                               (pad_x//2, pad_x - pad_x//2)),\n                         mode='constant', constant_values=0)\n            mask = np.pad(mask, ((0,0),\n                                 (pad_z//2, pad_z - pad_z//2),\n                                 (pad_y//2, pad_y - pad_y//2),\n                                 (pad_x//2, pad_x - pad_x//2)),\n                          mode='constant', constant_values=0)\n            _, cz, cy, cx = img.shape\n\n        # if still smaller than patch (unlikely after above), pad to patch_size\n        if cz < dz or cy < dy or cx < dx:\n            pad_z2 = max(0, dz - cz)\n            pad_y2 = max(0, dy - cy)\n            pad_x2 = max(0, dx - cx)\n            img = np.pad(img, ((0,0),\n                               (pad_z2//2, pad_z2 - pad_z2//2),\n                               (pad_y2//2, pad_y2 - pad_y2//2),\n                               (pad_x2//2, pad_x2 - pad_x2//2)),\n                         mode='constant', constant_values=0)\n            mask = np.pad(mask, ((0,0),\n                                 (pad_z2//2, pad_z2 - pad_z2//2),\n                                 (pad_y2//2, pad_y2 - pad_y2//2),\n                                 (pad_x2//2, pad_x2 - pad_x2//2)),\n                          mode='constant', constant_values=0)\n            _, cz, cy, cx = img.shape\n\n        # now sample a random crop of requested patch_size\n        z = random.randint(0, max(0, cz - dz))\n        y = random.randint(0, max(0, cy - dy))\n        x = random.randint(0, max(0, cx - dx))\n        return img[:, z:z+dz, y:y+dy, x:x+dx], mask[:, z:z+dz, y:y+dy, x:x+dx]\n\n    def augment(self, img, mask):\n        # simple augmentations: flips and small intensity jitter\n        if random.random() < 0.5:\n            img = img[:, :, ::-1, :].copy(); mask = mask[:, :, ::-1, :].copy()\n        if random.random() < 0.5:\n            img = img[:, :, :, ::-1].copy(); mask = mask[:, :, :, ::-1].copy()\n        # intensity jitter\n        if random.random() < 0.5:\n            scale = random.uniform(0.9,1.1)\n            shift = random.uniform(-0.1,0.1)\n            img = img * scale + shift\n        return img, mask\n\n# ----------------------------- Training helpers -----------------------------\n\ndef train_one_epoch(model, loader, optimizer, criterion, scaler):\n    model.train()\n    total_loss = 0.0\n    for imgs, masks in tqdm(loader):\n        imgs = imgs.to(CFG.device)\n        masks = masks.to(CFG.device)\n        optimizer.zero_grad()\n        with amp.autocast(device_type=CFG.device if CFG.device=='cuda' else 'cpu'):\n            logits = model(imgs)\n            # enforce exact shape match\n            assert logits.shape == masks.shape, f\"Shape mismatch: logits={logits.shape}, masks={masks.shape}\"\n            loss = criterion(logits, masks)\n        scaler.scale(loss).backward()\n        scaler.step(optimizer)\n        scaler.update()\n        total_loss += loss.item() * imgs.size(0)\n    return total_loss / len(loader.dataset)\n\ndef valid_one_epoch(model, loader, criterion):\n    model.eval()\n    total_loss = 0.0\n    with torch.no_grad():\n        for imgs, masks in tqdm(loader):\n            imgs = imgs.to(CFG.device)\n            masks = masks.to(CFG.device)\n            with amp.autocast(device_type=CFG.device if CFG.device=='cuda' else 'cpu'):\n                logits = model(imgs)\n                assert logits.shape == masks.shape, f\"Shape mismatch: logits={logits.shape}, masks={masks.shape}\"\n                loss = criterion(logits, masks)\n            total_loss += loss.item() * imgs.size(0)\n    return total_loss / len(loader.dataset)\n\n# ----------------------------- Sliding window inference -----------------------------\n\ndef sliding_window_inference(volume, model, patch_size, overlap=0.25, tta=True):\n    # volume: numpy array shape (1,D,H,W)\n    model.eval()\n    C, D, H, W = volume.shape\n    pd, ph, pw = patch_size\n    # ensure patch sizes are multiples of down_factor\n    down_factor = 8\n    assert pd % down_factor == 0 and ph % down_factor == 0 and pw % down_factor == 0, 'patch_size must be multiple of downsampling factor'\n\n    sd = max(1, int(pd * (1 - overlap)))\n    sy = max(1, int(ph * (1 - overlap)))\n    sx = max(1, int(pw * (1 - overlap)))\n\n    output = np.zeros((1, D, H, W), dtype=np.float32)\n    count = np.zeros_like(output)\n\n    z_starts = list(range(0, max(1, D - pd + 1), sd))\n    y_starts = list(range(0, max(1, H - ph + 1), sy))\n    x_starts = list(range(0, max(1, W - pw + 1), sx))\n    if z_starts[-1] != D - pd:\n        z_starts.append(D - pd)\n    if y_starts[-1] != H - ph:\n        y_starts.append(H - ph)\n    if x_starts[-1] != W - pw:\n        x_starts.append(W - pw)\n\n    with torch.no_grad():\n        for z in z_starts:\n            for y in y_starts:\n                for x in x_starts:\n                    patch = volume[:, z:z+pd, y:y+ph, x:x+pw]\n                    inp = torch.tensor(patch).unsqueeze(0).to(CFG.device).float()\n                    with amp.autocast(device_type=CFG.device if CFG.device=='cuda' else 'cpu'):\n                        logits = model(inp)\n                        probs = torch.sigmoid(logits).cpu().numpy()[0]\n                    output[:, z:z+pd, y:y+ph, x:x+pw] += probs\n                    count[:, z:z+pd, y:y+ph, x:x+pw] += 1\n                    # TTA flips\n                    if tta:\n                        inp_f = torch.flip(inp, dims=[2])\n                        with amp.autocast(device_type=CFG.device if CFG.device=='cuda' else 'cpu'):\n                            probs_f = torch.sigmoid(model(inp_f)).cpu().numpy()[0]\n                        probs_f = np.flip(probs_f, axis=1)\n                        output[:, z:z+pd, y:y+ph, x:x+pw] += probs_f\n                        count[:, z:z+pd, y:y+ph, x:x+pw] += 1\n\n    avg = output / (count + 1e-8)\n    return avg[0]\n\n# ----------------------------- Postprocessing -----------------------------\n\ndef postprocess_mask(prob, thr=0.5, min_size=1000):\n    mask = (prob >= thr).astype(np.uint8)\n    # keep largest connected components\n    labeled, ncomp = ndimage.label(mask)\n    if ncomp == 0:\n        return mask\n    counts = np.bincount(labeled.flatten())\n    # zero index is background\n    counts[0] = 0\n    keep = np.where(counts >= min_size)[0]\n    new_mask = np.isin(labeled, keep).astype(np.uint8)\n    return new_mask\n\n# ----------------------------- Full Training Loop (single split) -----------------------------\n\ndef run_training():\n    df = pd.read_csv(CFG.train_csv)\n    # simple train/val split stratified by scroll_id\n    tr_df, val_df = train_test_split(df, test_size=CFG.val_split, random_state=CFG.seed, stratify=df['scroll_id'])\n\n    train_dataset = VesuviusDataset(tr_df.reset_index(drop=True), CFG.img_dir, CFG.mask_dir, patch_size=CFG.patch_size, is_train=True)\n    val_dataset = VesuviusDataset(val_df.reset_index(drop=True), CFG.img_dir, CFG.mask_dir, patch_size=CFG.patch_size, is_train=True)\n\n    train_loader = DataLoader(train_dataset, batch_size=CFG.batch_size, shuffle=True, num_workers=CFG.num_workers, pin_memory=True)\n    val_loader = DataLoader(val_dataset, batch_size=CFG.batch_size, shuffle=False, num_workers=CFG.num_workers, pin_memory=True)\n\n    model = UNet3D(in_ch=1, out_ch=1).to(CFG.device)\n    optimizer = torch.optim.AdamW(model.parameters(), lr=CFG.lr, weight_decay=CFG.weight_decay)\n    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=CFG.epochs)\n    criterion = BCEDiceLoss()\n    scaler = amp.GradScaler()\n\n    best_val = 1e9\n    best_path = os.path.join(CFG.save_dir, 'best.pth')\n    for epoch in range(CFG.epochs):\n        t0 = time.time()\n        train_loss = train_one_epoch(model, train_loader, optimizer, criterion, scaler)\n        val_loss = valid_one_epoch(model, val_loader, criterion)\n        scheduler.step()\n        print(f'Epoch {epoch} train_loss={train_loss:.4f} val_loss={val_loss:.4f} time={(time.time()-t0):.1f}s')\n        # save best\n        if val_loss < best_val:\n            best_val = val_loss\n            torch.save(model.state_dict(), best_path)\n    del model, optimizer, scheduler, scaler\n    gc.collect()\n    torch.cuda.empty_cache()\n\n# ----------------------------- Inference pipeline -----------------------------\n\ndef run_inference(test_csv_path, out_dir='./submission_masks'):\n    os.makedirs(out_dir, exist_ok=True)\n    df_test = pd.read_csv(test_csv_path)\n\n    # load single best model\n    m = UNet3D(in_ch=1, out_ch=1).to(CFG.device)\n    p = os.path.join(CFG.save_dir, 'best.pth')\n    if not os.path.exists(p):\n        raise FileNotFoundError('best.pth not found. Run training first or adjust save path.')\n    m.load_state_dict(torch.load(p, map_location=CFG.device))\n    m.eval()\n\n    for _, row in tqdm(df_test.iterrows(), total=len(df_test)):\n        image_id = row['id']\n        img_path = os.path.join(CFG.img_dir, f'{image_id}.tif')\n        vol = read_tiff(img_path).astype(np.float32)\n        vol = (vol - vol.mean()) / (vol.std() + 1e-8)\n        vol = vol[np.newaxis, ...]\n        prob = sliding_window_inference(vol, m, CFG.patch_size, overlap=CFG.inference_overlap, tta=CFG.tta)\n        mask = postprocess_mask(prob, thr=0.5, min_size=1000)\n        save_tiff(mask, os.path.join(out_dir, f'{image_id}.tif'))\n\n    # zip\n    zipname = 'submission.zip'\n    with zipfile.ZipFile(zipname, 'w') as z:\n        for p in glob(os.path.join(out_dir, '*.tif')):\n            z.write(p, arcname=os.path.basename(p))\n    print('Saved', zipname)\n\n# ----------------------------- Example usage -----------------------------\nif __name__ == '__main__':\n    # 1) Train (single split)\n    run_training()\n\n    # 2) After training, run inference on test.csv and produce submission.zip\n    run_inference(CFG.test_csv, out_dir='./submission_masks')\n\n    print('This script contains a baseline single-split pipeline. Adjust CFG parameters for your GPU.')\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-11-21T03:38:40.607537Z","iopub.execute_input":"2025-11-21T03:38:40.607906Z","iopub.status.idle":"2025-11-21T04:29:04.075928Z","shell.execute_reply.started":"2025-11-21T03:38:40.607874Z","shell.execute_reply":"2025-11-21T04:29:04.069617Z"}},"outputs":[],"execution_count":null}]}