Anatomically Constrained Neural Networks (ACNN): Application to Cardiac Image Enhancement and Segmentation

In plain words

Image networks that label or sharpen medical scans usually judge each pixel on its own. When the scan is messy, for example because the patient moved between slices or parts of the heart are missing, the network can output shapes that no real heart would have.

This paper adds a shape check that is only used during training. A separate network learns a compact summary of what heart shapes look like. While the main network trains, its output is compressed with that summary and compared with the compressed form of the true answer. If the two summaries disagree, the main network is penalised, so it learns to produce outputs that look like plausible hearts.

Because the check is dropped once training ends, the final model runs as fast as the original. The authors apply the idea to cardiac MRI segmentation, cardiac ultrasound segmentation, and MRI super-resolution, and report better accuracy than the baselines they compare against.

TL;DR

• Adds a training-only loss term that compares low-dimensional codes of the predicted and true anatomy, learned with a convolutional autoencoder §2.2.
• For super-resolution, a predictor network maps intensity images to the same kind of code, trained in two stages with the autoencoder §2.4.
• Improves cardiac MR segmentation, cardiac 3D ultrasound segmentation and MR super-resolution over the paper's baselines (author-reported) tab. 1 tab. 2 tab. 3.
• Uses only a 64-d code and no extra inference cost; the regulariser is not used at test time §2.4 [supp. Implementation Details].
• Weak points: several training details (optimiser, schedule, layer tables of the main networks, loss reductions) are not stated; see Pitfalls.

Problem

Segmentation and super-resolution CNNs are trained with pixel-level losses such as cross-entropy or the smooth ℓ1 loss. These losses do not capture the structure or dependencies of the output, so the predicted shape is not constrained to be anatomically plausible §1. The authors argue this matters most for cardiac imaging, where 2D stacks acquired in separate breath-holds show motion artefacts and slice misalignment, and where the stacks are often too thick along the slice direction for accurate analysis §1.1 fig. 1.

Earlier shape-prior approaches include PCA-based and active shape models, deep Boltzmann-machine shape models, and cascaded networks. These either need explicit shape specification, add model complexity, or need post-processing such as conditional random fields with heuristic tuning §1 §1.2.

Method

Data and preprocessing §3.1 §3.3

• UK Digital Heart cohort: 1200 pairs of cine 2D short-axis stacks (low resolution, LR) and cine 3D high-resolution (HR) cardiac MR volumes, from healthy subjects. Voxel sizes are 1.25 × 1.25 × 10.00 mm (LR) and 1.25 × 1.25 × 2.00 mm (HR) §3.1.1.
• HR annotations (endocardium and myocardium) were made by correcting multi-atlas segmentations §3.1.1.
• Intensities are linearly normalised. Images are cropped around six automatically detected landmarks §3.3.
• Segmentation split: 900 training, 100 validation, 200 test §3.3. Super-resolution split: 1000 training, 200 test, disjoint §3.5.
• CETUS'14 3D ultrasound: 3D+time sequences from 15 healthy subjects and 30 patients, resampled to 1 mm isotropic. Testing uses 30 of 45 sequences, evaluated blindly on the challenge platform §3.1.2 §3.4.
• ACDC classification: 20 cine stacks per class, 60 sequences used in the classification experiment §3.1.3 §3.6.

Step 1. Baseline segmentation network §2.1 fig. 2

The network estimates P(y | x) by producing one feature map per class and applying softmax. It is trained with the average cross-entropy over classes using stochastic gradient descent §2.1. It has a VGG-like analysis path and a 3D U-Net-like reconstruction path. Features are computed at several image scales to enlarge the receptive field §2.1 fig. 2. The authors train the up-sampling layers against HR labels, so the output is on a sub-pixel grid and no separate up-sampling step is needed §2.1. The paper gives only a block diagram for this backbone, not a layer table.

Step 2. Baseline super-resolution network §2.3 fig. 2

The SR model uses the same architecture. Features are computed on the LR grid rather than up-sampled early, which the authors argue avoids redundant computation in HR space. Multi-scale features come from strides in the in-plane direction, which gives a larger receptive field than earlier patch-based SR models §2.3. The reconstruction loss is the smooth ℓ1 (Huber) loss:

\Psi_{\ell_1}(k) = \begin{cases} 0.5\,k^2 & |k| < 1 \\ |k| - 0.5 & \text{otherwise} \end{cases} §2.3

Step 3. Convolutional autoencoder (shape prior for segmentation) §2.2 tab. 5

A stacked convolutional autoencoder is trained on label maps. The input is a one-hot multi-label map with 3 input channels, the bottleneck is a 64-d code, and the output has 3 channels tab. 5. The encoder ends with a convolution with one output channel, then a fully connected layer to 64 units. The decoder starts from a fully connected layer to 125 units and uses transposed convolutions tab. 5. No non-linearity is applied to the 64-d code [supp. Network Structure].

Step 4. ACNN-Seg objective §2.2 eq. 1

The autoencoder encoder f is applied to the predicted label maps φ(x) and to the ground-truth labels y. The shape term is the squared distance between the two codes. Only the segmentation network's parameters θ_s are updated §2.2.

L_{h_e} = \left\lVert f(\phi(\boldsymbol{x};\boldsymbol{\theta}_s);\boldsymbol{\theta}_f) - f(\boldsymbol{y};\boldsymbol{\theta}_f) \right\rVert_2^2

\min_{\boldsymbol{\theta}_s}\; \Big( L_x(\phi(\boldsymbol{x};\boldsymbol{\theta}_s), \boldsymbol{y}) + \lambda_1 \cdot L_{h_e} + \frac{\lambda_2}{2} \lVert \boldsymbol{w} \rVert_2^2 \Big) eq. 1

The weight-decay term applies to convolution filter weights w. λ1 weights the shape term and λ2 the weight decay §2.2.

# ACNN-Seg training objective  [§2.2] [eq. 1]
# Input:  images (B, D, H, W) at LR; labels (B, D, H, W) on the HR grid
# Output: scalar loss; only segmentation weights receive gradients

def acnn_seg_loss(images, labels, seg_net, ae_encoder, lambda_1, lambda_2):
    # Class probability maps on the sub-pixel grid  [§2.1]
    predictions = seg_net(images)
    # Pixel-wise cross-entropy against the labels  [§2.1] [eq. 1]
    # (mean or sum over voxels not stated)
    ce_loss = cross_entropy(predictions, labels)
    # Low-dimensional codes of predicted and true label maps  [§2.2] [eq. 1]
    codes_pred = ae_encoder(predictions)              # (B, 64)
    codes_true = ae_encoder(labels)                   # (B, 64)
    # Squared Euclidean distance between the codes  [§2.2] [eq. 1]
    shape_loss = sum_of_squares(codes_pred - codes_true)
    # Weight decay on convolution filter weights  [§2.2] [eq. 1]
    decay = sum_of_squares(conv_filter_weights(seg_net))
    return ce_loss + lambda_1 * shape_loss + (lambda_2 / 2) * decay

Step 5. T-L network (predictor for intensity images) §2.4 tab. 4 fig. 3

For super-resolution the output is an intensity image, which the segmentation autoencoder cannot encode. The authors therefore add a predictor p that maps an intensity image to the same 64-d code space, forming a "T-L" network with the autoencoder §2.4 fig. 3. The predictor is a stack of 3D convolutions with ReLU, followed by a fully connected layer to 64 units tab. 4.

Training has two stages §2.4:

• Stage 1: train the autoencoder alone on ground-truth masks with cross-entropy L_x. Then train the predictor to match the autoencoder's codes, minimising the Euclidean distance L_h.
• Stage 2: train the two jointly. The encoder f receives two gradients, one from L_x and one from L_h. The two losses are scaled to similar ranges. The authors report that a single joint update converged better than alternating updates §2.4 §3.2.

The T-L network is used only during training, not at inference §2.4.

Step 6. ACNN-SR objective §2.4 eq. 2

The predictor codes of the synthesised HR image and of the true HR image are compared with a squared Euclidean distance. Only the SR network parameters θ_r are optimised in this equation §2.4.

L_{h_p} = \left\lVert p(\Phi(\boldsymbol{x};\boldsymbol{\theta}_r);\boldsymbol{\theta}_p) - p(\boldsymbol{y}_r;\boldsymbol{\theta}_p) \right\rVert_2^2

\min_{\boldsymbol{\theta}_r}\; \Big( \Psi_{\ell_1}\big(\Phi(\boldsymbol{x};\boldsymbol{\theta}_r) - \boldsymbol{y}_r\big) + \lambda_1 \cdot L_{h_p} + \frac{\lambda_2}{2} \lVert \boldsymbol{w} \rVert_2^2 \Big) eq. 2

# ACNN-SR training objective  [§2.3] [§2.4] [eq. 2]
# Input:  LR stacks (B, D_lr, H, W); HR targets (B, D, H, W)
# Output: scalar loss; the predictor is not used at inference

def acnn_sr_loss(lr_images, hr_targets, sr_net, predictor, lambda_1, lambda_2):
    # Super-resolved volume from the LR-grid features  [§2.3]
    sr_images = sr_net(lr_images)
    # Smooth L1 (Huber) reconstruction term, threshold 1  [§2.3] [eq. 2]
    recon_loss = smooth_l1(sr_images - hr_targets, threshold=1.0)
    # Codes of the synthesised and true HR images  [§2.4] [eq. 2]
    # (whether predictor weights also update here is not stated)
    codes_sr = predictor(sr_images)                   # (B, 64)
    codes_hr = predictor(hr_targets)                  # (B, 64)
    # Squared Euclidean distance between the codes  [§2.4] [eq. 2]
    shape_loss = sum_of_squares(codes_sr - codes_hr)
    # Weight decay on convolution filter weights  [§2.4] [eq. 2]
    decay = sum_of_squares(conv_filter_weights(sr_net))
    return recon_loss + lambda_1 * shape_loss + (lambda_2 / 2) * decay

Step 7. Predictor and autoencoder architectures tab. 4 tab. 5

Predictor (input 120 × 120 × 60 HR intensity volume) tab. 4:

Autoencoder (input and output 120 × 120 × 60 label map) tab. 5:

Batch normalisation follows every convolution except the last one. Residual connections are not used [supp. Network Structure].

Step 8. Training recipe §3.2 [supp. Implementation Details]

• Affine augmentation is applied to inputs for the segmentation and T-L models §3.2.
• For super-resolution, respiratory motion is simulated by in-plane rigid transforms applied independently to each slice. The HR targets are not transformed §3.2.
• Additive Gaussian noise is applied to the input intensity images §3.2.
• For the autoencoder, tissue labels are randomly swapped with probability 0.1 §3.2.
• Training uses full images, not patches, with mini-batches of 8 [supp. Implementation Details].
• Super-resolution training LR data are synthesised from HR by Gaussian blurring along the through-plane direction (σ = 4.0 mm) followed by decimation by K = 5 [supp. Implementation Details] §3.5.
• Dice loss was tested for segmentation and gave worse results than cross-entropy [supp. Implementation Details].

Hyperparameters

Results

Author-reported. The setups are described in each row's source.

Cardiac MR segmentation (200 stacks; LV endocardium and myocardium; Dice in fractions; distances in mm; trainable parameters) tab. 1 §3.3

The paper's text says Wilcoxon signed-rank tests compare the top two methods on each metric §3.3. For Hausdorff distance, the ACNN-Seg difference from the best competitor is not significant in this table.

Cardiac 3D ultrasound segmentation (CETUS'14, 30 blinded test sequences; LV cavity) tab. 2 §3.4

BEAS is the CETUS'14 challenge winner [barbosa2013fast]; FCN is [chen2016iterative]. The text describes the Dice gain as not significant, although the table gives p ≈ 0.05 §3.4.

Cardiac MR super-resolution (SR-CNN comparison, 200 test volumes; SSIM and mean opinion score from two cardiologists on a 1–5 scale; inference per 120 × 120 × 12 LR input on a GTX-1080) tab. 3 §3.5

Pathology classification from latent codes (10-fold cross-validation on 60 cardiac MR sequences; ensemble of decision trees) §3.6

Paper vs code

No official code.

Searched: the LaTeX source contains no code or project-page link; its only URLs are the UK Digital Heart, CETUS and ACDC challenge pages §3.1. GitHub repository searches for "anatomically constrained neural networks", "acnn" and "ACNN anatomical shape prior" returned no matching repository. Searches of the first author's GitHub user (ozan-oktay) and of the Imperial BioMedIA organisation (biomedia-mira) for the title terms also returned no match.

Pitfalls

• Autoencoder for segmentation is under-specified. §2.2 gives only the loss. Eq. 1 updates only θ_s, which implies a frozen encoder, but the paper does not say how or for how long the segmentation autoencoder is trained. The two-stage schedule in §2.4 is described for the T-L (super-resolution) model only.
• Input to the encoder is not stated. Eq. 1 applies f to φ(x). The paper does not say whether this is softmax probabilities, logits or argmax labels. Gradients reach θ_s only through a differentiable input, so argmax labels would block them. This is an inference from the equation, not a stated detail.
• Reductions and scale of the loss terms are not stated. Cross-entropy mean vs sum over voxels, whether the squared code distance is summed or averaged over the 64 dimensions, and the intensity range all matter. λ1 = 0.01 was "chosen experimentally" [supp. Implementation Details] and is only meaningful for the authors' normalisation.
• SR objective notation changes. §2.3 writes the SR objective as a sum over voxels or samples. Eq. 2 has one Ψ term. Sum vs mean is not stated. The Huber threshold of 1 depends on the intensity scale, which the paper does not give.
• Joint T-L update is under-specified. The two gradients into the encoder are "scaled to match their range", but the scale factors and the combination rule are not stated. The paper says a joint update converged better, without details.
• Predictor training in the SR stage is ambiguous. Eq. 2 optimises only θ_r. Whether the predictor's weights θ_p are also updated during SR training is not stated.
• Main segmentation and SR backbones have no layer table. Only a block diagram exists fig. 2. The number of scales, filter counts and the through-plane up-sampling layers are not given. The parameter counts in tab. 1 (1.60 × 10⁶ for 3D-Seg and ACNN-Seg) cannot be reproduced from this card.
• Padding and sizes are not stated for the T-L networks. The predictor tab. 4 has three stride-2 layers, which divide each dimension by 8. The depth of 60 slices gives 7.5, so the paper's padding or rounding rule must be chosen by the implementer, and it is not stated. The autoencoder encoder ends at 5 × 5 × 5 × 1 = 125 units, which matches the FC layer of 125 units only if reshaped to 5 × 5 × 5. That reshape is an inference, not a stated detail. The decoder output must be cropped or padded back to 120 × 120 × 60, and the rule is not stated.
• Capacity matching is not as described. §3.3 says trainable parameters were "kept fixed to avoid any bias", but tab. 1 lists 1.39 × 10⁶ to 1.91 × 10⁶ across models.
• Significance claims are mixed. The Hausdorff gain in tab. 1 is not significant (p ≈ 0.890 and p ≈ 0.071). The US Dice gain has p ≈ 0.05 in tab. 2, which the text calls not significant. Read the claims in §3.3 and §3.4 against the p-values.
• Speed-up wording. §3.5 says ACNN-SR is faster "by a factor of 5". tab. 3 gives 0.29 s vs 0.06 s, which is about 4.8×. Also, the 3D-UNet row runs at 0.07 s.
• ACDC class set is inconsistent. §3.1.3 names infarction, dilated and hypertrophic cardiomyopathy. §3.6 describes healthy vs dilated and hypertrophic cardiomyopathy, with 60 sequences. The paper does not state which classes were used. The decision-tree ensemble (type, number of trees, depth) is not stated, and the spatial-normalisation step for the PCA baseline is mentioned but not described.
• Training details missing. Optimiser variant, momentum, schedule, epochs, early stopping, Gaussian input-noise level, affine augmentation ranges and motion magnitudes are all not stated.
• Numbering of figures and tables. Figure numbers (1 to 9) and table numbers (1 to 5) in this card follow the order in the LaTeX source. The compiled IEEE PDF may label them differently, and the supplement's numbering may follow the same rule. Check against the PDF before quoting.
• Sub-pixel output. The segmentation model is said to reach sub-pixel accuracy by training up-sampling layers against HR labels §2.1, but the text does not describe the through-plane up-sampling operation beyond the block diagram.

Lineage

• Builds on: Oktay et al. 2016, the multi-input cardiac SR CNN the paper extends and compares against [oktay2016multi] §2.3. The T-L architecture comes from Girdhar et al. 2016 [girdhar2016learning] §2.4. The stacked convolutional autoencoder follows Masci et al. 2011 [masci2011stacked] §2.2.
• Compared against: 2D-FCN (Tran 2016) [tran2016fully]; 3D-UNet (Çiçek et al. 2016) [cicek20163d]; cascaded AE-Seg (Ravishankar et al. 2017) [ravishankarlearning]; FCN for ultrasound (Chen et al. 2016) [chen2016iterative]; BEAS (Barbosa et al. 2013) [barbosa2013fast]; SR-CNN [oktay2016multi].
• Superseded by: not stated in this card's sources.