Diffusion Based Shape-Aware Learning with Multi-Scale Context for Segmentation of Tibiofemoral Knee Joint Tissues: an End-to-End Approach

Knee Osteoarthritis (KOA) is commonly evaluated through an automated segmentation of the femur (FB), tibia (TB), and tibiofemoral cartilage (FC and TC) in Magnetic Resonance Imaging. But, in recent works, such automated segmentation (1) is possible only from the multistage framework, thus creating data handling challenges and requiring continuous manual intervention, and (2) lacks poor learning capability in multiclass segmentation problem. To solve these issues, Multi-Scale Attentive-Unet (MiSA-Unet) model is proposed that includes (1) Scale-aware Attentive Feature Enhancement (SAFE) module to learn the multi-contextual information by attending to varied receptive fields and (2) Diffusion-based Multiple Tissue Shape Reconstruction (DMulTiSR) loss to progressively refine the tissue details by contemplating the feature difference and utilizing discrete diffusion process. Unlike existing methods, the proposed framework is end-to-end and single-stage model that improved average DSC by 2.33% and segmentation speed by 82.6%, while with post-processing, it out-performed by 9.19% in FC, 33.82% in TC, and 13.31% in FB for Hausdorff distance. The code is available at: https://github.com/adaydar/MiSA-Unet.