{"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,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":13818749,"sourceType":"datasetVersion","datasetId":8800020},{"sourceId":284794863,"sourceType":"kernelVersion"},{"sourceId":668931,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":501042,"modelId":516237}],"dockerImageVersionId":31192,"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-package-scraper\" --no-deps -q pytorch_lightning monai albumentations imagecodecs # \"numpy==1.26.4\" \"scipy==1.15.3\"\n!pip uninstall -q -y tensorflow  # preventing AttributeError\n# !pip install --upgrade ipython ipython_genutils","metadata":{"trusted":true,"_kg_hide-output":true,"_kg_hide-input":true,"execution":{"iopub.status.busy":"2025-12-08T21:22:06.003286Z","iopub.execute_input":"2025-12-08T21:22:06.003595Z","iopub.status.idle":"2025-12-08T21:22:31.289997Z","shell.execute_reply.started":"2025-12-08T21:22:06.003571Z","shell.execute_reply":"2025-12-08T21:22:31.289231Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%load_ext autoreload\n%autoreload 2","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-08T21:22:31.291639Z","iopub.execute_input":"2025-12-08T21:22:31.291888Z","iopub.status.idle":"2025-12-08T21:22:31.324888Z","shell.execute_reply.started":"2025-12-08T21:22:31.291854Z","shell.execute_reply":"2025-12-08T21:22:31.324275Z"}},"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 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(\".\")\n# CACHE_DIR = Path(\"/kaggle/tmp/dataset_cache\")  # Cache for preprocessed slices\nCACHE_DIR = Path(\"/kaggle/input/vesuvius-surface-slices\")\nCHECKPOINT_DIR = \"/kaggle/input/vesuvius-surface-detection-2d-checkpoints/pytorch/default/6\"\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 = 28\nNUM_EPOCHS = 15\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-12-08T21:22:31.325555Z","iopub.execute_input":"2025-12-08T21:22:31.325754Z","iopub.status.idle":"2025-12-08T21:23:26.485563Z","shell.execute_reply.started":"2025-12-08T21:22:31.325723Z","shell.execute_reply":"2025-12-08T21:23:26.484939Z"}},"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":false}},{"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                try:\n                    num_slices = int(marker_path.read_text().strip())\n                except Exception:\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(\n        cache_dir: Path, volume_file: str, slice_axis: str, slice_idx: int\n    ) -> 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, volume_file: str, slice_axis: str, 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        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, volume_file: str, slice_axis: str, 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(\n                    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(\n                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        # 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        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-12-08T21:23:26.486397Z","iopub.execute_input":"2025-12-08T21:23:26.486970Z","iopub.status.idle":"2025-12-08T21:23:26.555196Z","shell.execute_reply.started":"2025-12-08T21:23:26.486950Z","shell.execute_reply":"2025-12-08T21:23:26.554487Z"}},"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-12-08T21:23:26.555977Z","iopub.execute_input":"2025-12-08T21:23:26.556345Z","iopub.status.idle":"2025-12-08T21:23:26.625685Z","shell.execute_reply.started":"2025-12-08T21:23:26.556325Z","shell.execute_reply":"2025-12-08T21:23:26.624950Z"}},"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-12-08T21:23:26.626564Z","iopub.execute_input":"2025-12-08T21:23:26.626772Z","iopub.status.idle":"2025-12-08T21:23:27.513247Z","shell.execute_reply.started":"2025-12-08T21:23:26.626753Z","shell.execute_reply":"2025-12-08T21:23:27.512559Z"}},"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-12-08T21:23:27.515633Z","iopub.execute_input":"2025-12-08T21:23:27.515915Z","iopub.status.idle":"2025-12-08T21:23:27.600589Z","shell.execute_reply.started":"2025-12-08T21:23:27.515893Z","shell.execute_reply":"2025-12-08T21:23:27.599872Z"}},"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-12-08T21:23:27.601354Z","iopub.execute_input":"2025-12-08T21:23:27.601590Z","iopub.status.idle":"2025-12-08T21:23:30.113549Z","shell.execute_reply.started":"2025-12-08T21:23:27.601550Z","shell.execute_reply":"2025-12-08T21:23:30.112628Z"}},"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-12-08T21:23:30.114440Z","iopub.execute_input":"2025-12-08T21:23:30.114728Z","iopub.status.idle":"2025-12-08T21:23:31.614358Z","shell.execute_reply.started":"2025-12-08T21:23:30.114700Z","shell.execute_reply":"2025-12-08T21:23:31.613546Z"}},"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\nfrom monai.losses import DiceCELoss, TverskyLoss\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau, CosineAnnealingWarmRestarts\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: nn.Module,\n        out_channels: int = 2,\n        learning_rate: float = 1e-3,\n    ):\n        super().__init__()\n        self.save_hyperparameters(ignore=[\"net\"])\n        self.learning_rate = learning_rate\n        # Class label to be ignored in loss calculation\n        self.ignore_index_val = 2\n        self.net_module = net\n        self.out_channels = out_channels\n\n        # Use TverskyLoss to enforce structure continuity for thin lines.\n        # Tversky Index = TP / (TP + alpha * FP + beta * FN)\n        # alpha: Weight for False Positives (FP). Controls penalty for over-segmentation.\n        # beta: Weight for False Negatives (FN). Controls penalty for missed pixels (gaps).\n        # Setting beta > alpha (e.g., beta=0.7, alpha=0.3) penalizes gaps more heavily,\n        # which helps preserve the continuity of thin structures.\n        self.criterion_tversky = TverskyLoss(\n            include_background=True,\n            to_onehot_y=False,\n            softmax=True,\n            alpha=0.3,\n            beta=0.7,\n        )\n\n        # Combine with Standard Smooth Segmentation Loss (DiceCELoss) for better convergence\n        self.criterion_dice_ce = DiceCELoss(\n            include_background=True,\n            to_onehot_y=False,\n            softmax=True,\n            lambda_dice=0.5,\n            lambda_ce=0.5,\n        )\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        return self.net_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        # 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            # 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            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            # 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            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        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        # 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        # 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        # 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        # 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 combined loss\n        loss_tversky = self.criterion_tversky(logits, target_one_hot)\n        loss_dice_ce = self.criterion_dice_ce(logits, target_one_hot)\n        loss = loss_tversky + loss_dice_ce\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\n        optimizer = torch.optim.AdamW(self.parameters(), lr=self.learning_rate, weight_decay=1e-2)\n        # CosineAnnealingWarmRestarts scheduler that restarts every epoch\n        scheduler = CosineAnnealingWarmRestarts(\n            optimizer,\n            T_0=1,      # Number of epochs for the first restart\n            T_mult=1,   # Factor to increase T_0 after each restart (1 means T_0 is constant)\n            eta_min=1e-7 # Minimum learning rate\n        )\n        return {\n            \"optimizer\": optimizer,\n            \"lr_scheduler\": {\n                \"scheduler\": scheduler,\n                \"monitor\": \"val_dice\", # Monitor validation Dice score for logging, not for scheduling\n                \"interval\": \"epoch\",\n            }\n        }","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-08T21:23:31.615232Z","iopub.execute_input":"2025-12-08T21:23:31.615522Z","iopub.status.idle":"2025-12-08T21:23:31.689995Z","shell.execute_reply.started":"2025-12-08T21:23:31.615497Z","shell.execute_reply":"2025-12-08T21:23:31.689190Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from monai.networks.nets import UNet, DynUNet, AttentionUnet, SegResNet, SwinUNETR, FlexibleUNet\n\n# Model architecture\n# net = UNet(\n#     spatial_dims=2,\n#     in_channels=1,\n#     out_channels=2,\n#     channels=(32, 64, 128, 256, 512),   # UNet channel progression\n#     strides=(2, 2, 2, 2),               # Downsampling strides\n#     num_res_units=4,\n#     dropout=0.2,\n# )\n\n# Initialize SwinUNETR model\n# net = SwinUNETR(\n#     in_channels=1,\n#     out_channels=2,\n#     feature_size=36,\n#     spatial_dims=2,\n#     use_v2=True,\n#     drop_rate=0.2,\n#     attn_drop_rate=0.2,\n#     dropout_path_rate=0.2,\n# )\n\n# Initialize FlexibleUNet with EfficientNet backbone\n# Backbone can be \"efficientnet-b0\" to \"efficientnet-b7\"\n# net = FlexibleUNet(\n#     in_channels=1,\n#     out_channels=2,\n#     backbone=\"efficientnet-b4\",\n#     pretrained=False,\n#     spatial_dims=2,\n#     dropout=0.2,\n# )\n\n# SegResNet is the optimized, efficient standard.\n# net = SegResNet(\n#     spatial_dims=2,\n#     in_channels=1,\n#     out_channels=2,\n#     init_filters=16,\n#     dropout_prob=0.2\n# )\n\n# DynUNet for advanced segmentation\nnet = DynUNet(\n    spatial_dims=2,\n    in_channels=1,\n    out_channels=2,\n    kernel_size=[3, 3, 3, 3, 3, 3],\n    strides=[1, 2, 2, 2, 2, 2],\n    upsample_kernel_size=[2, 2, 2, 2, 2],\n    filters=[32, 64, 128, 256, 512, 1024],\n    dropout=0.2,\n)\n\n# Create the Lightning model wrapper\nnet_name = net.__class__.__name__\nmodel = VesuviusSegmentationModel(\n    net=net, learning_rate=LEARNING_RATE,\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-08T21:23:31.690979Z","iopub.execute_input":"2025-12-08T21:23:31.691260Z","iopub.status.idle":"2025-12-08T21:23:32.209877Z","shell.execute_reply.started":"2025-12-08T21:23:31.691238Z","shell.execute_reply":"2025-12-08T21:23:32.209083Z"},"_kg_hide-input":false},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 🏋️ Lit Training\n\nMain training pipeline...","metadata":{}},{"cell_type":"code","source":"import re\nfrom pathlib import Path\nfrom typing import List, Union, Tuple, Optional\n\ndef get_best_checkpoint(\n    checkpoint_dirs: Union[str, Path, List[Union[str, Path]]],\n    name: str = \"\",\n) -> Tuple[str, float]:\n    \"\"\"Finds the checkpoint with the highest val_dice score across multiple directories.\"\"\"\n    # Normalize input to a list of Paths\n    if not isinstance(checkpoint_dirs, list):\n        checkpoint_dirs = [checkpoint_dirs]\n    checkpoint_dirs = [d for d in checkpoint_dirs if Path(d).exists()]\n    if not checkpoint_dirs:\n        print(\"No valid folder provided.\")\n        return None, None\n    \n    checkpoints = []\n    # Regex for val_dice\n    pattern = re.compile(r\"val_dice=?([0-9]+\\.[0-9]+)\")\n    # Iterate over all files in all valid directories\n    for path in [f for d in checkpoint_dirs for f in Path(d).glob(f\"{name}*.ckpt\")]:\n        match = pattern.search(path.name)\n        if not match:\n            continue\n        checkpoints.append((float(match.group(1)), str(path)))\n\n    if not checkpoints:\n        print(\"No valid checkpoints found.\")\n        return None, None\n\n    # Sort by score descending so the best is first\n    checkpoints.sort(key=lambda x: x[0], reverse=True)\n    best_score, best_path = checkpoints[0]\n    print(f\"Found {len(checkpoints)} checkpoints.\")\n    print(f\"Best  (Score={best_score}): {Path(best_path)}\")\n    return best_path, best_score\n\nckpt_path, ckpt_score = get_best_checkpoint(\n    [OUTPUT_DIR, CHECKPOINT_DIR], name=net_name)","metadata":{"trusted":true,"_kg_hide-input":true,"execution":{"iopub.status.busy":"2025-12-08T21:23:32.210740Z","iopub.execute_input":"2025-12-08T21:23:32.211878Z","iopub.status.idle":"2025-12-08T21:23:32.280670Z","shell.execute_reply.started":"2025-12-08T21:23:32.211859Z","shell.execute_reply":"2025-12-08T21:23:32.280068Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping, LearningRateMonitor\nfrom pytorch_lightning.loggers import CSVLogger\nfrom pytorch_lightning.utilities.exceptions import MisconfigurationException\n\n# Callbacks\ncheckpoint_callback = ModelCheckpoint(\n    dirpath=OUTPUT_DIR,\n    filename=net_name + \"-{epoch:02d}-{val_dice:.4f}\",\n    monitor=\"val_dice\",\n    mode=\"max\",\n    save_top_k=3,\n)\nearly_stop_callback = EarlyStopping(\n    monitor=\"val_dice\",\n    patience=10,\n    mode=\"max\",\n)\nlr_monitor = LearningRateMonitor(logging_interval=\"epoch\")\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, lr_monitor], # early_stop_callback\n    logger=csv_logger, # Add csv_logger to the Trainer\n    log_every_n_steps=20,\n    limit_train_batches=0.5,\n    # limit_val_batches=0.5,\n    accumulate_grad_batches=5,\n    val_check_interval=0.5, # Run validation four times per epoch\n    precision='16-mixed', # Enable mixed precision training\n    gradient_clip_val=1.0, # Clips gradient norm to 1.0 to prevent exploding gradients\n)\n\n# Train\ntry:\n    trainer.fit(model, datamodule=data_module, ckpt_path=ckpt_path)\nexcept MisconfigurationException as ex:\n    print(ex)","metadata":{"trusted":true,"_kg_hide-output":true,"_kg_hide-input":false,"execution":{"iopub.status.busy":"2025-12-08T21:23:32.281404Z","iopub.execute_input":"2025-12-08T21:23:32.281765Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nimport seaborn as sns\nfrom IPython.display import display\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) / 'lightning_logs'\n\n# Get the latest version directory\nmetrics_path = log_base_dir / f\"version_{trainer.logger._version}\" / 'metrics.csv'\nprint(f\"Loading metrics from: {metrics_path}\")\n\nif metrics_path.exists():\n    metrics = pd.read_csv(metrics_path)\n    # Remove any columns that are entirely NaN (e.g., from different logging frequencies)\n    display(metrics.dropna(axis=1, how=\"all\").head())\n    # Fill any NaN values by propagating the last valid observation forward (useful for sparse logging)\n    metrics.ffill(inplace=True)\n    # Melt the DataFrame to long-form for plotting\n    # We assume 'epoch' is a reliable identifier for x-axis\n    metrics_melted = metrics.reset_index().melt(\n        id_vars='epoch', var_name='metric', value_name='value')\n    # Define metric groups based on available metrics from VesuviusSegmentationModel\n    metric_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    # Plot metrics for each group in a separate chart\n    for 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        # Plitting...\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        # Apply log scale only for Loss, not for Dice/IoU which are typically 0-1\n        if title == 'Loss':\n            plt.yscale('log')\n        plt.tight_layout()\n    plt.show()","metadata":{"trusted":true,"_kg_hide-input":true},"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},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def resize_nearest_neighbor(\n    image_array: np.ndarray, target_height: int, target_width: int\n) -> 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},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from skimage.morphology import remove_small_objects, ball\nimport torch.nn.functional as F\n\ndef remove_edge_artifacts(mask: np.ndarray, margin: int = 20) -> np.ndarray:\n    \"\"\"Zeroes out margins of a 2D mask to remove edge artifacts.\"\"\"\n    h, w = mask.shape\n    # Handle Height (Top and Bottom)\n    if h > 2 * margin:\n        mask[:margin, :] = 0\n        mask[-margin:, :] = 0\n    # Handle Width (Left and Right)\n    if w > 2 * margin:\n        mask[:, :margin] = 0\n        mask[:, -margin:] = 0\n    return mask\n\ndef get_spherical_kernel(radius):\n    \"\"\"Generates a spherical kernel (structuring element) for 3D morphological operations.\"\"\"\n    # Generate boolean ball on CPU\n    kernel_np = ball(radius)\n    # Convert to float tensor: (1, 1, D, H, W)\n    kernel = torch.from_numpy(kernel_np.astype(np.float32)).unsqueeze(0).unsqueeze(0)\n    return kernel\n\ndef post_process_3d(\n    volume: np.ndarray,\n    min_size: int = 1000,\n    closing_radius: int = 5,\n    device: str = DEVICE,\n) -> np.ndarray:\n    \"\"\"Applies 3D morphological operations to clean up a segmentation volume using a spherical element.\n\n    Steps:\n    1. Performs morphological closing (Dilation -> Erosion) with a spherical kernel.\n    2. Removes small connected components.\n    \"\"\"\n    # Ensure input is boolean\n    binary_vol = volume > 0\n    clean_vol_np = binary_vol\n\n    # 1. Close gaps with Spherical Element (GPU accelerated)\n    if closing_radius > 0:\n        # print(f\"Closing gaps with spherical radius {closing_radius} (GPU accelerated)...\")\n        # Prepare Input: (1, 1, D, H, W)\n        input_tensor = torch.from_numpy(clean_vol_np.astype(np.float32)).unsqueeze(0).unsqueeze(0).to(device)\n        # Prepare Kernel\n        kernel = get_spherical_kernel(closing_radius).to(device)\n        # Dilation: (Input * Kernel) > 0\n        # We use padding=closing_radius to maintain the same spatial dimensions (same as 'same' padding)\n        dilated = (F.conv3d(input_tensor, kernel, padding=closing_radius) > 0).float()\n        # Erosion: 1 - ((1 - Dilated) * Kernel > 0)\n        # This relies on the duality: Erosion(A) = ~Dilation(~A)\n        eroded = 1.0 - (F.conv3d(1.0 - dilated, kernel, padding=closing_radius) > 0).float()\n        # Retrieve result\n        clean_vol_np = eroded.squeeze().cpu().numpy().astype(bool)\n\n    # 2. Remove small objects (CPU-bound)\n    # print(f\"Removing small objects < {min_size} voxels (CPU-bound)...\")\n    clean_vol_np = remove_small_objects(clean_vol_np, min_size=min_size)\n\n    return clean_vol_np.astype(np.uint8)","metadata":{"trusted":true,"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Load best model\nbest_checkpoint_path, _ = get_best_checkpoint(\n    [OUTPUT_DIR, CHECKPOINT_DIR], name=net_name)\n\nassert best_checkpoint_path, \"No checkpoint found in trainer, using current model state.\"\nprint(f\"Loading best checkpoint: {best_checkpoint_path}\")\n# We must pass the 'net' argument because it was ignored in save_hyperparameters\nmodel = VesuviusSegmentationModel.load_from_checkpoint(best_checkpoint_path, net=net)\n\nmodel.eval()\nmodel.to(DEVICE)","metadata":{"trusted":true,"_kg_hide-output":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- 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        # Store original dimensions before resizing\n        original_h, original_w = image_input_for_transform.shape\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        # --- Post-processing: Remove edge artifacts ---\n        pred_class_original_size = remove_edge_artifacts(pred_class_original_size, margin=20)\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    predictions_volume = post_process_3d(predictions_volume, min_size=25*25*25)\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},"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},"outputs":[],"execution_count":null}]}