[Paper Review] Pyramid-Focus-Augmentation: Medical Image Segmentation with Step-Wise Focus

Vajira Thambawita, Steven A. Hicks+2 more|arXiv (Cornell University)|Dec 14, 2020

Retinal Imaging and Analysis8 references6 citations

TL;DR

This paper introduces Pyramid-Focus-Augmentation (PYRA), a novel data augmentation technique for medical image segmentation that applies hierarchical grid-based transformations (from 256×256 down to 4×4) to enhance model performance on small datasets. By training a U-Net model on these multi-scale grid-converted masks, PYRA improves segmentation accuracy and enables step-wise attention to polyp regions, achieving state-of-the-art results on the Medico 2020 colon polyp segmentation challenge with a test mIoU of 0.6981 and Dice score of 0.7887.

ABSTRACT

Segmentation of findings in the gastrointestinal tract is a challenging but also an important task which is an important building stone for sufficient automatic decision support systems. In this work, we present our solution for the Medico 2020 task, which focused on the problem of colon polyp segmentation. We present our simple but efficient idea of using an augmentation method that uses grids in a pyramid-like manner (large to small) for segmentation. Our results show that the proposed methods work as indented and can also lead to comparable results when competing with other methods.

Motivation & Objective

• To address the challenge of limited annotated medical images in colon polyp segmentation, especially when positive findings are rare.
• To develop a data augmentation strategy that enhances deep learning model performance on small medical datasets without requiring complex architectures.
• To enable gradual, step-wise attention to polyp regions by progressively refining segmentation at multiple grid scales.
• To provide a flexible, adjustable augmentation method that can reduce computational cost by using lower-resolution grids when sufficient.
• To ensure compatibility with any segmentation model and task, making the method broadly applicable across medical imaging domains.

Proposed method

• PYRA generates checkerboard grids of sizes 4×4, 8×8, 16×16, 32×32, 64×64, 128×128, and 256×256 to augment input images, increasing training data from 800 to 6,400 samples.
• Ground truth masks are converted into grid-based representations by dividing them into cells and marking each cell as positive if it contains any true segmentation pixels.
• A U-Net model is trained on stacked inputs of original images and randomly selected grid images, predicting the corresponding grid-converted ground truth masks.
• Monte Carlo dropout inference is applied with 50 forward passes per input to estimate prediction uncertainty via mean and standard deviation outputs.
• Final predictions are generated by averaging 50 stochastic forward passes using the model in training mode, with outputs reconstructed at the original resolution (256×256).
• The method is evaluated using mIoU and Dice coefficient on both validation and test sets, with results submitted at the Medico 2020 challenge.

Experimental results

Research questions

• RQ1 Can pyramid-structured grid augmentation improve segmentation performance on small medical datasets with limited positive examples?
• RQ2 Does step-wise focus on polyp regions through multi-scale grids enhance model attention and prediction accuracy?
• RQ3 Can the proposed augmentation method be effectively combined with standard data augmentation and deep learning models like U-Net?
• RQ4 To what extent does the model’s uncertainty, as measured by standard deviation, correlate with accurate polyp region detection?
• RQ5 Can the method be adjusted to lower-resolution grids to reduce computational cost without sacrificing performance?

Key findings

• PYRA achieved the highest test mIoU of 0.6981 and Dice score of 0.7887 among all experiments, outperforming baseline methods.
• The best validation performance was achieved with PYRA alone (Exp-3), yielding mIoU of 0.7693 and Dice of 0.8447, indicating strong generalization.
• Model confidence, visualized via standard deviation maps, was highest in the central regions of polyps and lower at outer borders, confirming focused attention on key areas.
• The model successfully predicted polyp regions across all grid scales from 2×2 to 256×256, demonstrating progressive refinement from coarse to fine segmentation.
• Exp-4, which combined general augmentation with PYRA, underperformed compared to Exp-3, suggesting that standard augmentations may interfere with the pyramid-focus mechanism.
• The method is flexible and can be applied to any segmentation model, with results adjustable to lower-resolution grids for faster inference when needed.