K-SAM: A Prompting Method Using Pretrained U-Net to Improve Zero Shot Performance of SAM on Lung Segmentation in CXR Images

TL;DR

This paper introduces K-SAM, a method that uses pretrained U-Net models to automatically generate prompts, significantly enhancing SAM's zero-shot lung segmentation performance on chest X-ray images, especially in challenging cases with abnormalities.

Contribution

The study presents a fully automated prompt selection technique using pretrained U-Net models to improve SAM's zero-shot lung segmentation in CXR images, addressing challenges with abnormal and distorted lungs.

Findings

Achieved an average Dice score of 95.5% and 94.9% on two datasets.

Improved zero-shot segmentation performance with automated prompt selection.

SAM's predictions were less accurate on images with extreme abnormalities.

Abstract

In clinical procedures, precise localization of the target area is an essential step for clinical diagnosis and screening. For many diagnostic applications, lung segmentation of chest X-ray images is an essential first step that significantly reduces the image size to speed up the subsequent analysis. One of the primary difficulties with this task is segmenting the lung regions covered by dense abnormalities also known as opacities due to diseases like pneumonia and tuberculosis. SAM has astonishing generalization capabilities for category agnostic segmentation. In this study we propose an algorithm to improve zero shot performance of SAM on lung region segmentation task by automatic prompt selection. Two separate UNet models were trained, one for predicting lung segments and another for heart segment. Though these predictions lack fine details around the edges, they provide positive and negative points as prompt for SAM. Using proposed prompting method zero shot performance of SAM is evaluated on two benchmark datasets. ViT-l version of the model achieved slightly better performance compared to other two versions, ViTh and ViTb. It yields an average Dice score of 95.5 percent and 94.9 percent on hold out data for two datasets respectively. Though, for most of the images, SAM did outstanding segmentation, its prediction was way off for some of the images. After careful inspection it is found that all of these images either had extreme abnormality or distorted shape. Unlike most of the research performed so far on lung segmentation from CXR images using SAM, this study proposes a fully automated prompt selection process only from the input image. Our finding indicates that using pretrained models for prompt selection can utilize SAM impressive generalization capability to its full extent.