{"cells": [{"cell_type": "markdown", "metadata": {}, "source": "# Vesuvius Surface Detection -- skeleton-recall (700ep, FULL 786-case training data), NO postprocessing\n\nRuns inference with the skeleton-recall winner-extension checkpoint trained on **all 786\navailable cases** (fold=all -- see README.md step 8's \"train on whole of 786 data\"\nextension), packaging `submission.zip` straight from the raw predictions -- 1st-place\npostprocessing is still computed below (for parity/logging) but its output is NOT what gets\nzipped.\n\nThis checkpoint has no local held-out LOSO score by design (trained on every case) -- only\nreal Kaggle scores are meaningful for it. Submitted alongside the postprocessed variant\n(`vesuvius-skelrecall700-fulldata-submission`) so both the raw and postprocessed full-data\nresults are real, comparable Kaggle numbers -- same pattern as the fold_0 A2/A3 pair.\n\nAdapted from `vesuvius-skelrecall700-nopp-submission` (same mount/offline-install/staging/\ninference mechanics, unchanged) -- only CHECKPOINT_DATASET_SLUG and FOLD (0 -> \"all\") differ.\n"}, {"cell_type": "markdown", "metadata": {}, "source": "## Configuration"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "from __future__ import annotations\n\nimport json\nimport os\nimport shutil\nimport subprocess\nimport zipfile\nfrom pathlib import Path\nfrom typing import Optional, Tuple, Union\n\n# ---- Full-786-case checkpoint (fold=all, no held-out split). ----\nCHECKPOINT_DATASET_SLUG = \"vesuvius-skelrecall700-fulldata-checkpoint-v1\"  # REAL final checkpoint (epoch 700, fold=all, training complete)\nWHEELS_DATASET_SLUG = \"vesuvius-nnunet-wheels-v3\"\n\n# ---- Paths (Kaggle mounts) ----\nINPUT_DIR = Path(\"/kaggle/input/competitions/vesuvius-challenge-surface-detection\")\nKAGGLE_USERNAME = \"vigneshk96\"\n# Kaggle's /kaggle/input mount convention for custom (dataset_sources) datasets is NOT\n# consistent -- confirmed directly via two throwaway diagnostic kernel runs\n# (vigneshk96/vesuvius-diag-mount) that ran `mount` + `ls -R /kaggle/input`. Run 1: both\n# attached datasets mounted at the NEW nested path /kaggle/input/datasets/<owner>/<slug>/.\n# Run 2 (different dataset, same kernel/account): the wheels dataset again mounted at the\n# nested path, but the checkpoint dataset mounted at the OLD flat /kaggle/input/<slug>/\n# instead -- two datasets, two different conventions, same run. Real, observed platform\n# behavior, not something to assume away. CHECKPOINT_DIR/WHEELS_DIR are resolved for real\n# below (see resolve_dataset_mount in the offline-install cell) by probing both candidate\n# paths and using whichever one actually exists -- these are just the candidate list.\nCHECKPOINT_DIR_CANDIDATES = [\n    Path(f\"/kaggle/input/datasets/{KAGGLE_USERNAME}/{CHECKPOINT_DATASET_SLUG}\"),\n    Path(f\"/kaggle/input/{CHECKPOINT_DATASET_SLUG}\"),\n]\nWHEELS_DIR_CANDIDATES = [\n    Path(f\"/kaggle/input/datasets/{KAGGLE_USERNAME}/{WHEELS_DATASET_SLUG}\"),\n    Path(f\"/kaggle/input/{WHEELS_DATASET_SLUG}\"),\n]\nCHECKPOINT_DIR = CHECKPOINT_DIR_CANDIDATES[0]  # placeholder, re-resolved for real below\nWHEELS_DIR = WHEELS_DIR_CANDIDATES[0]  # placeholder, re-resolved for real below\nWORKING_DIR = Path(\"/kaggle/temp\")\nOUTPUT_DIR = Path(\"/kaggle/working\")\n\nNNUNET_RAW = WORKING_DIR / \"nnUNet_data\" / \"nnUNet_raw\"\nNNUNET_PREPROCESSED = WORKING_DIR / \"nnUNet_data\" / \"nnUNet_preprocessed\"\nNNUNET_RESULTS = WORKING_DIR / \"nnUNet_results\"\n\n# ---- Model identity (must match what was trained) ----\nDATASET_ID = 100\nDATASET_NAME = \"Dataset100_VesuviusSurface\"\nCONFIGURATION = \"3d_lowres\"\nPLANS_NAME = \"nnUNetResEncUNetMPlans\"\nTRAINER = \"nnUNetTrainerSkeletonRecall_700epochs\"  # stub trainer staged below -- inference only\n                                                     # needs build_network_architecture, which\n                                                     # skeleton-recall never overrides, so the\n                                                     # full loss/transform code isn't needed here\n                                                     # (verified: nnUNetv2_predict calls\n                                                     # build_network_architecture directly on the\n                                                     # class without instantiating it, __init__\n                                                     # never runs at inference time)\nCHECKPOINT_FILENAME = \"checkpoint_best.pth\"\nFOLD = \"all\"  # trained on all 786 cases, no held-out split (final-submission-style run)\n\nMODEL_DIR_NAME = f\"{TRAINER}__{PLANS_NAME}__{CONFIGURATION}\"\nEXT_TRAINER_DIR = WORKING_DIR / \"ext_trainers\"\n\nTEST_INPUT_DIR = WORKING_DIR / \"test_input\"\nPREDICTIONS_DIR = WORKING_DIR / \"predictions\"\nPREDICTIONS_TIFF_DIR = OUTPUT_DIR / \"predictions_tiff\"\nPOSTPROCESSED_DIR = OUTPUT_DIR / \"predictions_postprocessed\"\nSUBMISSION_ZIP = OUTPUT_DIR / \"submission.zip\"\n\nprint(\"INPUT_DIR:\", INPUT_DIR)\nprint(\"CHECKPOINT_DIR:\", CHECKPOINT_DIR)\nprint(\"WHEELS_DIR:\", WHEELS_DIR)\n"}, {"cell_type": "markdown", "metadata": {}, "source": "## Offline install\n\nReal internet is disabled during Kaggle submission scoring. Installs `nnunetv2==2.8.1`\n(exactly the version the checkpoint was trained with -- avoids any train/inference version\ndrift) plus its non-torch dependencies from the offline wheel bundle. torch/torchvision are\ndeliberately *not* bundled: Kaggle's GPU notebook image already ships a working torch, and\npulling a mismatched one from PyPI's default index during wheel-building actually resolved to\nthe wrong CUDA toolkit version entirely -- see `baselinerun/kaggle_submission/README.md`.\n\nThe wheel bundle is built for **Python 3.12** (`cp312`) -- Kaggle's actual kernel Python,\nconfirmed from a real failed run's traceback paths (`/usr/local/lib/python3.12/...`), not\nassumed. The first build targeted 3.11 to match the local dev env and failed on Kaggle with\n`ERROR: Could not find a version that satisfies the requirement nnunetv2==2.8.1 (from\nversions: none)` -- pip correctly refused to install `cp311`-tagged compiled wheels\n(numpy/scipy/scikit-image/etc.) into a 3.12 interpreter."}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "import time\n\n\ndef resolve_dataset_mount(candidates: list, min_files: int = 1, timeout_s: int = 1800, poll_s: int = 15) -> Path:\n    \"\"\"Polls every candidate path each round (Kaggle's mount convention for custom datasets\n    is inconsistent even within one account/kernel run -- see the cell-2 comment) and returns\n    the first one that actually has content. Raises only if NONE of the candidates mount\n    within timeout_s.\"\"\"\n    waited = 0\n    while waited <= timeout_s:\n        for c in candidates:\n            if c.exists() and len(list(c.iterdir())) >= min_files:\n                print(f\"{c} mounted ({len(list(c.iterdir()))} entries) after {waited}s\")\n                return c\n        time.sleep(poll_s)\n        waited += poll_s\n    raise RuntimeError(\n        f\"none of {candidates} mounted after {timeout_s}s \"\n        f\"(exists={[c.exists() for c in candidates]})\"\n    )\n\n\ndef wait_for_mount(path: Path, min_files: int = 1, timeout_s: int = 900, poll_s: int = 15) -> None:\n    \"\"\"Single-path convenience wrapper around resolve_dataset_mount, for mounts whose path\n    convention is already known/stable (e.g. competition data, which has consistently used\n    the /kaggle/input/competitions/<slug>/ nesting in every run so far -- unlike the custom\n    dataset_sources paths, which need the multi-candidate resolver above).\"\"\"\n    resolve_dataset_mount([path], min_files=min_files, timeout_s=timeout_s, poll_s=poll_s)\n\n\nWHEELS_DIR = resolve_dataset_mount(WHEELS_DIR_CANDIDATES, min_files=50, timeout_s=1800)\nCHECKPOINT_DIR = resolve_dataset_mount(CHECKPOINT_DIR_CANDIDATES, min_files=1, timeout_s=1800)\n\nresult = subprocess.run(\n    f\"pip install --no-index --find-links={WHEELS_DIR} nnunetv2==2.8.1 nibabel tifffile tqdm -q\",\n    shell=True, capture_output=True, text=True,\n)\nprint(result.stdout[-3000:])\nif result.returncode != 0:\n    print(\"STDERR:\", result.stderr[-3000:])\n    raise RuntimeError(\"Offline install failed\")\n\nimport nnunetv2\nprint(\"nnunetv2 installed OK\")\n"}, {"cell_type": "markdown", "metadata": {}, "source": "## Environment setup\n\nVerbatim from `baselinerun/src/training/environment.py::setup_environment`."}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "def setup_environment():\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    # baselinerun's own config uses \"true\" (torch.compile), but Kaggle's GPU pool can assign\n    # older cards (confirmed: Tesla P100, CUDA capability 6.0) that torch.compile's Triton\n    # backend does not support (requires >=7.0) -- a real run hit\n    # \"torch._inductor.exc.GPUTooOldForTriton\" here. Disabled for Kaggle specifically;\n    # arunodhayan's own notebook independently made the same call for the same reason.\n    os.environ[\"nnUNet_compile\"] = \"false\"\n\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\n\nsetup_environment()\n"}, {"cell_type": "markdown", "metadata": {}, "source": "## Stage the checkpoint and the custom trainer\n\nCopies the uploaded checkpoint dataset into the layout `nnUNetv2_predict` expects:\n`nnUNet_results/{DATASET_NAME}/{Trainer}__{Plans}__{Config}/{dataset.json,plans.json,fold_0/checkpoint_best.pth}`.\n\nThe checkpoint dataset is uploaded *flat* (dataset.json, plans.json, checkpoint_best.pth all\nat the dataset root, no subfolder) deliberately -- an earlier version uploaded a nested folder\nvia `kaggle datasets create -r zip`, which bundles the whole folder as a single `.zip` file\nrather than unpacking it, and that dataset consistently failed to mount into a running kernel\nat all (confirmed: still not mounted after a 15-minute wait, while flat, non-zipped datasets\nmounted instantly). Flat upload avoids the zip path entirely.\n\nAlso writes a minimal stub `nnUNetTrainerSkeletonRecall_700epochs.py` to disk and sets\n`nnUNet_extTrainer` so `nnUNetv2_predict` can resolve the trainer by name (the checkpoint\nembeds this trainer name internally). The real trainer adds a skeleton-recall loss + training\ntransforms that never run at inference -- see the code cell below for why a bare stub is\nsufficient (verified from nnU-Net's own inference source, not assumed)."}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "dst_model_dir = NNUNET_RESULTS / DATASET_NAME / MODEL_DIR_NAME\ndst_model_dir.mkdir(parents=True, exist_ok=True)\n(dst_model_dir / f\"fold_{FOLD}\").mkdir(parents=True, exist_ok=True)\n\nshutil.copy2(CHECKPOINT_DIR / \"dataset.json\", dst_model_dir / \"dataset.json\")\nshutil.copy2(CHECKPOINT_DIR / \"plans.json\", dst_model_dir / \"plans.json\")\nshutil.copy2(CHECKPOINT_DIR / CHECKPOINT_FILENAME, dst_model_dir / f\"fold_{FOLD}\" / CHECKPOINT_FILENAME)\n\ncheckpoint_path = dst_model_dir / f\"fold_{FOLD}\" / CHECKPOINT_FILENAME\nassert checkpoint_path.exists(), f\"Missing checkpoint at {checkpoint_path}\"\nprint(\"Staged checkpoint:\", checkpoint_path, f\"({checkpoint_path.stat().st_size / 1e6:.1f} MB)\")\nprint(\"dataset.json:\", (dst_model_dir / \"dataset.json\").exists())\nprint(\"plans.json:\", (dst_model_dir / \"plans.json\").exists())\n\n# ---- Stage a minimal stub trainer so nnUNetv2_predict can resolve\n# nnUNetTrainerSkeletonRecall_700epochs by name via the external-trainer mechanism.\n# The real trainer (vesuvius-surface/src/training/trainers/nnUNetTrainerSkeletonRecall*.py)\n# adds a skeleton-recall loss term and extra training-time transforms -- none of that runs at\n# inference. nnUNetv2_predict resolves the trainer class from the checkpoint's embedded\n# trainer_name and calls build_network_architecture DIRECTLY ON THE CLASS, never instantiating\n# it (verified via source read of predict_from_raw_data.py -- no trainer_class(...) call\n# anywhere in the inference path). Confirmed neither nnUNetTrainerSkeletonRecall nor\n# nnUNetTrainerSkeletonRecall_700epochs override build_network_architecture (source-read) --\n# it's inherited unchanged from stock nnUNetTrainer, so a bare passthrough subclass is\n# functionally identical for inference purposes. ----\nEXT_TRAINER_DIR.mkdir(parents=True, exist_ok=True)\n(EXT_TRAINER_DIR / \"nnUNetTrainerSkeletonRecall_700epochs.py\").write_text('''\\\n\"\"\"Inference-only stub for nnUNetTrainerSkeletonRecall_700epochs.\n\nThe real trainer (vesuvius-surface/src/training/trainers/nnUNetTrainerSkeletonRecall.py and\nits _700epochs subclass) adds a DC+CE+skeleton-recall loss and extra training-time transforms\non top of stock nnUNetTrainer. Neither overrides build_network_architecture, and\nnnUNetv2_predict calls build_network_architecture directly on the resolved class without ever\ninstantiating it -- so at inference time this bare subclass is functionally identical to the\nreal one. Loaded via nnU-Net's external-trainer mechanism (env var nnUNet_extTrainer); see\nnnunetv2.utilities.find_objects.recursive_find_trainer_class_by_name.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer\n\n\nclass nnUNetTrainerSkeletonRecall_700epochs(nnUNetTrainer):\n    pass\n''')\nos.environ[\"nnUNet_extTrainer\"] = str(EXT_TRAINER_DIR)\nprint(\"Staged custom trainer at:\", EXT_TRAINER_DIR / \"nnUNetTrainerSkeletonRecall_700epochs.py\")\nprint(\"nnUNet_extTrainer:\", os.environ[\"nnUNet_extTrainer\"])"}, {"cell_type": "markdown", "metadata": {}, "source": "## Prepare test data\n\nVerbatim from `baselinerun/src/data/prepare_training_data.py`\n(`create_spacing_json`, `prepare_single_case`, `prepare_test_data`) -- only the test-data\nsubset of that module is needed here, since this notebook only runs inference."}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "from tqdm.auto import tqdm\n\n\ndef create_spacing_json(output_path: Path, shape: tuple, spacing: tuple = (1.0, 1.0, 1.0)):\n    json_data = {\"spacing\": list(spacing)}\n    with open(output_path, \"w\") as f:\n        json.dump(json_data, f)\n\n\ndef prepare_single_case(src_path: Path, dest_path: Path, json_path: Path, use_symlinks: bool = True) -> bool:\n    try:\n        import tifffile\n        with tifffile.TiffFile(src_path) as tif:\n            shape = tif.pages[0].shape if len(tif.pages) == 1 else (len(tif.pages), *tif.pages[0].shape)\n\n        if use_symlinks:\n            if not dest_path.exists():\n                dest_path.symlink_to(src_path.resolve())\n        else:\n            shutil.copy2(src_path, dest_path)\n\n        create_spacing_json(json_path, shape)\n        return True\n    except Exception as e:\n        print(f\"Error processing {src_path.name}: {e}\")\n        return False\n\n\ndef prepare_test_data(input_dir: Path, output_dir: Path, use_symlinks: bool = True) -> Path:\n    output_dir.mkdir(parents=True, exist_ok=True)\n    test_images_dir = input_dir / \"test_images\"\n    if not test_images_dir.exists():\n        raise FileNotFoundError(f\"{test_images_dir} not found\")\n\n    test_files = sorted(test_images_dir.glob(\"*.tif\"))\n    print(f\"Found {len(test_files)} test cases\")\n\n    for img_path in tqdm(test_files, desc=\"Preparing test data\"):\n        case_id = img_path.stem\n        prepare_single_case(img_path, output_dir / f\"{case_id}_0000.tif\", output_dir / f\"{case_id}_0000.json\", use_symlinks)\n\n    return output_dir\n\n\n# The competition dataset (attached via competition_sources, not dataset_sources) previously\n# hit the exact same mount-lag issue as our own datasets -- confirmed via the Kaggle API that\n# test_images/1407735.tif genuinely exists, yet a run still saw test_images/ never mount even\n# after a full 15-minute wait (unlike our own datasets, which always mounted instantly once\n# uploaded correctly). Diagnose broadly this time instead of assuming the fix is identical:\n# wait for INPUT_DIR itself first, print what's actually there, then handle test_images\n# specifically with clear diagnostics either way.\nwait_for_mount(INPUT_DIR, min_files=1, timeout_s=300)\nprint(\"INPUT_DIR contents:\", sorted(p.name for p in INPUT_DIR.iterdir()))\n\ntest_images_dir = INPUT_DIR / \"test_images\"\ntry:\n    wait_for_mount(test_images_dir, min_files=1, timeout_s=600)\nexcept RuntimeError as e:\n    print(f\"WARNING: {e}\")\n    print(\"Trying case-insensitive / alternate-name search under INPUT_DIR...\")\n    candidates = [p for p in INPUT_DIR.rglob(\"*\") if p.is_dir() and \"test\" in p.name.lower()]\n    print(\"Directories with 'test' in the name:\", candidates)\n    for c in candidates:\n        try:\n            print(f\"  {c}: {sorted(p.name for p in c.iterdir())[:10]}\")\n        except Exception as inner_e:\n            print(f\"  {c}: could not list ({inner_e})\")\n    raise\n\nprepare_test_data(INPUT_DIR, TEST_INPUT_DIR)\n"}, {"cell_type": "markdown", "metadata": {}, "source": "## Run inference\n\n`_run_command` and `run_inference` verbatim from `baselinerun/src/training/commands.py`."}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "def _run_command(cmd: str, name: str = \"Command\", tail_lines: int = 30, timeout: Optional[int] = None) -> bool:\n    print(f\"Running: {cmd}\")\n    try:\n        result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=timeout)\n    except subprocess.TimeoutExpired:\n        print(f\"{name} TIMEOUT after {timeout}s!\")\n        return False\n\n    if result.returncode != 0:\n        print(f\"{name} FAILED!\")\n        print(f\"STDERR:\\n{result.stderr[-3000:]}\")\n        return False\n\n    print(f\"{name} complete!\")\n    if result.stdout.strip():\n        lines = result.stdout.strip().split(\"\\n\")\n        print(\"\\n\".join(lines[-tail_lines:]))\n    return True\n\n\ndef run_inference(\n    input_dir: Path, output_dir: Path, dataset_id: int, config: str, fold: Union[int, str],\n    plans: str, trainer: str, checkpoint_name: str = \"checkpoint_final.pth\", save_probabilities: bool = True,\n    num_processes_preprocessing: int = 2, num_processes_segmentation: int = 2,\n) -> bool:\n    output_dir.mkdir(parents=True, exist_ok=True)\n    cmd = f\"nnUNetv2_predict -d {dataset_id:03d} -c {config} -f {fold}\"\n    cmd += f\" -i {input_dir} -o {output_dir} -p {plans} -tr {trainer} -chk {checkpoint_name}\"\n    cmd += f\" -npp {num_processes_preprocessing} -nps {num_processes_segmentation}\"\n    cmd += \" --verbose\"\n    if save_probabilities:\n        cmd += \" --save_probabilities\"\n    return _run_command(cmd, \"Inference\")\n\n\nok = run_inference(\n    input_dir=TEST_INPUT_DIR, output_dir=PREDICTIONS_DIR,\n    dataset_id=DATASET_ID, config=CONFIGURATION, fold=FOLD,\n    plans=PLANS_NAME, trainer=TRAINER, checkpoint_name=CHECKPOINT_FILENAME,\n)\nassert ok, \"Inference failed -- see STDERR above\"\n"}, {"cell_type": "markdown", "metadata": {}, "source": "## Convert predictions to submission TIFFs\n\n`load_probabilities` and `predictions_to_tiff` verbatim from\n`baselinerun/src/training/postprocess.py` (the `.nii.gz` fallback branch is dropped here --\nit exists in the source only for a legacy prediction path this notebook never produces)."}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "import numpy as np\nimport tifffile\n\n\ndef load_probabilities(npz_path: Path) -> np.ndarray:\n    data = np.load(npz_path)\n    return data[\"probabilities\"]\n\n\ndef predictions_to_tiff(pred_dir: Path, output_dir: Path):\n    output_dir.mkdir(parents=True, exist_ok=True)\n    npz_files = list(pred_dir.glob(\"*.npz\"))\n    tif_files = list(pred_dir.glob(\"*.tif\"))\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            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            pred = tifffile.imread(str(tif_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\n\npredictions_to_tiff(PREDICTIONS_DIR, PREDICTIONS_TIFF_DIR)\n"}, {"cell_type": "markdown", "metadata": {}, "source": "## Apply 1st-place postprocessing\n\nInlined verbatim from `vesuvius-surface/src/postprocess/first_place.py` (the validated\ncontrol chain -- not the unmerge novelty layer, which is still under calibration and not part\nof this submission). Same defaults as the writeup / research_log.md: remove components <20k\nvoxels, per-sheet binary closing (radius 3, 26-connectivity), height-map gap patching (discard\nif it increases hole count), 1-voxel hole plugging via 2x2x2 LUT, then a global\n`binary_fill_holes`. Applied per predicted TIFF, `PREDICTIONS_TIFF_DIR` -> `POSTPROCESSED_DIR`."}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "from scipy import ndimage\nfrom scipy.ndimage import (\n    binary_closing,\n    binary_fill_holes,\n    distance_transform_edt,\n    find_objects,\n    generate_binary_structure,\n    label as cc_label,\n)\nfrom dataclasses import dataclass\n\n\n@dataclass\nclass PostprocessConfig:\n    min_component_size: int = 20_000\n    closing_radius: int = 3\n    connectivity: int = 26\n    enable_closing: bool = True\n    enable_patching: bool = True\n    enable_hole_plugging: bool = True\n    enable_fill_holes: bool = True\n    surface_label: int = 1\n    threshold: float = 0.5\n\n\ndef binarize_prediction(volume, *, threshold=0.5, surface_label=1):\n    arr = np.asarray(volume)\n    if arr.ndim == 4:\n        if arr.shape[0] > surface_label:\n            arr = arr[surface_label]\n        else:\n            arr = arr.argmax(axis=0)\n    if np.issubdtype(arr.dtype, np.floating):\n        return (arr >= threshold).astype(np.uint8)\n    uniq = set(np.unique(arr).tolist())\n    if uniq <= {0, 1}:\n        return (arr > 0).astype(np.uint8)\n    return (arr == surface_label).astype(np.uint8)\n\n\ndef _structure(connectivity):\n    if connectivity == 6:\n        return generate_binary_structure(3, 1)\n    if connectivity == 18:\n        return generate_binary_structure(3, 2)\n    if connectivity == 26:\n        return generate_binary_structure(3, 3)\n    raise ValueError(f\"connectivity must be 6, 18, or 26; got {connectivity}\")\n\n\ndef make_ball_footprint(radius):\n    zz, yy, xx = np.ogrid[-radius : radius + 1, -radius : radius + 1, -radius : radius + 1]\n    return (zz**2 + yy**2 + xx**2) <= radius**2\n\n\ndef _pad_slices(sl, shape, pad):\n    return tuple(\n        slice(max(0, s.start - pad), min(dim, s.stop + pad)) for s, dim in zip(sl, shape)\n    )\n\n\ndef remove_small_components(mask, min_size=20_000, connectivity=26):\n    struct = _structure(connectivity)\n    labeled, n = cc_label(mask.astype(np.uint8), structure=struct)\n    if n == 0:\n        return mask.astype(np.uint8)\n    sizes = ndimage.sum(mask, labeled, range(1, n + 1))\n    keep = np.zeros_like(mask, dtype=np.uint8)\n    for i, size in enumerate(sizes, 1):\n        if size >= min_size:\n            keep[labeled == i] = 1\n    return keep\n\n\ndef _count_internal_holes(mask):\n    inv = 1 - mask.astype(np.uint8)\n    labeled, n = cc_label(inv, structure=generate_binary_structure(3, 1))\n    if n == 0:\n        return 0\n    border = np.zeros(n + 1, dtype=bool)\n    border[labeled[0]] = True\n    border[labeled[-1]] = True\n    border[labeled[:, 0]] = True\n    border[labeled[:, -1]] = True\n    border[labeled[:, :, 0]] = True\n    border[labeled[:, :, -1]] = True\n    border[0] = True\n    return int(n - border[1:].sum())\n\n\ndef height_map_patch_crop(crop):\n    if crop.sum() == 0:\n        return crop\n\n    best_axis, best_area = 0, 0\n    for axis in range(3):\n        area = int(crop.max(axis=axis).sum())\n        if area > best_area:\n            best_area = area\n            best_axis = axis\n\n    crop_t = np.moveaxis(crop, best_axis, 0)\n    depth, height, width = crop_t.shape\n    depth_coords = np.arange(depth, dtype=np.float32).reshape(depth, 1, 1)\n    valid_3d = crop_t.astype(bool)\n    count_map = valid_3d.sum(axis=0)\n    has_voxels = count_map > 0\n\n    height_map = np.full((height, width), np.nan, dtype=np.float32)\n    thick_map = np.zeros((height, width), dtype=np.float32)\n    depth_sum = (depth_coords * valid_3d).sum(axis=0)\n    height_map[has_voxels] = depth_sum[has_voxels] / count_map[has_voxels]\n    thick_map[has_voxels] = count_map[has_voxels]\n\n    filled_proj = binary_fill_holes(has_voxels)\n    gap_mask = filled_proj & ~has_voxels\n    if not gap_mask.any():\n        return crop\n\n    holes_before = _count_internal_holes(crop)\n\n    fill_row_h = np.full((height, width), np.nan, dtype=np.float32)\n    fill_row_t = np.full((height, width), np.nan, dtype=np.float32)\n    for r in range(height):\n        valid_cols = np.where(has_voxels[r])[0]\n        gap_cols = np.where(gap_mask[r])[0]\n        if len(valid_cols) >= 2 and len(gap_cols) > 0:\n            fill_row_h[r, gap_cols] = np.interp(gap_cols, valid_cols, height_map[r, valid_cols])\n            fill_row_t[r, gap_cols] = np.interp(gap_cols, valid_cols, thick_map[r, valid_cols])\n\n    fill_col_h = np.full((height, width), np.nan, dtype=np.float32)\n    fill_col_t = np.full((height, width), np.nan, dtype=np.float32)\n    for c in range(width):\n        valid_rows = np.where(has_voxels[:, c])[0]\n        gap_rows = np.where(gap_mask[:, c])[0]\n        if len(valid_rows) >= 2 and len(gap_rows) > 0:\n            fill_col_h[gap_rows, c] = np.interp(gap_rows, valid_rows, height_map[valid_rows, c])\n            fill_col_t[gap_rows, c] = np.interp(gap_rows, valid_rows, thick_map[valid_rows, c])\n\n    not_valid = ~has_voxels\n    row_dist = distance_transform_edt(not_valid, sampling=[1e6, 1])\n    col_dist = distance_transform_edt(not_valid, sampling=[1, 1e6])\n\n    gap_r, gap_c = np.where(gap_mask)\n    hr, hc = fill_row_h[gap_r, gap_c], fill_col_h[gap_r, gap_c]\n    tr, tc = fill_row_t[gap_r, gap_c], fill_col_t[gap_r, gap_c]\n    dr = np.maximum(row_dist[gap_r, gap_c], 1e-6)\n    dc = np.maximum(col_dist[gap_r, gap_c], 1e-6)\n\n    valid_r, valid_c = ~np.isnan(hr), ~np.isnan(hc)\n    both = valid_r & valid_c\n    only_r, only_c = valid_r & ~valid_c, valid_c & ~valid_r\n    wr = np.where(both, 1.0 / dr, 0.0)\n    wc = np.where(both, 1.0 / dc, 0.0)\n    w_total = np.maximum(wr + wc, 1e-12)\n\n    h_avg = np.where(\n        both,\n        (np.nan_to_num(hr) * wr + np.nan_to_num(hc) * wc) / w_total,\n        np.where(only_r, hr, np.where(only_c, hc, np.nan)),\n    )\n    t_avg = np.where(\n        both,\n        (np.nan_to_num(tr) * wr + np.nan_to_num(tc) * wc) / w_total,\n        np.where(only_r, tr, np.where(only_c, tc, 0.0)),\n    )\n\n    patched_t = crop_t.copy()\n    for idx in range(len(gap_r)):\n        h_val = h_avg[idx]\n        if np.isnan(h_val):\n            continue\n        r, c = int(gap_r[idx]), int(gap_c[idx])\n        center = int(round(float(h_val)))\n        half = max(0, int(round(float(t_avg[idx]) / 2)))\n        z0 = max(0, center - half)\n        z1 = min(depth - 1, center + half)\n        patched_t[z0 : z1 + 1, r, c] = 1\n\n    patched_3d = np.moveaxis(patched_t, 0, best_axis)\n    if _count_internal_holes(patched_3d) > holes_before:\n        return crop\n    return patched_3d.astype(np.uint8)\n\n\n_HOLE_PLUG_LUT = None\n\n\ndef _build_hole_plug_lut():\n    face_diags = [\n        (0, 3, 1, 2), (1, 2, 0, 3), (4, 7, 5, 6), (5, 6, 4, 7),\n        (0, 5, 1, 4), (1, 4, 0, 5), (2, 7, 3, 6), (3, 6, 2, 7),\n        (0, 6, 2, 4), (2, 4, 0, 6), (1, 7, 3, 5), (3, 5, 1, 7),\n    ]\n    lut = np.zeros(256, dtype=np.uint8)\n    for pattern in range(256):\n        add = 0\n        for fa, fb, g1, g2 in face_diags:\n            if (\n                ((pattern >> fa) & 1)\n                and ((pattern >> fb) & 1)\n                and not ((pattern >> g1) & 1)\n                and not ((pattern >> g2) & 1)\n            ):\n                add |= 1 << g1\n        lut[pattern] = add\n    return lut\n\n\ndef plug_holes_lut(mask, max_iterations=5):\n    global _HOLE_PLUG_LUT\n    if _HOLE_PLUG_LUT is None:\n        _HOLE_PLUG_LUT = _build_hole_plug_lut()\n    lut = _HOLE_PLUG_LUT\n    result = mask.astype(np.uint8).copy()\n    depth, height, width = result.shape\n    if depth < 2 or height < 2 or width < 2:\n        return result\n\n    offsets = [(dz, dy, dx) for dz in range(2) for dy in range(2) for dx in range(2)]\n    for _ in range(max_iterations):\n        pattern = np.zeros((depth - 1, height - 1, width - 1), dtype=np.uint8)\n        for bit, (dz, dy, dx) in enumerate(offsets):\n            pattern |= result[dz : depth - 1 + dz, dy : height - 1 + dy, dx : width - 1 + dx] << bit\n        additions = lut[pattern]\n        if not additions.any():\n            break\n        for bit, (dz, dy, dx) in enumerate(offsets):\n            add_bit = ((additions >> bit) & 1).astype(np.uint8)\n            result[dz : depth - 1 + dz, dy : height - 1 + dy, dx : width - 1 + dx] |= add_bit\n    return result\n\n\ndef apply_first_place(prediction, config=None, *, through_stage=\"fill\"):\n    stages = (\"raw\", \"remove_small\", \"closing\", \"patch\", \"plug\", \"fill\")\n    cfg = config or PostprocessConfig()\n    stop_at = stages.index(through_stage)\n\n    mask = binarize_prediction(prediction, threshold=cfg.threshold, surface_label=cfg.surface_label)\n    if stop_at == 0:\n        return mask\n\n    mask = remove_small_components(mask, cfg.min_component_size, cfg.connectivity)\n    if stop_at == 1:\n        return mask\n\n    struct = _structure(cfg.connectivity)\n    labeled, n = cc_label(mask, structure=struct)\n    if n == 0:\n        return mask\n\n    slices = find_objects(labeled)\n    footprint = make_ball_footprint(cfg.closing_radius) if cfg.closing_radius > 0 else None\n    pad = cfg.closing_radius if cfg.enable_closing and cfg.closing_radius > 0 else 0\n    result = np.zeros_like(mask, dtype=np.uint8)\n\n    do_closing = cfg.enable_closing and stop_at >= 2\n    do_patch = cfg.enable_patching and stop_at >= 3\n    do_plug = cfg.enable_hole_plugging and stop_at >= 4\n\n    for comp_id, sl in enumerate(slices, 1):\n        if sl is None:\n            continue\n        padded_sl = _pad_slices(sl, mask.shape, pad)\n        crop = (labeled[padded_sl] == comp_id).astype(np.uint8)\n\n        if do_closing and footprint is not None:\n            crop = binary_closing(crop, structure=footprint).astype(np.uint8)\n        if do_patch:\n            crop = height_map_patch_crop(crop)\n        if do_plug:\n            crop = plug_holes_lut(crop)\n\n        result[padded_sl] |= crop\n\n    if stop_at < 5 or not cfg.enable_fill_holes:\n        return result\n    return binary_fill_holes(result).astype(np.uint8)\n\n\nPOSTPROCESSED_DIR.mkdir(parents=True, exist_ok=True)\npp_config = PostprocessConfig()\npred_tiffs = sorted(PREDICTIONS_TIFF_DIR.glob(\"*.tif\"))\nprint(f\"Postprocessing {len(pred_tiffs)} predictions (1st-place chain, config={pp_config})...\")\nfor tif_path in tqdm(pred_tiffs, desc=\"Postprocessing\"):\n    raw = tifffile.imread(str(tif_path))\n    processed = apply_first_place(raw, pp_config)\n    tifffile.imwrite(POSTPROCESSED_DIR / tif_path.name, processed.astype(np.uint8))\nprint(\"Postprocessing done ->\", POSTPROCESSED_DIR)"}, {"cell_type": "markdown", "metadata": {}, "source": "## Sanity check: dimensions and dtype\n\nNot part of the source pipeline -- added here because the competition rules are explicit\nthat each mask \"must match the dimensions of the source image exactly, and use the same data\ntype as the train mask\" (uint8). Cheap to check, expensive to get wrong."}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "train_labels_dir = INPUT_DIR / \"train_labels\"\nsample_train_label = next(train_labels_dir.glob(\"*.tif\"))\nexpected_dtype = tifffile.imread(str(sample_train_label)).dtype\nprint(f\"Expected dtype (from a train label): {expected_dtype}\")\n\nall_ok = True\n# A2: validate the raw predictions -- what actually gets zipped below, not the\n# unused postprocessed output.\nfor pred_path in sorted(PREDICTIONS_TIFF_DIR.glob(\"*.tif\")):\n    case_id = pred_path.stem\n    src_path = INPUT_DIR / \"test_images\" / f\"{case_id}.tif\"\n    pred_arr = tifffile.imread(str(pred_path))\n    src_arr = tifffile.imread(str(src_path))\n\n    shape_ok = pred_arr.shape == src_arr.shape\n    dtype_ok = pred_arr.dtype == expected_dtype\n    all_ok &= shape_ok and dtype_ok\n\n    print(f\"{case_id}: pred shape={pred_arr.shape} dtype={pred_arr.dtype} | \"\n          f\"src shape={src_arr.shape} | shape_ok={shape_ok} dtype_ok={dtype_ok}\")\n\nassert all_ok, \"Dimension/dtype mismatch detected -- fix before submitting\"\nprint(\"\\nAll predictions match source dimensions and expected dtype.\")\n"}, {"cell_type": "markdown", "metadata": {}, "source": "## Generate submission.zip\n\nVerbatim from `baselinerun/src/training/submission.py::generate_submission`."}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "def generate_submission(predictions_tiff_dir: Path, output_zip: Path, delete_after_zip: bool = True) -> Optional[Path]:\n    tiff_files = sorted(predictions_tiff_dir.glob(\"*.tif\"))\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    with zipfile.ZipFile(output_zip, \"w\", zipfile.ZIP_DEFLATED) as zipf:\n        for tiff_path in tqdm(tiff_files, desc=\"Zipping predictions\"):\n            zipf.write(tiff_path, tiff_path.name)\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    return output_zip\n\n\n# A2: zip the RAW predictions, not the postprocessed ones (postprocessing above still runs\n# for logging/parity but its output is intentionally unused here).\nsubmission_path = generate_submission(PREDICTIONS_TIFF_DIR, SUBMISSION_ZIP)\nassert submission_path is not None and submission_path.exists()\nprint(\"\\nDone:\", submission_path)\n"}], "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, "language_info": {"name": "python", "version": "3.11"}}, "nbformat": 4, "nbformat_minor": 5}