{"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":[{"sourceId":117682,"databundleVersionId":15062069,"sourceType":"competition"}],"dockerImageVersionId":31260,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install -q imagecodecs","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-02-03T10:14:50.213966Z","iopub.execute_input":"2026-02-03T10:14:50.214358Z","iopub.status.idle":"2026-02-03T10:14:53.607861Z","shell.execute_reply.started":"2026-02-03T10:14:50.214331Z","shell.execute_reply":"2026-02-03T10:14:53.607053Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nimport tifffile\nimport cv2\nimport random\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.optim import Adam\nfrom torch.amp import autocast, GradScaler\n\n# --- Configuration ---\nCFG = {\n    'data_dir': '/kaggle/input/vesuvius-challenge-surface-detection',\n    'device': 'cuda',\n    'img_size': (64, 128, 128),  # (Depth, Height, Width)\n    'batch_size': 1,             \n    'epochs': 15,\n    'lr': 1e-4,\n    # Restored to 2 to improve training speed (8s/it -> ~2s/it)\n    'num_workers': 2, \n    'threshold': 0.5,\n    'patch_overlap': 16\n}\n\n# --- Seed ---\ndef seed_everything(seed):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = True\nseed_everything(42)\n\n# --- 3D U-Net Model ---\nclass DoubleConv3d(nn.Module):\n    def __init__(self, in_channels, out_channels):\n        super().__init__()\n        self.double_conv = nn.Sequential(\n            nn.Conv3d(in_channels, out_channels, kernel_size=3, padding=1),\n            nn.BatchNorm3d(out_channels),\n            nn.ReLU(inplace=True),\n            nn.Conv3d(out_channels, out_channels, kernel_size=3, padding=1),\n            nn.BatchNorm3d(out_channels),\n            nn.ReLU(inplace=True)\n        )\n    def forward(self, x): return self.double_conv(x)\n\nclass Down3d(nn.Module):\n    def __init__(self, in_channels, out_channels):\n        super().__init__()\n        self.maxpool_conv = nn.Sequential(nn.MaxPool3d(2), DoubleConv3d(in_channels, out_channels))\n    def forward(self, x): return self.maxpool_conv(x)\n\nclass Up3d(nn.Module):\n    def __init__(self, in_channels, out_channels, trilinear=True):\n        super().__init__()\n        if trilinear:\n            self.up = nn.Upsample(scale_factor=2, mode='trilinear', align_corners=True)\n            self.conv = DoubleConv3d(in_channels + out_channels, out_channels)\n        else:\n            self.up = nn.ConvTranspose3d(in_channels, in_channels // 2, kernel_size=2, stride=2)\n            self.conv = DoubleConv3d(in_channels, out_channels)\n\n    def forward(self, x1, x2):\n        x1 = self.up(x1)\n        diffZ = x2.size()[2] - x1.size()[2]\n        diffY = x2.size()[3] - x1.size()[3]\n        diffX = x2.size()[4] - x1.size()[4]\n        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2, diffZ // 2, diffZ - diffZ // 2])\n        x = torch.cat([x2, x1], dim=1)\n        return self.conv(x)\n\nclass UNet3d(nn.Module):\n    def __init__(self, in_channels=1, n_classes=1):\n        super(UNet3d, self).__init__()\n        self.inc = DoubleConv3d(in_channels, 32)\n        self.down1 = Down3d(32, 64)\n        self.down2 = Down3d(64, 128)\n        self.down3 = Down3d(128, 256)\n        self.up1 = Up3d(256, 128)\n        self.up2 = Up3d(128, 64)\n        self.up3 = Up3d(64, 32)\n        self.outc = nn.Conv3d(32, n_classes, kernel_size=1)\n    def forward(self, x):\n        x1 = self.inc(x); x2 = self.down1(x1); x3 = self.down2(x2); x4 = self.down3(x3)\n        x = self.up1(x4, x3); x = self.up2(x, x2); x = self.up3(x, x1)\n        return self.outc(x)\n\n# --- Dataset ---\nclass VesuviusDataset(Dataset):\n    def __init__(self, df, mode='train', cfg=CFG):\n        self.df = df\n        self.mode = mode\n        self.cfg = cfg\n        self.crop_size = cfg['img_size']\n    def __len__(self): return len(self.df) * 10\n\n    def __getitem__(self, idx):\n        vol_idx = idx % len(self.df)\n        row = self.df.iloc[vol_idx]\n        img_path = os.path.join(self.cfg['data_dir'], f\"{self.mode}_images\", f\"{row['id']}.tif\")\n        \n        volume = tifffile.imread(img_path)\n        \n        if volume.ndim == 2: volume = np.expand_dims(volume, axis=-1)\n        volume = np.transpose(volume, (2, 0, 1)) # (D, H, W)\n        volume = volume.astype(np.float32) / 65535.0 \n        \n        D, H, W = volume.shape\n        \n        if self.mode == 'train':\n            mask_path = os.path.join(self.cfg['data_dir'], f\"{self.mode}_labels\", f\"{row['id']}.tif\")\n            mask = tifffile.imread(mask_path)\n            if mask.ndim == 3: mask = mask[:, :, 0]\n            \n            # Random Crop\n            d_start = random.randint(0, max(0, D - self.crop_size[0]))\n            h_start = random.randint(0, max(0, H - self.crop_size[1]))\n            w_start = random.randint(0, max(0, W - self.crop_size[2]))\n            d_end = min(d_start + self.crop_size[0], D)\n            h_end = min(h_start + self.crop_size[1], H)\n            w_end = min(w_start + self.crop_size[2], W)\n            \n            volume_crop = volume[d_start:d_end, h_start:h_end, w_start:w_end]\n            mask_crop = mask[h_start:h_end, w_start:w_end]\n            \n            # Padding\n            pad_d = self.crop_size[0] - (d_end - d_start)\n            pad_h = self.crop_size[1] - (h_end - h_start)\n            pad_w = self.crop_size[2] - (w_end - w_start)\n            if pad_d > 0 or pad_h > 0 or pad_w > 0:\n                volume_crop = np.pad(volume_crop, ((0, pad_d), (0, pad_h), (0, pad_w)), mode='constant')\n                mask_crop = np.pad(mask_crop, ((0, pad_h), (0, pad_w)), mode='constant', constant_values=0)\n\n            valid_mask = (mask_crop != 2).astype(np.float32) # 2 is unlabeled\n            target_mask = (mask_crop == 1).astype(np.float32)\n            \n            volume_crop = np.expand_dims(volume_crop, axis=0)\n            return torch.tensor(volume_crop, dtype=torch.float32), \\\n                   torch.tensor(target_mask, dtype=torch.float32).unsqueeze(0), \\\n                   torch.tensor(valid_mask, dtype=torch.float32).unsqueeze(0)\n        else:\n            return row['id'], volume.shape\n\n# --- Loss ---\ndef bce_loss(y_pred, y_true, valid_mask=1):\n    if y_pred.dim() == 5 and y_true.dim() == 4:\n        y_true = y_true.unsqueeze(2).expand_as(y_pred)\n        valid_mask = valid_mask.unsqueeze(2).expand_as(y_pred)\n    \n    loss = F.binary_cross_entropy_with_logits(y_pred, y_true, reduction='none')\n    loss = loss * valid_mask\n    return loss.sum() / (valid_mask.sum() + 1e-8)\n\ndef dice_loss(y_pred, y_true, valid_mask=1):\n    if y_pred.dim() == 5 and y_true.dim() == 4:\n        y_true = y_true.unsqueeze(2).expand_as(y_pred)\n        valid_mask = valid_mask.unsqueeze(2).expand_as(y_pred)\n\n    y_pred = torch.sigmoid(y_pred) * valid_mask\n    y_true = y_true * valid_mask\n    return 1 - (2. * (y_true * y_pred).sum() + 1) / (y_true.sum() + y_pred.sum() + 1)\n\n# --- Training ---\ndef train():\n    print(\"Loading Data...\")\n    train_df = pd.read_csv(os.path.join(CFG['data_dir'], 'train.csv'))\n    \n    # --- FIX: Filter dataset to remove IDs that are missing files ---\n    print(\"Checking for valid image files...\")\n    valid_ids = []\n    for _, row in train_df.iterrows():\n        img_path = os.path.join(CFG['data_dir'], f\"train_images\", f\"{row['id']}.tif\")\n        if os.path.exists(img_path):\n            valid_ids.append(row['id'])\n    \n    original_count = len(train_df)\n    train_df = train_df[train_df['id'].isin(valid_ids)].reset_index(drop=True)\n    print(f\"Filtered dataset: {len(train_df)}/{original_count} volumes have valid files.\")\n    # ---------------------------------------------------------------\n    \n    train_dataset = VesuviusDataset(train_df, mode='train')\n    train_loader = DataLoader(train_dataset, batch_size=CFG['batch_size'], shuffle=True, num_workers=CFG['num_workers'], pin_memory=True)\n    \n    model = UNet3d().to(CFG['device'])\n    optimizer = Adam(model.parameters(), lr=CFG['lr'])\n    scaler = GradScaler(\"cuda\") \n    \n    best_loss = float('inf')\n    print(f\"Starting Training...\")\n    \n    for epoch in range(CFG['epochs']):\n        model.train()\n        epoch_loss = 0\n        pbar = tqdm(train_loader, desc=f\"Epoch {epoch+1}/{CFG['epochs']}\")\n        \n        for images, masks, valid_masks in pbar:\n            images = images.to(CFG['device'])\n            masks = masks.to(CFG['device'])\n            valid_masks = valid_masks.to(CFG['device'])\n            \n            optimizer.zero_grad()\n            \n            with autocast(device_type='cuda'): \n                outputs = model(images)\n                loss = bce_loss(outputs, masks, valid_masks) + dice_loss(outputs, masks, valid_masks)\n            \n            scaler.scale(loss).backward()\n            scaler.step(optimizer)\n            scaler.update()\n            \n            epoch_loss += loss.item()\n            pbar.set_postfix({'loss': loss.item()})\n            \n        print(f\"Epoch {epoch+1} Loss: {epoch_loss / len(train_loader)}\")\n        if epoch_loss < best_loss:\n            torch.save(model.state_dict(), 'best_model.pth')\n            best_loss = epoch_loss\n\n# --- Inference ---\ndef predict(model, volume, cfg):\n    model.eval()\n    D, H, W = volume.shape\n    pd, ph, pw = cfg['img_size']\n    final_mask = np.zeros((H, W), dtype=np.float32)\n    count_buffer = np.zeros((H, W), dtype=np.float32)\n    \n    volume = volume.astype(np.float32) / 65535.0\n    \n    with torch.no_grad():\n        for h in range(0, H, ph - 16):\n            for w in range(0, W, pw - 16):\n                h_end = min(h + ph, H)\n                w_end = min(w + pw, W)\n                \n                vol_patch = volume[:pd, h:h_end, w:w_end] \n                curr_d = vol_patch.shape[0]\n                if curr_d < pd:\n                    vol_patch = np.pad(vol_patch, ((0, pd - curr_d), (0,0), (0,0)), mode='constant')\n\n                x = torch.from_numpy(vol_patch).unsqueeze(0).unsqueeze(0).to(cfg['device'])\n                with autocast(device_type='cuda'):\n                    out = model(x)\n                \n                prob = torch.sigmoid(out).squeeze().cpu().numpy()\n                prob_2d = np.mean(prob, axis=0) \n                \n                final_mask[h:h_end, w:w_end] += prob_2d[:h_end-h, :w_end-w]\n                count_buffer[h:h_end, w:w_end] += 1\n                \n    return final_mask / (count_buffer + 1e-8)\n\n# --- Main ---\ndef main():\n    train()\n    print(\"Inference...\")\n    test_df = pd.read_csv(os.path.join(CFG['data_dir'], 'test.csv'))\n    model = UNet3d().to(CFG['device'])\n    model.load_state_dict(torch.load('best_model.pth'))\n    \n    sub = []\n    for idx, row in tqdm(test_df.iterrows(), total=len(test_df)):\n        vol_id = row['id']\n        vol = tifffile.imread(os.path.join(CFG['data_dir'], f\"test_images/{vol_id}.tif\"))\n        if vol.ndim == 2: vol = np.expand_dims(vol, axis=-1)\n        vol = np.transpose(vol, (2, 0, 1))\n        \n        pred = predict(model, vol, CFG)\n        sub.append({'id': vol_id, 'predicted': \",\".join(map(str, pred.flatten()))})\n        cv2.imwrite(f\"{vol_id}_pred.png\", (pred * 255).astype(np.uint8))\n        \n    pd.DataFrame(sub).to_csv('submission.csv', index=False)\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-03T10:14:53.609646Z","iopub.execute_input":"2026-02-03T10:14:53.609923Z","iopub.status.idle":"2026-02-03T10:27:47.273091Z","shell.execute_reply.started":"2026-02-03T10:14:53.609894Z","shell.execute_reply":"2026-02-03T10:27:47.271805Z"}},"outputs":[],"execution_count":null}]}