CS-UNet: A generalizable and flexible segmentation algorithm

Alrfou, Khaled,Zhao, Tian,Kordijazi, Amir
DOI: https://doi.org/10.1007/s11042-024-19242-4
IF: 2.577
2024-04-27
Multimedia Tools and Applications
Abstract:This study introduces a novel U-shaped image-segmentation algorithm, CS-UNet, which contains parallel CNN and Transformer encoders. This algorithm leverages the relative strength of CNN and Transformers, and enables flexible combination of encoders pre-trained on different datasets to extract the maximum benefit of transfer-learning. CS-UNet is evaluated for its segmentation accuracy on microscopy images of materials science. The performance of CS-UNet is comparable or better than state-of-the-art algorithms based on CNN or Transformer encoders. As expected, the performance of CS-UNet is better when its encoders are pre-trained on microscopy images than when its encoders are pre-trained on natural images. However, the strength of in-domain pre-training is more significant in use cases such as out-of-distribution learning and one-shot learning. In particular, the Intersection over Union (IoU) accuracy of nickel-based super-alloy images is improved from 77.89% to 82.13% for out-of-distribution learning and IoU accuracy of environmental-barrier-coating images is improved from 65.9% to 70.45% for one-shot learning. CS-UNet also has surprisingly good performance on medical images. For Synapse multi-organ dataset, CS-UNet pre-trained on materials microscopy images has an average accuracy of 84.2% in Dice Similarity Coefficient (DSC), and 8.89 mm in 95% Hausdorff Distance (HD). In comparison, state-of-the-art segmentation algorithms pre-trained on ImageNet have an average DSC ranging from 76.5% to 80.39% and average HD ranging from 14.7 to 39.7 mm. Even when pre-trained on ImageNet, CS-UNet still has DSC of 83.27% and HD of 15.26 mm. This suggests that Transformer and CNN complement each other and pre-training on images with similar attributes is beneficial to the downstream tasks.
computer science, information systems, theory & methods,engineering, electrical & electronic, software engineering
What problem does this paper attempt to address?