Generalisable 3D Fabric Architecture for Streamlined Universal Multi-Dataset Medical Image Segmentation

TL;DR

FIRENet is a universal 3D architecture that effectively segments multiple medical datasets simultaneously and adapts well to transfer learning, handling diverse image sizes and features without hyper-parameter tuning.

Contribution

This work introduces FIRENet, a novel 3D fabric-based architecture enabling universal multi-dataset medical image segmentation and transfer learning with implicit multi-scale feature extraction.

Findings

FIRENet achieved excellent multi-dataset segmentation performance.

FIRENet demonstrated strong transfer learning capabilities.

FIRENet handled diverse image sizes effectively.

Abstract

Data scarcity is common in deep learning models for medical image segmentation. Previous works proposed multi-dataset learning, either simultaneously or via transfer learning to expand training sets. However, medical image datasets have diverse-sized images and features, and developing a model simultaneously for multiple datasets is challenging. This work proposes Fabric Image Representation Encoding Network (FIRENet), a universal architecture for simultaneous multi-dataset segmentation and transfer learning involving arbitrary numbers of dataset(s). To handle different-sized image and feature, a 3D fabric module is used to encapsulate many multi-scale sub-architectures. An optimal combination of these sub-architectures can be implicitly learnt to best suit the target dataset(s). For diverse-scale feature extraction, a 3D extension of atrous spatial pyramid pooling (ASPP3D) is used in each fabric node for a fine-grained coverage of rich-scale image features. In the first experiment, FIRENet performed 3D universal bone segmentation of multiple musculoskeletal datasets of the human knee, shoulder and hip joints and exhibited excellent simultaneous multi-dataset segmentation performance. When tested for transfer learning, FIRENet further exhibited excellent single dataset performance (when pre-training on a prostate dataset), as well as significantly improved universal bone segmentation performance. The following experiment involves the simultaneous segmentation of the 10 Medical Segmentation Decathlon (MSD) challenge datasets. FIRENet demonstrated good multi-dataset segmentation results and inter-dataset adaptability of highly diverse image sizes. In both experiments, FIRENet's streamlined multi-dataset learning with one unified network that requires no hyper-parameter tuning.