PaperlayerView exact revision on arXiv

Source-linked AI summary

Semi-supervised Task-driven Data Augmentation for Medical Image Segmentation

Krishna Chaitanya, Neerav Karani, Christian F. Baumgartner, Ertunc Erdil, Anton Becker, Olivio Donati, Ender Konukoglu

arXiv:2007.05363v2eess.IVcs.CVcs.LG

TL;DR

Medical segmentation requires many costly pixel-wise annotations, while small labeled sets may not capture the variation needed for generalization. The paper learns task-driven shape and intensity augmentations using labeled and unlabeled data. It reports that the method outperforms other augmentation and semi-supervised approaches across three medical datasets, while not addressing scanner-related intensity differences.

  • Problem

    Medical segmentation models need many costly annotations, and limited labeled data may not expose them to the shape and intensity variations present in unseen images.

    Source: §I · p. 1

  • Method

    The method learns deformation-field and intensity-field generators for task-driven augmentation, combining segmentation loss with regularization that leverages unlabeled images.

    Source: A. Semi-Supervised and Task-Driven… · p. 3

  • Results

    The proposed augmentation outperforms the other considered augmentation and semi-supervised methods on the cardiac, prostate, and pancreas datasets.

    Source: §V · p. 8

  • Takeaways & Limitations

    Generated images need not always have realistic anatomical shapes to yield improved segmentation performance.

    Source: §V · p. 8, §VI · p. 11

    The study does not address intensity differences arising from scanner differences.

    Source: §V · p. 9

Abstract

from arXiv · show

Supervised learning-based segmentation methods typically require a large number of annotated training data to generalize well at test time. In medical applications, curating such datasets is not a favourable option because acquiring a large number of annotated samples from experts is time-consuming and expensive. Consequently, numerous methods have been proposed in the literature for learning with limited annotated examples. Unfortunately, the proposed approaches in the literature have not yet yielded significant gains over random data augmentation for image segmentation, where random augmentations themselves do not yield high accuracy. In this work, we propose a novel task-driven data augmentation method for learning with limited labeled data where the synthetic data generator, is optimized for the segmentation task. The generator of the proposed method models intensity and shape variations using two sets of transformations, as additive intensity transformations and deformation fields. Both transformations are optimized using labeled as well as unlabeled examples in a semi-supervised framework. Our experiments on three medical datasets, namely cardic, prostate and pancreas, show that the proposed approach significantly outperforms standard augmentation and semi-supervised approaches for image segmentation in the limited annotation setting. The code is made publicly available at https://github.com/krishnabits001/task$\_$driven$\_$data$\_$augmentation.

I. INTRODUCTION

Medical segmentation models need annotated examples to learn shape and intensity variation, but pixel-wise annotation is expensive and small labeled sets leave important variations unseen. The paper proposes task-driven, semi-supervised augmentation and evaluates it against existing approaches on cardiac, prostate, and pancreas datasets.

I. INTRODUCTION

  • Small labeled datasets expose segmentation networks to too little anatomical, tissue, and scanner-related variation, limiting performance on unseen images.The introduction frames this as a gap between human generalization from few examples and current deep-learning methods.Source: §I · p. 1
  • The proposed method optimizes generators for segmentation while using unlabeled images through an adversarial term to model population shape and intensity variation.It uses deformation fields and additive intensity transformations to create augmented image-label pairs.Source: §I · p. 1
  • Unlike random elastic augmentation with a few tuned hyperparameters, the approach learns flexible transformation networks from labeled and unlabeled training images.The authors state that the iteration count is its only hyperparameter and is determined using validation performance.Source: §I · p. 1
  • The study evaluates the method on cardiac, prostate, and pancreas datasets, compares related methods, and includes ablations of regularization terms and joint generator-segmentation optimization.The extended version specifically analyzes adversarial and large-deviation loss terms and compares joint with independent optimization.Source: §I · p. 2

A. Related work

  • Prior augmentation methods transform image-label pairs or generate samples, but the surveyed methods do not jointly optimize generators for task performance and leverage unlabeled data.The paper contrasts its approach with random transformations, GAN-based generation, and MixUp [18].Source: A. Related work · p. 2
  • The closest cited task-optimized generator is fully supervised and synthesizes images directly, whereas this method models transformations and uses unlabeled data to obtain masks and improve performance.A concurrent method explores geometric variations near decision boundaries rather than optimizing augmentation for task performance.Source: A. Related work · p. 2

II. METHODS

The methods section sets out supervised segmentation and augmentation, then contrasts fixed model-based transformations with GAN generation. These provide the background for learning augmentation generators in the task-driven approach.

II. METHODS

  • Supervised segmentation minimizes disagreement between predicted masks and ground-truth labels by optimizing the segmentation network parameters.The labeled training images and corresponding labels define the loss.Source: §II · p. 3
  • With augmentation, training uses the union of labeled and generated image-label pairs, with a shared transformation applied to each image and its label.For model-based transformations, fixed parameters govern the transformation and a random component samples its variation.Source: §II · p. 3
  • In GAN-based generation, a neural-network generator and discriminator are trained adversarially, while conventional generators are predefined or trained without considering the downstream task.The discriminator distinguishes generated from real pairs, and the generator aims to make its output indistinguishable from real pairs.Source: §II · p. 3

A. Semi-Supervised and Task-Driven Data Augmentation

The method learns generators for shape and intensity transformations using task loss and unlabeled data. Its regularization combines adversarial matching with a preference for larger transformations.

A. Semi-Supervised and Task-Driven Data Augmentation

  • The proposed optimization combines segmentation task loss with a regularization term and uses unlabeled images in generator learning.The figure marks segmentation-loss inclusion in generator optimization with a dotted red line.Source: A. Semi-Supervised and Task-Driven… · p. 3, Figure 1 · p. 4
  • Conditional generators model shape with deformation fields and intensity with intensity fields, applying the same transformation to the image and its mask.This construction yields the augmented segmentation mask corresponding to each generated image.Source: A. Semi-Supervised and Task-Driven… · p. 3
  • Segmentation loss favors augmented pairs useful for the task, while adversarial and large-deformation losses encourage larger transformations that match the unlabeled-image distribution.The regularization discourages generators from producing only easy-to-segment images with small changes.Source: A. Semi-Supervised and Task-Driven… · pp. 3–4

1) Deformation Field Generator:

The method learns generators for shape and intensity transformations, using segmentation and regularization losses to produce task-driven augmented image-label pairs. Adversarial training incorporates unlabeled images, while a large-deviation term favors substantial transformations.

1) Deformation Field Generator:

  • The deformation generator warps a labeled image and its one-hot label with the same generated dense per-pixel field to create an augmented pair.Source: 1) Deformation Field Generator · p. 4

2) Additive Intensity Field Generator:

3) Regularization Loss:

  • The regularization combines adversarial and large-deviation losses: unlabeled images guide realistic variation, while negative L1 norms favor large fields that the adversarial loss constrains.Source: 3) Regularization Loss · p. 4

4) Optimization Sequence:

The segmentation network is first pretrained on labeled data, then optimized alongside each generator and discriminator. Validation performance selects generator parameters before the segmentation network is retrained using labeled and generated data.

4) Optimization Sequence:

  • The segmentation network is pretrained on labeled data before each generator-discriminator pair is optimized with the segmentation cost and its corresponding regularization.Source: 4) Optimization Sequence · p. 4
  • Validation Dice scores select generator parameters, after which the segmentation network is reinitialized and trained on original labeled data plus samples from the selected generators.Source: 4) Optimization Sequence · p. 5
  • The algorithm trains deformation and intensity generators using labeled and unlabeled sets, then samples augmented pairs for segmentation training.Source: 4) Optimization Sequence · p. 5
  • The final training data can include deformation-only, intensity-only, and sequentially combined shape-and-intensity image-label pairs alongside affine-transformed and original labeled data.Source: 4) Optimization Sequence · p. 5

III. DATASET AND NETWORK DETAILS

The experiments use cardiac, prostate, and pancreas datasets with dataset-specific preprocessing and network architectures. The paper describes cardiac end-systolic images as the main-article choice.

A. Dataset details

  • The cardiac data include 100 subjects and LV, myocardium, and RV masks, while prostate has 48 scans with two regions and pancreas has 282 CT scans with pancreas and tumor merged as foreground.Source: A. Dataset details · p. 5
  • The main article uses only cardiac end-systolic images, a random choice; end-diastolic evaluation is reported in an appendix table.Source: A. Dataset details · p. 5

B. Pre-processing

  • Cardiac and prostate volumes receive N4 bias correction; all datasets are intensity-normalized, resampled to dataset-specific in-plane resolutions, and cropped or padded to fixed sizes.Source: B. Pre-processing · p. 5

C. Network Architecture

  • The method uses segmentation, discriminator, and generator networks; separate generator variants produce deformation fields or intensity masks.Source: C. Network Architecture · pp. 5–6
  • The discriminator uses a DCGAN-like architecture, and the segmentation network uses an encoder-decoder architecture similar to U-Net.Source: C. Network Architecture · p. 6

D. Training Details

Training uses weighted segmentation loss and fixed optimization settings, while experiments compare augmentation and semi-supervised alternatives under limited-label conditions. Ablations examine regularization, segmentation-loss use, and labeled or unlabeled data amounts.

D. Training Details

  • Weighted cross entropy assigns greater total weight to foreground pixels, and the loss weights and regularization coefficients are empirically selected for training.The regularization coefficients were chosen using one cardiac training volume; the authors did not exhaustively grid-search combinations because of computational expense.Source: D. Training Details · p. 6

IV. EXPERIMENTS

  • Experiments use one or three labeled volumes, fixed unlabeled, test, and validation sets, and 15 runs per experiment across training-volume and network-initialization variations.Source: §IV · p. 6
  • Comparisons include no augmentation, affine augmentation, random elastic deformation, and random contrast and brightness changes, with affine augmentation also included by default in later methods.Source: §IV · pp. 6–7
  • Other comparisons evaluate MixUp, MixUp over generated and affine data, adversarial training, and self-training.Source: §IV · p. 7
  • Cardiac-only ablations vary adversarial and large-deviation loss terms and test generator optimization without segmentation loss.Source: §IV · p. 7

E. Different set of train, validation, test and unlabeled

The authors test whether the method generalizes across a different cardiac data split and describe a training setup without validation images. They evaluate segmentation with Dice’s similarity coefficient.

E. Different set of train, validation, test and unlabeled

  • A new cardiac training, validation, test, and unlabeled split is used to retrain the transformation fields with one labeled volume, testing sensitivity to the original split.The experiment is intended to assess whether the method overfits to a specific split or generalizes across splits.Source: E. Different set of train, validation,… · p. 7
  • Without validation images, generator training stops after a predefined number of iterations, and the resulting parameters generate augmentation images.Source: E. Different set of train, validation,… · p. 7
  • Segmentation performance is evaluated using Dice’s similarity coefficient on the test images for each dataset’s structures.Source: E. Different set of train, validation,… · p. 8

V. RESULTS AND DISCUSSION

Across cardiac, prostate, and pancreas experiments, the proposed GD+GI augmentation outperforms the compared augmentation and semi-supervised methods. The analyses examine its component losses, data requirements, and computational and scanner-related scope.

V. RESULTS AND DISCUSSION

  • Without augmentation performance is lowest; affine augmentation improves it, and adding random elastic deformation or intensity fluctuation yields further gains.Source: §V · p. 8
  • Learned deformation and intensity transformations outperform their random counterparts, and combining them substantially improves DSC, although generated images need not have realistic anatomy or intensity.Source: §V · pp. 8–9
  • GD+GI outperforms the compared augmentation and semi-supervised methods across cardiac, prostate, and pancreas; for most test subjects, it matches or exceeds random augmentation.Fig. 2 reports DSC boxplots, using one labeled volume for cardiac and prostate and three for pancreas.Source: Figure 2 · p. 8, §V · pp. 8–9
  • Mixup improves over affine and random elastic augmentation, with marginal further gains when applied to generator outputs; self-training improves over affine but generally trails GD+GI, as does adversarial semi-supervised training.Self-training is an exception for the prostate central gland; adversarial semi-supervised training yields marginal gains over affine augmentation.Source: §V · p. 9
  • Deployment on a new dataset is computationally expensive because it requires searching for regularization-loss hyperparameters; the study does not address scanner intensity differences and evaluates images from one scanner type.Source: §V · p. 9
  • Both regularization terms together give the highest reported boost: adversarial loss constrains extreme transformations, while large-deviation loss encourages larger fields.Disabling regularization gives the least gain; large-deviation loss alone deteriorates performance, whereas adversarial loss alone improves it.Source: §V · p. 9
  • Including segmentation loss in generator training yields higher DSC than independently optimizing the generators without that task loss.The authors contrast joint optimization with independent generator optimization used in earlier work, [17].Source: §V · pp. 9–10
  • 1–3 unlabeled cardiac volumes reduce performance; gains are similar from 5 to 25, with negligible improvement at 50 versus 25, while the augmentation gap narrows as labeled volumes increase.Improvements remain statistically significant through 10 labeled volumes; 10 proposed-method volumes perform similarly to 40 with random augmentation.Source: §V · p. 10, Figure 3 · p. 10

E. Different set of train, validation, test, and unlabeled

The reported findings indicate that the approach’s results persist across a different cardiac dataset split and change little when validation images are not used to select generator parameters.

E. Different set of train, validation, test, and unlabeled

  • The results hold for a randomly chosen cardiac dataset split, rather than overfitting to one particular split.The passage points to Table II for this finding but does not provide its numerical values.Source: E. Different set of train, validation,… · p. 10
  • Without validation images, performance does not vary much from performance obtained using validation images.The no-validation model parameters are selected after a predefined number of training iterations.Source: E. Different set of train, validation,… · p. 10

VI. CONCLUSION

The conclusion presents task-driven, semi-supervised augmentation using conditional generators for shape and intensity transformations. On three public datasets, the method reports substantial gains, although its generated images need not look realistic to improve segmentation.

VI. CONCLUSION

  • Fig. 4 presents qualitative comparisons of GD+GI with affine, random, adversarial-training, and Mixup approaches across three datasets.It shows two images per dataset alongside the input and ground truth.Source: Figure 4 · p. 11
  • The method uses task-driven optimization and unlabeled data with conditional generators for shape and intensity transformations.The generators produce transformations intended to capture variation in the population and optimize augmentation for segmentation performance.Source: §VI · p. 11
  • The approach reports substantial segmentation gains over existing methods on cardiac, prostate, and pancreas datasets with limited annotations.The datasets are described as publicly available.Source: §VI · p. 11
  • Generated augmentations were not necessarily realistic, yet yielded improved segmentation performance.The authors question whether generating realistic examples is optimal for segmentation.Source: §VI · p. 11

VIII. APPENDIX

The appendix describes supplementary comparisons and visualizations of augmentation settings, learned transformations, and segmentation performance across datasets and evaluation conditions.

VIII. APPENDIX

  • Fig. 6 evaluates random elastic augmentation across initial matrix sizes and sigma values against learned deformation-generator augmentation, reporting mean Dice for cardiac at NL=1.Random-field matrix values are sampled from a zero-mean Gaussian distribution and interpolated to image dimensions.Source: Figure 6 · p. 13
  • Fig. 7 compares mean Dice for cardiac at NL=1 across contrast and brightness settings against augmentation from the learned intensity-field generator.The caption identifies the intensity settings and learned-generator comparison, but provides no outcome values.Source: Figure 7 · p. 13
  • Fig. 8 compares learned deformation-generator performance without validation images across predefined training durations with a best-validation model and random deformations.The caption defines 1k, 2k, 4k, 6k, and 10k as training iterations in thousands; mean Dice is reported for cardiac at NL=1.Source: Figure 8 · p. 13
  • Table III compares GD+GI with RD+RI using mean DSC and standard deviation over 15 runs on 20 ACDC test ED images with one labeled ED volume.The caption marks Wilcoxon signed-rank significance at p-value threshold 0.05 but gives no cell values.Source: Table III · p. 13
  • Fig. 9 compares affine, random, and proposed augmentations per test subject, reporting mean Dice over 15 runs for cardiac and prostate at NL=1 and pancreas at NL=3.The caption identifies the comparison and evaluation conditions but supplies no per-subject values.Source: Figure 9 · p. 14
  • Figs. 10–11 show generated deformation and additive intensity fields and their transformed cardiac images; the accompanying schematic labels an intensity-transformed image as X + ∆I.Fig. 10 also identifies shape-transformed images as X ◦ v.Source: Figure 10 · p. 14, Figure 11 · p. 15, §VIII · p. 15

Related papers

Loading 2007.05363v2…