Teaching AI the Anatomy Behind the Scan: Addressing Anatomical Flaws in Medical Image Segmentation with Learnable Prior (2403.18878v2)
Abstract: Imposing key anatomical features, such as the number of organs, their shapes and relative positions, is crucial for building a robust multi-organ segmentation model. Current attempts to incorporate anatomical features include broadening the effective receptive field (ERF) size with data-intensive modules, or introducing anatomical constraints that scales poorly to multi-organ segmentation. We introduce a novel architecture called the Anatomy-Informed Cascaded Segmentation Network (AIC-Net). AIC-Net incorporates a learnable input termed "Anatomical Prior", which can be adapted to patient-specific anatomy using a differentiable spatial deformation. The deformed prior later guides decoder layers towards more anatomy-informed predictions. We repeat this process at a local patch level to enhance the representation of intricate objects, resulting in a cascaded network structure. AIC-Net is a general method that enhances any existing segmentation models to be more anatomy-aware. We have validated the performance of AIC-Net, with various backbones, on two multi-organ segmentation tasks: abdominal organs and vertebrae. For each respective task, our benchmarks demonstrate improved dice score and Hausdorff distance.
- Fabian Bongratz, Anne-Marie Rickmann and Christian Wachinger “Abdominal organ segmentation via deep diffeomorphic mesh deformations” In Scientific Reports 13.1 Nature Publishing Group UK London, 2023, pp. 18270
- Fred L. Bookstein “Principal warps: Thin-plate splines and the decomposition of deformations” In IEEE Transactions on pattern analysis and machine intelligence 11.6 IEEE, 1989, pp. 567–585
- “Transunet: Transformers make strong encoders for medical image segmentation” In arXiv preprint arXiv:2102.04306, 2021
- “Encoder-decoder with atrous separable convolution for semantic image segmentation” In Proceedings of the European conference on computer vision (ECCV), 2018, pp. 801–818
- “3D U-Net: learning dense volumetric segmentation from sparse annotation” In Medical Image Computing and Computer-Assisted Intervention–MICCAI 2016: 19th International Conference, Athens, Greece, October 17-21, 2016, Proceedings, Part II 19, 2016, pp. 424–432 Springer
- “An image is worth 16x16 words: Transformers for image recognition at scale” In arXiv preprint arXiv:2010.11929, 2020
- “Swin unetr: Swin transformers for semantic segmentation of brain tumors in mri images” In International MICCAI Brainlesion Workshop, 2021, pp. 272–284 Springer
- “Topology-preserving deep image segmentation” In Advances in neural information processing systems 32, 2019
- “ViG-UNet: Vision Graph Neural Networks for Medical Image Segmentation” In arXiv preprint arXiv:2306.04905, 2023
- Fanwei Kong, Nathan Wilson and Shawn Shadden “A deep-learning approach for direct whole-heart mesh reconstruction” In Medical image analysis 74 Elsevier, 2021, pp. 102222
- “Decoupled weight decay regularization” In arXiv preprint arXiv:1711.05101, 2017
- “Word: A large scale dataset, benchmark and clinical applicable study for abdominal organ segmentation from ct image” In arXiv preprint arXiv:2111.02403, 2021
- “U-net transformer: Self and cross attention for medical image segmentation” In Machine Learning in Medical Imaging: 12th International Workshop, MLMI 2021, Held in Conjunction with MICCAI 2021, Strasbourg, France, September 27, 2021, Proceedings 12, 2021, pp. 267–276 Springer
- “Topology Preserving Compositionality for Robust Medical Image Segmentation” In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2023, pp. 543–552
- “The graph neural network model” In IEEE transactions on neural networks 20.1 IEEE, 2008, pp. 61–80
- “Robust deep learning-based semantic organ segmentation in hyperspectral images” In Medical Image Analysis 80 Elsevier, 2022, pp. 102488
- “Attention is all you need” In Advances in neural information processing systems 30, 2017
- “Progressive deep segmentation of coronary artery via hierarchical topology learning” In International Conference on Medical Image Computing and Computer-Assisted Intervention, 2022, pp. 391–400 Springer