Self-Supervised Learning for Cardiac MR Image Segmentation by Anatomical Position Prediction

In plain words

Labelling heart scans for segmentation is slow and needs experts. This paper uses a signal that comes with the scan for free. Each cardiac MRI has known view planes, and its header says which direction is left, right, front and back. The method marks nine spots on each image at places defined by those planes, and trains a network to say which spot each pixel belongs to. No human labels are needed. To do this well, the network has to learn what the heart's chambers look like and what sits next to them.

The trained network is then retrained to outline the heart's ventricles, using a small set of labelled scans. The gain is largest when only a few subjects are labelled. With many labelled subjects, the gap to a network trained from scratch mostly closes. All results come from the authors.

TL;DR

• Pretext task: a U-net predicts nine anatomical positions, defined along the cardiac view lines, with no extra labels §2.
• Three ways to transfer it to ventricle segmentation; the best on average in short-axis is SSL+MultiTask, which keeps the pretext head and adds the segmentation head with weight β = 10 §2 §3 tab. 1.
• With five labelled short-axis subjects, mean Dice goes from 0.811 (U-net from scratch) to 0.852 (SSL+MultiTask), author-reported tab. 1.
• The gain shrinks with more labels and is not uniform: in long-axis with one subject, SSL+MultiTask (0.600) is below the scratch baseline (0.678) tab. 2.
• No official code was found. The method needs UK Biobank-scale data and several details the paper does not state (see Pitfalls).

Problem

Cardiac MR segmentation supports diagnosis and monitoring, but deep networks usually need many manual annotations, which are expensive and limited by the time of expert analysts §1. Most published cardiac segmentation methods are fully supervised and treat short-axis and long-axis views separately, training each from random weights §1 §2. The authors assume that each scan carries view-plane geometry and DICOM orientation headers that can define positions without extra annotation §1 §2.

Method

The method has two stages. Task 1 is the self-supervised pretext task, which predicts anatomical positions from unlabelled scans. Task 2 is the target task, which is segmentation of the LV cavity, myocardium and RV cavity (short-axis) or additionally LA and RA (long-axis 4Ch) §2 §3.

Anatomical positions, short-axis image. The 2Ch and 4Ch view lines are drawn on a short-axis slice. Their intersection lies in the LV. Along these lines the authors define nine positions, shown as bounding boxes: the intersection, two on the left, two on the right, two anterior and two posterior. Left/right and posterior/anterior orientation comes from the DICOM header §2 fig. 1.

Anatomical positions, long-axis 4Ch image. The 2Ch view line and the mid short-axis line are drawn on a 4Ch image. Nine positions follow the same pattern: the intersection, two left, two right, two superior, two inferior §2 fig. 1.

Pretext network. A standard U-net is trained to output 10 classes: the nine boxes plus background §2. The encoder, decoder and skip connections are kept; the last convolution is the task head. The loss is cross-entropy §2. The network depth, width, normalisation and pooling are not stated.

Transfer to segmentation. The pretext head is replaced by a new head: a 1×1 convolution with K output channels, where K is the number of classes §2. Three transfer variants are compared §2 fig. 2:

1. SSL+Decoder: the encoder is frozen; the decoder and head are trained on the annotated subjects §2 fig. 2(b).
2. SSL+All: all weights, encoder, decoder and head, are fine-tuned §2 fig. 2(c).
3. SSL+MultiTask: the network has two heads, the pretext head and the new segmentation head, trained together, so the pretext task is not forgotten §2 fig. 2(d).

Multi-task loss. For SSL+MultiTask the paper writes the loss as shown below (the equation has no number in the source). x is the image, y1 the 9-position label map, y2 the expert segmentation, θ the network parameters and β the weight §2.

L(\theta) = L_{task_1}(x, y_1 \mid \theta) + \beta \cdot L_{task_2}(x, y_2 \mid \theta) §2

Training schedule. Stochastic optimisation alternates the two tasks, because the two datasets differ in size and subjects §3:

# Pretext pre-training on unlabelled scans  [§2] [§3]
# Input:  slices (B, H, W); position maps (B, H, W), values 0..9
# Output: scalar loss for one optimiser step

def pretext_loss(slices, position_maps):
    # Predict the 10-way map: 9 position boxes plus background  [§2]
    logits = unet(slices, head="pretext")          # (B, 10, H, W)
    # Cross-entropy against the box label map  [§2]
    return cross_entropy(logits, position_maps)

# Multi-task fine-tuning, one outer iteration  [§3] [§2]
# Weight beta = 10 for the new task  [§3]

def multitask_iteration(pretext_batch, labelled_batch, beta):
    # One sub-iteration on task 1 (unlabelled pool, 3,825 subjects)  [§3]
    optimiser_step(pretext_loss(*pretext_batch))
    # beta sub-iterations on task 2 (labelled pool)  [§3]
    for _ in range(beta):
        optimiser_step(segmentation_loss(*labelled_batch))

Baseline. U-net-scratch uses the same network and settings, but starts from random weights §3.

Hyperparameters

Results

Author-reported. Setup: UK Biobank scans, 100 test subjects for transfer, Dice averaged over the labelled structures, baseline U-net trained from random weights with the same settings §3.

Short-axis (LV, myocardium, RV averaged), mean (standard deviation) of Dice tab. 1:

Bold marks the best value per row, as in the paper. The paper states that SSL+MultiTask is best on average across rows §3. Averaging the five rows gives 0.838 for SSL+MultiTask, 0.820 for SSL+All and 0.759 for U-net-scratch.

Long-axis (LV, myocardium, RV, LA, RA; Dice), mean (standard deviation) tab. 2:

Figures (described, not tabulated in the source): fig. 3 shows short-axis segmentations for U-net-scratch and SSL+MultiTask at 1, 5, 10, 50 and 100 subjects, against manual labels; the text says the scratch network fails completely at one subject. fig. 4 plots Dice and mean contour distance per structure for the two methods; the values are only plotted, so no numbers are given in the text. fig. 5 shows long-axis segmentations with few subjects.

Paper vs code

No official code. The LaTeX source has no code link, and the arXiv metadata records none. GitHub searches (api.github.com repository search) for "cardiac self-supervised position", "anatomical position prediction", "self-supervised cardiac segmentation", "wenjia bai cardiac", "BioMedIA", and the UK Biobank pretext wording returned only unrelated third-party repositories, none by the authors or their lab. Nothing else was checked, such as a project page.

Pitfalls

• Position geometry is under-specified. The paper does not give how box centres are computed from the view lines, how the 30 px spacing is measured, or what happens when boxes or view lines are near the image border. It also depends on DICOM orientation fields, so data without them cannot be used as is.
• Long-axis result contradicts the prose. The text says SSL+MultiTask "generally produces better segmentations" with limited data §3, but tab. 2 shows it at 0.600 with one subject, below U-net-scratch (0.678). It is above scratch at 5, 10, 50 and 100 subjects, so the gain is real from five subjects up, but the one-subject result is not stated in the prose.
• Short-axis claim is averaged. "SSL+MultiTask performs the best" §3 is a mean over five rows. It is best in four of five rows (not at 50 subjects, where SSL+All is higher). "Comparable" to scratch at large data is a small margin in both directions.
• Frozen encoder is not free. SSL+Decoder is worse than scratch at 100 subjects (0.867 vs 0.886) tab. 1. The "comparable" claim §3 holds only at some sizes.
• Multi-task loss vs schedule. The paper writes a summed loss with weight β §2, but implements the new task with β separate sub-iterations per task-1 step §3, and calls this "approximately equivalent" to a weight. The optimiser is Adam §3, which normalises gradient scale per parameter, so the effect of β is not obvious. No ablation over β is reported. Implement the schedule as stated in §3.
• Which network gets transferred to long-axis? Pretext data include both SA and LA images §3. It is not stated whether one encoder serves both tasks or whether the LA model was pretrained separately.
• Slice counts. The "(slices)" column counts annotated slices. The text says annotation was done at ED and ES §3, but does not say whether each subject contributes both frames. Dice units (per slice or per 3D volume) and the spread behind the ± values are not stated.
• Subsets and split. The 100/100 split, the subset choice for 1 to 50 subjects and the checkpoint used at test time are not stated. Results come from single runs with no seeds reported, so differences of about 0.01 may be noise.
• Data access. The method needs about 3,825 subjects with view-plane metadata, and the cohort is UK Biobank, which requires access approval.
• Augmentation and preprocessing. Rotation and scaling ranges, intensity normalisation, cropping and resampling are not stated.
• Only one pretext design. The paper does not test pretext variants, fewer or more positions, or pretext-only transfer with other heads. The nine-position design is the only one reported.

Lineage

• Builds on: the U-net architecture, cited as Ronneberger et al. §2.
• Related pretext tasks named in the introduction: relative position prediction, image rotation, colourisation, inpainting, subject identification (siamese MR pairs), recolouring surgical video and rotation prediction for lung CT §1. No arXiv ids are given in the source.
• Compared against: U-net trained from scratch with the same settings §3.
• Superseded by: not known from this wiki.