/LUS-multitask-learning

Multi-Task Learning for Segmentation and Classification of Lung Ultrasound Images.

Primary LanguageJupyter NotebookGNU General Public License v3.0GPL-3.0

A Multi-Task Deep Learning Approach for the Assessment of COVID-19 in Lung Ultrasound

PyPI - Python Version Tensorflow UNIVPM License


πŸ“ Table of Contents

πŸ“‹ Abstract

Deep Learning (DL) has established itself as a highly effective approach in medical imaging, a fact that became particularly evident during the COVID-19 pandemic when the need for advanced diagnostic tools became critical. Although much of the existing research has centered around computed tomography (CT) scans, this study focuses on the application of DL in lung ultrasound (LUS), emphasizing LUS as a non-invasive imaging technique that is safer and more accessible than CT. This paper introduces a multi-task learning framework, meticulously crafted for both the classification and segmentation of lung damage induced by COVID-19. The approach aims to evaluate whether integrating shared features between classification and segmentation tasks enhances predictions and boosts the overall effectiveness of the model, even in domains characterized by highly variable, complex images lacking geometric patterns such as LUS images. The model's performance was thoroughly assessed using the publicly available ICLUS dataset, showing promising results. For the task of classification, the model achieved an accuracy of 66%, which is an improvement over the 63% accuracy rate achieved by single-task models. In segmentation, the model attained a Dice similarity coefficient of 49%, surpassing the 47% obtained with established techniques. This integrated approach ultimately leads to a more precise and visually clearer assessment of patients' clinical states, enhancing the diagnostic process.

πŸ—‚οΈ Dataset ICLUS

The ICLUS dataset, provided by the US Laboratory of the University of Trento (ULTRA), includes 277 pulmonary ultrasound (US) videos with 58,924 frames collected in Italian hospitals during the pandemic. It covers 35 patients, including 17 COVID-19 positive, 4 suspected, and 14 healthy individuals. The dataset features both linear and convex images, with a thorough four-tier evaluation system for pleural line conditions. Additionally, the dataset supports segmentation tasks with expert-annotated truth masks for about 2000 frames (4% of the dataset). Among these frames, 701 were classified as 0 (33%), 425 as 1 (19%), 769 as 2 (36%) and 259 as 3 (12%).

πŸ•ΈοΈ Multitask architecture: MT-UNet++

MT-UNet++

Our proposed network, MT-Unet++, integrates two shared branches to support both segmentation and classification tasks, utilizing a UNet++ architecture for enhanced feature extraction across various resolutions. For the encoder path, we employ ResNet18 with convolutional blocks followed by batch normalization and ReLU activation. The decoder reconstructs the original resolution using feature maps from the ResBlocks, organized in nodes denoted as $X^{i,j}$. The UNet++ structure is divided into stages, with stages 1-4 for the encoder, stage 5 as the bottleneck, and stages 6-9 for the decoder.

A classification branch is added to the bottom of the UNet++, utilizing features from stages 4, 5, and 6. This branch employs Global Average Pooling to standardize feature map sizes, followed by a two-dense layer configuration. The first dense layer has 256 units with ReLU activation, and the final layer has 4 units with softmax activation for classification tasks.

πŸ“Š Results

Quantitative Results

Qualitative Results