{"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":"gpu","dataSources":[{"sourceId":117682,"databundleVersionId":14443416,"sourceType":"competition"},{"sourceId":278922471,"sourceType":"kernelVersion"}],"dockerImageVersionId":31193,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# 🏛️ Vesuvius Challenge - Surface Detection Baseline\n\n**Simple 2D Segmentation Approach**\n\nThis notebook implements a baseline for the Vesuvius Challenge Surface Detection competition using:\n- 2D slicing of 3D TIFF volumes along configurable axis (X, Y, or Z)\n- MONAI UNet for 2D segmentation\n- PyTorch Lightning for clean training workflow\n\n### 📊 Data Structure\n\n```\nvesuvius-challenge-surface-detection/\n├── train_images/       # 3D TIFF volumes\n│   ├── 1004283650.tif\n│   └── ...\n├── train_labels/       # 3D mask annotations (same filenames)\n│   ├── 1004283650.tif\n│   └── ...\n└── test_images/        # Test volumes (no labels)\n    └── ...\n```","metadata":{}},{"cell_type":"code","source":"!pip install --no-index --find-links=\"/kaggle/input/surface-detect-package-scraper\" -q monai albumentations imagecodecs --no-deps # \"numpy==1.26.4\" \"scipy==1.15.3\"\n!pip uninstall -q -y tensorflow  # preventing AttributeError","metadata":{"trusted":true,"_kg_hide-output":true,"_kg_hide-input":true,"execution":{"iopub.status.busy":"2025-11-19T11:08:30.016706Z","iopub.execute_input":"2025-11-19T11:08:30.017325Z","iopub.status.idle":"2025-11-19T11:08:36.403574Z","shell.execute_reply.started":"2025-11-19T11:08:30.017303Z","shell.execute_reply":"2025-11-19T11:08:36.401952Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport imagecodecs\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport pytorch_lightning as pl\nimport tifffile\nimport albumentations as A\nimport matplotlib.pyplot as plt\nfrom pathlib import Path\nfrom typing import Tuple, List, Literal, Optional\nfrom torch.utils.data import Dataset, DataLoader\nfrom pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping\nfrom monai.losses import DiceCELoss\nfrom sklearn.model_selection import train_test_split\nfrom albumentations.pytorch import ToTensorV2\nfrom tqdm.auto import tqdm\nimport warnings\n\nwarnings.simplefilter(action='ignore', category=FutureWarning)\n# Data paths\nDATA_DIR = Path(\"/kaggle/input/vesuvius-challenge-surface-detection\")\nTRAIN_IMAGES_DIR = DATA_DIR / \"train_images\"\nTRAIN_LABELS_DIR = DATA_DIR / \"train_labels\"\nTEST_IMAGES_DIR = DATA_DIR / \"test_images\"\nOUTPUT_DIR = Path(\".\")\nCACHE_DIR = Path(\"/kaggle/tmp/dataset_cache\")  # Cache for preprocessed slices\n\n# Dataset configuration\nSLICE_AXIS: Literal[\"x\", \"y\", \"z\"] = \"z\"  # Axis to slice along\nIMAGE_SIZE = (320, 320)                   # Resize slices to this size\nUSE_CACHE = True                          # Enable caching for faster data loading\n\n# Training configuration\nBATCH_SIZE = 36\nNUM_EPOCHS = 10\nLEARNING_RATE = 1e-3\nVAL_SPLIT = 0.1\nSEED = 42\n\n# Device\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n# Set seed\npl.seed_everything(SEED)\n# Create output directory\nOUTPUT_DIR.mkdir(parents=True, exist_ok=True)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2025-11-19T11:08:36.405899Z","iopub.execute_input":"2025-11-19T11:08:36.406262Z","iopub.status.idle":"2025-11-19T11:08:36.749563Z","shell.execute_reply.started":"2025-11-19T11:08:36.406233Z","shell.execute_reply":"2025-11-19T11:08:36.74861Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 📊 Dataset\n\nThe dataset loads 3D TIFF volumes and slices them into 2D samples:\n1. Load entire 3D volume (Z, Y, X)\n2. Extract slice along chosen axis\n3. Stack N adjacent slices as channels (provides 3D context)\n4. **Cache slices as .npy files for fast loading**\n5. Apply transformations\n\n### 🚀 Caching System\n\n**First run:** Creates cache by extracting all slices from 3D volumes\n- Loads each volume once\n- Extracts all valid slices with multi-channel context\n- Saves as .npy files (fast to load)\n- Can take 5-15 minutes depending on data size\n\n**Subsequent runs:** Loads directly from cache\n- 10-100x faster than loading full 3D volumes\n- No repeated TIFF decompression\n\n**Cache invalidation:** Automatic when you change:\n- `SLICE_AXIS`\n\nSet `USE_CACHE = False` to disable caching.","metadata":{"_kg_hide-output":true,"_kg_hide-input":true}},{"cell_type":"code","source":"%%writefile vesuvius_dataset.py\n\nimport numpy as np\nimport torch\nimport tifffile\nimport multiprocessing\nimport albumentations as A\nfrom torch.utils.data import Dataset\nfrom tqdm.auto import tqdm\nfrom functools import lru_cache, partial\nfrom pathlib import Path\nfrom typing import Tuple, List, Literal, Optional\n\nclass VesuviusSliceDataset(Dataset):\n    \"\"\"Dataset for slicing 3D volumes into 2D segmentation samples.\n\n    Key features:\n    - Caches extracted slices as .npy files for fast loading (lazy caching)\n    - Can also be prepopulated with a static method.\n    - Configurable slice axis (x, y, or z)\n    \"\"\"\n\n    def __init__(\n        self,\n        images_dir: Path,\n        labels_dir: Optional[Path],\n        volume_files: List[str],\n        slice_axis: Literal[\"x\", \"y\", \"z\"] = \"z\",\n        cache_dir: Optional[Path] = None,\n        use_cache: bool = True,\n        transform: Optional[A.Compose] = None,\n    ):\n        super().__init__()\n        self.images_dir = images_dir\n        self.labels_dir = labels_dir\n        self.volume_files = volume_files\n        self.slice_axis = slice_axis\n        self.cache_dir = cache_dir\n        self.use_cache = use_cache\n        self.transform = transform\n\n        if self.cache_dir is not None:\n            self.cache_dir.mkdir(parents=True, exist_ok=True)\n\n        # The slice_index is now always built in-memory and cached\n        print(\"Building slice index in-memory (using LRU cache)...\")\n        self.slice_index = self._build_slice_index(\n            self.images_dir,\n            tuple(self.volume_files), # Convert list to tuple for hashability with lru_cache\n            self.slice_axis,\n            self.cache_dir # Pass cache_dir here\n        )\n        print(f\"Dataset initialized with {len(self.slice_index)} slices.\")\n\n    @staticmethod\n    def _map_axis_name_to_index(slice_axis: str) -> int:\n        axis_map = {\"z\": 0, \"y\": 1, \"x\": 2}\n        return axis_map[slice_axis]\n\n    @staticmethod\n    @lru_cache(maxsize=9) # Cache only the most recent set of parameters\n    def _build_slice_index(\n        images_dir: Path,\n        volume_files: Tuple[str, ...],\n        slice_axis: Literal[\"x\", \"y\", \"z\"],\n        cache_dir: Optional[Path] = None, # Added cache_dir argument\n    ) -> List[Tuple[str, str, int]]:\n        \"\"\"Build slice index for all volumes (cached static method).\"\"\"\n        slice_index = []\n\n        for volume_file in tqdm(list(volume_files), desc=\"Building slice index\"):\n            num_slices = 0\n\n            if cache_dir is not None:\n                volume_stem = Path(volume_file).stem\n                marker_dir = cache_dir / volume_stem\n                marker_path = marker_dir / f\"axis_{slice_axis}.done\"\n\n                if marker_path.exists():\n                    try:\n                        num_slices = int(marker_path.read_text().strip())\n                    except ValueError:\n                        print(f\"Warning: Could not read num_slices from marker file {marker_path}. Falling back to reading TIFF.\")\n                \n            if not num_slices:\n                # If cache_dir was None, or marker file didn't exist, or parsing failed\n                volume_tiff = tifffile.imread(str(images_dir / volume_file))\n                if len(volume_tiff.shape) == 2:\n                    num_slices = 1\n                else:\n                    axis_idx_tiff = VesuviusSliceDataset._map_axis_name_to_index(slice_axis)\n                    num_slices = volume_tiff.shape[axis_idx_tiff]\n\n            slice_index += [(volume_file, slice_axis, i) for i in range(num_slices)]\n\n        return slice_index\n\n    @staticmethod\n    def _extract_slice_from_volume(\n        volume: np.ndarray,\n        slice_idx: int,\n        axis: Literal[\"x\", \"y\", \"z\"]\n    ) -> np.ndarray:\n        \"\"\"Extract slice from 3D volume.\"\"\"\n        if len(volume.shape) == 2:\n            return volume[np.newaxis, ...]\n\n        if axis == \"z\":\n            slice_2d = volume[slice_idx, :, :]\n        elif axis == \"y\":\n            slice_2d = volume[:, slice_idx, :]\n        else:  # x\n            slice_2d = volume[:, :, slice_idx]\n\n        # Stack as channels (C, H, W)\n        return slice_2d[np.newaxis, ...]\n\n    @staticmethod\n    def _get_cache_path(cache_dir: Path, volume_file: str, slice_axis: str, slice_idx: int) -> Path:\n        \"\"\"Get path to cached slice file containing both image and mask.\"\"\"\n        volume_stem = Path(volume_file).stem\n        filename = f\"{slice_axis}_slice-{slice_idx:04d}.npz\"\n        return cache_dir / volume_stem / filename\n\n    @staticmethod\n    def _get_image_and_mask_from_volume(\n        image_volume: np.ndarray,\n        label_volume: Optional[np.ndarray],\n        slice_idx: int,\n        slice_axis: str,\n    ) -> Tuple[np.ndarray, np.ndarray]:\n        \"\"\"Extract and return image and mask slices (image: uint8, 0-255 range; mask: uint8 integer).\"\"\"\n        image_slice = VesuviusSliceDataset._extract_slice_from_volume(\n            image_volume, slice_idx, slice_axis)\n\n        if label_volume is not None:\n            label_slice = VesuviusSliceDataset._extract_slice_from_volume(\n                label_volume, slice_idx, slice_axis)\n        else:\n            # If no label, return a zero mask of the same shape, as integer type\n            label_slice = np.zeros(image_slice.shape, dtype=np.uint8)\n\n        # Ensure image_slice is uint8 for caching as per user request\n        image_slice = image_slice.astype(np.uint8)\n\n        # Ensure label_slice is an integer type and within expected range (e.g., 0 or 1 or 2)\n        label_slice = label_slice.astype(np.uint8)\n        assert label_slice.max() <= 2, \\\n            f\"Label values should only be 0, 1, or 2 after processing, found max: {label_slice.max()}\"\n        return image_slice, label_slice\n\n    def _load_from_raw(\n        self,\n        volume_file: str,\n        slice_axis: str,\n        slice_idx: int,\n    ) -> Tuple[np.ndarray, np.ndarray]:\n        \"\"\"Helper to load image and mask from raw TIFF files.\"\"\"\n        image_path = self.images_dir / volume_file\n        image_volume = tifffile.imread(str(image_path))\n\n        label_volume = None\n        if self.labels_dir is not None:\n            label_path = self.labels_dir / volume_file\n            if label_path.exists():\n                label_volume = tifffile.imread(str(label_path))\n\n        image, label_slice = VesuviusSliceDataset._get_image_and_mask_from_volume(\n            image_volume, label_volume, slice_idx, slice_axis)\n        return image, label_slice\n\n    @staticmethod\n    def _save_cache(\n        cache_dir: Path,\n        volume_file: str,\n        slice_axis: str,\n        slice_idx: int,\n        image_data: np.ndarray,\n        mask_data: np.ndarray\n    ):\n        \"\"\"Saves image and mask data for a slice to cache.\"\"\"\n        cache_path = VesuviusSliceDataset._get_cache_path(\n            cache_dir, volume_file, slice_axis, slice_idx)\n        cache_path.parent.mkdir(parents=True, exist_ok=True)\n\n        # Convert to appropriate types for saving (image as uint8, mask as uint8)\n        image_to_save = image_data.astype(np.uint8)\n        label_to_save = mask_data.astype(np.uint8)\n\n        np.savez_compressed(str(cache_path), image=image_to_save, mask=label_to_save)\n\n    @staticmethod\n    def _load_cache(\n        cache_dir: Path,\n        volume_file: str,\n        slice_axis: str,\n        slice_idx: int,\n    ) -> Tuple[np.ndarray, np.ndarray]:\n        \"\"\"Loads image and mask data for a slice from cache.\"\"\"\n        cache_path = VesuviusSliceDataset._get_cache_path(\n            cache_dir, volume_file, slice_axis, slice_idx)\n        data = np.load(str(cache_path), allow_pickle=True)\n        image = data['image'].astype(np.uint8) # Load image as uint8\n        mask = data['mask'].astype(np.uint8)\n        return image, mask\n\n    def __len__(self) -> int:\n        return len(self.slice_index)\n\n    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:\n        volume_file, slice_axis, slice_idx = self.slice_index[idx]\n\n        image = None\n        label_slice = None\n\n        if self.use_cache and self.cache_dir is not None:\n            cache_path = VesuviusSliceDataset._get_cache_path(\n                self.cache_dir, volume_file, slice_axis, slice_idx)\n\n            if cache_path.exists():\n                image, label_slice = VesuviusSliceDataset._load_cache(\n                    self.cache_dir, volume_file, slice_axis, slice_idx)\n            else:\n                image, label_slice = self._load_from_raw(volume_file, slice_axis, slice_idx)\n                VesuviusSliceDataset._save_cache(\n                    self.cache_dir, volume_file, slice_axis, slice_idx, image, label_slice)\n        else:\n            image, label_slice = self._load_from_raw(volume_file, slice_axis, slice_idx)\n\n        # Convert image to float32 and normalize to 0-1 range here\n        image = image.astype(np.float32) / 255.0\n\n        image = np.transpose(image, (1, 2, 0))\n        label_slice = label_slice.squeeze(0)\n\n        if self.transform is not None:\n            transformed = self.transform(image=image, mask=label_slice)\n            image = transformed[\"image\"]\n            label_slice = transformed[\"mask\"]\n        else:\n            # Convert numpy array to PyTorch tensor and reorder dimensions for PyTorch (C, H, W)\n            image = torch.from_numpy(image).permute(2, 0, 1)\n            label_slice = torch.from_numpy(label_slice).unsqueeze(0)\n\n        # The `label_slice` passed to `DiceCELoss` will now contain 0, 1, or 2.\n        # `DiceCELoss` will handle the `ignore_index`.\n        return image.float(), label_slice.float()\n\n    @staticmethod\n    def _process_volume_for_cache(\n        volume_file: str,\n        images_dir: Path,\n        labels_dir: Optional[Path],\n        slice_axis: str,\n        cache_dir: Path\n    ) -> List[Tuple[str, str, int]]:\n        \"\"\"Helper function to process a single volume for cache creation (static).\"\"\"\n\n        volume_stem = Path(volume_file).stem\n        marker_dir = cache_dir / volume_stem\n        marker_path = marker_dir / f\"axis_{slice_axis}.done\"\n\n        if marker_path.exists():\n            # If the marker exists, we can optionally read the number of slices from it\n            # but for now, we just skip re-processing.\n            print(f\"Cache for {volume_file} along axis {slice_axis} already exists. Skipping.\")\n            return [] # Return empty list if already cached\n\n        image_path = images_dir / volume_file\n        image_volume = tifffile.imread(str(image_path))\n\n        label_volume = None\n        if labels_dir is not None:\n            label_path = labels_dir / volume_file\n            if label_path.exists():\n                label_volume = tifffile.imread(str(label_path))\n\n        if len(image_volume.shape) == 2:\n            num_slices = 1\n        else:\n            axis_idx = VesuviusSliceDataset._map_axis_name_to_index(slice_axis)\n            num_slices = image_volume.shape[axis_idx]\n\n        slice_entries = []\n        for slice_idx in range(num_slices):\n            image_slice, label_slice = VesuviusSliceDataset._get_image_and_mask_from_volume(\n                image_volume, label_volume, slice_idx, slice_axis)\n\n            VesuviusSliceDataset._save_cache(\n                cache_dir, volume_file, slice_axis, slice_idx, image_slice, label_slice)\n            slice_entries.append((volume_file, slice_axis, slice_idx))\n\n        # Create a marker file to indicate that caching for this volume and axis is done\n        marker_dir.mkdir(parents=True, exist_ok=True) # Ensure directory exists\n        marker_path.write_text(str(num_slices)) # Store the number of slices in the marker file\n\n        return slice_entries\n\n    @staticmethod\n    def clear_done_markers(\n        cache_dir: Path,\n        slice_axis: Literal[\"x\", \"y\", \"z\"]\n    ) -> None:\n        \"\"\"Deletes all '.done' marker files for a specific slice_axis from cache subfolders.\"\"\"\n        print(f\"Clearing '.done' markers for slice_axis='{slice_axis}' in {cache_dir}...\")\n        if not cache_dir.exists():\n            print(\"Cache directory does not exist, nothing to clear.\")\n            return\n\n        count = 0\n        for subdir in cache_dir.iterdir():\n            marker_path = subdir / f\"axis_{slice_axis}.done\"\n            if not marker_path.exists():\n                continue\n            marker_path.unlink() # Delete the file\n            count += 1\n        print(f\"Cleared {count} marker files for slice_axis='{slice_axis}'.\")\n\n    @staticmethod\n    def prepopulate_cache(\n        images_dir: Path,\n        labels_dir: Optional[Path],\n        slice_axis: Literal[\"x\", \"y\", \"z\"],\n        cache_dir: Path,\n        volume_files: Optional[List[str]] = None,\n        num_processes: int = -1,\n    ) -> None:\n        \"\"\"Statically prepopulate the entire slice cache using multiprocessing.\"\"\"\n        if not cache_dir:\n            raise ValueError(\"cache_dir must be provided for prepopulating cache.\")\n        if not images_dir.exists():\n            raise ValueError(f\"Images directory {images_dir} does not exist.\")\n\n        cache_dir.mkdir(parents=True, exist_ok=True)\n\n        # Clear existing done markers for the current slice_axis to ensure a fresh cache for this axis\n        VesuviusSliceDataset.clear_done_markers(cache_dir, slice_axis)\n\n        if num_processes < 0:\n            num_processes = multiprocessing.cpu_count()\n\n        if volume_files is None:\n            # If volume_files are not provided, discover all .tif files in the directory\n            volume_files = sorted([f.name for f in images_dir.glob(\"*.tif\")])\n        if not volume_files:\n            raise ValueError(\n                f\"No TIFF files found in {images_dir} or `volume_files` was not provided.\")\n\n        # Use functools.partial to fix common arguments for the multiprocessing function\n        process_func = partial(\n            VesuviusSliceDataset._process_volume_for_cache,\n            images_dir=images_dir,\n            labels_dir=labels_dir,\n            slice_axis=slice_axis,\n            cache_dir=cache_dir\n        )\n\n        print(f\"Prepopulating cache for {len(volume_files)} volumes using {num_processes} processes...\")\n        all_slice_index = []\n        with multiprocessing.Pool(processes=num_processes) as pool:\n            # Now, imap_unordered only needs to iterate over the unique argument: volume_file\n            for volume_slices in tqdm(\n                pool.imap_unordered(process_func, volume_files),\n                total=len(volume_files),\n                desc=\"Prepopulating cache\",\n            ):\n                all_slice_index.extend(volume_slices)\n\n        print(f\"Cache prepopulated with {len(all_slice_index)} slices. No master index file saved.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-19T11:08:36.75088Z","iopub.execute_input":"2025-11-19T11:08:36.751145Z","iopub.status.idle":"2025-11-19T11:08:36.764184Z","shell.execute_reply.started":"2025-11-19T11:08:36.75112Z","shell.execute_reply":"2025-11-19T11:08:36.762351Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def visualize_sample(\n    dataset,  # VesuviusSliceDataset\n    idx: int,\n    figsize: Tuple[int, int] = (18, 5) # Increased figsize for colorbar\n):\n    \"\"\"Visualize a sample from the dataset.\"\"\"\n    # Import mark_boundaries here as requested\n    from skimage.segmentation import mark_boundaries\n\n    image, mask = dataset[idx]\n\n    # Convert to numpy\n    image = image.cpu().numpy()\n    mask = mask.cpu().numpy()\n\n    fig, axes = plt.subplots(1, 3, figsize=figsize)\n\n    # Calculate mid_channel for display\n    mid_channel = image.shape[0] // 2\n\n    axes[0].imshow(image[mid_channel], cmap='gray')\n    axes[0].set_title(f'Image (channel {mid_channel}/{image.shape[0]})')\n    axes[0].axis('off')\n\n    # Display image with mask contours overlaid\n    # Ensure the image for mark_boundaries is 2D and normalized for display\n    img4contour = image[mid_channel].copy()\n    if img4contour.max() > 1.0:\n        img4contour = img4contour / img4contour.max()\n\n    # mark_boundaries expects a binary mask (0 or 1)\n    # It will treat any non-zero value as foreground. So, '2' will be treated as foreground.\n    binary_mask = (mask > 0).astype(int)\n    contoured_image = mark_boundaries(\n        img4contour, binary_mask, outline_color=(1, 0, 0))\n    axes[1].imshow(contoured_image)\n    axes[1].set_title('Image with Mask Contours (2 treated as 1)')\n    axes[1].axis('off')\n\n    # Show mask with clarification for label-to-color mapping and add colorbar\n    # Setting vmin=0, vmax=2 to visualize all 3 possible labels (0, 1, 2)\n    im = axes[2].imshow(mask, cmap='hot', vmin=0, vmax=2)\n    axes[2].set_title('Ground Truth Mask (0: Bg, 1: Fg, 2: Ignored)')\n    axes[2].axis('off')\n\n    # Add colorbar\n    cbar = fig.colorbar(im, ax=axes[2], fraction=0.046, pad=0.04)\n    cbar.set_ticks([0, 1, 2])\n    cbar.set_ticklabels(['0 (Background)', '1 (Foreground)', '2 (Ignored)'])\n    cbar.set_label('Label Value')\n\n    plt.tight_layout()\n    plt.show()","metadata":{"trusted":true,"_kg_hide-input":true,"execution":{"iopub.status.busy":"2025-11-19T11:08:36.765226Z","iopub.execute_input":"2025-11-19T11:08:36.765544Z","iopub.status.idle":"2025-11-19T11:08:36.792543Z","shell.execute_reply.started":"2025-11-19T11:08:36.765525Z","shell.execute_reply":"2025-11-19T11:08:36.791728Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from vesuvius_dataset import VesuviusSliceDataset\n\nVesuviusSliceDataset.prepopulate_cache(\n    images_dir=TRAIN_IMAGES_DIR,\n    labels_dir=TRAIN_LABELS_DIR,\n    slice_axis=SLICE_AXIS,\n    cache_dir=CACHE_DIR,\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-19T11:08:36.794441Z","iopub.execute_input":"2025-11-19T11:08:36.794681Z","iopub.status.idle":"2025-11-19T11:20:55.893655Z","shell.execute_reply.started":"2025-11-19T11:08:36.794663Z","shell.execute_reply":"2025-11-19T11:20:55.892608Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 🔄 DataModule\n\nPyTorch Lightning DataModule handles data loading and splitting:","metadata":{}},{"cell_type":"code","source":"class VesuviusDataModule(pl.LightningDataModule):\n    \"\"\"DataModule for train/val split and data loading.\"\"\"\n\n    def __init__(\n        self,\n        images_dir: Path,\n        labels_dir: Path,\n        slice_axis: Literal[\"x\", \"y\", \"z\"],\n        image_size: Tuple[int, int],\n        cache_dir: Optional[Path],\n        use_cache: bool,\n        batch_size: int,\n        val_split: float,\n        num_workers: int = -1,\n        seed: int = 42,\n    ):\n        super().__init__()\n        self.images_dir = images_dir\n        self.labels_dir = labels_dir\n        self.slice_axis = slice_axis\n        self.image_size = image_size\n        self.cache_dir = cache_dir\n        self.use_cache = use_cache\n        self.batch_size = batch_size\n        if num_workers <= 0:\n            num_workers = min(8, os.cpu_count())\n        self.num_workers = num_workers\n        self.val_split = val_split\n        self.seed = seed\n\n    def setup(self, stage: Optional[str] = None):\n        \"\"\"Setup datasets.\"\"\"\n        # Get all TIFF files\n        volume_files = sorted([f.name for f in self.images_dir.glob(\"*.tif\")])\n\n        if not volume_files:\n            raise ValueError(f\"No TIFF files found in {self.images_dir}\")\n\n        print(f\"Found {len(volume_files)} volume files\")\n\n        # Split into train/val\n        train_files, val_files = train_test_split(\n            volume_files, test_size=self.val_split, random_state=self.seed\n        )\n\n        print(f\"Train: {len(train_files)} volumes, Val: {len(val_files)} volumes\")\n\n        # Create transforms\n        self.train_transform = A.Compose([\n            A.Resize(height=self.image_size[0], width=self.image_size[1]),\n            A.HorizontalFlip(p=0.5),\n            A.VerticalFlip(p=0.5),\n            # A.RandomRotate90(p=0.5),\n            # A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.1, rotate_limit=15, p=0.5),\n            A.MultiplicativeNoise(multiplier=(0.9, 1.1), elementwise=True, p=0.5),\n            A.RandomBrightnessContrast(brightness_limit=0.1, contrast_limit=0.1, p=0.5),\n            A.Normalize(mean=[0.5], std=[0.5]),\n            ToTensorV2(),\n        ])\n        self.val_transform = A.Compose([\n            A.Resize(height=self.image_size[0], width=self.image_size[1]),\n            A.Normalize(mean=[0.5], std=[0.5]),\n            ToTensorV2(),\n        ])\n\n        # Create datasets\n        self.train_dataset = VesuviusSliceDataset(\n            images_dir=self.images_dir,\n            labels_dir=self.labels_dir,\n            volume_files=train_files,\n            slice_axis=self.slice_axis,\n            cache_dir=self.cache_dir,\n            use_cache=self.use_cache,\n            transform=self.train_transform,\n        )\n        self.val_dataset = VesuviusSliceDataset(\n            images_dir=self.images_dir,\n            labels_dir=self.labels_dir,\n            volume_files=val_files,\n            slice_axis=self.slice_axis,\n            cache_dir=self.cache_dir,\n            use_cache=self.use_cache,\n            transform=self.val_transform,\n        )\n\n        print(f\"Train dataset: {len(self.train_dataset)} slices\")\n        print(f\"Val dataset: {len(self.val_dataset)} slices\")\n\n    def train_dataloader(self) -> DataLoader:\n        return DataLoader(\n            self.train_dataset,\n            batch_size=self.batch_size,\n            shuffle=True,\n            num_workers=self.num_workers,\n            pin_memory=True,\n        )\n\n    def val_dataloader(self) -> DataLoader:\n        return DataLoader(\n            self.val_dataset,\n            batch_size=self.batch_size,\n            shuffle=False,\n            num_workers=self.num_workers,\n            pin_memory=True,\n        )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-19T11:20:55.894821Z","iopub.execute_input":"2025-11-19T11:20:55.895102Z","iopub.status.idle":"2025-11-19T11:20:55.915276Z","shell.execute_reply.started":"2025-11-19T11:20:55.895078Z","shell.execute_reply":"2025-11-19T11:20:55.913862Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Create data module\ndata_module = VesuviusDataModule(\n    images_dir=TRAIN_IMAGES_DIR,\n    labels_dir=TRAIN_LABELS_DIR,\n    slice_axis=SLICE_AXIS,\n    image_size=IMAGE_SIZE,\n    cache_dir=CACHE_DIR,\n    use_cache=USE_CACHE,\n    batch_size=BATCH_SIZE,\n    val_split=VAL_SPLIT,\n    seed=SEED,\n)\ndata_module.setup()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-19T11:20:55.91641Z","iopub.execute_input":"2025-11-19T11:20:55.916697Z","iopub.status.idle":"2025-11-19T11:20:56.234564Z","shell.execute_reply.started":"2025-11-19T11:20:55.916669Z","shell.execute_reply":"2025-11-19T11:20:56.23382Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Visualize samples\nfor i in range(3):\n    visualize_sample(data_module.train_dataset, i)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-19T11:20:56.236576Z","iopub.execute_input":"2025-11-19T11:20:56.236936Z","iopub.status.idle":"2025-11-19T11:20:57.497684Z","shell.execute_reply.started":"2025-11-19T11:20:56.236911Z","shell.execute_reply":"2025-11-19T11:20:57.496368Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 🧠 Model - MONAI\n\nMONAI UNet wrapped in PyTorch Lightning module:","metadata":{}},{"cell_type":"code","source":"import torch.nn.functional as F\n\nclass VesuviusSegmentationModel(pl.LightningModule):\n    \"\"\"2D segmentation model with MONAI UNet.\n\n    Architecture:\n    - Multi-level encoder-decoder (UNet)\n    - Skip connections for detail preservation\n    - Residual blocks at each level\n    \"\"\"\n\n    def __init__(\n        self,\n        net_module: nn.Module,\n        out_channels: int = 2,\n        learning_rate: float = 1e-3,\n    ):\n        super().__init__()\n        self.save_hyperparameters()\n        self.learning_rate = learning_rate\n        # Class label to be ignored in loss calculation\n        self.ignore_index_val = 2\n        self.module = net_module\n        self.out_channels = out_channels\n\n        # DiceCELoss configured for binary segmentation (2 classes) without `ignore_index` argument.\n        # We will manually handle ignored labels by manipulating the target masks.\n        self.criterion = DiceCELoss(\n            softmax=True,\n            to_onehot_y=False, # We will manually create one-hot targets\n            lambda_dice=0.5,\n            lambda_ce=0.5,\n            include_background=True,\n        )\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        return self.module(x)\n\n    def _compute_metrics(self, preds_logits: torch.Tensor, targets_class_indices: torch.Tensor) -> dict:\n        \"\"\"Compute multi-class (now binary) Dice and IoU metrics.\n        preds_logits: (B, C, H, W) raw logits from the module.\n        targets_class_indices: (B, 1, H, W) with class labels (0, 1, or 2).\n        \"\"\"\n        # Convert logits to probabilities (B, C, H, W)\n        preds_proba = torch.softmax(preds_logits, dim=1)\n        # Get hard predictions (class indices) (B, 1, H, W)\n        preds_hard = torch.argmax(preds_proba, dim=1, keepdim=True)\n\n        # Create a mask to ignore pixels with the ignore_index (2)\n        valid_mask = (targets_class_indices != self.ignore_index_val).float() # (B, 1, H, W)\n\n        num_classes = preds_logits.shape[1] # This will be 2 (background, foreground)\n        dice_scores_per_class = []\n        iou_scores_per_class = []\n\n        for i in range(num_classes):\n            # Create binary masks for the current class 'i'\n            pred_class_i = (preds_hard == i).float() # (B, 1, H, W)\n            target_class_i = (targets_class_indices == i).float() # (B, 1, H, W)\n\n            # Apply valid_mask to only consider pixels that are not ignored\n            pred_class_i_valid = pred_class_i * valid_mask\n            target_class_i_valid = target_class_i * valid_mask\n\n            intersection = (pred_class_i_valid * target_class_i_valid).sum()\n            union_sum_dice = pred_class_i_valid.sum() + target_class_i_valid.sum()\n            union_sum_iou = pred_class_i_valid.sum() + target_class_i_valid.sum() - intersection\n\n            # Handle cases where union might be zero for a class (e.g., no foreground in valid pixels)\n            dice = (2 * intersection + 1e-8) / (union_sum_dice + 1e-8)\n            iou = (intersection + 1e-8) / (union_sum_iou + 1e-8)\n\n            dice_scores_per_class.append(dice)\n            iou_scores_per_class.append(iou)\n\n        # Average over both classes (background and foreground) where valid pixels exist\n        mean_dice = torch.mean(torch.stack(dice_scores_per_class))\n        mean_iou = torch.mean(torch.stack(iou_scores_per_class))\n\n        return {\"dice\": mean_dice, \"iou\": mean_iou}\n\n    def _step(self, batch: Tuple[torch.Tensor, torch.Tensor], prefix: str) -> torch.Tensor:\n        \"\"\"Common step logic for training and validation.\"\"\"\n        images, masks = batch # masks are (B, 1, H, W) with 0, 1, or 2\n        logits = self(images) # logits are (B, 2, H, W)\n\n        if masks.ndim == 3:\n            masks = masks.unsqueeze(1) # ensure (B, 1, H, W)\n\n        # Prepare target for DiceCELoss: manually create one-hot tensor and zero out ignored regions.\n        masks_long_for_onehot = masks.long().squeeze(1) # (B, H, W) with 0, 1, 2\n\n        # Temporarily replace ignore_index_val (2) with 0 for F.one_hot to avoid error\n        # The effect of ignoring will come from zeroing out the one-hot vectors later.\n        temp_masks_for_onehot = masks_long_for_onehot.clone()\n        temp_masks_for_onehot[temp_masks_for_onehot == self.ignore_index_val] = 0\n\n        # Create one-hot target for the 2 module output channels (0 and 1)\n        target_one_hot = F.one_hot(temp_masks_for_onehot, num_classes=self.out_channels).permute(0, 3, 1, 2).float()\n\n        # Create a mask for the actual ignored pixels from the original masks_long_for_onehot\n        # (B, 1, H, W) -> (B, 1, H, W) float\n        ignore_mask_for_target = (masks_long_for_onehot == self.ignore_index_val).unsqueeze(1).float()\n\n        # Zero out the one-hot target entries where pixels are ignored\n        target_one_hot = target_one_hot * (1 - ignore_mask_for_target) # (B, 2, H, W)\n\n        # Compute loss using the manually prepared one-hot target\n        loss = self.criterion(logits, target_one_hot)\n\n        # Metrics are computed using the original masks, which still contain the ignore_index_val\n        metrics = self._compute_metrics(logits, masks)\n\n        self.log(f\"{prefix}_loss\", loss, prog_bar=True)\n        self.log(f\"{prefix}_dice\", metrics[\"dice\"], prog_bar=True)\n        self.log(f\"{prefix}_iou\", metrics[\"iou\"])\n        return loss\n\n    def training_step(self, batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int) -> torch.Tensor:\n        return self._step(batch, \"train\")\n\n    def validation_step(self, batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int) -> torch.Tensor:\n        return self._step(batch, \"val\")\n\n    def configure_optimizers(self):\n        # Using AdamW with ReduceLROnPlateau for adaptive learning rate adjustment\n        optimizer = torch.optim.AdamW(self.parameters(), lr=self.learning_rate, weight_decay=1e-2)\n        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n            optimizer,\n            mode=\"max\",\n            factor=0.5,\n            patience=2,\n            min_lr=1e-6,\n        )\n        return {\n            \"optimizer\": optimizer,\n            \"lr_scheduler\": {\n                \"scheduler\": scheduler,\n                \"interval\": \"epoch\",\n                \"monitor\": \"val_dice\", # ReduceLROnPlateau needs a metric to monitor\n            }\n        }","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-19T11:20:57.499176Z","iopub.execute_input":"2025-11-19T11:20:57.500431Z","iopub.status.idle":"2025-11-19T11:20:57.512998Z","shell.execute_reply.started":"2025-11-19T11:20:57.500379Z","shell.execute_reply":"2025-11-19T11:20:57.512229Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from monai.networks.nets import UNet, AttentionUnet, SegResNet\n\n\n# Model architecture\n# MODEL_CHANNELS = (32, 64, 128, 256, 512)   # UNet channel progression\n# MODEL_STRIDES = (2, 2, 2, 2)              # Downsampling strides\n# net = UNet(\n#     spatial_dims=2,\n#     in_channels=1,\n#     out_channels=2,\n#     channels=MODEL_CHANNELS,\n#     strides=MODEL_STRIDES,\n#     num_res_units=6,  # for Unet only\n#     dropout=0.2,\n# )\n\nnet = SegResNet(\n    spatial_dims=2,  # 2D segmentation\n    in_channels=1,  # Single 2D slice as input\n    out_channels=2,  # Binary segmentation (foreground/background)\n    init_filters=32,  # Start with 32 filters (can go 16/32/48)\n    blocks_down=(1, 2, 2, 4),  # Encoder depth\n    blocks_up=(1, 1, 1),  # Decoder depth\n    dropout_prob=0.2,  # Regularization\n    norm=(\"GROUP\", {\"num_groups\": 8}),\n    act=(\"RELU\", {\"inplace\": True}),\n    use_conv_final=True\n)\n\n# Create model\nmodel = VesuviusSegmentationModel(\n    net_module=net,\n    learning_rate=LEARNING_RATE,\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-19T11:20:57.513804Z","iopub.execute_input":"2025-11-19T11:20:57.514022Z","iopub.status.idle":"2025-11-19T11:20:57.712665Z","shell.execute_reply.started":"2025-11-19T11:20:57.514005Z","shell.execute_reply":"2025-11-19T11:20:57.711309Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 🏋️ Lit Training\n\nMain training pipeline:","metadata":{}},{"cell_type":"code","source":"from pytorch_lightning.loggers import CSVLogger\n\n# Callbacks\ncheckpoint_callback = ModelCheckpoint(\n    dirpath=OUTPUT_DIR,\n    filename=\"vesuvius-{epoch:02d}-{val_dice:.4f}\",\n    monitor=\"val_dice\",\n    mode=\"max\",\n    save_top_k=3,\n)\n\nearly_stop_callback = EarlyStopping(\n    monitor=\"val_dice\",\n    patience=8,\n    mode=\"max\",\n)\n\n# Instantiate CSVLogger\ncsv_logger = CSVLogger(save_dir=OUTPUT_DIR)\n\n# Trainer\ntrainer = pl.Trainer(\n    max_epochs=NUM_EPOCHS,\n    accelerator=\"auto\",\n    callbacks=[checkpoint_callback, early_stop_callback],\n    logger=csv_logger, # Add csv_logger to the Trainer\n    log_every_n_steps=5,\n    limit_train_batches=0.4,\n    limit_val_batches=0.5,\n    accumulate_grad_batches=3,\n    val_check_interval=0.5, # Run validation four times per epoch\n    precision='16-mixed', # Enable mixed precision training\n)\n\n# Train\ntrainer.fit(model, data_module)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-19T11:20:57.713525Z","iopub.execute_input":"2025-11-19T11:20:57.713785Z","execution_failed":"2025-11-19T12:05:07.415Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nfrom IPython.display import display # Ensure display is available\nsns.set()\n\n# Read the metrics.csv using the trainer's logger directory\n# We need to find the latest version directory within the logger's save_dir\nlog_base_dir = Path(trainer.logger.save_dir)\n# Correct the path to include the 'lightning_logs' subdirectory\nlog_version_parent_dir = log_base_dir / 'lightning_logs'\nversion_dirs = sorted([\n    d for d in log_version_parent_dir.iterdir()\n    if d.is_dir() and d.name.startswith('version_')])\n\n# Get the latest version directory\nlatest_version_log_dir = version_dirs[-1]\nmetrics_path = latest_version_log_dir / 'metrics.csv'\nprint(f\"Loading metrics from: {metrics_path}\")\nmetrics = pd.read_csv(metrics_path)\n\n# Remove any columns that are entirely NaN (e.g., from different logging frequencies)\ndisplay(metrics.dropna(axis=1, how=\"all\").head())\n\n# Fill any NaN values by propagating the last valid observation forward (useful for sparse logging)\nmetrics.ffill(inplace=True)\n\n# Melt the DataFrame to long-form for plotting\n# We assume 'epoch' is a reliable identifier for x-axis\nmetrics_melted = metrics.reset_index().melt(\n    id_vars='epoch', var_name='metric', value_name='value')\n\n# Define metric groups based on available metrics from VesuviusSegmentationModel\nmetric_groups = {\n    'Loss': [c for c in metrics.columns if '_loss' in c],\n    'Dice Score': [c for c in metrics.columns if '_dice' in c],\n    'IoU Score': [c for c in metrics.columns if '_iou' in c],\n}\n\n# Plot metrics for each group in a separate chart\nfor title, metric_list in metric_groups.items():\n    # Filter melted DataFrame for the current group\n    group_metrics = metrics_melted[metrics_melted['metric'].isin(metric_list)]\n\n    plt.figure(figsize=(10, 5))\n    sns.lineplot(data=group_metrics, x='epoch', y='value', hue='metric')\n    plt.title(f'{title} over Epochs', fontsize=14, fontweight='bold')\n    plt.xlabel('Epoch', fontsize=12)\n    plt.ylabel(title, fontsize=12)\n    plt.grid(True, alpha=0.3)\n\n    # Apply log scale only for Loss, not for Dice/IoU which are typically 0-1\n    if title == 'Loss':\n        plt.yscale('log')\n\n    plt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"execution_failed":"2025-11-19T12:05:07.415Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 🔮 Inference\n\nPredict on new 3D volumes:","metadata":{}},{"cell_type":"code","source":"def visualize_prediction(\n    image: np.ndarray,\n    pred: np.ndarray,\n    target: np.ndarray = None,\n    figsize: Tuple[int, int] = (18, 6)\n):\n    \"\"\"Visualize prediction vs ground truth.\"\"\"\n    n_plots = 3 if target is not None else 2\n    fig, axes = plt.subplots(1, n_plots, figsize=figsize)\n    \n    # Image\n    if len(image.shape) == 3:\n        mid_channel = image.shape[0] // 2\n        axes[0].imshow(image[mid_channel], cmap='gray')\n    else:\n        axes[0].imshow(image, cmap='gray')\n    axes[0].set_title('Input Image')\n    axes[0].axis('off')\n    \n    # Prediction\n    axes[1].imshow(pred, cmap='hot', vmin=0, vmax=1)\n    axes[1].set_title('Prediction')\n    axes[1].axis('off')\n    \n    # Ground truth (if available)\n    if target is not None:\n        axes[2].imshow(target, cmap='hot', vmin=0, vmax=1)\n        axes[2].set_title('Ground Truth')\n        axes[2].axis('off')\n    \n    plt.tight_layout()\n    plt.show()","metadata":{"trusted":true,"execution":{"execution_failed":"2025-11-19T12:05:07.415Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Helper function for nearest neighbor resizing without cv2\ndef resize_nearest_neighbor(image_array: np.ndarray, target_height: int, target_width: int) -> np.ndarray:\n    \"\"\"Resizes a 2D numpy array using nearest neighbor interpolation.\"\"\"\n    original_h, original_w = image_array.shape\n    \n    # Create coordinate grids for the target image by sampling original indices\n    y_indices = np.linspace(0, original_h - 1, target_height, dtype=np.int32)\n    x_indices = np.linspace(0, original_w - 1, target_width, dtype=np.int32)\n\n    # Use outer product for broadcasting to create the full 2D grid and sample from the image\n    resized_image = image_array[np.ix_(y_indices, x_indices)]\n    return resized_image","metadata":{"trusted":true,"execution":{"execution_failed":"2025-11-19T12:05:07.415Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# import cv2 # Removed: Not needed with custom resize\n\nmodel.eval()\nmodel.to(DEVICE)\n\n# --- Setup volume_path (assuming it's a test image, for demonstration) ---\n# In a real scenario, this block would typically be part of an inference function\n# where the path to the volume is passed as an argument.\n# For this notebook's flow, we'll get the first test image.\ntest_volume_files = sorted([f.name for f in TEST_IMAGES_DIR.glob(\"*.tif\")])\nif not test_volume_files:\n    raise FileNotFoundError(f\"No TIFF files found in {TEST_IMAGES_DIR} for inference.\")\n\ntest_filenames = []\nfor test_volume_file in tqdm(test_volume_files):\n    volume_path = TEST_IMAGES_DIR / test_volume_file\n    # Load volume\n    volume = tifffile.imread(str(volume_path))\n\n    # Setup inference parameters\n    # Use VesuviusSliceDataset's static method for axis mapping\n    num_slices = volume.shape[VesuviusSliceDataset._map_axis_name_to_index(SLICE_AXIS)]\n    # Use the validation transform from the data module\n    transform = data_module.val_transform\n\n    predictions_list_2d = [] # Store 2D slices here\n    # Iterate directly over slices, as NUM_ADJACENT_SLICES = 1\n    for i in tqdm(range(num_slices)):\n        # Extract a single slice directly\n        # _extract_slice_from_volume returns (1, H, W), squeeze to (H, W)\n        image_input_for_transform = \\\n            VesuviusSliceDataset._extract_slice_from_volume(volume, i, SLICE_AXIS).squeeze(0)\n\n        # Store original dimensions before resizing\n        original_h, original_w = image_input_for_transform.shape\n\n        transformed = transform(image=image_input_for_transform)\n        image_tensor = transformed[\"image\"].unsqueeze(0).to(DEVICE) # Add batch dimension (B, C, H, W)\n\n        # Predict\n        with torch.no_grad():\n            logits = model(image_tensor) # Model now outputs logits (B, C, H, W) for multi-class\n            # Apply softmax to get probabilities, then argmax to get predicted class for each pixel\n            pred_proba = torch.softmax(logits, dim=1) # (B, C, H, W)\n            pred_class = torch.argmax(pred_proba, dim=1).cpu().numpy()[0] # (H, W) predicted class labels\n\n        # Resize prediction back to original slice dimensions using nearest neighbor interpolation\n        # Using custom numpy resize instead of Albumentations with cv2\n        pred_class_original_size = resize_nearest_neighbor(pred_class, original_h, original_w)\n\n        predictions_list_2d.append(pred_class_original_size)\n\n    # Stack the 2D slices into a 3D volume. This volume has its first dimension as the slice axis.\n    # If SLICE_AXIS='z', predictions will be (Z, Y, X)\n    # If SLICE_AXIS='y', predictions will be (Y, Z, X)\n    # If SLICE_AXIS='x', predictions will be (X, Z, Y)\n    predictions_volume = np.stack(predictions_list_2d, axis=0).astype(np.uint8)\n\n    print(f\"Inference for {volume.shape} generated {predictions_volume.shape} slices\"\n          f\" of predictions with values {np.unique(predictions_volume)}.\")\n\n    # Save predictions as a TIFF volume within the new subfolder\n    # This saves the volume in the orientation (SLICE_AXIS_DIM, OtherDim1, OtherDim2)\n    test_filename = f\"{Path(volume_path).stem}.tif\"\n    tifffile.imwrite(test_filename, predictions_volume)\n    test_filenames.append(test_filename)\n\n# Optional: Visualize a sample prediction\nif len(predictions_list_2d) > 0:\n    # Take a slice from the original volume corresponding to the first prediction\n    original_slice_for_display = VesuviusSliceDataset._extract_slice_from_volume(volume, 0, SLICE_AXIS).squeeze(0)\n    # For visualization, ensure the prediction is a single channel (class labels)\n    visualize_prediction(original_slice_for_display, predictions_list_2d[0])","metadata":{"trusted":true,"execution":{"execution_failed":"2025-11-19T12:05:07.415Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Export submission","metadata":{}},{"cell_type":"code","source":"import zipfile\n\nwith zipfile.ZipFile('submission.zip', 'w', zipfile.ZIP_DEFLATED) as zipf:\n    for filename in tqdm(test_filenames, desc=\"Zipping files\"):\n        if not os.path.exists(filename):\n            print(f\"Missing <> {filename}\")\n            continue\n        zipf.write(filename)\n        os.remove(filename)","metadata":{"trusted":true,"execution":{"execution_failed":"2025-11-19T12:05:07.415Z"}},"outputs":[],"execution_count":null}]}