Cascaded V-Net using ROI masks for brain tumor segmentation

In plain words

Brain tumours are a small part of a scan, and most voxels are healthy tissue or background. A network trained on whole scans can learn to answer "no tumour" almost everywhere and still score well, so the authors restrict training to the region that matters.

They split the job in two. A first network looks at the whole brain and marks where the tumour is, as one blob. The area around that blob is cut out as a box. A second network looks only inside that box and separates the tumour's parts: swelling, bright enhancing tissue and the dead or non-enhancing core. Both networks are 3D versions of a known segmentation design, with small changes to how layers are connected.

The paper's own results show the trade-off. The whole tumour is found well, but the smaller parts are often missed, which the authors list as future work.

TL;DR

• Two 3D networks, each a modified V-Net. Network 1 segments the whole tumour; network 2 segments its sub-regions inside a box around network 1's output §2.
• ROI masks zero out voxels outside the brain (training) or outside the box (stage 2), and block gradients from them §2.1.
• Author-reported BraTS 2017 validation Dice: 0.877 whole tumour, 0.714 enhancing tumour, 0.637 tumour core tab. 1.
• Many settings are not stated: optimiser, learning rate, epochs, threshold and morphological operator, and how the second network's outputs form the three regions.
• No official code was found.

Problem

Tumours cover only a few percent of the image and are split among several sub-regions of very different size §2. Training on whole images or on random patches is then biased towards predicting background. The authors want to train densely on whole volumes but only on voxels near the tumour §1 §2.

Method

The pipeline has two networks fig. 1. Both are 3D networks built from the same V-Net-style encoder-decoder §2.1 fig. 3. They are trained independently §2.2. At inference, the second network reads the first network's output §2.3.

Changes to V-Net §2.1 fig. 2 fig. 3.

• Convolutions use 3x3x3 kernels instead of 5x5x5, which reduces the number of parameters §2.1.
• Activation is ReLU instead of PReLU §2.1.
• Batch normalisation is applied before each non-linearity §2.1.
• Short residual connections are identity mappings, following pre-activation residual units §2.1 fig. 2.
• Where the residual branch and the shortcut differ in spatial size or channel count, the shortcut uses max-pooling and repeated up-sampling for size and a 1x1x1 convolution for channels §2.1 fig. 2.
• The figure gives the widths: encoder levels with 8, 16, 32, 64 and 128 channels, each entered by a 2x2x2 stride-2 convolution, and decoder levels with 32, 16, 8 channels; the output is a 2-channel softmax fig. 3. The text does not state the number of levels, and the figure is the only source for widths.

ROI masks §2.1. An ROI mask fixes voxels outside it to background. Those voxels are multiplied by 0 and their background probability is set to 1, so they add no loss and pass no gradient. The mask is applied before the final prediction in both training and inference §2.1. Network 1 uses the brain mask. Network 2 uses a rectangular box §2.2.

Network 1, whole tumour §2.2. Input is the four modalities T1, T1c, T2 and FLAIR §2.2 §3.1. FLAIR is concatenated with the feature maps of the last decoder level, so deeper layers see the FLAIR intensities §2.2 fig. 3. It is a binary tumour/non-tumour problem, trained with the soft Dice score in eq. 1 §2.2, computed over brain voxels only.

Network 2, sub-regions §2.2. Input is the same four modalities. The training box is the smallest rectangle covering the ground-truth tumour, which limits training to the tumour's vicinity §2.2. There are four classes: non-tumour, edema, enhancing core and non-enhancing core §2.2. The loss is cross-entropy plus 0.5 times the sum of three Dice terms, for whole tumour, enhancing tumour and tumour core eq. 2 §2.2. The 0.5 weight was chosen empirically §2.2.

Inference §2.3. Network 1 gives a whole-tumour prediction. Morphological filtering removes small spurious detections. The smallest rectangle covering what remains becomes the ROI mask for network 2, which reduces false positives outside the tumour. The two stages must run in sequence at inference, although the mask step can be decoupled from the network computation §2.3.

Training regime §2.2 §3.1. Dense training with one subject per batch, with no patch sampling and no class-balancing sampler §2.2. Each scan is normalised by its own mean and standard deviation. Augmentation is reflection across the sagittal plane §3.1.

Pseudocode (the paper gives the components, not an algorithm box; the structure below is ours, built from the cited sections):

# Two-stage inference  [§2.3] [fig. 1]
# Input:  four volumes (T1, T1c, T2, FLAIR), brain mask
# Output: label map with the three tumour sub-regions

def segment_volume(volumes, brain_mask):
    # Stage 1: whole-tumour probability inside the brain  [§2.2]
    wt_prob = network_1(volumes, brain_mask)     # (D, H, W)
    # Binarise and remove small spurious detections  [§2.3]
    # Operator and threshold are not stated in the paper.
    wt_mask = morphological_filter(binarise(wt_prob))
    # Smallest rectangular box covering the detection  [§2.3]
    roi = smallest_box_covering(wt_mask)
    # Stage 2: sub-region probabilities, one channel per class  [§2.2]
    probs = network_2(volumes, roi)              # (C, D, H, W)
    # Outside the ROI, force background with full probability  [§2.1]
    probs = where(roi, probs, background_one_hot)
    return argmax(probs, axis=0)

# Stage 1 loss: soft Dice over brain voxels  [§2.2] [eq. 1]
# tumour_prob: softmax tumour probability (N,); labels: 0 or 1 (N,)

def loss_stage_1(tumour_prob, labels, brain_mask):
    # Voxels outside the brain carry no gradient  [§2.1]
    p = tumour_prob * brain_mask
    l = labels * brain_mask
    overlap = sum(p * l)
    # No smoothing constant is stated in [eq. 1]
    return overlap / (sum(p) + sum(l))

# Stage 2 loss as written in the paper  [§2.2] [eq. 2]
# The paper does not state how the four softmax channels are merged
# into the three regions; that step is left to the implementer.

def loss_stage_2(probs, labels, roi_mask):
    # Cross-entropy over the ROI  [§2.2] [eq. 2]
    ce = cross_entropy(probs, labels, mask=roi_mask)
    # Dice per region: whole tumour, enhancing tumour, tumour core  [§2.2]
    d_wt = dice_wt(probs, labels, roi_mask)
    d_et = dice_et(probs, labels, roi_mask)
    d_tc = dice_tc(probs, labels, roi_mask)
    return ce + 0.5 * (d_wt + d_et + d_tc)

Key equations:

L_1 = \frac{\sum_{i=1}^{N} p_i \, l_i}{\sum_{i=1}^{N} p_i + \sum_{i=1}^{N} l_i} eq. 1

where N is the number of voxels, p_i the softmax output for voxel i and l_i ∈ {0, 1} its label §2.2.

L_2 = X_E + 0.5\,(D_{WT} + D_{ET} + D_{TC}) eq. 2

where X_E is cross-entropy and the D terms are Dice scores for the three regions §2.2. The paper does not define the D terms by formula.

Hyperparameters

Results

Author-reported. Setup: BraTS 2017. Development is a 30% hold-out of the training set; validation is the BraTS 2017 validation set (46 scans), scored on the challenge leaderboard §3.2. The paper compares with other participants only in text; it gives no numbers for them §3.2. Each row is one set, each column one region.

The Hausdorff unit is not stated. The authors say the whole-tumour Dice is close to the top challenge entries, and that the Hausdorff distances are above the best methods, most of all for the whole tumour §3.2. They also say enhancing and core regions are underrepresented in predictions, judged from sensitivity §3.2 tab. 2, and that over-prediction of edema reduces the smaller classes in a failure case §3.3 fig. 5. Training curves show the tumour core Dice starts to overfit from epoch 18 §3.2 fig. 4.

Paper vs code

No official code. The paper has no code link, and its only URL is the BraTS leaderboard footnote §3.2. The arXiv listing has no code link. We searched GitHub for "cascaded vnet brats", "casamitjana brain tumor" and "brats2017 cascade" (no results), and the author profile of Adrià Casamitjana, which lists 24 repositories, none about this method. The imatge-upc organisation has one BraTS repository, mri-braintumor-segmentation. Its README is for a 2020 V-Net and 3D-UNet uncertainty method and cites arXiv 2012.15294. It is not the code for this paper, so it is not used.

Pitfalls

• Label mapping conflict. The data section gives ET as label 4, but the evaluation section maps ET to label 1 after merging §3.1 §3.2. The paper does not say whether training uses the merged or the original labels.
• Loss sign. eq. 1 is a Dice overlap, which rises with agreement, but it is called a loss. Whether to minimise it or its negative is not stated. eq. 2 adds Dice terms to cross-entropy with a positive weight, with the same ambiguity.
• Missing optimiser and schedule. No optimiser, learning rate, decay or number of epochs is given. The only training length is the mention of epoch 18 at iteration 3600 §3.2. That implies about 200 iterations per epoch, so one iteration per training subject, if the 70% split is about 200 subjects (derived by us, not stated).
• Parallel or sequential. Method §2 says the two tasks run "in parallel", while inference needs the first network's output before the second runs §2.3. Training is independent; inference is sequential.
• Small regions. The authors state that sub-regions are underrepresented, and suggest up-weighting them as future work §4. The loss in eq. 2 does not weight classes.
• Binarisation and morphology. Stage 1 needs a threshold and a morphological operator to make the box §2.3. Neither is stated. A box built from a noisy stage-1 output can cut off tumour, and the paper gives no measure of that.
• Dice implementation. eq. 1 has no smoothing constant. An empty mask gives a zero denominator. The per-region definitions behind D_WT, D_ET and D_TC are not formulas in the paper.
• Dense training cost. Full volumes per batch and a second network are heavier than patch training, and the authors give no timing or memory figures. §4 mentions a cost trade-off but no numbers.

Lineage

• Builds on: V-Net, which uses 5x5x5 kernels and PReLU §2.1; pre-activation residual units §2.1.
• Compared against: DeepMedic-style 3D CNNs and other BraTS entries, named in the introduction §1. No numbers are given.
• Superseded by: not known from content already in this wiki.