AI muscle segmentation of the leg, hip, shoulder and arm using CNN

Automated muscle segmentation

muscle segmentation based on UNET architecture

The SegmentData function within the QMRITools package automatically segments the muscles of the lower extremities, hip and shoulder, and (3D only) the arm, and can segment several of these regions in one dataset. This function is an essential tool in the SegmentationTools package, leveraging neural networks for precise muscle segmentation. It is the segmentation engine used in the Muscle-BIDS pipeline.

Function Overview for muscle segmentation

SegmentData works in two steps:

  • Side and location recognition: one classification network (Body_Pos_Side) identifies the body location (lower leg, knee, upper leg, hip, torso, shoulder, head and neck) and the side (left/right) of the data. If the data contains multiple locations, it is cut in regions with overlap.
  • Muscle segmentation: for each region a dedicated network, available in 2D and 3D, performs the segmentation (shoulder, hip, upper leg and lower leg; the arm network is 3D only). The results are merged into one segmentation with standardized muscle labels. An overview is given on the networks page.

Full leg automated muscle segmentation using a convolutional neural network with UNET architecture.

The networks work on gradient echo and turbo spin echo Dixon data (in-phase, out-phase, raw echo, water and fat images) and diffusion data. The implementation supports any data size, and if the voxel size is given the data is rescaled to the resolution of the network. 3D continuous data with the 3D networks gives the best results; the 2D networks can be used for data with slice gaps or very thick slices. The trained networks are available as .wlnet files within the QMRITools package, and .ONNX files can be requested.

seg = SegmentData[data];                        (* full body default *)
seg = SegmentData[{data, vox}, "UpperLegsBody"]; (* a specific region, with voxel size *)

With the what argument the region is selected, e.g. "Body" (default), "LegsBody", "HipBody", "UpperLegsBody", "LowerLegsBody" or "ShoulderBody". The legacy options "Legs", "LegsHip", "UpperLegs", "LowerLegs", "Hip", "Shoulder" and "Arm" return the older label sets. The options SegmentationDimension ("2D" or "3D"), TargetDevice ("CPU" or "GPU") and NetworkOutput (e.g. "Both" to also return the confidence) control the segmentation. SegmentDataGUI[] opens a graphical interface, and MuscleMap can be run with RunMuscleMap.

Architecture and Training

The segmentation networks are based on the UNET architecture with ResNet convolution blocks that contain short skip connections. Other architectures (UNET++, strided-convolution downscaling) and block types are available. The networks are built with the NeuralNetworkTools. The block diagrams of the encoding and decoding blocks used in our UNET architecture are shown below.

Schematic representation of the unet architecture and the unet encoding and decoding blocks.

Data Augmentation and Training

For training, data is heavily augmented using various techniques such as scale, skew, rotation, translation, noise, sharpen, contrast, and brightness adjustments. Patches are selected after data augmentation. The default training settings are a batch size of 2, a patch size of 32x96x96 voxels and 256 patches seen per training round, see the methods page for details and for training your own network. Below are examples of 36 different patches with augmentation selected from the same dataset. The resulting segmentation from both networks are displayed. The training for the upper leg is shown below.

Example of data augmentation used for training the upper and lower leg segmentation UNET.

Animation of 3D UNET learning to segment muscle data.

Using SegmentData as a Script

No Mathematica license? Not a problem it also works without, using the free Wolfram Engine and the script Segment_Nii.wls, which also lets you select the region and the GPU. How this works you can read in the script tutorial. The README in the scripts folder on GitHub has the same instructions.

Example segmentations of the thigh and leg automated segmentation UNET neural networks.

Features Include

  • Automated region recognition
  • Multiple regions and contrasts, in 2D and 3D
  • Memory optimized segmentation
  • CPU and GPU support
  • Automated patching and merging
  • Standardized muscle labels
  • Confidence maps and segmentation quality checks
  • Train your own networks, with multi-contrast training and transfer learning

We gratefully acknowledge the support of NVIDIA with the donation of the GPU used for this work.