Source-linked AI summary
Retina U-Net: Embarrassingly Simple Exploitation of Segmentation Supervision for Medical Object Detection
Paul F. Jaeger, Simon A. A. Kohl, Sebastian Bickelhaupt, Fabian Isensee, Tristan Anselm Kuder, Heinz-Peter Schlemmer, Klaus H. Maier-Hein
TL;DR
Medical object detection often either relies on pixel-wise segmentation with ad-hoc object scoring or uses end-to-end detectors that discard valuable segmentation supervision, especially with small datasets. The paper proposes Retina U-Net, which combines a RetinaNet one-stage detector with U-Net-inspired auxiliary semantic segmentation, and evaluates it on two medical datasets and toy experiments. Retina U-Net outperforms models without full segmentation supervision and reaches the performance of more complex two-stage counterparts.
Problem
Medical object detection lacks a simple approach that preserves end-to-end object scoring while exploiting full pixel-wise supervision from commonly available annotations.
Method
Retina U-Net augments a RetinaNet one-stage detector with high-resolution U-Net-like feature levels and an auxiliary semantic segmentation loss.
Results
Retina U-Net outperforms models without full segmentation supervision on two medical datasets and matches the performance level of more complex two-stage counterparts.
Takeaways & Limitations
Full segmentation supervision is important for medical object detection, while Retina U-Net provides a simple architecture for leveraging it.
Takeaways & Limitations
Compared two-stage mask supervision is constrained because mask loss is evaluated on cropped proposal regions and masks are resampled to a fixed-sized grid.
Abstract
from arXiv · showhide
The task of localizing and categorizing objects in medical images often remains formulated as a semantic segmentation problem. This approach, however, only indirectly solves the coarse localization task by predicting pixel-level scores, requiring ad-hoc heuristics when mapping back to object-level scores. State-of-the-art object detectors on the other hand, allow for individual object scoring in an end-to-end fashion, while ironically trading in the ability to exploit the full pixel-wise supervision signal. This can be particularly disadvantageous in the setting of medical image analysis, where data sets are notoriously small. In this paper, we propose Retina U-Net, a simple architecture, which naturally fuses the Retina Net one-stage detector with the U-Net architecture widely used for semantic segmentation in medical images. The proposed architecture recaptures discarded supervision signals by complementing object detection with an auxiliary task in the form of semantic segmentation without introducing the additional complexity of previously proposed two-stage detectors. We evaluate the importance of full segmentation supervision on two medical data sets, provide an in-depth analysis on a series of toy experiments and show how the corresponding performance gain grows in the limit of small data sets. Retina U-Net yields strong detection performance only reached by its more complex two-staged counterparts. Our framework including all methods implemented for operation on 2D and 3D images is available at github.com/pfjaeger/medicaldetectiontoolkit.
1. Introduction
Medical object detection must reconcile pixel-wise annotations and object-level decisions: segmentation uses available supervision well, while detectors provide end-to-end object scoring but discard pixel-level signals. Retina U-Net addresses this trade-off by combining a RetinaNet-style detector with U-Net-inspired segmentation supervision and evaluates the approach across medical datasets and detector settings.
- Motivation: Pixel-wise annotations make semantic segmentation a natural fit for medical images, although many applications primarily require lesion localization or presence.Medical imaging also provides spatial separation between objects, reducing the need for overlapping-instance discrimination in many settings.
- Motivation: One-stage and two-stage detectors provide object-level predictions from coarse feature representations, but coarse levels lose pixel-wise information.Two-stage detectors generate proposals before categorization, whereas one-stage detectors classify directly on pyramid representations.
- Approach: Retina U-Net adds high-resolution pyramid levels and skip connections for auxiliary semantic segmentation while retaining RetinaNet-style detection.Its segmentation loss is computed from the highest-resolution logits using cross-entropy and soft Dice losses.
- Evaluation: The study evaluates lesion detection and categorization on Lung-CT and Breast-Diffusion-MRI datasets, supported by toy experiments analyzing observed performance gains.The framework covers comparative studies of object detectors operating on both 2D and 3D images.
- Contributions: The paper contributes weighted box clustering to consolidate detections across predictions of the same 2D or 3D image.It also provides modular implementations of the explored models and an efficient clustering implementation.
2. Related Work
Prior approaches add segmentation to object detection, but proposal-based designs restrict how fully segmentation supervision is used. Retina U-Net instead fuses a top-down FPN pathway with segmentation in a simple one-stage detector.
- Proposal-based detectors evaluate mask loss only on cropped regions, excluding gradients from surrounding context.
- RoIAlign resamples proposal regions and ground-truth masks to a fixed-sized grid, potentially limiting their original resolution.
- Mask supervision in proposal-based methods uses only positive matched proposals, making it dependent on region proposal performance.
- Mask-loss gradients in these detectors flow only through the corresponding pyramid level and higher layers, rather than the entire model.
- Retina U-Net naturally fuses object detection and segmentation through a top-down FPN pathway with additional high-resolution levels.
3. Methods
The method combines a RetinaNet-style one-stage detector with U-Net-like semantic segmentation supervision. It also adapts detection to small medical objects and combines patch-level predictions using weighted box clustering.
- RetinaNet detector: RetinaNet uses an FPN with classification and bounding-box regression subnetworks operating across pyramid levels.The described implementation shifts these operations toward P2–P5 to address small object sizes in medical images.
- Segmentation supervision: Retina U-Net adds full semantic segmentation supervision through pyramid levels P1 and P0 with skip connections, producing a U-FPN resembling a symmetric U-Net.Detection subnetworks do not operate on P1 and P0, limiting their parameter impact at inference.
- Segmentation loss: Dice scores are computed over a pseudovolume containing all images in a batch to stabilize classes with zero foreground pixels.Per-class scores can otherwise become unstable and fail to penalize false-positive predictions when a class has no foreground pixels.
- Patch-based processing: Large medical images are processed as patch crops, creating a trade-off between patch size and batch size under GPU-memory limits.This constraint applies to high-resolution images and 3D volumes, and can also affect inference when a single image exceeds available memory.
- Weighted box clustering: Weighted box clustering combines clustered detections using confidence, overlap, box area, and patch-center factors.The area factor favors larger boxes, while the patch-center factor down-weights boxes farther from the patch center.
- Weighted box clustering: Figure 4 provides an example application of weighted box clustering.
4.1. Backbone & Baselines
The study compares Retina U-Net with one- and two-stage detectors and a U-Net-like segmentation baseline under a shared FPN-based feature-extraction framework.
- Shared Backbone: All methods use an FPN with a ResNet50 backbone and are evaluated in both 2D slice-based and 3D volumetric-patch settings.Anchor sizes are reduced for smaller medical objects to {42, 82, 162, 322} across P2-P5.
- Baselines: The baseline set includes Retina Net, Mask R-CNN, Faster R-CNN+, U-Faster R-CNN+, and DetU-Net.Faster R-CNN+ disables the mask loss in toy-data ablations, while U-Faster R-CNN+ places the detector on U-FPN.
- Architecture Comparison: Figure 3 distinguishes regular FPN and symmetric U-FPN extractors while showing detector heads, decoder scales, and auxiliary segmentation logits.The upper panel contains regular-FPN baselines; the lower panel contains U-FPN baselines.
- Baselines: DetU-Net converts segmentation outputs from P0 into connected components, then assigns bounding boxes and component-level class scores.Only the five largest components per image, or fifteen in 3D, are retained to reduce noise.
4.2. Training & Inference Setup
Training and inference compare 2D, context-augmented 2D, and 3D processing under shared augmentation, oversampling, ensembling, and prediction-consolidation procedures.
- Processing Modes: The experiments compare slice-wise 2D, 2D with three neighboring slices as additional channels, and 3D volumetric convolution processing.Foreground regions are oversampled during training on patch crops.
- Reported Results: Table 1 reports lung lesion detection results on CT.The table corresponds to the lung CT task evaluated in the experiments.
- Training: Extensive 2D and 3D data augmentation is used to account for overfitting on small data sets.Results are reported on aggregated inner-loop test sets because small data sets produce unstable statistics.
- Inference: Inference ensembles the five highest-scoring validation epochs and applies test-time mirroring.Predictions from ensemble members and overlapping tiles are consolidated by clustering and weighted averaging.
- Reported Results: Table 2 reports breast lesion detection results on Diffusion-MRI.The table corresponds to the breast lesion task evaluated in the experiments.
4.3. Evaluation
Evaluation uses 3D mean average precision at a low IoU threshold, while a worked example contrasts weighted box clustering with non-maximum suppression.
- Evaluation Metric: Performance is measured by mean average precision at IoU = 0.1, reflecting coarse localization and non-overlapping objects in 3D.Evaluation and matching are performed in 3D for every model and processing setup.
- Patient-level Evaluation: Patient-level scores are computed as the maximum predicted score per class and patient, with average precision calculated over those scores.The paper reports patient-level scores for comparability but cautions that this metric should be interpreted carefully in the study setting.
- Prediction Consolidation: In the Figure 4 example, weighted box clustering leaves one benign-lesion prediction at 93% confidence and one malignant-class false positive.Non-maximum suppression leaves multiple false positives with confidence scores as high as 83%.
4.4. Lung nodule detection and categorization
Retina U-Net performs strongly on lung lesion detection, especially in 2D with context slices, while nearly matching the best two-stage 3D model.
- Retina U-Net and U-FRCNN+ maintain a clear performance margin over the remaining baselines.The comparison suggests that full segmentation supervision is important for this task.
- DetU-Net performs worse, reaching 0.41 mAP with context slices in 2D and 0.37 mAP in 3D.The authors associate this with high-confidence false positives from ad-hoc score aggregation.
- Using neighbouring slices provides important 3D context, whereas fully 3D operation yields no benefit over a 2D network.The 2D network receives ± 3 neighbouring slices.
4.5. Breast lesion detection and categorization
For breast lesion detection, Retina U-Net achieves the best results with context slices and in 3D, while plain 2D is best served by Mask R-CNN.
- 0.33 mAP is Retina U-Net’s best 2D-with-context result, and 0.36 mAP is its best 3D result.Mask R-CNN performs best on plain 2D with 0.34 mAP.
- Retina U-Net and U-FRCNN+ achieve the overall best results, indicating the value of full segmentation information.
- Three-dimensional context is less important for this breast MRI task than for the lung CT task.The authors attribute this to the data’s low z-resolution.
- Patient-level scores correlate poorly with actual model performance and should not guide clinical model selection.The authors attribute this to ambiguities in score aggregation.
4.6. Toy Experiments
Toy experiments test when full segmentation supervision helps object categorization under limited data. It helps when discriminative information involves shape or patterns, but not scale alone.
- Experiment design: Three tasks examine shape, discriminative patterns unrelated to shape, and object scale while progressively reducing training data.The experiments target the importance of full segmentation supervision with limited training data.
- Data generation: The toy data use 2D noisy images with 1000 training, 500 validation, and 1000 held-out test images.Images are 320 × 320, with foreground intensity increases of 0.2 before uniform noise is added.
- Shape discrimination: In the shape task, Retina U-Net and U-FRCNN perform best, with their margin increasing as training data decreases.The segmentation mask explicitly contains the discriminative hole distinguishing circles from donuts.
- Pattern discrimination: In the pattern task, Retina U-Net and U-FRCNN+ retain similar margins over other models despite the masks omitting the discriminative hole.This setup tests implicit use of segmentation supervision.
- Scale discrimination: In the scale task, semantic supervision provides no gain for small training sets because class information is entirely encoded in target-box scale.Two-stage detectors perform better in this setting, contrary to the expected effect of RoIAlign scale invariance.
- Baseline comparison: Mask R-CNN versus Faster R-CNN+ shows no detection gain from sub-optimal mask supervision with limited training data.
5. Conclusion
The paper presents Retina U-Net as a simple one-stage detector that leverages semantic segmentation supervision. Across medical datasets and toy settings, it improves detection over models lacking full supervision and is especially advantageous with scarce data.
- Retina U-Net is a simple one-stage detection model that leverages semantic segmentation supervision.
- Across lung CT and breast MRI datasets, Retina U-Net outperforms models without full segmentation supervision and matches performance reached by a more complex two-stage counterpart.
- Toy experiments show benefits when discriminative power lies beyond mere object scale.
- The approach is designed to exploit scarce labelled data, with experiments highlighting an advantage on small datasets.