⚡ electrical engineering

OTLesMix: Wasserstein Barycenter and Optimal Transport Map for Synthetic Lesion Generation with Diverse Shapes and Locations

This paper introduces OTLesMix, a novel data augmentation method that leverages Wasserstein barycenters and optimal transport maps to generate realistic synthetic brain lesions with diverse shapes and locations, significantly improving segmentation model performance compared to existing mix-based approaches.

Original authors: Robin Trombetta, Carole Lartizien

Published 2026-08-07

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Robin Trombetta, Carole Lartizien

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to recognize a specific type of cloud in the sky. You show it a thousand pictures of fluffy white cumulus clouds, and it gets pretty good. But then you show it a picture of a cloud that looks like a dragon, or a cloud that is shaped like a teapot, and the robot gets confused. It hasn't seen those shapes before. This is a common problem in the world of medical imaging, where computers need to learn to spot diseases like tumors or strokes in brain scans. The problem is that real medical data is hard to get; it takes hours for a human expert to carefully draw a map of exactly where a disease is on a single 3D brain scan. Because there aren't enough examples, the computer often only learns to recognize the "standard" shapes and locations it has already seen.

To fix this, scientists use a trick called "data augmentation." Think of it like a chef trying to teach a student to bake. If the student only ever sees one recipe, they might fail if the ingredients change slightly. So, the chef gives them a few extra practice batches. In computer science, this means taking the existing images and tweaking them—rotating them, changing the brightness, or even cutting a piece of one image and pasting it onto another. The goal is to create a massive library of "fake" but realistic examples so the computer learns to handle any shape or location a disease might take. However, the old ways of doing this often result in weird, blurry, or impossible-looking images that don't teach the computer much about new shapes.

This is where a new method called OTLesMix comes in, proposed by researchers Robin Trombetta and Carole Lartizien. They wanted to create a smarter way to mix medical images that could invent brand-new, realistic-looking disease shapes and locations, rather than just shuffling the old ones around.

The Magic of the "Perfect Blend"

The researchers' big idea relies on a branch of math called Optimal Transport. To understand this, imagine you have two piles of sand. One pile is shaped like a star, and the other is shaped like a circle. You want to move the sand from both piles to create a new pile that is a perfect blend of the two—maybe a star-circle hybrid. The "Optimal Transport" problem asks: What is the most efficient way to move every grain of sand so that you use the least amount of energy?

In the world of medical images, the "sand" is the pixels that make up a tumor or a lesion. The researchers use a specific type of optimal transport called a Wasserstein barycenter. If you think of two brain scans with tumors as two different "maps," this math tool finds a "middle ground" map. It doesn't just average the pictures (which would make a blurry mess); instead, it calculates a new, unique shape that sits perfectly between the two original shapes. It's like taking a photo of a cat and a photo of a dog, and mathematically generating a brand-new animal that has the ears of the cat and the tail of the dog, but looks like a real, living creature in between.

How OTLesMix Works

The process happens in two main steps, like a two-stage cooking recipe:

Cooking the Shape: First, the computer picks two real brain scans that have lesions (diseased areas). It uses the "Wasserstein barycenter" math to calculate a brand-new lesion shape and location. This new shape is a smooth, realistic interpolation between the two originals. It might be a tumor that is slightly larger, slightly shifted, or shaped differently than anything in the original training set.
Painting the Picture: Once the new shape is decided, the computer needs to figure out what the inside of this new lesion should look like. It uses an "optimal transport plan" to figure out exactly which pixels from the original images should be moved to fill this new shape. It's like taking the texture and color from the original tumors and carefully "inpainting" them into the new shape. Finally, this new, synthetic lesion is pasted onto a healthy brain scan (or one of the original ones) to create a completely new training example.

Does It Actually Work?

The researchers tested this method on three different brain imaging datasets involving different types of brain issues: glioblastoma (a type of brain cancer), stroke, and ischemic stroke lesions. They compared their new method against older, popular mixing techniques like MixUp (which just blends images linearly) and CutMix (which cuts and pastes rectangles).

The results were quite promising. When they trained a computer model to segment (identify) these brain lesions, the model that used the OTLesMix synthetic data performed significantly better than models trained without it.

On the brain cancer dataset (BraTS 2020), the model's accuracy improved by 2.9 to 6.6 points on a standard scoring metric called the Dice score compared to a model that didn't use synthetic data.
It also beat the other mixing methods (MixUp, CutMix, and CarveMix) on all three datasets. For example, on the BraTS dataset, OTLesMix achieved a score of 83.3% for the whole tumor, while the next best method (CarveMix) scored 82.5%.

The authors note that the difference wasn't just a tiny fluke; statistical tests showed that OTLesMix was consistently better than the others. Visually, the new method also produced lesions that looked more diverse in shape and location, whereas the older methods tended to produce shapes that looked a bit too similar to the originals or just awkwardly cut-and-pasted.

The Catch

While the method is a success, the authors are honest about its current limits. Generating these new images takes time. For a 2D slice of a brain, it takes about 1 second. But for full 3D brain scans, it can take anywhere from 1 to 3 minutes per image. This is much slower than just cutting and pasting a rectangle.

Also, the method relies on a simple linear interpolation of pixel colors. If the two original brain scans came from very different machines or had very different lighting conditions, the new "blended" lesion might look a bit unrealistic, like a photo that was poorly edited. The authors suggest that in the future, they could combine this shape-generation math with powerful AI image generators to fix those texture issues.

In short, OTLesMix is a clever new tool that uses advanced math to invent realistic, diverse examples of brain diseases. By doing so, it helps teach medical AI to recognize diseases in shapes and places it has never seen before, potentially leading to better diagnosis tools in the future.

Technical Summary of OTLesMix: Wasserstein Barycenter and Optimal Transport Map for Synthetic Lesion Generation

Problem Statement
Deep learning has revolutionized medical image segmentation, yet supervised training remains constrained by the high cost and time required for expert annotations, particularly for dense 3D segmentation tasks. While data augmentation is a standard technique to improve model training, traditional methods (spatial/intensity transformations) and advanced generative models (GANs, Diffusion Models) present specific limitations. Generative models require large, diverse datasets for their own training, which may not be available in medical contexts. Conversely, existing "mix-based" augmentation strategies (e.g., MixUp, CutMix, CarveMix, LesionMix) often fail to generate sufficient variability in lesion shapes and locations. They typically rely on simple random sampling or rigid geometric transformations of existing masks, resulting in synthetic samples that lack the diversity necessary to robustly train segmentation models.

Methodology: OTLesMix
The authors propose OTLesMix, a novel image synthesis method that leverages Optimal Transport (OT) theory to generate realistic and diverse synthetic lesions. The method operates in two primary stages to create a new sample $(\tilde{X}, \tilde{Y})$ from a base dataset $D = \{(X_i, Y_i)\}$ :

Lesion Mask Generation via Wasserstein Barycenter:
- Two base samples $(X_i, Y_i)$ and $(X_j, Y_j)$ are randomly selected.
- The coordinates of the non-zero pixels in the ground truth masks $Y_i$ and $Y_j$ are treated as probabilistic distributions.
- A new lesion mask $\tilde{Y}_{les}$ is generated by computing the Wasserstein barycenter of these two distributions. This involves optimizing the coordinates of the new mask points to minimize the weighted sum of transport costs to the source masks:
  $B(\lambda, \alpha_i, \alpha_j, \beta) = \lambda OT(\alpha_i, \beta) + (1-\lambda)OT(\alpha_j, \beta)$
  where $\lambda \sim U(0,1)$ is a random weighting parameter.
- This process interpolates between the shapes and locations of the two original masks, allowing the new lesion to exist anywhere between the original locations with a shape that is a smooth interpolation of the originals.
- To ensure computational efficiency, the authors utilize the de-biased Sinkhorn divergence and the Sinkhorn algorithm, enabling the processing of up to $10^6$ points in seconds on modern GPUs.
Image and Label Generation via Optimal Transport Map:
- Once the new mask $\tilde{Y}_{les}$ is defined, the pixel intensities and labels for the synthetic lesion must be determined.
- An optimal transport plan is computed between the generated mask $\tilde{Y}_{les}$ and one of the source masks (e.g., $Y_j$ ).
- For each pixel in the new mask, the transport map $T$ determines its corresponding location in the source image $X_j$ .
- Intensity Synthesis: The intensity of the synthetic pixel is obtained via linear interpolation of the source image intensities at the mapped location.
- Label Synthesis: For multi-class tasks, the ground truth label is assigned via nearest-neighbor interpolation of the source label map.
- The synthesized lesion is then inpainted onto a base image (either $X_i$ , $X_j$ , or a third image from the dataset). Post-processing steps remove intersections with existing lesions or unrealistic anatomical locations (e.g., skull, CSF).

Key Contributions

Novel Synthesis Framework: Introduction of OTLesMix, which uniquely combines Wasserstein barycenters for shape/location interpolation and optimal transport maps for intensity/label transfer.
Enhanced Diversity: The method explicitly addresses the limitation of existing mix-based methods by generating lesions with diverse shapes and locations that are not merely copies or simple rotations of existing data.
State-of-the-Art Performance: The method demonstrates superior performance in improving downstream segmentation tasks compared to standard mix-based augmentation techniques.

Experimental Results
The authors evaluated OTLesMix on three brain lesion segmentation datasets using the nnUNet backbone:

BraTS 2020: Glioblastoma and low-grade glioma segmentation (2D slices).
ATLAS v2.0: Stroke neuroimaging (3D volumes).
ISLES 2022: Ischemic stroke lesion segmentation (3D volumes).

Quantitative Findings:

OTLesMix outperformed state-of-the-art mix-based methods (MixUp, CutMix, CarveMix) and the baseline (nnUNet with only traditional data augmentation) across all datasets.
Improvement over Baseline: The Dice Similarity Coefficient (DSC) improved by 2.9 to 6.6 points compared to models trained without synthetic data.
Improvement over Competitors: OTLesMix achieved a 0.8 to 1.8 point increase in DSC over the second-best method.
Statistical Significance: The improvements were statistically significant ( $p \le 0.05$ ) when compared to other generative methods, with the largest p-value being 0.012 (vs. CarveMix on ATLAS).

Qualitative Findings:
Visual inspection confirmed that OTLesMix produces lesions with greater diversity in shape and location compared to other methods. The method successfully handled multi-class segmentation (BraTS) and lesions with multiple connected components (ATLAS, ISLES), maintaining consistency in the interpolated shapes.

Significance and Limitations
The paper claims that OTLesMix provides a robust solution for data augmentation in medical imaging where data is scarce, offering a middle ground between the simplicity of mixing methods and the complexity of training full generative models. By leveraging optimal transport, it ensures that synthetic samples are geometrically and intensity-wise consistent with the underlying data distribution.

However, the authors acknowledge specific limitations:

Computational Cost: Generating a single sample takes approximately 1 second for 2D images and 1–3 minutes for 3D volumes, depending on resolution.
Realism Constraints: The method relies on linear interpolation of pixel intensities. This can lead to unrealistic inpainted lesions if there is a significant domain gap (e.g., different scanners or acquisition protocols) between the two source images used for mixing.
Future Direction: The authors suggest that future work could combine the OTLesMix mask generation procedure with mask-conditioned deep learning generation models to further improve the realism of the synthesized intensities.

The Magic of the "Perfect Blend"

How OTLesMix Works

Does It Actually Work?

The Catch

Technical Summary of OTLesMix: Wasserstein Barycenter and Optimal Transport Map for Synthetic Lesion Generation

More like this