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

Abstract

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.