{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":117682,"databundleVersionId":15062069,"sourceType":"competition"},{"sourceId":14330461,"sourceType":"datasetVersion","datasetId":8944423},{"sourceId":289294474,"sourceType":"kernelVersion"},{"sourceId":724325,"sourceType":"modelInstanceVersion","modelInstanceId":551208,"modelId":563810},{"sourceId":738346,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":545966,"modelId":558965}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# DONT CARE THIS","metadata":{}},{"cell_type":"markdown","source":"```\n# Summary\nInput: .tif\nOutput: binary segmentation mask\nFramework: nnUnetv2\n\n# Autoconfig by nnUnet\n- Patch size\n- Batch size\n- Depth, Norm, Data aug, lr scheduler\n\n# Manualconfig\n- epochs (chooses from list, cannot choose outside the predefined list)\n- panner\n- fold\n- model config\n- gpu\n```","metadata":{}},{"cell_type":"code","source":"import os\nimport json\nimport shutil\nimport subprocess\nfrom functools import partial\nfrom multiprocessing import Pool\nfrom pathlib import Path\nfrom typing import Optional, Tuple, List, Literal, Union\n\nEpochs = Literal[1, 5, 10, 20, 50, 100, 250, 500, 750, 1000, 2000, 4000, 8000]\nDATA_DIR = Path(\"/kaggle/input/vesuvius-challenge-surface-detection\")\nPREPARED_DATA_DIR = Path(\"/kaggle/input/vesuvius-surface-nnunet-preprocessed\")\nWORKING_DIR = Path(\"/kaggle/temp\") # large file storage\nOUTPUT_DIR = Path(\"/kaggle/working\") # persisted output\n\nNNUNET_BASE = WORKING_DIR / \"nnUNet_data\"\nNNUNET_RAW = NNUNET_BASE / \"nnUNet_raw\"\nNNUNET_PREPROCESSED = NNUNET_BASE / \"nnUNet_preprocessed\"\nNNUNET_RESULTS = OUTPUT_DIR / \"nnUNet_results\"\n\nDS_ID = 100\nDS_NAME = f\"Dataset{DS_ID:03d}_VesuviusSurface\" # DatasetXXXX_{CustomName}\n\nFOLD: Union[int, str] = \"all\"\nCONFIGURATION = \"3d_fullres\"\nPLANNER = \"nnUNetPlannerResEncM\"\nPLANS_NAME = \"nnUNetResEncUNetMPlans\"\n\nNUM_WORKERS = os.cpu_count() or 4\nEPOCHS: Optional[Epochs] = 1000\nCOMMAND_TIMEOUT: Optional[int] = 18000\n\nimport torch\nNUM_GPUS: int = torch.cuda.device_count() ","metadata":{"execution":{"iopub.status.busy":"2026-01-27T13:26:20.312482Z","iopub.execute_input":"2026-01-27T13:26:20.312748Z","iopub.status.idle":"2026-01-27T13:26:24.544664Z","shell.execute_reply.started":"2026-01-27T13:26:20.312718Z","shell.execute_reply":"2026-01-27T13:26:24.543907Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def _get_trainer_name_simple(epochs: Optional[int]) -> str:\n    \"\"\"Get trainer class name based on epochs (simple version for path construction).\"\"\"\n    if epochs is None or epochs == 1000:\n        return \"nnUNetTrainer\"\n    elif epochs == 1:\n        return \"nnUNetTrainer_1epoch\"  # Special case: singular form\n    else:\n        return f\"nnUNetTrainer_{epochs}epochs\"\n    \ndef get_training_output_dir(\n    epochs: Optional[Epochs] = None,\n    plans: str = PLANS_NAME,\n    config: str = CONFIGURATION,\n    fold: Union[int, str] = FOLD\n) -> Path:\n    \"\"\"\n    Get the training output directory path based on configuration.\n    \n    nnUNet creates this folder structure:\n    NNUNET_RESULTS/Dataset100_VesuviusSurface/nnUNetTrainer__nnUNetResEncUNetMPlans__3d_lowres/fold_all/\n    \n    Use this to find checkpoints, logs, and progress.png\n    \"\"\"\n    _epochs = epochs if epochs is not None else EPOCHS\n    trainer = _get_trainer_name_simple(_epochs)\n    return NNUNET_RESULTS / DS_NAME / f\"{trainer}__{plans}__{config}\" / f\"fold_{fold}\"\n\ndef get_progress_image_path(\n    epochs: Optional[Epochs] = None,\n    plans: str = PLANS_NAME,\n    config: str = CONFIGURATION,\n    fold: Union[int, str] = FOLD\n) -> Path:\n    \"\"\"Get path to training progress image (loss curves, metrics over epochs).\"\"\"\n    return get_training_output_dir(epochs, plans, config, fold) / \"progress.png\"","metadata":{"execution":{"iopub.status.busy":"2026-01-27T13:26:24.546162Z","iopub.execute_input":"2026-01-27T13:26:24.546561Z","iopub.status.idle":"2026-01-27T13:26:24.553652Z","shell.execute_reply.started":"2026-01-27T13:26:24.546536Z","shell.execute_reply":"2026-01-27T13:26:24.552745Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Environment\n\n- nnUNet_raw: Where nnUNet looks for raw dataset\n- nnUNet_preprocessed: Where preprocessed data is stored\n- nnUNet_results: Where trained models are saved\n- nnUNet_compile: Disable torch.compile (can cause issues)","metadata":{}},{"cell_type":"code","source":"def setup_environment():\n    \"\"\"Set up nnUNet environment variables and directories.\"\"\"\n    for d in [NNUNET_RAW, NNUNET_PREPROCESSED, NNUNET_RESULTS, OUTPUT_DIR]:\n        d.mkdir(parents=True, exist_ok=True)\n    \n    os.environ[\"nnUNet_raw\"] = str(NNUNET_RAW)\n    os.environ[\"nnUNet_preprocessed\"] = str(NNUNET_PREPROCESSED)\n    os.environ[\"nnUNet_results\"] = str(NNUNET_RESULTS)\n    os.environ[\"nnUNet_compile\"] = \"true\"\n    print(f\"nnUNet_raw: {NNUNET_RAW}\")\n    print(f\"nnUNet_preprocessed: {NNUNET_PREPROCESSED}\")\n    print(f\"nnUNet_results: {NNUNET_RESULTS}\")\n    print(f\"nnUNet_USE_BLOSC2: {os.environ.get('nnUNet_USE_BLOSC2', 'not set')} (0=NPZ, 1=blosc2)\")\n    print(f\"NUM_WORKERS: {NUM_WORKERS}\")\n\n\nsetup_environment()","metadata":{"execution":{"iopub.status.busy":"2026-01-27T13:26:24.554507Z","iopub.execute_input":"2026-01-27T13:26:24.554832Z","iopub.status.idle":"2026-01-27T13:26:24.574825Z","shell.execute_reply.started":"2026-01-27T13:26:24.554796Z","shell.execute_reply":"2026-01-27T13:26:24.574071Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def _link_prepared_preprocessed() -> bool:\n    \"\"\"\n    Link/copy pre-prepared preprocessed data if available.\n    \n    Copies the folder structure and metadata files (which nnUNet may need to modify),\n    but symlinks the heavy .npz/.b2nd data files to save space.\n    \n    Handles two possible structures:\n    1. PREPARED_DATA_DIR points directly to Dataset100_* folder\n    2. PREPARED_DATA_DIR contains Dataset100_* as subfolder\n    \n    Returns True if linked/copied successfully.\n    \"\"\"\n    if not PREPARED_DATA_DIR.exists():\n        return False\n    \n    source_dir = PREPARED_DATA_DIR\n    if not (source_dir / \"dataset.json\").exists():\n        dataset_folders = list(PREPARED_DATA_DIR.glob(f\"Dataset*_{DS_NAME.split('_')[1]}*\"))\n        if not dataset_folders:\n            dataset_folders = list(PREPARED_DATA_DIR.glob(\"Dataset*\"))\n        if dataset_folders:\n            source_dir = dataset_folders[0]\n        else:\n            print(f\"No dataset folder found in {PREPARED_DATA_DIR}\")\n            return False\n    \n    dest_dir = NNUNET_PREPROCESSED / DS_NAME\n    if dest_dir.exists():\n        print(f\"Preprocessed data already exists at {dest_dir}\")\n        return True\n    \n    dest_dir.mkdir(parents=True, exist_ok=True)\n    # copy_patterns = [\"*.json\", \"*.pkl\", \"*.txt\"]\n    symlink_patterns = [\"*.npz\", \"*.npy\", \"*.b2nd\"]\n\n    cnt_copy = 0   \n    cnt_symlink = 0\n    for src_path in source_dir.rglob(\"*\"): # recursive\n        if src_path.is_dir():\n            continue\n        # Compute relative path and create target path\n        rel_path = src_path.relative_to(source_dir)\n        dst_path = dest_dir / rel_path\n        dst_path.parent.mkdir(parents=True, exist_ok=True)\n        \n        # Check if this is a heavy data file (symlink) or metadata (copy)\n        is_data_file = any(src_path.match(pat) for pat in symlink_patterns)\n        \n        if is_data_file:\n            if not dst_path.exists():\n                dst_path.symlink_to(src_path.resolve())\n                cnt_symlink += 1\n        else:\n            if not dst_path.exists():\n                shutil.copy2(src_path, dst_path)\n                cnt_copy += 1\n    print(f\"Prepared preprocessed data: {cnt_copy} files copied, {cnt_symlink} files symlinked\")\n    print(f\"Location: {dest_dir}\")\n    return True\n","metadata":{"execution":{"iopub.status.busy":"2026-01-27T13:26:24.575835Z","iopub.execute_input":"2026-01-27T13:26:24.576121Z","iopub.status.idle":"2026-01-27T13:26:24.589876Z","shell.execute_reply.started":"2026-01-27T13:26:24.576092Z","shell.execute_reply":"2026-01-27T13:26:24.589157Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Install nnUNetv2, tifffile","metadata":{}},{"cell_type":"code","source":"!mkdir -p /kaggle/temp\n!mkdir predictions_tiff\n# !pip install nnunetv2 nibabel tifffile tqdm -q --no-index -f \"/kaggle/input/surface-package-scraper\"\n!pip install nnunetv2 nibabel tifffile tqdm\n\n# IMPORTANT: Set this BEFORE importing nnunetv2\n# Blosc2 is a newer compression format but can cause compatibility issues\nos.environ[\"nnUNet_USE_BLOSC2\"] = \"1\"  # Use blosc2 format (faster, smaller files)\n\nimport nibabel as nib\nimport numpy as np\nimport pandas as pd\nimport tifffile\nimport matplotlib.pyplot as plt\nfrom tqdm.auto import tqdm\n\n# Show GPU configuration\nprint(f\"Available GPUs: {torch.cuda.device_count()}\")\nprint(f\"Using NUM_GPUS={NUM_GPUS}\")","metadata":{"execution":{"iopub.status.busy":"2026-01-27T13:26:24.591073Z","iopub.execute_input":"2026-01-27T13:26:24.591317Z","iopub.status.idle":"2026-01-27T13:26:54.675607Z","shell.execute_reply.started":"2026-01-27T13:26:24.591296Z","shell.execute_reply":"2026-01-27T13:26:54.674792Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## DATA","metadata":{}},{"cell_type":"markdown","source":"### Utilities","metadata":{}},{"cell_type":"code","source":"def create_spacing_json(output_path: Path, shape: tuple, spacing: tuple = (1.0, 1.0, 1.0)):\n    \"\"\"Create JSON sidecar with spacing info for TIFF files.\"\"\"\n    json_data = {\"spacing\": list(spacing)}\n    output_path.parent.mkdir(parents=True, exist_ok=True)  # Ensure directory exists\n    with open(output_path, \"w\") as f:\n        json.dump(json_data, f, indent=2)\n        f.flush()  # Ensure data is written to disk\n    # Verify file was created\n    if not output_path.exists():\n        raise IOError(f\"Failed to create JSON file: {output_path}\")\n\ndef load_nifti(path: Path) -> np.ndarray:\n    \"\"\"Load NIfTI file (used for loading nnUNet predictions).\"\"\"\n    return nib.load(str(path)).get_fdata()\n\ndef create_dataset_json(output_dir: Path, num_training: int, file_ending: str = \".tif\") -> dict:\n    \"\"\"Create dataset.json with ignore label support and 3D TIFF reader.\"\"\"\n    \n    dataset_json = {\n        \"channel_names\": {\"0\": \"CT\"},\n        \"labels\": {\"background\": 0, \"surface\": 1, \"ignore\": 2},\n        \"numTraining\": num_training,\n        \"file_ending\": file_ending,\n        \"overwrite_image_reader_writer\": \"SimpleTiffIO\" # custom reader \n    }\n    \n    json_path = output_dir / \"dataset.json\"\n    with open(json_path, \"w\") as f:\n        json.dump(dataset_json, f, indent=4)\n    \n    print(f\"Created {json_path}\")\n    print(f\"  - {num_training} training cases\")\n    print(f\"  - Labels: background(0), surface(1), ignore(2)\")\n    print(f\"  - Reader: SimpleTiffIO (3D TIFF)\")\n    \n    return dataset_json","metadata":{"execution":{"iopub.status.busy":"2026-01-27T13:26:54.676738Z","iopub.execute_input":"2026-01-27T13:26:54.677514Z","iopub.status.idle":"2026-01-27T13:26:54.684425Z","shell.execute_reply.started":"2026-01-27T13:26:54.677464Z","shell.execute_reply":"2026-01-27T13:26:54.683667Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Prepare","metadata":{}},{"cell_type":"markdown","source":"nnUNet Dataset Structure:\n```\nnnUNet_raw/Dataset100_VesuviusSurface/\n├── imagesTr/           # Training images\n│   ├── case001_0000.tif  # _0000 suffix = channel 0 (only one for CT)\n│   └── case001_0000.json # Spacing information\n├── labelsTr/           # Training labels\n│   ├── case001.tif\n│   └── case001.json\n└── dataset.json        # Dataset configuration\n```","metadata":{}},{"cell_type":"code","source":"def prepare_single_case(\n    src_path: Path, \n    dest_path: Path, \n    json_path: Path, \n    use_symlinks: bool = True\n) -> bool:\n    \"\"\"\n    Prepare a single TIFF file for nnUNet: create symlink/copy and JSON sidecar.\n    Returns True on success, False on failure.\n    \"\"\"\n    try: \n        # Ensure destination directory exists\n        dest_path.parent.mkdir(parents=True, exist_ok=True)\n        \n        # Read TIFF shape\n        with tifffile.TiffFile(src_path) as tif:\n            shape = tif.pages[0].shape if len(tif.pages) == 1 \\\n                    else (len(tif.pages), *tif.pages[0].shape)\n        \n        # Link or copy TIFF file\n        if use_symlinks:\n            if not dest_path.exists():\n                dest_path.symlink_to(src_path.resolve())\n        else:\n            if not dest_path.exists():\n                shutil.copy2(src_path, dest_path)\n        \n        # Create JSON sidecar\n        create_spacing_json(json_path, shape)\n        \n        # Verify both files exist\n        if not dest_path.exists():\n            raise IOError(f\"TIFF file not created: {dest_path}\")\n        if not json_path.exists():\n            raise IOError(f\"JSON file not created: {json_path}\")\n            \n        return True\n\n    except Exception as e:\n        print(f\"Error processing {src_path.name}: {e}\")\n        return False\n    \ndef prepare_training_case(\n    img_path: Path,\n    train_labels_dir: Path,\n    images_dir: Path,\n    labels_dir: Path,\n    use_symlinks: bool\n) -> bool:\n    \"\"\"Worker function for parallel dataset preparation.\"\"\"\n    case_id = img_path.stem\n    label_path = train_labels_dir / img_path.name\n    \n    if not label_path.exists():\n        return False\n    \n    img_ok = prepare_single_case(\n        src_path=img_path,\n        dest_path=images_dir / f\"{case_id}_0000.tif\",\n        json_path=images_dir / f\"{case_id}_0000.json\",\n        use_symlinks=use_symlinks\n    )\n\n    label_ok = prepare_single_case(\n        src_path=label_path,\n        dest_path=labels_dir / f\"{case_id}.tif\",\n        json_path=labels_dir / f\"{case_id}.json\",\n        use_symlinks=use_symlinks\n    )\n    return img_ok and label_ok\n\ndef prepare_dataset(input_dir: Path, max_cases: Optional[int] = None, use_symlinks: bool = True):\n    \"\"\"\n    Convert competition data to nnUNet format using TIFF directly (no NIfTI).\n    Uses multiprocessing for faster preparation.\n    \n    Competition structure:\n    - train_images/*.tif  (3D volumes)\n    - train_labels/*.tif  (3D labels: 0=bg, 1=surface, 2=ignore)\n    \"\"\"\n    dataset_dir = NNUNET_RAW / DS_NAME\n    images_dir = dataset_dir / \"imagesTr\"\n    labels_dir = dataset_dir / \"labelsTr\"\n\n    images_dir.mkdir(parents=True, exist_ok=True)\n    labels_dir.mkdir(parents=True, exist_ok=True)\n\n    train_images_dir = input_dir / \"train_images\"\n    train_labels_dir = input_dir / \"train_labels\"\n    if not train_images_dir.exists():\n        print(f\"ERROR: {train_images_dir} not found!\")\n        return None\n    \n    image_files = sorted(train_images_dir.glob(\"*.tif\"))\n    if max_cases:\n        image_files = image_files[:max_cases]\n    \n    print(f\"Found {len(image_files)} training cases\")\n    print(f\"Using {'symlinks' if use_symlinks else 'copy'}\")\n    print(f\"Processing with {NUM_WORKERS} workers...\")\n    \n    # Create worker function with fixed arguments\n    worker = partial(\n        prepare_training_case,\n        train_labels_dir=train_labels_dir,\n        images_dir=images_dir,\n        labels_dir=labels_dir,\n        use_symlinks=use_symlinks\n    )\n    # Process in parallel with progress bar\n    with Pool(NUM_WORKERS) as pool:\n        results = list(tqdm(\n            pool.imap(worker, image_files),\n            total=len(image_files),\n            desc=\"Preparing dataset\"\n        ))\n    \n    num_converted = sum(results)\n    create_dataset_json(dataset_dir, num_converted, file_ending=\".tif\")\n    \n    # Verify JSON files were created\n    json_count = len(list(images_dir.glob(\"*.json\")))\n    print(f\"\\nDataset prepared: {num_converted} cases\")\n    print(f\"JSON files created: {json_count} (expected: {num_converted * 2})\")\n    print(f\"Location: {dataset_dir}\")\n    \n    if json_count < num_converted:\n        print(f\"WARNING: Missing JSON files! Expected {num_converted}, found {json_count}\")\n    \n    return dataset_dir","metadata":{"execution":{"iopub.status.busy":"2026-01-27T13:26:54.687331Z","iopub.execute_input":"2026-01-27T13:26:54.687722Z","iopub.status.idle":"2026-01-27T13:26:54.712091Z","shell.execute_reply.started":"2026-01-27T13:26:54.687684Z","shell.execute_reply":"2026-01-27T13:26:54.711385Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## CMD\n1. run_preprocessing() → nnUNetv2_plan_and_preprocess\n2. run_training() → nnUNetv2_train\n3. run_inference() → nnUNetv2_predict\n```\n# Preprocessing\nnnUNetv2_plan_and_preprocess -d 100 -np 4 -pl nnUNetPlannerResEncM\n\n# Training (uses ignore label automatically)\nnnUNetv2_train 100 3d_fullres 0 -p nnUNetResEncUNetMPlans\n\n# Training with fewer epochs (faster)\nnnUNetv2_train 100 3d_fullres 0 -p nnUNetResEncUNetMPlans -tr nnUNetTrainer_250epochs\n\n# Training all folds in parallel (different GPUs)\n\nCUDA_VISIBLE_DEVICES=0 nnUNetv2_train 100 3d_fullres 0 ... &\n\nCUDA_VISIBLE_DEVICES=1 nnUNetv2_train 100 3d_fullres 1 ... &\n\n# Multi-GPU DDP training (single fold, multiple GPUs)\nnnUNetv2_train 100 3d_fullres 0 ... -num_gpus 2\n\n# Inference\nnnUNetv2_predict -d 100 -c 3d_fullres -f 0 -i INPUT -o OUTPUT -p nnUNetResEncUNetMPlans\n```","metadata":{}},{"cell_type":"code","source":"def _run_command(\n    cmd: str,\n    name: str = \"Command\",\n    tail_lines: int = 20,\n    timeout: Optional[int] = COMMAND_TIMEOUT\n) -> bool:\n    \"\"\"\n    Execute shell command and handle output parsing.\n    \n    Args:\n        cmd: Shell command to execute\n        name: Display name for logging\n        tail_lines: Number of stdout lines to show on success\n        timeout: Timeout in seconds (None for no timeout)\n    \n    Returns:\n        True if command succeeded, False otherwise\n    \"\"\"\n    print(f\"Running: {cmd}\")\n    if timeout:\n        print(f\"Timeout: {timeout}s ({timeout/3600:.1f}h)\")\n    \n    try:\n        result = subprocess.run(\n            cmd, shell=True,\n            capture_output=True,\n            text=True,\n            timeout=timeout\n        )\n    except subprocess.TimeoutExpired:\n        print(f\"{name} timed out after {timeout} seconds.\")\n        return False\n    if result.returncode != 0:\n        print(f\"{name} failed\")\n        print(f\"STDERR:\\n{result.stderr[-3000:]}\")\n        return False\n    print(f\"{name} succeeded\")\n    if result.stdout.strip():\n        lines = result.stdout.strip().split('\\n')\n        print('\\n'.join(lines[-tail_lines:]))\n    \n    return True","metadata":{"execution":{"iopub.status.busy":"2026-01-27T13:26:54.712991Z","iopub.execute_input":"2026-01-27T13:26:54.713344Z","iopub.status.idle":"2026-01-27T13:26:54.730766Z","shell.execute_reply.started":"2026-01-27T13:26:54.713298Z","shell.execute_reply":"2026-01-27T13:26:54.729996Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Run preprocessing","metadata":{}},{"cell_type":"code","source":"def run_preprocessing(\n    dataset_id: int = DS_ID, \n    planner: str = PLANNER,\n    num_workers: int = NUM_WORKERS,\n    configurations: Optional[List[str]] = None,\n    timeout: Optional[int] = COMMAND_TIMEOUT\n) -> bool:\n    \"\"\"\n    Run nnUNet preprocessing.\n    \n    Args:\n        dataset_id: nnUNet dataset ID\n        planner: Planner class name\n        num_workers: Number of CPU workers for parallel processing\n        configurations: List of configs to preprocess (e.g., [\"3d_fullres\"])\n        timeout: Timeout in seconds (None for no timeout)\n    \n    Returns:\n        True if preprocessing succeeded\n    \"\"\"\n    if configurations is None:\n        configurations = [CONFIGURATION]\n\n    cmd = f\"nnUNetv2_plan_and_preprocess -d {dataset_id:03d} \" \\\n          f\"-np {num_workers} \" \\\n          f\"-pl {planner} \" \\\n          f\"-c {' '.join(configurations)}\"\n    return _run_command(\n        cmd,\n        name=\"Preprocessing\",\n        timeout=timeout\n    )","metadata":{"execution":{"iopub.status.busy":"2026-01-27T13:26:54.731691Z","iopub.execute_input":"2026-01-27T13:26:54.731930Z","iopub.status.idle":"2026-01-27T13:26:54.747149Z","shell.execute_reply.started":"2026-01-27T13:26:54.731899Z","shell.execute_reply":"2026-01-27T13:26:54.746512Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Run trainer","metadata":{}},{"cell_type":"code","source":"def _get_trainer_name(epochs: Optional[Epochs]) -> str:\n    \"\"\"Get trainer class name based on epochs.\"\"\"\n    if epochs is None or epochs == 1000:\n        return \"nnUNetTrainer\"\n    elif epochs == 1:\n        return \"nnUNetTrainer_1epoch\"\n    else:\n        return f\"nnUNetTrainer_{epochs}epochs\"\n\n\ndef run_training(\n    dataset_id: int = DS_ID,\n    config: str = CONFIGURATION,\n    fold: Union[int, str] = FOLD,\n    plans: str = PLANS_NAME,\n    epochs: Optional[Epochs] = EPOCHS,\n    pretrained_weights: Optional[Path] = None,\n    continue_training: bool = False,\n    only_run_validation: bool = False,\n    disable_checkpointing: bool = False,\n    npz: bool = False,\n    num_gpus: int = NUM_GPUS,\n    timeout: Optional[int] = COMMAND_TIMEOUT\n) -> bool:\n    \"\"\"\n    Run nnUNet training.\n    \n    Args:\n        dataset_id: nnUNet dataset ID\n        config: Configuration name (3d_fullres, 2d, etc.)\n        fold: Fold number (0-4) or \"all\" for training on all data\n        plans: Plans name matching the planner used\n        epochs: Number of epochs. Available: 1, 5, 10, 20, 50, 100, 150, 200, 250, \n                300, 400, 500, 750, 1000, 2000, 4000, 8000. None = 1000 (default)\n        pretrained_weights: Optional path to pretrained checkpoint for fine-tuning\n        continue_training: Continue from last checkpoint (add -c flag)\n        only_run_validation: Only run validation, skip training\n        disable_checkpointing: Disable saving checkpoints (saves disk space)\n        npz: Save softmax outputs during validation (needed for ensembling)\n        num_gpus: Number of GPUs for DDP training (default: auto-detected)\n        timeout: Timeout in seconds (None for no timeout)\n    \n    Returns:\n        True if training succeeded\n    \n    Note:\n        For multi-GPU (DDP) training, batch size should be divisible by num_gpus.\n        The first run extracts preprocessed data - wait for GPU usage before \n        starting additional folds on other GPUs.\n        \n        When using fold=\"all\", nnUNet will run validation on ALL training data\n        after training completes, which can be slow. This is normal behavior.\n    \"\"\"\n    trainer = _get_trainer_name(epochs)\n    cmd = f\"nnUNetv2_train {dataset_id:03d} {config} {fold} -p {plans} -tr {trainer}\"\n    \n    if pretrained_weights:\n        cmd += f\" -pretrained_weights {pretrained_weights}\"\n    if continue_training:\n        # Find checkpoint to resume from: prefer checkpoint_final.pth, fallback to checkpoint_best.pth\n        model_dir = get_training_output_dir(epochs=epochs, plans=plans, config=config, fold=fold)\n        checkpoint_final = model_dir / \"checkpoint_final.pth\"\n        checkpoint_best = model_dir / \"checkpoint_best.pth\"\n        if checkpoint_final.exists():\n            print(f\"Resuming from: {checkpoint_final}\")\n            cmd += \" --c\"\n        elif checkpoint_best.exists():\n            print(f\"Resuming from: {checkpoint_best}\")\n            cmd += \" --c\"\n        else:\n            print(f\"WARNING: No checkpoint found in {model_dir}, starting fresh\")\n    if only_run_validation:\n        cmd += \" --val\"\n    if disable_checkpointing:\n        cmd += \" --disable_checkpointing\"\n    if npz:\n        cmd += \" --npz\"\n    if num_gpus > 1:\n        cmd += f\" -num_gpus {num_gpus}\"\n    \n    epochs_str = epochs if epochs else 1000\n    gpu_str = f\", {num_gpus} GPUs\" if num_gpus > 1 else \"\"\n    return _run_command(cmd, f\"Training ({epochs_str} epochs{gpu_str})\", tail_lines=30, timeout=timeout)\n","metadata":{"execution":{"iopub.status.busy":"2026-01-27T13:26:54.747909Z","iopub.execute_input":"2026-01-27T13:26:54.748140Z","iopub.status.idle":"2026-01-27T13:26:54.767478Z","shell.execute_reply.started":"2026-01-27T13:26:54.748119Z","shell.execute_reply":"2026-01-27T13:26:54.766876Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Run inference","metadata":{}},{"cell_type":"code","source":"def run_inference(\n    input_dir: Path,\n    output_dir: Path,\n    dataset_id: int = DS_ID,\n    config: str = CONFIGURATION,\n    fold: Union[int, str] = FOLD,\n    plans: str = PLANS_NAME,\n    epochs: Optional[Epochs] = EPOCHS,\n    save_probabilities: bool = True,\n    num_processes_preprocessing: int = 2,\n    num_processes_segmentation: int = 2,\n    timeout: Optional[int] = COMMAND_TIMEOUT\n) -> bool:\n    \"\"\"\n    Run inference with trained model.\n    \n    Args:\n        input_dir: Directory with test images (must have _0000 suffix)\n        output_dir: Directory to save predictions\n        dataset_id: nnUNet dataset ID\n        config: Configuration name\n        fold: Fold number used for training (or \"all\" or tuple like \"0,1,2\")\n        plans: Plans name\n        epochs: Epochs used during training (must match trained model)\n        save_probabilities: Whether to save probability maps (.npz files)\n        num_processes_preprocessing: Parallel processes for preprocessing\n        num_processes_segmentation: Parallel processes for segmentation\n        timeout: Timeout in seconds (None for no timeout)\n    \n    Returns:\n        True if inference succeeded\n    \"\"\"\n    output_dir.mkdir(parents=True, exist_ok=True)\n    \n    trainer = _get_trainer_name(epochs)\n    \n    cmd = f\"nnUNetv2_predict -d {dataset_id:03d} \" \\\n          f\"-c {config} -f {fold} \" \\\n          f\"-i {input_dir} -o {output_dir} \" \\\n          f\"-p {plans} -tr {trainer} \" \\\n          f\"-npp {num_processes_preprocessing} -nps {num_processes_segmentation} \" \\\n          \"--verbose\"\n    \n    if save_probabilities:\n        cmd += \" --save_probabilities\"\n    \n    return _run_command(cmd, \"Inference\", timeout=timeout)","metadata":{"execution":{"iopub.status.busy":"2026-01-27T13:26:54.768290Z","iopub.execute_input":"2026-01-27T13:26:54.768587Z","iopub.status.idle":"2026-01-27T13:26:54.783025Z","shell.execute_reply.started":"2026-01-27T13:26:54.768558Z","shell.execute_reply":"2026-01-27T13:26:54.782508Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def prepare_test_data(input_dir: Path, output_dir: Path, use_symlinks: bool = True) -> Path:\n    \"\"\"Prepare test TIFF images for nnUNet inference.\"\"\"\n    \n    output_dir.mkdir(parents=True, exist_ok=True)\n    \n    test_images_dir = input_dir / \"test_images\"\n    \n    if not test_images_dir.exists():\n        print(f\"ERROR: {test_images_dir} not found!\")\n        return output_dir\n    \n    test_files = sorted(test_images_dir.glob(\"*.tif\"))\n    print(f\"Found {len(test_files)} test cases\")\n    print(f\"Using {'symlinks' if use_symlinks else 'copy'}\")\n    \n    for img_path in tqdm(test_files, desc=\"Preparing test data\"):\n        case_id = img_path.stem\n        prepare_single_case(\n            img_path,\n            output_dir / f\"{case_id}_0000.tif\",\n            output_dir / f\"{case_id}_0000.json\",\n            use_symlinks\n        )\n    \n    return output_dir","metadata":{"execution":{"iopub.status.busy":"2026-01-27T13:26:54.783850Z","iopub.execute_input":"2026-01-27T13:26:54.784124Z","iopub.status.idle":"2026-01-27T13:26:54.798581Z","shell.execute_reply.started":"2026-01-27T13:26:54.784080Z","shell.execute_reply":"2026-01-27T13:26:54.797922Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Inference utilities","metadata":{}},{"cell_type":"code","source":"def load_probabilities(npz_path: Path) -> np.ndarray:\n    \"\"\"\n    Load probability maps from nnUNet inference.\n    \n    Only available if inference was run with save_probabilities=True.\n    Shape: (num_classes, D, H, W) with float32 values in [0, 1].\n    \"\"\"\n    data = np.load(npz_path)\n    return data['probabilities']\n\n\ndef predictions_to_tiff(pred_dir: Path, output_dir: Path):\n    \"\"\"\n    Convert nnUNet predictions to 3D TIFF files.\n    \n    nnUNet outputs:\n    - .npz files with probability maps (if save_probabilities=True)\n    - .tif files with predictions (our SimpleTiffIO format)\n    - .pkl files with metadata\n    \n    This function:\n    1. First tries to load .npz files and convert probabilities to binary predictions\n    2. Falls back to .tif files if .npz not found\n    3. Saves as uint8 TIFF (0=background, 1=surface)\n    \"\"\"\n    output_dir.mkdir(parents=True, exist_ok=True)\n    \n    # Try NPZ files first (probability maps)\n    npz_files = list(pred_dir.glob(\"*.npz\"))\n    tif_files = list(pred_dir.glob(\"*.tif\"))\n    nii_files = list(pred_dir.glob(\"*.nii.gz\"))\n    \n    if npz_files:\n        print(f\"Converting {len(npz_files)} NPZ probability files to TIFF...\")\n        for npz_path in tqdm(npz_files, desc=\"Converting to TIFF\"):\n            case_id = npz_path.stem\n            # Load probabilities and take argmax to get class predictions\n            probs = load_probabilities(npz_path)\n            pred = np.argmax(probs, axis=0).astype(np.uint8)\n            tifffile.imwrite(output_dir / f\"{case_id}.tif\", pred)\n    elif tif_files:\n        print(f\"Copying {len(tif_files)} TIFF prediction files...\")\n        for tif_path in tqdm(tif_files, desc=\"Copying TIFF\"):\n            case_id = tif_path.stem\n            # Load and ensure uint8\n            pred = tifffile.imread(str(tif_path)).astype(np.uint8)\n            tifffile.imwrite(output_dir / f\"{case_id}.tif\", pred)\n    # Try NIfTI as last resort (legacy format)\n    elif nii_files:\n        print(f\"Converting {len(nii_files)} NIfTI files to TIFF...\")\n        for nii_path in tqdm(nii_files, desc=\"Converting to TIFF\"):\n            case_id = nii_path.stem.replace(\".nii\", \"\")\n            pred = load_nifti(nii_path).astype(np.uint8)\n            tifffile.imwrite(output_dir / f\"{case_id}.tif\", pred)\n    else:\n        print(f\"WARNING: No prediction files found in {pred_dir}\")\n        print(f\"  Checked for: *.npz, *.tif, *.nii.gz\")","metadata":{"execution":{"iopub.status.busy":"2026-01-27T13:26:54.799428Z","iopub.execute_input":"2026-01-27T13:26:54.800005Z","iopub.status.idle":"2026-01-27T13:26:54.814856Z","shell.execute_reply.started":"2026-01-27T13:26:54.799983Z","shell.execute_reply":"2026-01-27T13:26:54.814089Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Pipeline","metadata":{}},{"cell_type":"markdown","source":"```\n# Full pipeline with defaults\nfull_pipeline()\n\n# Quick experiment\nfull_pipeline(epochs=50, config=\"2d\")\n\n# Resume training\nfull_pipeline(continue_training=True, epochs=250)\n\n# Inference only (trained model exists)\nfull_pipeline(do_preprocess=False, do_train=False)\n```","metadata":{}},{"cell_type":"code","source":"def _parse_model_path(model_path: Union[str, Path, None]) -> Tuple[Optional[Path], Optional[int], Optional[str], Optional[str], Optional[str]]:\n    \"\"\"\n    Parse model path to extract configuration parameters.\n    \n    Model path format: .../DatasetXXX_Name/TrainerName__PlansName__Config/fold_X/checkpoint.pth\n    \n    Returns:\n        (model_dir, epochs, plans, config, fold) - extracted from path, None if not found\n    \n    Note: Parses path structure even if file doesn't exist (for path validation).\n    \"\"\"\n    if model_path is None:\n        return None, None, None, None, None\n    \n    model_path = Path(model_path)\n    \n    # Determine model directory from path (even if doesn't exist)\n    # If path ends with .pth, use parent directory\n    if model_path.suffix == \".pth\" or (model_path.exists() and model_path.is_file()):\n        model_dir = model_path.parent\n    else:\n        model_dir = model_path\n    \n    # Try to parse folder structure\n    # Expected: .../fold_X or .../TrainerName__PlansName__Config/fold_X\n    try:\n        parts = model_dir.parts\n        \n        # Find fold\n        fold = None\n        for part in reversed(parts):\n            if part.startswith(\"fold_\"):\n                fold = part.replace(\"fold_\", \"\")\n                break\n        \n        # Find trainer__plans__config\n        epochs = None\n        plans = None\n        config = None\n        for part in parts:\n            if \"__\" not in part:\n                continue\n            segments = part.split(\"__\")\n            if len(segments) >= 3:\n                trainer_name = segments[0]\n                plans = segments[1]\n                config = segments[2]\n                # Extract epochs from trainer name\n                if \"epochs\" in trainer_name:\n                    import re\n                    match = re.search(r'(\\d+)epochs?', trainer_name)\n                    if match:\n                        epochs = int(match.group(1))\n                elif trainer_name == \"nnUNetTrainer\":\n                    epochs = 1000  # Default\n                break\n        \n        return model_dir, epochs, plans, config, fold\n    except Exception:\n        return model_dir, None, None, None, None","metadata":{"execution":{"iopub.status.busy":"2026-01-27T13:26:54.815739Z","iopub.execute_input":"2026-01-27T13:26:54.816033Z","iopub.status.idle":"2026-01-27T13:26:54.832155Z","shell.execute_reply.started":"2026-01-27T13:26:54.816002Z","shell.execute_reply":"2026-01-27T13:26:54.831310Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def full_pipeline_1(\n    # Data options\n    max_cases: Optional[int] = None,\n    # Stage control\n    do_preprocess: bool = True,\n    do_train: bool = True,\n    do_inference: bool = True,\n    # Training options (all tunable parameters)\n    config: str = CONFIGURATION,\n    fold: Union[int, str] = FOLD,\n    planner: str = PLANNER,\n    plans: str = PLANS_NAME,\n    epochs: Optional[Epochs] = EPOCHS,\n    pretrained_weights: Optional[Path] = None,\n    continue_training: bool = False,\n    num_gpus: int = NUM_GPUS,\n    # Inference options\n    save_probabilities: bool = True,\n    # External model (for inference without training)\n    model_path: Optional[Union[str, Path]] = None,\n    # Timeout\n    timeout: Optional[int] = COMMAND_TIMEOUT,\n):\n    \"\"\"\n    Run complete pipeline: setup -> data prep -> preprocess -> train -> predict.\n    \n    Args:\n        max_cases: Limit number of training cases (None = use all)\n        do_preprocess: Run preprocessing step\n        do_train: Run training step\n        do_inference: Run inference step\n        \n        config: nnUNet configuration (3d_fullres, 2d, 3d_lowres, 3d_cascade_fullres)\n        fold: Fold number (0-4) or \"all\" for training on all data\n        planner: Planner class name\n        plans: Plans name matching the planner\n        epochs: Number of training epochs (1, 5, 10, 20, 50, 100, 150, 200, 250, \n                300, 400, 500, 750, 1000, 2000, 4000, 8000). None = 1000\n        pretrained_weights: Path to pretrained weights for fine-tuning\n        continue_training: Continue from last checkpoint\n        num_gpus: Number of GPUs for DDP training\n        \n        save_probabilities: Save probability maps during inference\n        \n        model_path: Path to trained model checkpoint or directory (str or Path).\n                    When provided with do_train=False, parameters (epochs, plans, config, fold)\n                    are auto-extracted from the path if possible.\n                    Example: \"/path/to/nnUNetTrainer_5epochs__nnUNetResEncUNetMPlans__3d_fullres/fold_all/checkpoint_best.pth\"\n        timeout: Command timeout in seconds (None for no timeout)\n    \n    Returns:\n        True if pipeline completed successfully\n    \"\"\"\n    \n    print(\"=\" * 60)\n    print(\"Vesuvius Surface Detection - nnUNet Pipeline\")\n    print(\"=\" * 60)\n    print(f\"Stages: preprocess={do_preprocess}, train={do_train}, inference={do_inference}\")\n    print(f\"Config: {config}, Fold: {fold}, Epochs: {epochs or 1000}, GPUs: {num_gpus}\")\n    \n    # 1. Setup\n    print(\"\\n[1/5] Environment setup...\")\n    setup_environment()\n    \n    # 2. Prepare raw data (always - uses symlinks, fast)\n    print(\"\\n[2/5] Preparing raw dataset (symlinks)...\")\n    raw_dataset_dir = NNUNET_RAW / DS_NAME\n    if not raw_dataset_dir.exists():\n        prepare_dataset(DATA_DIR, max_cases=max_cases)\n    else:\n        print(f\"Raw dataset already exists: {raw_dataset_dir}\")\n    \n    # 3. Preprocessing\n    if do_preprocess:\n        # Check if pre-prepared preprocessed data exists\n        if _link_prepared_preprocessed():\n            print(\"\\n[3/5] Using pre-prepared preprocessed data...\")\n        else:\n            print(\"\\n[3/5] Preprocessing...\")\n            success = run_preprocessing(planner=planner, configurations=[config], timeout=timeout)\n            if not success:\n                print(\"Preprocessing failed!\")\n                return False\n    else:\n        print(\"\\n[3/5] Skipping preprocessing...\")\n        _link_prepared_preprocessed()  # Still link if available\n    \n    \n    return True\n\ndef full_pipeline_2(\n    # Data options\n    max_cases: Optional[int] = None,\n    # Stage control\n    do_preprocess: bool = True,\n    do_train: bool = True,\n    do_inference: bool = True,\n    # Training options (all tunable parameters)\n    config: str = CONFIGURATION,\n    fold: Union[int, str] = FOLD,\n    planner: str = PLANNER,\n    plans: str = PLANS_NAME,\n    epochs: Optional[Epochs] = EPOCHS,\n    pretrained_weights: Optional[Path] = None,\n    continue_training: bool = False,\n    num_gpus: int = NUM_GPUS,\n    # Inference options\n    save_probabilities: bool = True,\n    # External model (for inference without training)\n    model_path: Optional[Union[str, Path]] = None,\n    # Timeout\n    timeout: Optional[int] = COMMAND_TIMEOUT,\n):\n    \n    # 4. Training\n    if do_train:\n        print(\"\\n[4/5] Training...\")\n        success = run_training(\n            config=config,\n            fold=fold,\n            plans=plans,\n            epochs=epochs,\n            pretrained_weights=pretrained_weights,\n            continue_training=continue_training,\n            num_gpus=num_gpus,\n            timeout=timeout\n        )\n        if not success:\n            print(\"Training failed!\")\n            return False\n    else:\n        print(\"\\n[4/5] Skipping training...\")\n        \n        # Parse model_path to extract configuration if provided\n        if model_path is not None:\n            model_path = Path(model_path) if isinstance(model_path, str) else model_path\n            _, parsed_epochs, parsed_plans, parsed_config, parsed_fold = _parse_model_path(model_path)\n            \n            # Override parameters with parsed values\n            if parsed_epochs is not None:\n                epochs = parsed_epochs\n                print(f\"  Detected epochs={epochs} from model path\")\n            if parsed_plans is not None:\n                plans = parsed_plans\n                print(f\"  Detected plans={plans} from model path\")\n            if parsed_config is not None:\n                config = parsed_config\n                print(f\"  Detected config={config} from model path\")\n            if parsed_fold is not None:\n                fold = parsed_fold\n                print(f\"  Detected fold={fold} from model path\")\n            \n            if model_path.exists():\n                print(f\"Using model: {model_path}\")\n            else:\n                print(f\"WARNING: model_path does not exist: {model_path}\")\n        \n        # Verify model exists\n        expected_model_dir = get_training_output_dir(epochs=epochs, plans=plans, config=config, fold=fold)\n        expected_model_dir.mkdir(exist_ok=True, parents=True)\n        checkpoint_final = expected_model_dir / \"checkpoint_final.pth\"\n        checkpoint_best = expected_model_dir / \"checkpoint_best.pth\"\n\n        if model_path and Path(model_path).exists():\n            # Symlink external -> final so nnUNet finds it\n            checkpoint_final.symlink_to(model_path)\n            print(f\"Map model: {checkpoint_final}\")\n        elif checkpoint_final.exists():\n            print(f\"Found model: {checkpoint_final}\")\n        elif checkpoint_best.exists():\n            # Symlink best -> final so nnUNet finds it\n            checkpoint_final.symlink_to(checkpoint_best)\n            print(f\"Found model: {checkpoint_best} (symlinked to checkpoint_final.pth)\")\n        elif do_inference:\n            print(f\"WARNING: No model found at {expected_model_dir}\")\n            print(\"  Provide model_path to a valid nnUNet checkpoint file\")\n    \n    # 5. Inference\n    if do_inference:\n        print(\"\\n[5/5] Running inference...\")\n        \n        # Prepare test data in temp location\n        test_input_dir = WORKING_DIR / \"test_input\"\n        prepare_test_data(DATA_DIR, test_input_dir)\n        \n        # Run inference\n        predictions_dir = WORKING_DIR / \"predictions\"\n        success = run_inference(\n            test_input_dir, \n            predictions_dir,\n            config=config,\n            fold=fold,\n            plans=plans,\n            epochs=epochs,\n            save_probabilities=save_probabilities,\n            timeout=timeout\n        )\n        if not success:\n            print(\"Inference failed!\")\n            return False\n        \n        # Convert predictions to TIFF\n        print(\"\\nConverting predictions to TIFF...\")\n        tiff_output_dir = OUTPUT_DIR / \"predictions_tiff\"\n        predictions_to_tiff(predictions_dir, tiff_output_dir)\n        \n        print(f\"\\nPredictions saved to: {tiff_output_dir}\")\n    else:\n        print(\"\\n[5/5] Skipping inference...\")\n    \n    print(\"\\n\" + \"=\" * 60)\n    print(\"Pipeline complete!\")\n    print(\"=\" * 60)\n    \n    # Show training progress if training was done\n    if do_train:\n        print(\"\\n[Visualization] Training progress:\")\n        show_progress(epochs=epochs, plans=plans, config=config, fold=fold)\n    \n    # Visualize predictions if inference was done\n    if do_inference:\n        print(\"\\n[Visualization] Sample prediction:\")\n        visualize_predictions(num_samples=1)\n    return True","metadata":{"execution":{"iopub.status.busy":"2026-01-27T13:26:54.833187Z","iopub.execute_input":"2026-01-27T13:26:54.833598Z","iopub.status.idle":"2026-01-27T13:26:54.853354Z","shell.execute_reply.started":"2026-01-27T13:26:54.833565Z","shell.execute_reply":"2026-01-27T13:26:54.852765Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def show_progress(\n    epochs: Optional[Epochs] = None,\n    plans: str = PLANS_NAME,\n    config: str = CONFIGURATION,\n    fold: Union[int, str] = FOLD\n):\n    \"\"\"\n    Display training progress image from nnUNet.\n    \n    Args:\n        epochs: Number of epochs used for training (to find correct folder)\n        plans: Plans name\n        config: Configuration name\n        fold: Fold number or \"all\"\n    \"\"\"\n    progress_path = get_progress_image_path(epochs, plans, config, fold)\n    \n    if not progress_path.exists():\n        print(f\"Progress image not found: {progress_path}\")\n        print(\"Training may not have started or completed yet.\")\n        return\n    \n    from IPython.display import Image, display\n    print(f\"Training progress: {progress_path}\")\n    display(Image(filename=str(progress_path)))\n\n\ndef plot_three_axis_cuts(\n    image_vol_path: Path,\n    mask_vol_path: Path,\n    figsize: tuple = (12, 15)\n):\n    \"\"\"\n    Plot middle slices of XY, XZ, and YZ planes for image volume and predicted mask.\n    \n    Args:\n        image_vol_path: Path to input image TIFF\n        mask_vol_path: Path to prediction mask TIFF\n        figsize: Figure size (width, height)\n    \"\"\"\n    print(f\"Visualizing: {image_vol_path.name}\")\n    \n    # Load volumes\n    image_vol = tifffile.imread(str(image_vol_path))\n    mask_vol = tifffile.imread(str(mask_vol_path)).astype(np.uint8)\n    \n    # Get dimensions\n    d, h, w = image_vol.shape\n    z_mid, y_mid, x_mid = d // 2, h // 2, w // 2\n    \n    # Extract slices\n    slices = {\n        'XY Plane (Z-axis)': (image_vol[z_mid, :, :], mask_vol[z_mid, :, :]),\n        'XZ Plane (Y-axis)': (image_vol[:, y_mid, :], mask_vol[:, y_mid, :]),\n        'YZ Plane (X-axis)': (image_vol[:, :, x_mid], mask_vol[:, :, x_mid])\n    }\n    \n    fig, axes = plt.subplots(3, 2, figsize=figsize)\n    for i, (plane_name, (img_slice, mask_slice)) in enumerate(slices.items()):\n        # Image Volume\n        axes[i, 0].imshow(img_slice, cmap='gray')\n        axes[i, 0].set_title(f\"{plane_name} - Image Volume\")\n        axes[i, 0].axis('off')\n        \n        # Mask\n        axes[i, 1].imshow(mask_slice, cmap='gray')\n        axes[i, 1].set_title(f\"{plane_name} - Predicted Mask\")\n        axes[i, 1].axis('off')\n    \n    plt.tight_layout()\n    plt.show()\n\n\ndef visualize_predictions(\n    predictions_dir: Path = OUTPUT_DIR / \"predictions_tiff\",\n    test_images_dir: Path = DATA_DIR / \"test_images\",\n    num_samples: int = 1\n):\n    \"\"\"\n    Visualize prediction results by showing image/mask pairs.\n    \n    Args:\n        predictions_dir: Directory containing prediction TIFFs\n        test_images_dir: Directory containing input test images\n        num_samples: Number of samples to visualize\n    \"\"\"\n    if not predictions_dir.exists():\n        print(f\"Predictions directory not found: {predictions_dir}\")\n        return\n    \n    predictions = sorted(predictions_dir.glob(\"*.tif\"))\n    if not predictions:\n        print(f\"No TIFF predictions found in {predictions_dir}\")\n        return\n    \n    print(f\"Found {len(predictions)} predictions\")\n    \n    for pred_path in predictions[:num_samples]:\n        image_path = test_images_dir / pred_path.name\n        if image_path.exists():\n            plot_three_axis_cuts(image_path, pred_path)\n        else:\n            print(f\"Warning: Input image not found: {image_path}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-27T13:26:54.854254Z","iopub.execute_input":"2026-01-27T13:26:54.854540Z","iopub.status.idle":"2026-01-27T13:26:54.871637Z","shell.execute_reply.started":"2026-01-27T13:26:54.854516Z","shell.execute_reply":"2026-01-27T13:26:54.870858Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Submission","metadata":{}},{"cell_type":"code","source":"def generate_submission(\n    predictions_tiff_dir: Path = OUTPUT_DIR / \"predictions_tiff\",\n    output_zip: Path = OUTPUT_DIR / \"submission.zip\",\n    delete_after_zip: bool = True  # Default True to save space on Kaggle\n) -> Optional[Path]:\n    \"\"\"\n    Create submission ZIP from TIFF predictions.\n    \n    Args:\n        predictions_tiff_dir: Directory containing predicted TIFF files\n        output_zip: Output ZIP file path\n        delete_after_zip: Delete TIFF files after adding to ZIP (saves space)\n    \n    Returns:\n        Path to submission ZIP if successful, None otherwise\n    \"\"\"\n    import zipfile\n    \n    if not predictions_tiff_dir.exists():\n        print(f\"ERROR: Predictions directory not found: {predictions_tiff_dir}\")\n        print(\"Run inference first!\")\n        return None\n    \n    tiff_files = sorted(predictions_tiff_dir.glob(\"*.tif\"))\n    \n    if not tiff_files:\n        print(f\"No TIFF files found in {predictions_tiff_dir}\")\n        return None\n    \n    print(f\"Creating submission ZIP with {len(tiff_files)} files...\")\n    \n    with zipfile.ZipFile(output_zip, 'w', zipfile.ZIP_DEFLATED) as zipf:\n        for tiff_path in tqdm(tiff_files, desc=\"Zipping predictions\"):\n            # Add file with just the filename (no directory structure)\n            zipf.write(tiff_path, tiff_path.name)\n            \n            if delete_after_zip:\n                tiff_path.unlink()\n    \n    zip_size_mb = output_zip.stat().st_size / (1024 * 1024)\n    print(f\"Submission saved: {output_zip} ({zip_size_mb:.1f} MB)\")\n    \n    return output_zip\n\n# generate_submission()\n# generate_submission(delete_after_zip=True)  # Delete TIFFs after zipping to save space","metadata":{"execution":{"iopub.status.busy":"2026-01-27T13:26:54.872617Z","iopub.execute_input":"2026-01-27T13:26:54.872958Z","iopub.status.idle":"2026-01-27T13:26:54.887656Z","shell.execute_reply.started":"2026-01-27T13:26:54.872935Z","shell.execute_reply":"2026-01-27T13:26:54.886820Z"},"jupyter":{"source_hidden":true},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# RUN THIS","metadata":{}},{"cell_type":"markdown","source":"Default Configuration:\n\n- fold=\"all\" - Train on all data (no cross-validation)\n- config=\"3d_fullres\" - Best quality single-stage\n- epochs=1000 - Full training (can reduce to 250-500)\n- num_gpus=auto - Uses all available GPUs\n\nAfter Running:\n- Training progress displayed (progress.png)\n- Sample prediction visualized\n- Submission ZIP created at /kaggle/working/submission.zip","metadata":{}},{"cell_type":"code","source":"# Sample pipeline with 50 epochs and time limit\n# full_pipeline(epochs=100)\n\n# Then resume when training finishes and run inference with the best model\nEPOCHS = 1000\n\nfull_pipeline_1(\n    do_preprocess=False, epochs=EPOCHS\n)\n\n# full_pipeline_2(\n#     do_preprocess=False, epochs=EPOCHS, \n#     pretrained_weights=\"/kaggle/input/vesuvius-baseline-1000e/pytorch/default/1/nnUNet_results/Dataset100_VesuviusSurface/nnUNetTrainer__nnUNetResEncUNetMPlans__3d_fullres/fold_all/checkpoint_best.pth\"\n# )\n\n\nKAGGLE_INPUT_MODEL = Path(\"/kaggle/input/vesuvius-baseline/pytorch/default/15/nnUNet_results\")\nif KAGGLE_INPUT_MODEL.exists():\n    print(f\"\\nCopying pretrained model from {KAGGLE_INPUT_MODEL}\")\n    print(f\"Destination: {NNUNET_RESULTS}\")\n    \n    # Remove existing results if present\n    if NNUNET_RESULTS.exists():\n        print(f\"Removing existing nnUNet_results at {NNUNET_RESULTS}\")\n        shutil.rmtree(NNUNET_RESULTS)\n    \n    # Copy the entire nnUNet_results folder\n    shutil.copytree(KAGGLE_INPUT_MODEL, NNUNET_RESULTS)\n    print(f\"Model copied successfully!\")\n    \n    # Verify checkpoint exists\n    expected_model_dir = get_training_output_dir(epochs=EPOCHS)\n    checkpoint_files = list(expected_model_dir.glob(\"*.pth\"))\n    if checkpoint_files:\n        print(f\"Found checkpoints: {[f.name for f in checkpoint_files]}\")\n    else:\n        print(f\"WARNING: No checkpoint found in {expected_model_dir}\")\nelse:\n    print(f\"\\nKaggle input model not found at {KAGGLE_INPUT_MODEL}\")\n    print(\"Starting training from scratch or using existing local checkpoint\")\n\nfull_pipeline_2(\n    do_preprocess=False, epochs=EPOCHS,\n    continue_training=True\n)\n# full_pipeline_2(\n#     do_preprocess=False, do_train=False,\n#     model_path=\"/kaggle/input/vesuvius-baseline/pytorch/default/3/nnUNet_results/Dataset100_VesuviusSurface/nnUNetTrainer__nnUNetResEncUNetMPlans__3d_fullres/fold_all/checkpoint_best.pth\"\n# )","metadata":{"execution":{"iopub.status.busy":"2026-01-27T13:26:54.888373Z","iopub.execute_input":"2026-01-27T13:26:54.888621Z","execution_failed":"2026-01-27T13:30:10.872Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Generate submission ZIP\ngenerate_submission()","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-27T13:30:10.873Z"}},"outputs":[],"execution_count":null}]}