Skip to content

ARCHITECTURE.md

Latest commit

 

History

History
128 lines (102 loc) · 8.63 KB

ARCHITECTURE.md

File metadata and controls

128 lines (102 loc) · 8.63 KB
Raw

Architecture

This document describes the two-pipeline structure of Vesuvius and how a Kaggle submission flows from raw CT volumes to a submission.zip.

Overview

                     ┌───────────────────────────────┐
                     │  data/                          │
                     │  ├── download.py                │   Kaggle CLI wrapper with retry/resume
                     │  ├── download_vesuvius.py       │   Vesuvius competition data download
                     │  └── download_model.py         │   kagglehub model download (nnUNet ckpt)
                     └──────────────┬──────────────────┘
                                    │ produces DATA_ROOT/
                                    │   ├── train_images/*.tif
                                    │   ├── train_labels/*.tif
                                    │   ├── test_images/*.tif
                                    │   └── test.csv
                                    ▼
       ┌────────────────────────────┴────────────────────────────┐
       │                                                       │
       ▼                                                       ▼
┌──────────────────────────────┐              ┌──────────────────────────────────┐
│  Pipeline A: nnunet/         │              │  Pipeline B: segm3d/              │
│  nnUNetv2 CLI wrapper        │              │  MONAI + PyTorch Lightning        │
│                              │              │                                   │
│  data_process.py             │              │  train.py       (resize to cube)   │
│   → nnUNetv2_plan_and_       │              │  train_v1.py    (random patch)    │
│     preprocess               │              │  infer.py       (resize back)     │
│  train.py                    │              │  infer_v1.py    (sliding window)  │
│   → nnUNetv2_train (DDP)     │              │                                   │
│  infer.py                    │              │  Networks: SegResNet, SwinUNETR   │
│   → nnUNetv2_predict +       │              │  Loss: DiceCE / Tversky           │
│     binarize + zip           │              │  (ignore=2 skipped)               │
└──────────────┬───────────────┘              └────────────────┬─────────────────┘
               │                                                │
               │  optional predictions-style postprocess:        │
               │  3D hysteresis + anisotropic closing +          │
               │  small-CC removal (built into both infer_*)     │
               └────────────────────┬───────────────────────────┘
                                    ▼
                            ┌──────────────────┐
                            │  submission.zip   │
                            │  (1-bit TIFF       │
                            │   per test id)     │
                            └──────────────────┘

Pipeline A — nnunet/ (nnUNetv2 CLI wrapper)

nnunet/ is a thin Python wrapper around the nnUNetv2_* CLIs. It does not reimplement segmentation; it prepares inputs, launches training/inference as subprocesses, and converts nnUNet output back into the binary TIFFs the competition expects.

  • nnunet/data_process.py — Step 1/2. Converts DATA_ROOT/train_images + train_labels into nnUNet_raw/Dataset100_Vesuvius/ (TIFF + JSON spacing sidecar), then invokes nnUNetv2_plan_and_preprocess to produce nnUNet_preprocessed/.
  • nnunet/train.py — Step 2/2. Sets the nnUNet_raw / nnUNet_preprocessed / nnUNet_results env vars and invokes nnUNetv2_train for the requested fold(s). Handles the DDP constraint (global batch_size >= num_gpus) and offers a parallel-folds mode to fill 8 GPUs when a single DDP run is capped at 2.
  • nnunet/infer.py — Renames test_images/*.tif to the *_0000.tif naming nnUNet expects, writes spacing sidecars, invokes nnUNetv2_predict, then converts the predicted .npz/.tif/.nii.gz to 1-bit TIFFs and packs submission.zip. Optional 3D hysteresis + closing + small-CC removal postprocess.

Pipeline B — segm3d/ (MONAI + PyTorch Lightning)

segm3d/ is a self-contained MONAI + Lightning implementation. The network is built and trained in-process (no external CLI). Two generations live side-by-side:

  • segm3d/train.py (base) — Resizes each volume to --model-input-size (default 160³), trains a SegResNet / SwinUNETR with DiceCELoss (ignore=2 excluded), checkpoints on val_dice.
  • segm3d/infer.py (base) — Same resize-back path: predicts on the resized cube and resamples the mask to the original spacing. Faster, but resize-then-resample introduces topological artifacts at the surface boundary.
  • segm3d/train_v1.py — Random-patch training at --patch-size 192³ with pos/neg guided sampling (foreground-center probability --pos-fraction 0.7). EMA of weights, monitored on val_fg_dice (foreground class=1 only). Strict ignore mask in both loss and metrics.
  • segm3d/infer_v1.py — Native-resolution sliding_window_inference with --overlap 0.5~0.7 and optional light TTA (noise / scale jitter, no flips). Avoids the resize-then-resample topological hard-failure of base infer.py.

When to use which

Concern nnunet/ segm3d/ base segm3d/ v1
Reproducible SOTA baseline ✅ — —
Topology-correct surface mask ✅ (with postprocess) — ✅
Easy to extend (new loss, TTA, sampler) — (CLI) ✅ ✅
Multi-GPU on 8 cards parallel folds DDP / ddp_spawn DDP / ddp_spawn
Lowest effort to retrain from scratch medium (need plan_and_preprocess) low medium

Data flow: from raw TIFFs to a Kaggle submission

Step 1 — Acquire data

python data/download_vesuvius.py --out ./kaggle_data/vesuvius_surface_detection

Step 2a — Train with nnUNet (Pipeline A)

python nnunet/data_process.py --input-dir DATA_ROOT --work-dir ./work --dataset-id 100 \
    --configuration 3d_fullres --planner nnUNetPlannerResEncM
python nnunet/train.py --work-dir ./work --dataset-id 100 \
    --configuration 3d_fullres --plans-name nnUNetResEncUNetMPlans \
    --fold all --epochs 250 --num-gpus 1

Step 2b — Train with 3D-SegM (Pipeline B, v1)

python segm3d/train_v1.py --train-images-dir DATA_ROOT/train_images \
    --train-labels-dir DATA_ROOT/train_labels --output-dir ./work_3d_segm_v1 \
    --arch segresnet --patch-size 192 192 192 --samples-per-volume 16 \
    --pos-fraction 0.7 --batch-size 1 --max-epochs 20 --ema 1 --ema-decay 0.999

Step 3 — Infer + pack submission.zip

# Pipeline A
python nnunet/infer.py --root-dir DATA_ROOT --work-dir ./work_infer \
    --dataset-id 100 --configuration 3d_fullres --plans-name nnUNetResEncUNetMPlans \
    --trainer nnUNetTrainer_250epochs --fold all --checkpoint checkpoint_final.pth \
    --save-probabilities 1 --postprocess 1

# Pipeline B v1
python segm3d/infer_v1.py --root-dir DATA_ROOT --checkpoint-dir ./work_3d_segm_v1 \
    --work-dir ./work_3d_segm_v1_infer --roi-size 192 192 192 --overlap 0.6 \
    --sw-batch-size 1 --tta 1 --postprocess 1

Both produce work-*/submission.zip containing one 1-bit TIFF per id in test.csv.

Postprocessing (shared by both infer scripts)

--postprocess 1 enables the "predictions-style" 3D postprocess originally prototyped in notebooks/vesuvius-predictions-1feb.ipynb:

  1. 3D hysteresis thresholding — low threshold --T-low 0.30, high threshold --T-high 0.80 (voxels above high are kept; voxels above low are kept only if 6-connected to a high one).
  2. Anisotropic closing — structure element extended along z (--z-radius 3) vs xy (--xy-radius 2) to match the CT volume's anisotropic spacing.
  3. Small-CC removal — --dust-min-size 100 drops tiny isolated components likely to be false positives.

Requires scipy + scikit-image (already in requirements.txt).