Random forest-based out-of-distribution detection for robust lung cancer segmentation

TL;DR

This paper introduces RF-Deep, a random forest classifier that uses deep features from a pretrained transformer to detect out-of-distribution CT scans, improving the reliability of lung cancer segmentation in both in-distribution and OOD datasets.

Contribution

The paper presents RF-Deep, a novel OOD detection method leveraging deep features from a transformer encoder to enhance segmentation robustness in medical imaging.

Findings

RF-Deep outperforms existing OOD detection methods on multiple datasets.

The transformer-based segmentation model achieves high accuracy with self-supervised pretraining.

RF-Deep effectively detects OOD cases with low false positive rates.

Abstract

Accurate detection and segmentation of cancerous lesions from computed tomography (CT) scans is essential for automated treatment planning and cancer treatment response assessment. Transformer-based models with self-supervised pretraining can produce reliably accurate segmentation from in-distribution (ID) data but degrade when applied to out-of-distribution (OOD) datasets. We address this challenge with RF-Deep, a random forest classifier that utilizes deep features from a pretrained transformer encoder of the segmentation model to detect OOD scans and enhance segmentation reliability. The segmentation model comprises a Swin Transformer encoder, pretrained with masked image modeling (SimMIM) on 10,432 unlabeled 3D CT scans covering cancerous and non-cancerous conditions, with a convolution decoder, trained to segment lung cancers in 317 3D scans. Independent testing was performed on 603 3D CT public datasets that included one ID dataset and four OOD datasets comprising chest CTs with pulmonary embolism (PE) and COVID-19, and abdominal CTs with kidney cancers and healthy volunteers. RF-Deep detected OOD cases with a FPR95 of 18.26%, 27.66%, and less than 0.1% on PE, COVID-19, and abdominal CTs, consistently outperforming established OOD approaches. The RF-Deep classifier provides a simple and effective approach to enhance reliability of cancer segmentation in ID and OOD scenarios.