{"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":"gpu","dataSources":[{"sourceId":117682,"databundleVersionId":15062069,"sourceType":"competition"},{"sourceId":14677261,"sourceType":"datasetVersion","datasetId":9376767}],"dockerImageVersionId":31259,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install /kaggle/input/imagecodecs/imagecodecs-2025.3.30-cp312-cp312-manylinux*.whl --no-index","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport imagecodecs\nimport numpy as np\nimport tifffile as tiff\nimport torch\nimport torch.nn as nn\nimport zipfile","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def normalize(img):\n    img = img.astype(np.float32)\n    img = (img - img.mean()) / (img.std() + 1e-6)\n    return img","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def extract_patches(img, mask, patch_size=64, stride=64):\n    patches = []\n    masks = []\n\n    z, y, x = img.shape\n    for i in range(0, z - patch_size, stride):\n        for j in range(0, y - patch_size, stride):\n            for k in range(0, x - patch_size, stride):\n                patches.append(img[i:i+patch_size, j:j+patch_size, k:k+patch_size])\n                masks.append(mask[i:i+patch_size, j:j+patch_size, k:k+patch_size])\n\n    return patches, masks","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def filter_empty(patches, masks, min_voxels=50):\n    new_p, new_m = [], []\n    for p, m in zip(patches, masks):\n        if m.sum() > min_voxels:\n            new_p.append(p)\n            new_m.append(m)\n    return new_p, new_m","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def patch_generator(img, mask, ps=64, stride=64):\n    z, y, x = img.shape\n    for i in range(0, z-ps, stride):\n        for j in range(0, y-ps, stride):\n            for k in range(0, x-ps, stride):\n                patch = img[i:i+ps, j:j+ps, k:k+ps]\n                lbl   = mask[i:i+ps, j:j+ps, k:k+ps]\n                if lbl.sum() > 50:\n                    yield patch, lbl","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class DoubleConv(nn.Module):\n    def __init__(self, in_c, out_c):\n        super().__init__()\n        self.conv = nn.Sequential(\n            nn.Conv3d(in_c, out_c, 3, padding=1),\n            nn.BatchNorm3d(out_c),\n            nn.ReLU(inplace=True),\n            nn.Conv3d(out_c, out_c, 3, padding=1),\n            nn.BatchNorm3d(out_c),\n            nn.ReLU(inplace=True)\n        )\n\n    def forward(self, x):\n        return self.conv(x)\n\n\nclass UNet3D(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.enc1 = DoubleConv(1, 32)\n        self.enc2 = DoubleConv(32, 64)\n\n        self.pool = nn.MaxPool3d(2)\n\n        self.dec1 = DoubleConv(64, 32)\n        self.out = nn.Conv3d(32, 1, 1)\n\n    def forward(self, x):\n        x1 = self.enc1(x)\n        x2 = self.enc2(self.pool(x1))\n        x = nn.functional.interpolate(x2, scale_factor=2, mode=\"trilinear\")\n        x = self.dec1(x)\n        return self.out(x)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\nmodel = UNet3D().to(device)\n\ncriterion = nn.BCEWithLogitsLoss()\noptimizer = torch.optim.Adam(model.parameters(), lr=1e-4)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_img_dir = \"/kaggle/input/vesuvius-challenge-surface-detection/train_images\"\ntrain_lbl_dir = \"/kaggle/input/vesuvius-challenge-surface-detection/train_labels\"\n\nimage_ids = os.listdir(train_img_dir)\n\nmodel.train()\n\nfor img_id in image_ids:\n    img = tiff.imread(os.path.join(train_img_dir, img_id))\n    mask = tiff.imread(os.path.join(train_lbl_dir, img_id))\n\n    img = normalize(img)\n\n    for p, m in patch_generator(img, mask):\n        p = torch.tensor(p)[None, None].float().to(device)\n        m = torch.tensor(m)[None, None].float().to(device)\n\n        pred = model(p)\n        loss = criterion(pred, m)\n\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n\n        del p, m, pred  # 🔥 IMPORTANT\n        torch.cuda.empty_cache()\n\n    print(img_id, \"loss:\", loss.item())","metadata":{"trusted":true,"scrolled":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model.eval()\nmodel.to(device)\nCHUNK = 8","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_dir = \"/kaggle/input/vesuvius-challenge-surface-detection/test_images\"\ntest_dir_imgs = os.listdir(test_dir)\nos.makedirs(\"submission\", exist_ok=True)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"with torch.no_grad():\n    for test_dir_img in test_dir_imgs:\n        vol = tiff.imread(os.path.join(test_dir, test_dir_img))\n        vol = normalize(vol)\n\n        H, W, D = vol.shape\n        pred = np.zeros((H, W, D), dtype=np.float32)\n\n        for z in range(0, D, CHUNK):\n            chunk = vol[:, :, z:z+CHUNK]          # (H, W, d)\n\n            # 👉 Make it 5D: (B=1, C=1, D, H, W)\n            x = torch.from_numpy(chunk).float()\n            x = x.permute(2, 0, 1)                # (d, H, W)\n            x = x.unsqueeze(0).unsqueeze(0)       # (1,1,d,H,W)\n            x = x.to(device)\n\n            p = model(x)                          # (1,1,d,H,W)\n            p = p.squeeze().cpu().numpy()         # (d,H,W)\n            p = p.transpose(1, 2, 0)              # (H,W,d)\n\n            pred[:, :, z:z+CHUNK] = p\n\n            del x, p\n            torch.cuda.empty_cache()\n\n        binary = (pred > 0.5).astype(np.uint8)\n        \n        out_path = f\"submission/{test_dir_img}\"\n        tiff.imwrite(out_path, binary)\n\n        print(\"Saved:\", test_dir_img)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"zip_path = \"/kaggle/working/submission.zip\"\nwith zipfile.ZipFile(zip_path, \"w\", zipfile.ZIP_DEFLATED) as z:\n    for f in os.listdir(\"submission\"):\n        z.write(f\"submission/{f}\", arcname=f)\n\nprint(\"Created:\", zip_path)","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}