SDMT: Spatial Dependence Multi-Task Transformer Network for 3D Knee MRI Segmentation and Landmark Localization

• Published in: IEEE transactions on medical imaging 2023-08, Vol.42 (8), p.2274-2285
• Main Authors: Li, Xiang, Lv, Songcen, Li, Minglei, Zhang, Jiusi, Jiang, Yuchen, Qin, Yong, Luo, Hao, Yin, Shen
• Format: Article
• Language: English
• Subjects: Artificial neural networks; Attention task; Bones; Coders; convolutional neural network; Deep learning; Feature extraction; Image processing; Image Processing, Computer-Assisted; Image segmentation; Knee; Knee Joint - diagnostic imaging; Knee landmark localization; Knee segmentation; Ligaments; Localization; Location awareness; Machine learning; Magnetic Resonance Imaging; multi-task learning; Neural networks; Neural Networks, Computer; Task analysis; Three-dimensional displays; transformer; Transformers
• ISSN: 0278-0062; 1558-254X
• DOI: 10.1109/TMI.2023.3247543

Knee segmentation and landmark localization from 3D MRI are two significant tasks for diagnosis and treatment of knee diseases. With the development of deep learning, Convolutional Neural Network (CNN) based methods have become the mainstream. However, the existing CNN methods are mostly single-task methods. Due to the complex structure of bone, cartilage and ligament in the knee, it is challenging to complete the segmentation or landmark localization alone. And establishing independent models for all tasks will bring difficulties for surgeon's clinical using. In this paper, a Spatial Dependence Multi-task Transformer (SDMT) network is proposed for 3D knee MRI segmentation and landmark localization. We use a shared encoder for feature extraction, then SDMT utilizes the spatial dependence of segmentation results and landmark position to mutually promote the two tasks. Specifically, SDMT adds spatial encoding to the features, and a task hybrided multi-head attention mechanism is designed, in which the attention heads are divided into the inter-task attention head and the intra-task attention head. The two attention head deal with the spatial dependence between two tasks and correlation within the single task, respectively. Finally, we design a dynamic weight multi-task loss function to balance the training process of two task. The proposed method is validated on our 3D knee MRI multi-task datasets. Dice can reach 83.91% in the segmentation task, and MRE can reach 2.12 mm in the landmark localization task, it is competitive and superior over other state-of-the-art single-task methods.