{"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":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":117682,"databundleVersionId":15062069},{"sourceType":"datasetVersion","sourceId":14905580,"datasetId":9537256,"databundleVersionId":15770730},{"sourceType":"datasetVersion","sourceId":14943987,"datasetId":9555226,"databundleVersionId":15813346},{"sourceType":"datasetVersion","sourceId":14873403,"datasetId":9515025,"databundleVersionId":15735766},{"sourceType":"datasetVersion","sourceId":14778603,"datasetId":9440035,"databundleVersionId":15631478},{"sourceType":"datasetVersion","sourceId":14868235,"datasetId":9500850,"databundleVersionId":15730107},{"sourceType":"datasetVersion","sourceId":14758444,"datasetId":9092253,"databundleVersionId":15609338},{"sourceType":"datasetVersion","sourceId":14457580,"datasetId":9234403,"databundleVersionId":15278306},{"sourceType":"datasetVersion","sourceId":14066012,"datasetId":8953198,"databundleVersionId":14846463}],"dockerImageVersionId":31193,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# ============================================================\n# Vesuvius nnUNet MULTI-MODE ENSEMBLE (modes 1–10 registry)\n# 2-GPU PARALLEL BATCHING + DISK-SAFE (/tmp) + DEADLOCK-SAFE LOG STREAMING\n#\n# BASE:  LB 0.613 \"nnUNet MULTI-MODEL chunk ensemble\" notebook\n#\n# CHANGES (per your request):\n#   - ACTIVE ensemble: (7 + 5 + 2 + 1)\n#   - DIRECT_WEIGHTED weights: M7=0.42, M5=0.18, M2=0.28, M1=0.12 (sums to 1)\n#   - Threshold: PROB_T = 0.26\n#   - Postprocess CHANGED:\n#        from: Patch+Plug+Fill + binary_closing(5)\n#        to:   closing(7,7) + check holes (PatchAndPlugAndFillCavities-style)\n#\n# Peak-disk reduction controls preserved:\n# ✅ Chunked test processing (CASES_PER_CHUNK)\n# ✅ For each chunk: predict -> per-case load prob -> float16-quantize -> fuse+PP -> write->zip->delete\n# ✅ No \"store all cases float16 maps\" phase\n# ✅ No \"store all TIFFs then zip\" phase\n# ============================================================\n\nimport os, json, shutil, subprocess, re, zipfile, threading\nfrom pathlib import Path\nfrom typing import Optional, Union, List, Dict\n\nimport numpy as np\nimport tifffile\nfrom tqdm.auto import tqdm\n\n# ----------------------------\n# Install (offline) - no pip cache\n# ----------------------------\n!pip install -q --no-index --no-cache-dir -f /kaggle/input/surface-packages-offline nnunetv2\n\n# ============================================================\n# Debug / logging options\n# ============================================================\nDEBUG_DISK_REPORT = False\nDISK_REPORT_INCLUDE_WORKING = True\n\n# ============================================================\n# Peak-disk reduction controls\n# ============================================================\nPROCESS_IN_CHUNKS = True\nCASES_PER_CHUNK   = 32\nZIP_TIFF_IMMEDIATELY = True\nEMULATE_FLOAT16_MAPS = True  # keep True to match previous behavior\n\n# ============================================================\n# Paths / Settings\n# ============================================================\nINPUT_DIR   = Path(\"/kaggle/input/vesuvius-challenge-surface-detection\")\n\n# /tmp scratch for ALL big intermediates\nWORKING_DIR = Path(\"/tmp/vesuvius_multi_scratch\")\nOUTPUT_DIR  = Path(\"/kaggle/working\")\n\nWORKING_DIR.mkdir(parents=True, exist_ok=True)\nOUTPUT_DIR.mkdir(parents=True, exist_ok=True)\n\n# ---- Cache / temp redirection ----\nCACHE_ROOT = WORKING_DIR / \"_cache\"\nTMP_ROOT   = WORKING_DIR / \"_tmp\"\nCACHE_ROOT.mkdir(parents=True, exist_ok=True)\nTMP_ROOT.mkdir(parents=True, exist_ok=True)\n\nos.environ[\"XDG_CACHE_HOME\"]  = str(CACHE_ROOT / \"xdg\")\nos.environ[\"PIP_CACHE_DIR\"]   = str(CACHE_ROOT / \"pip\")\nos.environ[\"NUMBA_CACHE_DIR\"] = str(CACHE_ROOT / \"numba\")\nos.environ[\"CUPY_CACHE_DIR\"]  = str(CACHE_ROOT / \"cupy\")\nos.environ[\"TORCH_HOME\"]      = str(CACHE_ROOT / \"torch\")\nos.environ[\"MPLCONFIGDIR\"]    = str(CACHE_ROOT / \"mpl\")\n\nos.environ[\"TMPDIR\"] = str(TMP_ROOT)\nos.environ[\"TEMP\"]   = str(TMP_ROOT)\nos.environ[\"TMP\"]    = str(TMP_ROOT)\n\nos.environ[\"PIP_NO_CACHE_DIR\"] = \"1\"\n\n# nnUNet knobs\nos.environ[\"nnUNet_USE_BLOSC2\"] = \"1\"\nos.environ[\"nnUNet_compile\"] = \"false\"\n\n# ---- Per-run cleanup toggles ----\nCLEAN_RUN_CACHES = True\nCLEAN_NNUNET_PREPROCESSED = True\nCLEAN_NNUNET_RESULTS = True\n\nFOLD: Union[int, str] = \"all\"\nCONFIGURATION = \"3d_fullres\"\n\n# ============================================================\n# ACTIVE MODES + FUSION (UPDATED)\n# ============================================================\nACTIVE_MODES = [7, 5, 2, 1]\nFUSION_SCHEME = \"DIRECT_WEIGHTED\"\nPROB_T = 0.26\n\n# Requested weights (sum to 1)\nW1  = 0.12\nW2  = 0.28\nW3  = 0.0\nW4  = 0.0\nW5  = 0.18\nW6  = 0.0\nW7  = 0.42\nW8  = 0.0\nW9  = 0.0\nW10 = 0.0\nWEIGHTS: Dict[int, float] = {1: W1, 2: W2, 3: W3, 4: W4, 5: W5, 6: W6, 7: W7, 8: W8, 9: W9, 10: W10}\n\n# Keep these defined for compatibility (not used when DIRECT_WEIGHTED)\nP54_W5 = 0.7\nP54_W4 = 0.3\nP21_W2 = 0.7\nP21_W1 = 0.3\nFINAL_W54 = 0.6\nFINAL_W21 = 0.4\n\n# Parallel CPU workers per process\nNPP = 1\nNPS = 1\n\n# ============================================================\n# nnUNet env roots (keep in /tmp)\n# ============================================================\nNNUNET_BASE         = WORKING_DIR / \"nnUNet_data\"\nNNUNET_RAW          = NNUNET_BASE / \"nnUNet_raw\"\nNNUNET_PREPROCESSED = NNUNET_BASE / \"nnUNet_preprocessed\"\nNNUNET_RESULTS      = WORKING_DIR / \"nnUNet_results\"\n\nTEST_INPUT_DIR = WORKING_DIR / \"test_input\"\nFINAL_TIFF_DIR = WORKING_DIR / \"predictions_tiff\"\nSUB_ZIP        = OUTPUT_DIR / \"submission.zip\"\n\nNNUNET_OUT: Dict[int, Path] = {m: WORKING_DIR / f\"nnunet_out_m{m}\" for m in range(1, 11)}\n\n# ============================================================\n# Utilities (disk reporting + safe IO)\n# ============================================================\ndef _sh(cmd: str) -> str:\n    p = subprocess.run(cmd, shell=True, capture_output=True, text=True)\n    out = (p.stdout or \"\") + (p.stderr or \"\")\n    return out.strip()\n\ndef disk_report(tag: str):\n    if not DEBUG_DISK_REPORT:\n        return\n    print(f\"\\n==================== DISK REPORT: {tag} ====================\")\n    try:\n        du = shutil.disk_usage(\"/tmp\")\n        print(f\"/tmp disk_usage: total={du.total/1024**3:.2f}G used={du.used/1024**3:.2f}G free={du.free/1024**3:.2f}G\")\n    except Exception as e:\n        print(\"disk_usage(/tmp) failed:\", repr(e))\n\n    if DISK_REPORT_INCLUDE_WORKING:\n        try:\n            duw = shutil.disk_usage(str(OUTPUT_DIR))\n            print(f\"{OUTPUT_DIR} disk_usage: total={duw.total/1024**3:.2f}G used={duw.used/1024**3:.2f}G free={duw.free/1024**3:.2f}G\")\n        except Exception as e:\n            print(f\"disk_usage({OUTPUT_DIR}) failed:\", repr(e))\n\n    print(_sh(\"df -h / /tmp || true\"))\n    if DISK_REPORT_INCLUDE_WORKING:\n        print(_sh(f\"df -h {OUTPUT_DIR} || true\"))\n\n    print(_sh(f\"du -sh {WORKING_DIR} 2>/dev/null || true\"))\n    print(_sh(f\"du -sh {NNUNET_PREPROCESSED} 2>/dev/null || true\"))\n    print(_sh(f\"du -sh {NNUNET_RESULTS} 2>/dev/null || true\"))\n    print(_sh(f\"du -sh {CACHE_ROOT} 2>/dev/null || true\"))\n    if DISK_REPORT_INCLUDE_WORKING:\n        print(_sh(f\"du -sh {OUTPUT_DIR} 2>/dev/null || true\"))\n        print(_sh(f\"ls -lh {OUTPUT_DIR} | head -n 50 || true\"))\n    print(\"============================================================\\n\")\n\ndef safe_rmtree(p: Path):\n    if p.exists():\n        shutil.rmtree(p, ignore_errors=True)\n\ndef safe_mkdir_clean(p: Path):\n    if p.exists():\n        safe_rmtree(p)\n    p.mkdir(parents=True, exist_ok=True)\n\ndef popen(cmd: str, env: Optional[dict] = None) -> subprocess.Popen:\n    return subprocess.Popen(\n        cmd, shell=True,\n        stdout=subprocess.PIPE, stderr=subprocess.STDOUT,\n        text=True, env=env,\n        cwd=str(WORKING_DIR),\n    )\n\ndef run(cmd: str, env: Optional[dict] = None, timeout: Optional[int] = None) -> subprocess.CompletedProcess:\n    return subprocess.run(cmd, shell=True, capture_output=True, text=True, env=env, timeout=timeout, cwd=str(WORKING_DIR))\n\ndef stream_process_output_async(p: subprocess.Popen, prefix: str) -> threading.Thread:\n    assert p.stdout is not None\n    def _worker():\n        for line in p.stdout:\n            print(f\"[{prefix}] {line.rstrip()}\")\n    t = threading.Thread(target=_worker, daemon=True)\n    t.start()\n    return t\n\ndef get_gpu_ids() -> List[str]:\n    try:\n        import torch\n        n = int(torch.cuda.device_count())\n    except Exception:\n        n = 0\n    if n >= 2:\n        return [\"0\", \"1\"]\n    if n == 1:\n        return [\"0\"]\n    return []\n\ndef setup_environment():\n    for d in [NNUNET_RAW, NNUNET_PREPROCESSED, NNUNET_RESULTS, OUTPUT_DIR, TEST_INPUT_DIR, FINAL_TIFF_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\n    print(\"WORKING_DIR (/tmp scratch):\", WORKING_DIR)\n    print(\"OUTPUT_DIR  (/kaggle/working):\", OUTPUT_DIR)\n    print(\"nnUNet_raw:\", NNUNET_RAW)\n    print(\"nnUNet_preprocessed:\", NNUNET_PREPROCESSED)\n    print(\"nnUNet_results:\", NNUNET_RESULTS)\n    print(\"CACHE_ROOT:\", CACHE_ROOT)\n    print(\"TMP_ROOT:\", TMP_ROOT)\n    print(\"nnUNet_USE_BLOSC2:\", os.environ.get(\"nnUNet_USE_BLOSC2\"))\n    print(\"CUDA_VISIBLE_DEVICES (parent):\", os.environ.get(\"CUDA_VISIBLE_DEVICES\", \"ALL\"))\n\ndef chunked(seq: List[Path], k: int) -> List[List[Path]]:\n    return [seq[i:i+k] for i in range(0, len(seq), k)]\n\n# ============================================================\n# MODELS[1..10] registry\n# ============================================================\ndef _resolve_ckpt(meta_src: Path, ckpt_name: Optional[str] = None) -> Path:\n    meta_src = Path(meta_src)\n    candidates: List[Path] = []\n    fold_all = meta_src / \"fold_all\"\n\n    if ckpt_name:\n        candidates += [fold_all / ckpt_name, meta_src / ckpt_name]\n    else:\n        for n in (\"checkpoint_final.pth\", \"checkpoint_best.pth\"):\n            candidates += [fold_all / n, meta_src / n]\n\n    for p in candidates:\n        if p.exists():\n            return p\n\n    hits = sorted(meta_src.rglob(\"checkpoint_*.pth\"))\n    if hits:\n        for p in hits:\n            if p.name == \"checkpoint_final.pth\":\n                return p\n        return hits[0]\n\n    raise FileNotFoundError(f\"No checkpoint found under {meta_src} (ckpt_name={ckpt_name})\")\n\ndef make_model(mode: int, tag: str, plans: str, trainer: str, meta_src: str,\n               ckpt: Optional[str] = None, ckpt_name: Optional[str] = None) -> dict:\n    ds_id = 100 + mode\n    ds_name = f\"Dataset{ds_id:03d}_VesuviusSurface_M{mode}\"\n    meta_src_p = Path(meta_src)\n    ckpt_p = Path(ckpt) if ckpt is not None else _resolve_ckpt(meta_src_p, ckpt_name=ckpt_name)\n    return dict(mode=mode, name=f\"M{mode}_{tag}\", ds_id=ds_id, ds_name=ds_name,\n                plans=plans, trainer=trainer, ckpt=ckpt_p, meta_src=meta_src_p)\n\nMODELS: Dict[int, dict] = {\n    1: make_model(\n        mode=1, tag=\"MPlans_4000_patch128\",\n        plans=\"nnUNetResEncUNetMPlans\", trainer=\"nnUNetTrainer_4000epochs\",\n        meta_src=\"/kaggle/input/nnunet-ep4000-all\",\n        ckpt=\"/kaggle/input/nnunet-ep4000-all/checkpoint_final.pth\",\n    ),\n    2: make_model(\n        mode=2, tag=\"XLPlans_250_patch192_best\",\n        plans=\"nnUNetResEncUNetXLPlans\", trainer=\"nnUNetTrainer_250epochs\",\n        meta_src=\"/kaggle/input/kaggle-vesuvius-dataset/nnUNetTrainer_250epochs__nnUNetResEncUNetXLPlans__3d_fullres\",\n        ckpt_name=\"checkpoint_best.pth\",\n    ),\n    3: make_model(\n        mode=3, tag=\"XLPlans_250_patch192_final\",\n        plans=\"nnUNetResEncUNetXLPlans\", trainer=\"nnUNetTrainer_250epochs\",\n        meta_src=\"/kaggle/input/kaggle-vesuvius-dataset/nnUNetTrainer_250epochs__nnUNetResEncUNetXLPlans__3d_fullres\",\n        ckpt_name=\"checkpoint_final.pth\",\n    ),\n    4: make_model(\n        mode=4, tag=\"XLPlans_250_patch256_best\",\n        plans=\"nnUNetResEncUNetXLPlans\", trainer=\"nnUNetTrainer_250epochs\",\n        meta_src=\"/kaggle/input/kaggle-vesuvius-dataset-nnunet-256\",\n        ckpt_name=\"checkpoint_best.pth\",\n    ),\n    5: make_model(\n        mode=5, tag=\"XLPlans_250_patch256_final\",\n        plans=\"nnUNetResEncUNetXLPlans\", trainer=\"nnUNetTrainer_250epochs\",\n        meta_src=\"/kaggle/input/kaggle-vesuvius-dataset-nnunet-256\",\n        ckpt_name=\"checkpoint_final.pth\",\n    ),\n\n    # ---- ADDED 6–10 (as provided) ----\n    6: make_model(\n        mode=6,\n        tag=\"06_XLPlans_500_patch288_E478\",\n        plans=\"nnUNetResEncUNetXLPlans\",\n        trainer=\"nnUNetTrainer_500epochs\",\n        meta_src=\"/kaggle/input/datasets/yiheng/kaggle-vesuvius-dataset-nnunet-288\",\n        ckpt_name=\"checkpoint_478e.pth\",\n    ),\n    7: make_model(\n        mode=7,\n        tag=\"07_LPlans_4000_patch192_fromscratch_final\",\n        plans=\"nnUNetResEncUNetLPlans\",\n        trainer=\"nnUNetTrainer_4000epochs\",\n        meta_src=\"/kaggle/input/datasets/yiheng/kaggle-vesuvius-dataset-nnunet-192-from-scratch\",\n        ckpt_name=\"checkpoint_final.pth\",\n    ),\n    8: make_model(\n        mode=8,\n        tag=\"8_LPlans_4000_patch224_fromscratch_final\",\n        plans=\"nnUNetResEncUNetLPlans\",\n        trainer=\"nnUNetTrainer_4000epochs\",\n        meta_src=\"/kaggle/input/datasets/yiheng/kaggle-vesuvius-dataset-nnunet-224-from-scratch\",\n        ckpt_name=\"checkpoint_best.pth\",\n    ),\n    9: make_model(\n        mode=9,\n        tag=\"09_LPlans_4000_patch160_fromscratch_best\",\n        plans=\"nnUNetResEncUNetLPlans\",\n        trainer=\"nnUNetTrainer_4000epochs\",\n        meta_src=\"/kaggle/input/datasets/yiheng/kaggle-vesuvius-dataset-nnunet-160-from-scratch\",\n        ckpt_name=\"checkpoint_best.pth\",\n    ),\n    10: make_model(\n        mode=10,\n        tag=\"10_LPlans_4000_patch224_fromscratch_final\",\n        plans=\"nnUNetResEncUNetLPlans\",\n        trainer=\"nnUNetTrainer_4000epochs\",\n        meta_src=\"/kaggle/input/datasets/yiheng/kaggle-vesuvius-dataset-nnunet-224-from-scratch\",\n        ckpt_name=\"checkpoint_final.pth\",\n    ),\n}\n\ndef ensure_model_registered(model: dict) -> Path:\n    ds_name = model[\"ds_name\"]\n    ds_id   = int(model[\"ds_id\"])\n    if not ds_name.startswith(f\"Dataset{ds_id:03d}_\"):\n        raise ValueError(f\"ds_name must start with Dataset{ds_id:03d}_ ; got {ds_name}\")\n\n    trainer = model[\"trainer\"]\n    plans   = model[\"plans\"]\n\n    model_root = NNUNET_RESULTS / ds_name / f\"{trainer}__{plans}__{CONFIGURATION}\"\n    fold_dir   = model_root / \"fold_all\"\n    fold_dir.mkdir(parents=True, exist_ok=True)\n\n    ckpt_src = Path(model[\"ckpt\"])\n    if not ckpt_src.exists():\n        raise FileNotFoundError(f\"Checkpoint not found: {ckpt_src}\")\n\n    ckpt_dst = fold_dir / \"checkpoint_final.pth\"\n    if ckpt_dst.exists() or ckpt_dst.is_symlink():\n        ckpt_dst.unlink()\n    ckpt_dst.symlink_to(ckpt_src)\n\n    meta_src = Path(model[\"meta_src\"])\n    for fname in [\"dataset.json\", \"plans.json\", \"dataset_fingerprint.json\"]:\n        src = meta_src / fname\n        dst = model_root / fname\n        if not src.exists():\n            raise FileNotFoundError(f\"Missing {fname} at {src}\")\n        if dst.exists():\n            dst.unlink()\n        shutil.copy2(src, dst)\n\n    resolved = ckpt_dst.resolve()\n    print(f\"Registered {model['name']}\")\n    print(\"  model_root:\", model_root)\n    print(\"  ckpt_src :\", ckpt_src, f\"(size={ckpt_src.stat().st_size/1024/1024:.1f} MB)\")\n    print(\"  ckpt_link:\", ckpt_dst, \"->\", resolved)\n    assert resolved == ckpt_src.resolve(), f\"Symlink sanity failed: {resolved} != {ckpt_src.resolve()}\"\n    return model_root\n\n# ============================================================\n# Test data prep (subset-aware)\n# ============================================================\ndef create_spacing_json(path: Path, spacing=(1.0, 1.0, 1.0)):\n    with open(path, \"w\") as f:\n        json.dump({\"spacing\": list(spacing)}, f)\n\ndef list_all_test_tifs() -> List[Path]:\n    test_images_dir = INPUT_DIR / \"test_images\"\n    tifs = sorted(test_images_dir.glob(\"*.tif\"))\n    print(\"Found test cases:\", len(tifs))\n    if not tifs:\n        raise RuntimeError(f\"No test tif found under {test_images_dir}\")\n    return tifs\n\ndef prepare_test_data_subset(tifs: List[Path]) -> List[str]:\n    safe_mkdir_clean(TEST_INPUT_DIR)\n    for p in tqdm(tifs, desc=\"Preparing test data (chunk)\"):\n        case_id = p.stem\n        img_dst = TEST_INPUT_DIR / f\"{case_id}_0000.tif\"\n        if img_dst.exists() or img_dst.is_symlink():\n            img_dst.unlink()\n        img_dst.symlink_to(p.resolve())\n        create_spacing_json(TEST_INPUT_DIR / f\"{case_id}_0000.json\")\n        create_spacing_json(TEST_INPUT_DIR / f\"{case_id}.json\")\n    return [p.stem for p in tifs]\n\n# ============================================================\n# nnUNet predict helpers\n# ============================================================\ndef detect_save_prob_flag() -> str:\n    help_out = run(\"nnUNetv2_predict -h\")\n    txt = (help_out.stdout or \"\") + \"\\n\" + (help_out.stderr or \"\")\n\n    candidates = [\n        \"--save_probabilities\",\n        \"--save_probabilities_in_separate_folder\",\n        \"--save_probabilities_in_npz\",\n        \"--export_probabilities\",\n    ]\n    for c in candidates:\n        if c in txt:\n            print(\"Detected probability flag:\", c)\n            return c\n\n    m = re.search(r\"--save[_-]prob\\w+\", txt)\n    if m:\n        print(\"Detected probability flag (regex):\", m.group(0))\n        return m.group(0)\n\n    print(\"nnUNetv2_predict -h tail:\\n\", txt[-2500:])\n    raise RuntimeError(\"Could not find probability-saving flag in nnUNetv2_predict help.\")\n\ndef build_predict_cmd(model: dict, out_dir: Path, save_prob_flag: str) -> str:\n    return (\n        f\"nnUNetv2_predict \"\n        f\"-d {int(model['ds_id']):03d} -c {CONFIGURATION} -f {FOLD} \"\n        f\"-i {TEST_INPUT_DIR} -o {out_dir} \"\n        f\"-p {model['plans']} -tr {model['trainer']} \"\n        f\"-npp {NPP} -nps {NPS} --verbose \"\n        f\"{save_prob_flag}\"\n    )\n\ndef run_two_models_parallel(model1: dict, out1: Path, gpu1: str,\n                            model2: dict, out2: Path, gpu2: str,\n                            save_prob_flag: str,\n                            stage: int):\n    disk_report(f\"BEFORE Stage {stage} (parallel)\")\n\n    cmd1 = build_predict_cmd(model1, out1, save_prob_flag)\n    cmd2 = build_predict_cmd(model2, out2, save_prob_flag)\n\n    env1 = os.environ.copy(); env1[\"CUDA_VISIBLE_DEVICES\"] = gpu1\n    env2 = os.environ.copy(); env2[\"CUDA_VISIBLE_DEVICES\"] = gpu2\n\n    print(f\"\\n=== Stage {stage}: parallel inference ===\")\n    print(f\"GPU{gpu1}:\", cmd1)\n    print(f\"GPU{gpu2}:\", cmd2)\n\n    p1 = popen(cmd1, env=env1)\n    p2 = popen(cmd2, env=env2)\n\n    t1 = stream_process_output_async(p1, model1[\"name\"])\n    t2 = stream_process_output_async(p2, model2[\"name\"])\n\n    rc1 = p1.wait()\n    rc2 = p2.wait()\n    t1.join(timeout=2)\n    t2.join(timeout=2)\n\n    if rc1 != 0 or rc2 != 0:\n        disk_report(f\"AFTER Stage {stage} (parallel) - FAILED\")\n        raise RuntimeError(f\"nnUNetv2_predict failed: rc1={rc1}, rc2={rc2}\")\n\n    disk_report(f\"AFTER Stage {stage} (parallel)\")\n\ndef run_one_model(model: dict, out_dir: Path, gpu: str, save_prob_flag: str, stage: int):\n    disk_report(f\"BEFORE Stage {stage} (single)\")\n\n    cmd = build_predict_cmd(model, out_dir, save_prob_flag)\n    env = os.environ.copy(); env[\"CUDA_VISIBLE_DEVICES\"] = gpu\n\n    print(f\"\\n=== Stage {stage}: single inference ===\")\n    print(f\"GPU{gpu}:\", cmd)\n\n    p = popen(cmd, env=env)\n    t = stream_process_output_async(p, model[\"name\"])\n    rc = p.wait()\n    t.join(timeout=2)\n\n    if rc != 0:\n        disk_report(f\"AFTER Stage {stage} (single) - FAILED\")\n        raise RuntimeError(f\"nnUNetv2_predict failed for {model['name']}: rc={rc}\")\n\n    disk_report(f\"AFTER Stage {stage} (single)\")\n\ndef run_active_models_in_stages(active_modes: List[int], save_prob_flag: str):\n    gpu_ids = get_gpu_ids()\n    if not gpu_ids:\n        raise RuntimeError(\"No CUDA GPU detected. This pipeline expects at least 1 GPU.\")\n\n    modes = list(active_modes)\n    stage = 1\n    i = 0\n    while i < len(modes):\n        if len(gpu_ids) >= 2 and (i + 1) < len(modes):\n            m1, m2 = modes[i], modes[i + 1]\n            run_two_models_parallel(MODELS[m1], NNUNET_OUT[m1], gpu_ids[0],\n                                    MODELS[m2], NNUNET_OUT[m2], gpu_ids[1],\n                                    save_prob_flag, stage=stage)\n            i += 2\n        else:\n            m1 = modes[i]\n            run_one_model(MODELS[m1], NNUNET_OUT[m1], gpu_ids[0], save_prob_flag, stage=stage)\n            i += 1\n        stage += 1\n\n# ============================================================\n# Load class-1 prob DIRECTLY from nnUNet npz (per-case)\n# ============================================================\ndef find_case_npz(pred_dir: Path, case_id: str) -> Path:\n    direct = pred_dir / f\"{case_id}.npz\"\n    if direct.exists():\n        return direct\n\n    probs = pred_dir / \"probabilities\" / f\"{case_id}.npz\"\n    if probs.exists():\n        return probs\n\n    alt = pred_dir / f\"{case_id}_0000.npz\"\n    if alt.exists():\n        return alt\n\n    probs_alt = pred_dir / \"probabilities\" / f\"{case_id}_0000.npz\"\n    if probs_alt.exists():\n        return probs_alt\n\n    hits = sorted(pred_dir.rglob(f\"{case_id}*.npz\"))\n    if hits:\n        return hits[0]\n\n    raise FileNotFoundError(f\"No NPZ found for {case_id} under {pred_dir}\")\n\ndef load_surface_prob_from_npz(npz_path: Path) -> np.ndarray:\n    with np.load(npz_path) as data:\n        keys = list(data.keys())\n        arr = None\n        for k in [\"probabilities\", \"softmax\", \"predicted_probabilities\", \"probs\"]:\n            if k in data:\n                arr = data[k]\n                break\n        if arr is None:\n            arr = data[keys[0]]\n\n    arr = np.asarray(arr)\n    if arr.ndim != 4:\n        raise ValueError(f\"Unexpected prob tensor shape {arr.shape} in {npz_path}\")\n\n    # try to infer channel axis\n    if arr.shape[0] <= 10 and arr.shape[1] > 10:\n        p = arr[1].astype(np.float32)\n    elif arr.shape[-1] <= 10 and arr.shape[0] > 10:\n        p = arr[..., 1].astype(np.float32)\n    else:\n        p = arr[1].astype(np.float32)\n\n    if EMULATE_FLOAT16_MAPS:\n        p = p.astype(np.float16).astype(np.float32)\n\n    return p\n\ndef load_p_from_nnunet(mode: int, cid: str) -> np.ndarray:\n    npz = find_case_npz(NNUNET_OUT[mode], cid)\n    return load_surface_prob_from_npz(npz)\n\n# ============================================================\n# Fusion helpers\n# ============================================================\ndef normalize_weights(active_modes: List[int], weights: Dict[int, float]) -> Dict[int, float]:\n    w = {m: float(weights.get(m, 0.0)) for m in active_modes}\n    w = {m: v for m, v in w.items() if v > 0}\n    s = sum(w.values())\n    if s <= 0:\n        raise ValueError(f\"All weights are zero/non-positive for ACTIVE_MODES={active_modes}.\")\n    return {m: v / s for m, v in w.items()}\n\ndef fuse_probability_from_nnunet(cid: str, active_modes: List[int]) -> np.ndarray:\n    if FUSION_SCHEME == \"DIRECT_WEIGHTED\":\n        w_norm = normalize_weights(active_modes, WEIGHTS)\n        p = None\n        for m, w in w_norm.items():\n            pm = load_p_from_nnunet(m, cid)\n            if p is None:\n                p = w * pm\n            else:\n                p += w * pm\n        assert p is not None\n        return p\n\n    if FUSION_SCHEME == \"PAIR_ENSEMBLE\":\n        required = {1, 2, 4, 5}\n        if not required.issubset(set(active_modes)):\n            raise ValueError(f\"PAIR_ENSEMBLE requires ACTIVE_MODES to include {sorted(required)}; got {active_modes}\")\n\n        p5 = load_p_from_nnunet(5, cid)\n        p4 = load_p_from_nnunet(4, cid)\n        p2 = load_p_from_nnunet(2, cid)\n        p1 = load_p_from_nnunet(1, cid)\n\n        p54 = (P54_W5 * p5) + (P54_W4 * p4)\n        p21 = (P21_W2 * p2) + (P21_W1 * p1)\n        p   = (FINAL_W54 * p54) + (FINAL_W21 * p21)\n        return p\n\n    raise ValueError(f\"Unknown FUSION_SCHEME={FUSION_SCHEME}\")\n\n# ============================================================\n# Postprocess (REPLACED): closing(7,7) + check holes (PatchAndPlugAndFillCavities-style)\n# ============================================================\nfrom numba import jit\nfrom skimage import morphology as sk_morphology\nimport cv2\nimport cupy\nimport cupyx\nfrom cupyx.scipy import ndimage  # kept\n\nEULER_COEFS3D_26 = cupy.array([\n    0, 1, 1, 0, 1, 0, -2, -1, 1, -2, 0, -1, 0, -1, -1, 0, 1, 0, -2, -1, -2, -1, -1, -2, -6, -3, -3, -2, -3, -2, 0, -1, 1, -2,\n    0, -1, -6, -3, -3, -2, -2, -1, -1, -2, -3, 0, -2, -1, 0, -1, -1, 0, -3, -2, 0, -1, -3, 0, -2, -1, 0, 1, 1, 0, 1, -2, -6, -3, 0, -1, -3, -2, -2, -1,\n    -3, 0, -1, -2, -2, -1, 0, -1, -3, -2, -1, 0, 0, -1, -3, 0, 0, 1, -2, -1, 1, 0, -2, -1, -3, 0, -3, 0, 0, 1, -1, 4, 0, 3, 0, 3, 1, 2, -1, -2, -2, -1,\n    -2, -1, 1, 0, 0, 3, 1, 2, 1, 2, 2, 1, 1, -6, -2, -3, -2, -3, -1, 0, 0, -3, -1, -2, -1, -2, -2, -1, -2, -3, -1, 0, -1, 0, 4, 3, -3, 0, 0, 1, 0, 1,\n    3, 2, 0, -3, -1, -2, -3, 0, 0, 1, -1, 0, 0, -1, -2, 1, -1, 0, -1, -2, -2, -1, 0, 1, 3, 2, -2, 1, -1, 0, 1, 2, 2, 1, 0, -3, -3, 0, -1, -2, 0, 1,\n    -1, 0, -2, 1, 0, -1, -1, 0, -1, -2, 0, 1, -2, -1, 3, 2, -2, 1, 1, 2, -1, 0, 2, 1, -1, 0, -2, 1, -2, 1, 1, 2, -2, 3, -1, 2, -1, 2, 0, 1, 0, -1,\n    -1, 0, -1, 0, 2, 1, -1, 2, 0, 1, 0, 1, 1, 0\n], dtype=cupy.int32)\n\ndef cupy_euler_number(image, connectivity: int):\n    image = cupy.asarray(image)\n    image = (image > 0).astype(int)\n    image = cupy.pad(image, pad_width=1, mode='constant')\n\n    if image.ndim != 3:\n        raise ValueError('Input image is not 3D.')\n\n    if connectivity == 2:\n        raise NotImplementedError('For 3D images, Euler number is implemented for connectivities 1 and 3 only')\n\n    config = cupy.array([\n        [[0, 0, 0], [0, 0, 0], [0, 0, 0]],\n        [[0, 0, 0], [0, 1, 4], [0, 2, 8]],\n        [[0, 0, 0], [0, 16, 64], [0, 32, 128]],\n    ])\n\n    if connectivity == 1:\n        coefs = EULER_COEFS3D_26[::-1]\n    else:\n        coefs = EULER_COEFS3D_26\n\n    XF = cupyx.scipy.ndimage.convolve(image, config, mode='constant', cval=0)\n    h = cupy.bincount(XF.ravel(), minlength=256)\n    return int((0.125 * coefs @ h).item())\n\ndef get_best_projection(sheet_vol):\n    best_axis, best_area, best_mask = 0, 0, None\n    for axis in range(3):\n        proj_mask = cupy.max(sheet_vol, axis)\n        proj_area = cupy.sum(proj_mask)\n        if proj_area > best_area:\n            best_axis = axis\n            best_area = proj_area\n            best_mask = proj_mask\n    return best_axis, best_mask\n\ndef get_hole_patches(sheet_vol, proj_axis, proj_mask, params):\n    proj_mask = cupy.asnumpy(proj_mask)\n    heightmap = cupy.asnumpy(cupy.argmax(sheet_vol, proj_axis))\n\n    bdr = params['border_width']\n    if bdr > 0:\n        heightmap = heightmap[bdr:-bdr, bdr:-bdr]\n        proj_mask = proj_mask[bdr:-bdr, bdr:-bdr]\n\n    num_holes, hole_labels_img, hole_stats, _ = (\n        cv2.connectedComponentsWithStats(1 - proj_mask, connectivity=8, ltype=cv2.CV_16U)\n    )\n\n    patches = []\n    for i in range(1, num_holes):\n        x0, y0, hole_w, hole_h, _ = hole_stats[i]\n\n        if (x0 == 0) or (y0 == 0) or (y0 + hole_h >= heightmap.shape[0]) or (x0 + hole_w >= heightmap.shape[1]):\n            continue\n\n        ph, pw = hole_h + 2, hole_w + 2\n        interp_hm_h = np.full((ph, pw), -1, dtype=np.float32)\n        interp_hm_v = np.full((ph, pw), -1, dtype=np.float32)\n        wts_h = np.full((ph, pw), 0, dtype=np.float32)\n        wts_v = np.full((ph, pw), 0, dtype=np.float32)\n\n        hole_mask = (hole_labels_img[y0-1:y0+ph-1, x0-1:x0+pw-1] == i)\n        not_hole_mask = ~hole_mask\n\n        # horizontal interpolation\n        for y in range(0, ph):\n            hm_row = heightmap[y + y0 - 1, x0 - 1:]\n            mask_row = hole_mask[y]\n            wts_h_row = wts_h[y]\n            interp_hm_h_row = interp_hm_h[y]\n            ramping = False\n            x = 0\n            while x < pw:\n                if not ramping:\n                    if (not mask_row[x]) and (x < pw - 1) and mask_row[x + 1]:\n                        ramping = True\n                        ramp_start = (x, hm_row[x])\n                else:\n                    if not mask_row[x]:\n                        ramping = False\n                        ramp_end = (x, hm_row[x])\n                        slope = (ramp_end[1] - ramp_start[1]) / (ramp_end[0] - ramp_start[0])\n                        ht = ramp_start[1] + slope\n                        for xx in range(ramp_start[0] + 1, ramp_end[0]):\n                            interp_hm_h_row[xx] = ht\n                            ht += slope\n                            wts_h_row[xx] = 1 / min(xx - ramp_start[0], ramp_end[0] - xx)\n                        x -= 1\n                x += 1\n\n        # vertical interpolation\n        for x in range(0, pw):\n            hm_col = heightmap[y0 - 1:, x + x0 - 1]\n            mask_col = hole_mask[:, x]\n            wts_v_col = wts_v[:, x]\n            interp_hm_v_col = interp_hm_v[:, x]\n            ramping = False\n            y = 0\n            while y < ph:\n                if not ramping:\n                    if (not mask_col[y]) and (y < ph - 1) and mask_col[y + 1]:\n                        ramping = True\n                        ramp_start = (y, hm_col[y])\n                else:\n                    if not mask_col[y]:\n                        ramping = False\n                        ramp_end = (y, hm_col[y])\n                        slope = (ramp_end[1] - ramp_start[1]) / (ramp_end[0] - ramp_start[0])\n                        ht = ramp_start[1] + slope\n                        for yy in range(ramp_start[0] + 1, ramp_end[0]):\n                            interp_hm_v_col[yy] = ht\n                            ht += slope\n                            wts_v_col[yy] = 1 / min(yy - ramp_start[0], ramp_end[0] - yy)\n                        y -= 1\n                y += 1\n\n        interp_hts = np.round(\n            (wts_h * interp_hm_h + wts_v * interp_hm_v) / (wts_h + wts_v + 1e-8)\n        ).astype(np.uint16)\n        interp_hts[not_hole_mask] = heightmap[y0-1:y0-1+ph, x0-1:x0-1+pw][not_hole_mask]\n\n        patch = np.zeros((ph, pw, 2), dtype=np.int32)\n        for yy in range(1, hole_h + 1):\n            for xx in range(1, hole_w + 1):\n                nbr_hts = interp_hts[yy-1:yy+2, xx-1:xx+2]\n                patch[yy, xx, 0] = np.min(nbr_hts)\n                patch[yy, xx, 1] = np.max(nbr_hts)\n\n        patch[:, :, 0][not_hole_mask] = -1\n        patch[:, :, 1][not_hole_mask] = -1\n        patch = np.ascontiguousarray(patch[1:-1, 1:-1, :])\n        patches.append((patch, x0 + bdr, y0 + bdr, hole_w, hole_h))\n\n    return patches, proj_axis\n\n@jit(nopython=True)\ndef insert_patches_in_volume(patches, proj_axis, vol):\n    pz_max = vol.shape[proj_axis]\n    axes = (0,1,2) if proj_axis == 0 else (1,0,2) if proj_axis == 1 else (2,0,1)\n    vol_t = np.transpose(vol, axes)\n\n    for patch_info in patches:\n        patch, px0, py0, pw, ph = patch_info\n        for py in range(py0, py0+ph):\n            for px in range(px0, px0+pw):\n                for pz in range(patch[py-py0, px-px0, 0], patch[py-py0, px-px0, 1] + 1):\n                    if (pz > 0) and (pz < pz_max):\n                        vol_t[pz][py][px] = 1\n    return vol\n\ndef plug_small_holes(in_vol, lut, params):\n    dilation_struct = cupy.array(sk_morphology.footprint_rectangle((3,3,3)))\n    dilated_vol = cupyx.scipy.ndimage.binary_dilation(in_vol, dilation_struct)\n\n    voxel_locs = [cupy.asnumpy(a) for a in cupy.nonzero(dilated_vol)]\n    out_vol = cupy.asnumpy(in_vol)\n    count = plug_small_holes_numba(out_vol, voxel_locs, lut)\n    count += plug_small_holes_numba(out_vol, voxel_locs, lut)\n    return out_vol, count\n\n@jit(nopython=True)\ndef plug_small_holes_numba(in_vol, voxel_locs, lut):\n    count = 0\n    for z, y, x in zip(voxel_locs[0], voxel_locs[1], voxel_locs[2]):\n        lut_code = in_vol[z,y,x] +    2*in_vol[z,y,x+1] +    4*in_vol[z,y+1,x] +     8*in_vol[z,y+1,x+1] + \\\n                16*in_vol[z+1,y,x] + 32*in_vol[z+1,y,x+1] + 64*in_vol[z+1,y+1,x] + 128*in_vol[z+1,y+1,x+1]\n        lut_val = lut[lut_code]\n        if lut_val[0,0,0] < 256:\n            in_vol[z:z+2, y:y+2, x:x+2] = lut_val\n            count += 1\n    return count\n\ndef code_to_array(code):\n    result = np.zeros((8,), dtype=np.uint32)\n    b = f\"{code:08b}\"\n    for i in range(8):\n        result[7-i] = int(b[i])\n    return result.reshape((2,2,2))\n\ndef create_plug_lut():\n    lut = [None] * 256\n    for i in range(256):\n        lut[i] = np.full((8,), 512, dtype=np.uint32).reshape((2,2,2))\n\n    lut[6]   = code_to_array(7)\n    lut[9]   = code_to_array(11)\n    lut[18]  = code_to_array(19)\n    lut[20]  = code_to_array(21)\n    lut[22]  = code_to_array(23)\n    lut[24]  = code_to_array(27)\n    lut[25]  = code_to_array(27)\n    lut[26]  = code_to_array(27)\n    lut[28]  = code_to_array(29)\n    lut[30]  = code_to_array(31)\n\n    lut[33] = code_to_array(35)\n    lut[36] = code_to_array(39)\n    lut[37] = code_to_array(39)\n    lut[38] = code_to_array(39)\n    lut[40] = code_to_array(42)\n    lut[41] = code_to_array(43)\n    lut[44] = code_to_array(46)\n    lut[45] = code_to_array(47)\n\n    lut[52] = code_to_array(53)\n    lut[54] = code_to_array(55)\n    lut[56] = code_to_array(58)\n    lut[57] = code_to_array(59)\n    lut[60] = code_to_array(63)\n    lut[61] = code_to_array(63)\n    lut[62] = code_to_array(63)\n\n    lut[65] = code_to_array(69)\n    lut[66] = code_to_array(71)\n    lut[67] = code_to_array(71)\n    lut[70] = code_to_array(71)\n    lut[72] = code_to_array(76)\n    lut[73] = code_to_array(77)\n    lut[74] = code_to_array(78)\n    lut[75] = code_to_array(79)\n    lut[82] = code_to_array(83)\n    lut[86] = code_to_array(87)\n    lut[88] = code_to_array(92)\n    lut[89] = code_to_array(93)\n    lut[90] = code_to_array(95)\n    lut[91] = code_to_array(95)\n    lut[94] = code_to_array(95)\n\n    lut[96] = code_to_array(112)\n    lut[97] = code_to_array(113)\n    lut[98] = code_to_array(114)\n    lut[99] = code_to_array(115)\n    lut[100] = code_to_array(116)\n    lut[101] = code_to_array(117)\n    lut[102] = code_to_array(119)\n    lut[103] = code_to_array(119)\n    lut[104] = code_to_array(232)\n    lut[105] = code_to_array(251)\n    lut[106] = code_to_array(234)\n    lut[108] = code_to_array(236)\n    lut[107] = code_to_array(251)\n    lut[109] = code_to_array(253)\n    lut[110] = code_to_array(238)\n    lut[111] = code_to_array(239)\n    lut[118] = code_to_array(119)\n    lut[120] = code_to_array(248)\n    lut[121] = code_to_array(127)\n    lut[122] = code_to_array(127)\n    lut[123] = code_to_array(127)\n    lut[124] = code_to_array(127)\n    lut[125] = code_to_array(127)\n    lut[126] = code_to_array(127)\n\n    lut[129] = code_to_array(139)\n    lut[130] = code_to_array(138)\n    lut[131] = code_to_array(139)\n    lut[132] = code_to_array(140)\n    lut[133] = code_to_array(141)\n    lut[134] = code_to_array(142)\n    lut[135] = code_to_array(143)\n    lut[137] = code_to_array(139)\n    lut[144] = code_to_array(176)\n    lut[145] = code_to_array(177)\n    lut[146] = code_to_array(178)\n    lut[147] = code_to_array(179)\n    lut[148] = code_to_array(212)\n    lut[149] = code_to_array(213)\n    lut[150] = code_to_array(223)\n    lut[151] = code_to_array(223)\n    lut[152] = code_to_array(184)\n    lut[153] = code_to_array(187)\n    lut[154] = code_to_array(186)\n    lut[155] = code_to_array(187)\n    lut[156] = code_to_array(220)\n    lut[157] = code_to_array(221)\n    lut[158] = code_to_array(191)\n    lut[159] = code_to_array(191)\n\n    lut[161] = code_to_array(163)\n    lut[164] = code_to_array(172)\n    lut[165] = code_to_array(175)\n    lut[166] = code_to_array(174)\n    lut[167] = code_to_array(175)\n    lut[169] = code_to_array(171)\n    lut[173] = code_to_array(175)\n    lut[180] = code_to_array(244)\n    lut[181] = code_to_array(245)\n    lut[182] = code_to_array(247)\n    lut[183] = code_to_array(247)\n    lut[185] = code_to_array(187)\n    lut[188] = code_to_array(252)\n    lut[189] = code_to_array(253)\n    lut[190] = code_to_array(254)\n\n    lut[193] = code_to_array(197)\n    lut[194] = code_to_array(202)\n    lut[195] = code_to_array(207)\n    lut[198] = code_to_array(206)\n    lut[199] = code_to_array(207)\n    lut[201] = code_to_array(205)\n    lut[203] = code_to_array(207)\n    lut[210] = code_to_array(242)\n    lut[211] = code_to_array(243)\n    lut[214] = code_to_array(223)\n    lut[215] = code_to_array(223)\n    lut[217] = code_to_array(221)\n    lut[218] = code_to_array(223)\n    lut[219] = code_to_array(223)\n    lut[222] = code_to_array(223)\n\n    lut[225] = code_to_array(241)\n    lut[227] = code_to_array(243)\n    lut[229] = code_to_array(245)\n    lut[230] = code_to_array(238)\n    lut[231] = code_to_array(247)\n    lut[233] = code_to_array(239)\n    lut[235] = code_to_array(251)\n    lut[237] = code_to_array(253)\n    lut[246] = code_to_array(247)\n    lut[249] = code_to_array(253)\n    return lut\n\ndef postprocess(in_vol, params):\n    # Zero border\n    bw = params['border_width']\n    if bw > 0:\n        in_vol[0:bw, :, :] = in_vol[:, 0:bw, :] = in_vol[:, :, 0:bw] = 0\n        in_vol[-bw:, :, :] = in_vol[:, -bw:, :] = in_vol[:, :, -bw:] = 0\n\n    # Remove non-sheets\n    pruned_vol = sk_morphology.remove_small_objects(in_vol.astype(bool), min_size=params['small_obj_thresh'])\n\n    # Label sheets (26-connected)\n    struct_element333 = cupy.array(sk_morphology.footprint_rectangle((3,3,3)))\n    sheet_labels_vol, num_sheets = cupyx.scipy.ndimage.label(cupy.array(pruned_vol), struct_element333)\n\n    out_vol = cupy.zeros(in_vol.shape, dtype=np.uint8)\n    sheet_vol = cupy.empty(sheet_labels_vol.shape, dtype=np.uint8)\n\n    lut = create_plug_lut()\n\n    # Build oriented ellipsoidal closing elements\n    close_par  = int(params.get('close_par', 7))\n    close_perp = int(params.get('close_perp', 7))\n    sz = close_par  # source notebook assumption; with 7/7 this is correct\n    ctr = (sz - 1) // 2\n    r_par  = (close_par  - 1) // 2\n    r_perp = (close_perp - 1) // 2\n\n    closing_element = np.zeros((sz, sz, sz), dtype=np.uint8)\n    for z in range(sz):\n        for y in range(sz):\n            for x in range(sz):\n                if ((z-ctr)/(r_perp+0.25))**2 + ((y-ctr)/(r_par+0.25))**2 + ((x-ctr)/(r_par+0.25))**2 <= 1:\n                    closing_element[z, y, x] = 1\n\n    closing_elements = []\n    closing_elements.append(cupy.asarray(np.copy(np.transpose(closing_element, (0, 1, 2)))))\n    closing_elements.append(cupy.asarray(np.copy(np.transpose(closing_element, (1, 0, 2)))))\n    closing_elements.append(cupy.asarray(np.copy(np.transpose(closing_element, (1, 2, 0)))))\n\n    max_patch_size = int(params.get('max_patch_size', 64))\n\n    for sheet_idx in range(1, min(num_sheets, params['max_sheets']) + 1):\n        # sheet mask\n        sheet_vol.fill(0)\n        sheet_vol[cupy.where(sheet_labels_vol == sheet_idx)] = 1\n\n        # projection for orientation\n        proj_axis, proj_mask = get_best_projection(sheet_vol)\n        close_elem = closing_elements[proj_axis]\n\n        # oriented closing, keep original too\n        sheet_vol_closed = cupy.maximum(\n            sheet_vol,\n            cupyx.scipy.ndimage.binary_closing(sheet_vol, close_elem).astype(np.uint8)\n        )\n\n        # patches for holes (based on closed sheet)\n        patches, _ = get_hole_patches(sheet_vol_closed, proj_axis, proj_mask, params)\n\n        # apply only \"small\" patches\n        small_patches = [p for p in patches if p[3] * p[4] < max_patch_size**2]\n        sheet_vol_patched = sheet_vol_closed\n        if len(small_patches) > 0:\n            sheet_vol_patched = cupy.asarray(\n                insert_patches_in_volume(small_patches, proj_axis, cupy.asnumpy(sheet_vol_closed))\n            )\n            sheet_vol_patched = cupy.maximum(\n                sheet_vol_patched,\n                cupyx.scipy.ndimage.binary_closing(sheet_vol_patched, close_elem).astype(np.uint8)\n            )\n\n        # plug small holes\n        sheet_vol_fixed, _ = plug_small_holes(sheet_vol_patched, lut, params)\n        sheet_vol_fixed = cupy.asarray(sheet_vol_fixed)\n\n        # \"check holes\": discard patching if it made holes worse\n        num_holes_after = 1 - cupy_euler_number(sheet_vol_fixed, connectivity=1)\n        if num_holes_after > 0:\n            sheet_vol_nopatch, _ = plug_small_holes(sheet_vol_closed, lut, params)\n            sheet_vol_nopatch = cupy.asarray(sheet_vol_nopatch)\n            num_holes_before = 1 - cupy_euler_number(sheet_vol_nopatch, connectivity=1)\n            if num_holes_after > num_holes_before:\n                sheet_vol_fixed = sheet_vol_nopatch\n\n        occupied = cupyx.scipy.ndimage.binary_dilation((out_vol != 0), struct_element333)\n        out_vol[cupy.where((sheet_vol_fixed != 0) & (occupied == 0))] = sheet_idx\n\n    # fill cavities\n    out_vol = cupyx.scipy.ndimage.binary_fill_holes(out_vol).astype(cupy.uint8)\n\n    # zero border again\n    bw = params['border_width']\n    out_vol[0:bw,:,:] = out_vol[:,0:bw,:] = out_vol[:,:,0:bw] = 0\n    out_vol[-bw:,:,:] = out_vol[:,-bw:,:] = out_vol[:,:,-bw:] = 0\n\n    return cupy.asnumpy((out_vol != 0)).astype(np.uint8)\n\n# ============================================================\n# Main pipeline (chunked)\n# ============================================================\ndef main():\n    if any(m not in MODELS for m in ACTIVE_MODES):\n        missing = [m for m in ACTIVE_MODES if m not in MODELS]\n        raise ValueError(f\"ACTIVE_MODES contains undefined modes: {missing}\")\n\n    disk_report(\"RUN START (pre-clean)\")\n    setup_environment()\n\n    if CLEAN_RUN_CACHES:\n        safe_rmtree(CACHE_ROOT)\n        CACHE_ROOT.mkdir(parents=True, exist_ok=True)\n\n    if CLEAN_NNUNET_PREPROCESSED:\n        safe_rmtree(NNUNET_PREPROCESSED)\n\n    if CLEAN_NNUNET_RESULTS:\n        safe_rmtree(NNUNET_RESULTS)\n\n    safe_mkdir_clean(FINAL_TIFF_DIR)\n    safe_mkdir_clean(TEST_INPUT_DIR)\n    for m in ACTIVE_MODES:\n        safe_mkdir_clean(NNUNET_OUT[m])\n\n    setup_environment()\n    disk_report(\"AFTER CLEANUP\")\n\n    print(\"\\n=== Registering ACTIVE_MODES ===\")\n    for m in ACTIVE_MODES:\n        ensure_model_registered(MODELS[m])\n\n    all_tifs = list_all_test_tifs()\n    save_prob_flag = detect_save_prob_flag()\n\n    print(\"\\n=== FUSION CONFIG ===\")\n    print(\"ACTIVE_MODES:\", ACTIVE_MODES)\n    print(\"FUSION_SCHEME:\", FUSION_SCHEME)\n    w_norm = normalize_weights(ACTIVE_MODES, WEIGHTS)\n    print(\"WEIGHTS (normalized over active):\", {k: round(v, 6) for k, v in w_norm.items()})\n    print(\"PROB_T:\", PROB_T)\n    print(\"PROCESS_IN_CHUNKS:\", PROCESS_IN_CHUNKS, \"| CASES_PER_CHUNK:\", CASES_PER_CHUNK)\n\n    if SUB_ZIP.exists():\n        SUB_ZIP.unlink()\n\n    # UPDATED pp_params for closing(7,7)+check holes\n    pp_params = {\n        'small_obj_thresh': 20000,\n        'border_width': 3,\n        'max_sheets': 40,\n        'max_patch_size': 64,\n        'close_par': 7,\n        'close_perp': 7,\n    }\n\n    chunks = [all_tifs]\n    if PROCESS_IN_CHUNKS:\n        chunks = chunked(all_tifs, CASES_PER_CHUNK)\n\n    disk_report(\"BEFORE ALL CHUNKS\")\n\n    with zipfile.ZipFile(SUB_ZIP, \"w\", compression=zipfile.ZIP_DEFLATED) as z:\n        for ci, chunk_tifs in enumerate(chunks, start=1):\n            print(f\"\\n==================== CHUNK {ci}/{len(chunks)} ====================\")\n            print(f\"Chunk cases: {len(chunk_tifs)}\")\n\n            case_ids = prepare_test_data_subset(chunk_tifs)\n\n            for m in ACTIVE_MODES:\n                safe_mkdir_clean(NNUNET_OUT[m])\n\n            disk_report(f\"CHUNK {ci}: BEFORE PREDICT\")\n            run_active_models_in_stages(ACTIVE_MODES, save_prob_flag)\n            disk_report(f\"CHUNK {ci}: AFTER PREDICT\")\n\n            for cid in tqdm(case_ids, desc=f\"Chunk {ci}: Fuse+PP+Zip\"):\n                p = fuse_probability_from_nnunet(cid, ACTIVE_MODES)\n                mask = (p >= PROB_T).astype(np.uint8)\n                pp_mask = postprocess(mask, pp_params)\n\n                tif_path = FINAL_TIFF_DIR / f\"{cid}.tif\"\n                tifffile.imwrite(tif_path, pp_mask)\n\n                if ZIP_TIFF_IMMEDIATELY:\n                    z.write(tif_path, tif_path.name)\n                    tif_path.unlink()\n\n                # delete per-case npz immediately\n                for m in ACTIVE_MODES:\n                    try:\n                        npz = find_case_npz(NNUNET_OUT[m], cid)\n                        npz.unlink()\n                    except Exception:\n                        pass\n\n            for m in ACTIVE_MODES:\n                safe_rmtree(NNUNET_OUT[m])\n            safe_rmtree(TEST_INPUT_DIR)\n\n            disk_report(f\"CHUNK {ci}: AFTER CLEANUP\")\n\n        if not ZIP_TIFF_IMMEDIATELY:\n            tiffs = sorted(FINAL_TIFF_DIR.glob(\"*.tif\"))\n            for tp in tqdm(tiffs, desc=\"Zipping (final pass)\"):\n                z.write(tp, tp.name)\n                tp.unlink()\n\n    n_test = len(all_tifs)\n    with zipfile.ZipFile(SUB_ZIP, \"r\") as z:\n        n_zip = sum(1 for n in z.namelist() if n.endswith(\".tif\"))\n    print(f\"\\nSANITY: test_cases={n_test} | zip_tifs={n_zip}\")\n\n    safe_rmtree(FINAL_TIFF_DIR)\n    safe_rmtree(TEST_INPUT_DIR)\n\n    disk_report(\"RUN END\")\n    print(\"\\nDONE.\")\n    print(\"submission:\", SUB_ZIP, f\"({SUB_ZIP.stat().st_size/1024/1024:.1f} MB)\")\n\nmain()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-06T09:43:03.396778Z","iopub.execute_input":"2026-02-06T09:43:03.397245Z","iopub.status.idle":"2026-02-06T09:46:53.371535Z","shell.execute_reply.started":"2026-02-06T09:43:03.397221Z","shell.execute_reply":"2026-02-06T09:46:53.370806Z"}},"outputs":[],"execution_count":null}]}