Medical Image Synthesis with Context-Aware Generative Adversarial Networks

In plain words

Radiotherapy planning needs CT scans, because CT intensities tell the planning software how tissue absorbs radiation. But CT exposes the patient to radiation, while MRI does not. If you could compute a convincing CT from an MRI of the same patient, you could skip one scan. The hard part is that the two look very different: MRI shows rich soft-tissue texture, CT mostly shows density.

The paper trains a 3D convolutional network that takes a small cube of MRI and outputs the matching, smaller cube of CT. Training it only to minimise squared error tends to give blurry output, so two extra pressures are added. A second network, the discriminator, tries to tell real CT cubes from generated ones, and the generator is rewarded for fooling it. A gradient-matching loss also asks the edges in the output to be as strong as the edges in the real CT.

Because each cube only sees a small neighbourhood, the authors stack several such networks. The first predicts a whole CT volume; the second gets both the MRI and that first CT guess as input and produces a better one, and so on. Two rounds of this "auto-context" refinement were used.

On a 16-subject brain set and a 22-subject pelvis set, the authors report lower error and higher PSNR than an atlas method, a sparse-coding method and a structured random forest.

TL;DR

A 3D FCN generator maps 32³ MRI patches to 16³ CT patches and is trained as a GAN generator with loss = adversarial BCE + L2 + 3D gradient difference loss (weights 0.5, 1, 1). Several GANs are chained with the Auto-Context Model (ACM): each stage gets the MRI patch plus the previous stage's CT estimate as a second channel. Full volumes are predicted from overlapping patches whose outputs are averaged. Relevant if you need paired MR→CT (or other paired 3D modality) synthesis; the evaluation is small (16 and 22 subjects, leave-one-out), and the released 3D code differs from the paper in architecture and does not implement ACM.

Problem

CT is needed for dose calculation in radiotherapy and for PET attenuation correction, but it irradiates the patient; MRI is safe yet cannot be used directly for those purposes §1. Earlier MR→CT approaches segment MRI into tissue classes (dependent on segmentation accuracy and manual work), register an atlas (dependent on registration accuracy), or learn a regression on hand-crafted features (bounded by feature quality) §1. 2D CNNs process volumes slice by slice and can give results that are discontinuous across slices, and an L2 loss tends to produce blurry outputs §1 §2. The method assumes paired, voxel-aligned MR and CT of the same subject for training §3.

Method

Generator. A 3D fully convolutional network with no pooling, using batch normalization and ReLU §2.1. The text describes 8 stages of conv + BN + ReLU with 32, 32, 32, 64, 64, 64, 32, 32 filters, followed by a last convolutional layer whose output is the CT estimate §2.1. The architecture figure shows kernel sizes 9, 3, 3, 3, 9, 3, 3, 7 for the conv + BN + ReLU layers and 3 for the final plain conv, followed by a Euclidean loss block fig. 2. Input is a 32×32×32 MRI patch and output a 16×16×16 CT patch with the same centre §4. The paper does not state the padding; the code reaches 32→16 by using VALID padding on the two 9³ convolutions and SAME on all others g_model.py.

Discriminator. A CNN with three stages of conv + BN + ReLU + max pooling, then a convolution and three fully connected layers (ReLU, ReLU, sigmoid) §2.1. Conv filters are 5×5×5 with 32, 64, 128 and 256 filters; FC layers have 512, 128 and 1 output nodes §2.1. Its input is a real or generated CT patch fig. 2.

Losses. The standard GAN objectives are eq. 1 eq. 2:

\max_{D} \; \log(D(x)) + \log(1 - D(G(z))) eq. 1

\max_{G} \; \log(D(G(z))) eq. 2

Following Mathieu et al., the generator gets the MRI patch instead of noise §2.1. Discriminator loss eq. 3 with binary cross-entropy eq. 4:

L_{D} = L_{bce}(D(x), 1) + L_{bce}(D(G(x)), 0) eq. 3

L_{bce}(\hat{y}, y) = -\frac{1}{N}\sum_{i=1}^{N} y_i \log\hat{y}_i + (1 - y_i)\log(1 - \hat{y}_i) eq. 4

N is the minibatch size, label 1 is real CT and 0 is generated §2.1. (In eq. 3 as printed, the real term is written D(x) with x called the input MRI; in practice D receives the real CT, see Pitfalls.)

Generator adversarial + reconstruction loss, with Y the ground-truth CT eq. 5:

L_{G} = \lambda_1 L_{bce}(D(G(x)), 1) + \lambda_2 \lVert Y - G(x) \rVert_2^{2} eq. 5

3D gradient difference loss, implemented with finite differences eq. 6 §2.1:

L_{gdl} = \left| |\nabla Y_x| - |\nabla \hat{Y}_x| \right|^2 + \left| |\nabla Y_y| - |\nabla \hat{Y}_y| \right|^2 + \left| |\nabla Y_z| - |\nabla \hat{Y}_z| \right|^2 eq. 6

Total generator loss eq. 7:

L(X, Y) = L_G(X, Y) + \lambda_3 L_{gdl}(X, Y) eq. 7

Training. Alternating updates: first D on a minibatch of real CT and a minibatch of generated CT, then G on another minibatch of MRI/CT pairs §2.1.

Auto-Context Model (ACM). Several GANs are trained one after another. Stage 1 takes MRI patches. Its patch outputs are aggregated into a full predicted CT, patches are extracted from that CT and concatenated to the MRI patches as a second channel, and the result is the input for training the next GAN §2.2 fig. 3. Two ACM iterations are used for both datasets §4.2.

Inference. The MRI is split into overlapping patches, each is passed through the generator, and overlapping CT predictions are averaged §2.

# Context-aware GAN for MR -> CT  [§2] [§2.1] [§2.2] [fig. 3]
# Input:  mr_volumes, ct_volumes (paired, voxel-aligned); num_stages
# Output: list of trained generators, one per ACM stage

def train_context_aware_gan(mr_volumes, ct_volumes, num_stages):
    generators = []
    # stage 1 sees only MRI; later stages also see the previous CT estimate
    context_volumes = None
    for stage in range(num_stages):
        # MRI channel(s) at 32^3, CT target at 16^3, same centre  [§4]
        # previous estimate is added as a second channel  [§2.2]
        samples = sample_patch_pairs(
            mr_volumes, context_volumes, ct_volumes,
            input_size=32, target_size=16,
        )
        generator = new_fcn_generator(in_channels=1 if stage == 0 else 2)
        discriminator = new_cnn_discriminator()
        train_gan(generator, discriminator, samples)
        generators.append(generator)
        # predict full CT volumes to give context to the next stage  [fig. 3]
        context_volumes = [
            predict_volume(generator, mr, ctx)
            for mr, ctx in zip_with_optional(mr_volumes, context_volumes)
        ]
    return generators

def train_gan(generator, discriminator, samples):
    # Adam, lr 1e-6, beta1 0.5, batch 10  [§4]
    for step in range(num_steps):  # number of steps not stated
        # discriminator update: real CT -> 1, generated CT -> 0  [eq. 3]
        mr_batch, ct_batch = samples.next_batch(10)
        fake_ct = generator(mr_batch)
        loss_d = bce(discriminator(ct_batch), 1)
        loss_d = loss_d + bce(discriminator(stop_grad(fake_ct)), 0)
        update(discriminator, loss_d)
        # generator update on another batch  [§2.1] [eq. 5] [eq. 7]
        mr_batch, ct_batch = samples.next_batch(10)
        fake_ct = generator(mr_batch)
        loss_g = 0.5 * bce(discriminator(fake_ct), 1)
        loss_g = loss_g + 1.0 * sum_sq(ct_batch - fake_ct)
        loss_g = loss_g + 1.0 * gradient_difference_loss(ct_batch, fake_ct)
        update(generator, loss_g)

def gradient_difference_loss(real, fake):
    # squared difference of absolute finite differences per axis  [eq. 6]
    total = 0
    for axis in (x, y, z):
        real_grad = abs(finite_difference(real, axis))
        fake_grad = abs(finite_difference(fake, axis))
        total = total + sum_sq(real_grad - fake_grad)
    return total

def predict_volume(generator, mr, context):
    # overlapping patches, average overlapping outputs  [§2]
    # stride not stated in paper; code uses (2, 5, 5)
    # [code: 3dganversion/g_model.py#L178]
    output = zeros_like(mr)
    counts = zeros_like(mr)
    for centre in sliding_window_centres(mr, stride):
        inputs = crop(stack(mr, context), centre, size=32)
        output[crop16(centre)] += generator(inputs)
        counts[crop16(centre)] += 1
    return output / counts

Hyperparameters

Results

Author-reported. Leave-one-out evaluation; metrics are MAE and PSNR §4. Baselines: Atlas (multi-atlas registration and CT averaging) and SR (sparse representation on warped atlases), both re-implemented by the authors; SRF+ (structured random forest + auto-context), whose numbers are taken from its paper rather than re-run §4.

Brain: 16 ADNI subjects; MR voxels 1.2×1.2×1 mm³, CT voxels 0.59×0.59×3 mm³ §3. Mean (std) and median over subjects tab. 1:

Pelvis: 22 subjects, CT manually aligned to MR, cropped to 153×193×50 at 1×1×1 mm³ §3 tab. 2:

Ablations:

Qualitatively, the GAN output shows fewer artifacts than the plain FCN fig. 4, and the proposed method's brain and pelvis predictions look closer to the ground-truth CT than the baselines' fig. 6 fig. 7.

Paper vs code

The repository is the authors' (README cites this work under its MICCAI 2017 title; LICENSE copyright Dong Nie). It was released after the arXiv version and contains a 2D GAN at the root and a 3D GAN in 3dganversion/; the comparison below uses the 3D GAN.

• ACM is not in the code. No auto-context stage or second input channel exists; the generator input placeholder has a single channel g_model.py. Paper uses 2 ACM stages §4.2.
• Generator layers. Paper: filters 32, 32, 32, 64, 64, 64, 32, 32 §2.1, kernels 9, 3, 3, 3, 9, 3, 3, 7, 3 fig. 2. Code: filters 32, 32, 32, 32, 64, 64, 64, 32 then 1, kernels 9, 3, 3, 3, 9, 3, 3, 3, 3; the 7³ layer is a 3³ layer g_model.py. Final layer has 1 output channel in code versus "32 convolutional filters" in the text §2.1.
• Discriminator. Paper: three conv+BN+ReLU+maxpool stages, 5³ kernels, 32/64/128/256 filters, FC 512/128/1 §2.1. Code: one 5³ conv (32, no BN) then three 3³ conv+BN+ReLU layers (64, 64, 32), no pooling, FC 128 then 1; the 512 layer is commented out g_model.py.
• ReLU before sigmoid in D. The last FC layer is built with activation=True, which applies ReLU utils.py, so the "logits" fed to the BCE are non-negative and D's probability is never below 0.5 g_model.py.
• Learning rate. Paper 1e-6 §4; code flag 2e-7 main.py; class default 2e-8 with an lr_step of 30000 suggesting step decay g_model.py.
• Same batch for D and G. Paper updates G on another minibatch §2.1; code feeds the same X, y to both updates g_model.py.
• Weight decay 0.0005 on all weights via an L2 term in a losses collection utils.py; not in the paper.
• L2 term scaling. Code uses sum of squared errors divided by 2×batch size loss_functions.py. lp_loss also returns the sum of the whole losses collection (weight-decay terms included), and combined_loss_G adds its result to that collection and sums it again, so the L2 and weight-decay terms are counted more than once loss_functions.py g_model.py.
• GDL. Code raises the per-axis gradient-magnitude difference to a power alpha and sums over voxels, using [-1, 1] difference filters with SAME padding loss_functions.py; eq. 6 fixes the power at 2.
• Batch norm uses center=False, scale=False with ReLU fused in utils.py; the paper only says BN is used.
• Code is incomplete as released. combined_loss_G references self.alpha, self.lam_lp, self.lam_gdl and self.lam_adv, which are never set, and train reads self.learning_rate_tensor, which is never defined g_model.py g_model.py. The variable filters 'd_' in var.name and 'g_' in var.name g_model.py do not appear to match the layer names used (conv1_a, conv_dis_1_a, fc2, ...). Treat it as a reference for layer choices, not a runnable reproduction.
• PSNR in code is 20·log10(max_I / RMSE) with max_I the larger of the maxima of prediction and ground truth utils.py; the paper does not define PSNR or its peak value.

Pitfalls

• Patch geometry depends on padding. 32³ in, 16³ out §4 only works if exactly 16 voxels per axis are lost; the paper does not say where. The code gets this from VALID padding on the two 9³ convs (32→24→16) and SAME elsewhere g_model.py. With the 7³ kernel from fig. 2 and the same padding scheme the output size is unchanged, but any other mix of VALID layers gives a different size.
• Eq. 3 notation. x is defined as the input MRI, yet D(x) is the "real" term; D must see the real CT patch (16³), as the text and code do §2.1 g_model.py.
• Final generator layer. The text's "32 convolutional filters" for the last layer §2.1 cannot produce a single-channel CT; use 1 output channel as in the code g_model.py.
• Inconsistent discriminator description. "Three stages ... and a combination of convolution with three fully connected layers" lists four conv filter counts (32, 64, 128, 256) §2.1; with three 2× poolings on a 16³ input the feature map is 2³ before the fourth conv. Pooling size and stride are not stated.
• Unstated preprocessing. Intensity normalization, the units of MAE, how brain MR and CT with different voxel sizes were resampled and registered, and any masking for metrics are all not stated §3. MAE and PSNR are not defined in the paper, so numbers are not directly comparable to other work.
• ACM training cost. Each stage needs a full-volume prediction of every training subject before the next GAN can be trained §2.2 fig. 3; with leave-one-out, this repeats per fold. Whether stage k's training-set context comes from a model that saw those subjects (risking optimistic context) is not stated.
• Inference cost. Dense overlapping 3D patches with averaging §2; the stride is not stated, and the code's (2, 5, 5) g_model.py means many forward passes per volume.
• Small, unequal comparison. 16 and 22 subjects; SRF+ results are copied from its paper rather than re-run in the same protocol §4.
• FCN vs GAN ablation (24.7 vs 25.9 PSNR) does not say which dataset it uses §4.1, so it cannot be matched to a row of tab. 1 or tab. 2.
• Very low learning rate. 1e-6 with Adam §4 (and lower in code) implies long training; the number of iterations is not stated.

Lineage

• Builds on: GAN training of Goodfellow et al. (arXiv:1406.2661) §2.1; conditional, frame-to-frame GAN with gradient difference loss from Mathieu et al. (arXiv:1511.05440), here extended to 3D §2.1; architectural choices (no pooling in G, batch norm) motivated by DCGAN (arXiv:1511.06434) §2.1; Auto-Context Model of Tu and Bai, moved from segmentation classifiers to GAN regressors §2.2.
• Compared against: multi-atlas registration (Atlas), sparse representation (SR), and structured random forest with auto-context (SRF+, Huynh et al.) §4.
• Superseded by: the authors' code README points to an updated PyTorch implementation (github.com/ginobilinie/medSynthesisV1); not assessed here.