Comparison of atlas-based and neural-network-based semantic segmentation for DENSE MRI images

This study uses 256x256 pixel DENSE MR image data and associated masks for 63 patients. We randomly split the data into a training set, containing 51 examples used in the creation of the neural-network-based and atlas-based methods, and and a test set, consisting of 12 examples which we use to gauge the generalization of the methods. As is typical in MR imaging, the DENSE MRI data is stored as a complex-valued array.

The images typically visualized in MRI are based on magnitude, where the grey value of each pixel is created from the magnitude of the complex value in that pixel’s position (shown in Figure 2(1)). The color of each of the pixels in these magnitude images is controlled by the chemical makeup of the material that exists in that position in the brain, and therefore represent the brain’s anatomy.

Since the brightness and contrast levels varied across the magnitude images in a way that made segmentation more difficult, we pre-processed the data using histogram normalization. To this end, we used the MATLAB function histeq, which approximately linerarizes the cumulative distribution function and helps create a more uniform set of images. Figure 3 compares a sample of original and normalized images. For more explanation of this process, see [5].

DENSE imaging encodes displacement information in the phase of each complex number corresponding to the intensity of movement associated with each pixel position. For each subject many DENSE MRI scans were created corresponding to gates that occur over the patient’s cardiac cycle (shown in Figure 2(3)). The number of DENSE images varies per patient based on their heartbeat. The visual interpretation of these phase images is perhaps less intuitive than the magnitude images, mapping movement rather than anatomy.

Generally, areas of large movement in brain tissue appear lighter than those of small movement, but in areas of especially high movement – especially outside the brain where cerebrosprinal fluid (CSF) moves much more quickly than the micrometer deformations in the brain tissue – phase wrapping occurs. Because there is no distinction between the gray scale color of a pixel stored with a phase $\phi$ and $\phi+2\pi$, for example, they appear the same. These regions with high movement appear as random noise because the movement is beyond the smaller scale for which the phase encoding was calibrated. For more on phase wrapping, see [8].

In this study, we follow [11] and remove the temporal dimension of the DENSE MRI data by computing a mean DENSE image, showing the mean value per pixel over time (Figure 2(4)), and a peak DENSE image, showing the maximum value per pixel over time (Figure 2(5)), for each subject. This representation of the DENSE MRI data over time also helps reduce the regions of random noise and summarize over the patient’s whole cardiac cycle. For more discussion of this choice and of other ways to represent this data, see [11].

The masks associated with the magnitude images classify each pixel as being part of the background, the cerebellum, or the brain stem, and were drawn by researchers in Emory’s Department of Radiology and Imaging Science (shown in Figure 2(2)) as part of the study [11].

The neural-network-based and atlas-based approaches use image data in different ways. The first keeps the images as MRI discrete data while the latter uses these data to build a continuous model through interpolation.

In the neural-network-based method, the input, $I\in\mathds{R}^{256\times 256\times 3}$, is made of 3 channels: the normalized magnitude, peak DENSE, and mean DENSE images. These are each 256x256 matrices with entries corresponding to the gray values.

The atlas-based approach treats an image as an interpolated function, ${\cal I}:\Omega\to\mathds{R}$ where the domain $\Omega\subseteq\mathds{R}^{2}$ corresponds to coordinates of pixels in the image grid. The multilevel processes, described in Section 3.2, rely on this continuous model of our data. The range of this function, called on a specific 256x256 grid, creates the discrete 256x256 pixel images. To distinguish these understandings, calligraphic letters have been used to denote continuous functions while regular script denotes discrete. This notation, and the notation of other equations related to the atlas-based method, are formatted in the style of FAIR [9]. This MATLAB toolbox and its associated textbook [9] can be referred to for a more complete discussion of image registration.

These two understandings of images lend also themselves to two conceptions of the outputted masks. In the neural-network-based method, the output mask is a discrete 256x256 matrix which predicts a class $k\in\{0,1,2\}$ for each pixel where 0 corresponds to background, 1 to cerebellum, and 2 to brain stem. $M$ symbolizes the known target masks while $M_{p}$ corresponds to the model predicted mask.

Analogously, in the continuous function sense, let $C$, $B$, and $T$ be the sets representing the cerebellum, brain stem, and the total union between them respectively. Then let $\chi_{C},\chi_{B}:\Omega\to\{0,1\}$ be the characteristic functions for the cerebellum and brain stem respectively. Masks that display these regions will hence be defined as ${\mathcal{M}}:\Omega\to\{0,1,2\}$ where

$${\mathcal{M}}(x)=\chi_{C}(x)+2\chi_{B}(x).$$

(1)

In addition to this, we define the total (T) mask characteristic function to be $\chi_{T}:\Omega\to\{0,1\}$ where $\chi_{T}=\chi_{C\cup B}$.

Using this notation, the template mask will be referred to as ${\mathcal{M}}_{\mathcal{T}}$ and the reference mask as ${\mathcal{M}}_{\mathcal{R}}$.

Our ultimate goal is to automatically segment the brain stem and cerebellum so that we could produce peak-displacement biomarkers for those regions that match radiologist produced ones. The cerebellum’s average peak-displacement, for example, is calculated by averaging the peak-displacement of each pixel in the DENSE MRI that is classified as being part of the cerebellum (and likewise for the brain stem). In addition to producing accurate biomarkers we also quantify the accuracy of the segmentations to ensure the plausibility of the biomarker.

Radiologist drawn masks that identify what displacement data to average to obtain the biomarkers are treated as the gold standard. It may be that these algorithmic methods can produce better, more consistent masks, but that lies outside the scope of this study. To balance the dual concerns of producing masks and biomarkers that match manually produced ones, Dice similarity indices and biomarkers were considered at each step.

Given a new image without a mask, our method compares the new image to a set of 51 labeled images in order to predict a mask for the new image. To allow for a pixel-to-pixel comparison, we solve a set of registration problems. Following the language used in FAIR, the new image is called the reference, while the known labeled images are templates.

Here, we discuss how templates are chosen to be compared to new references along with an averaging method to reduce outlying results, and in the next section we will describe the registration performed in each of these comparisons.

The goal of the image registration is to find some transformation $y:\Omega\to\mathds{R}^{2}$ such that

$${\mathcal{T}}[y](x)\approx{\mathcal{R}}(x),\quad\text{ for all }\quad x\in\Omega.$$

(3)

Note that ${\mathcal{T}}[y](x)={\mathcal{T}}(y(x))$ is written this way to emphasize that the function $y$ is an input as well. Section 2.2 provides insight into other notational choices.

We find that function $y$ by minimizing the nonconvex objective functional

$${\cal{J}}_{\rm{atlas}}[y]=D_{\rm{SSD}}[{\mathcal{T}}[y],{\mathcal{R}}]+\alpha{\cal S}[y],$$

(4)

where $D_{\rm{SSD}}$ is the sum of squares distance measure, $\alpha$ is the regularization parameter, and $\cal S$ is the regularizing functional. These parameters along with some fundamental design decisions are described in more detail below:

Distance: When performing the image registration, the algorithm tries to minimize the distance measure. There are many options for how to pick this distance measure, but the perhaps simplest and most common is the Sum of Squares Distance (SSD) given by

$$D_{\rm{SSD}}[{\mathcal{T}}[y],{\mathcal{R}}]=\frac{1}{2}\int_{\Omega}({\mathcal{T}}[y](x)-{\mathcal{R}}(x))^{2}dx.$$

(5)

The SSD measure penalizes different brightness levels. As a result, altering the contrast of the images through normalization can make the SSD measure perform much better.

Regularizer: When only focusing on the distance measure, the image registration program can come up with some unrealistic scenarios (for further discussion, see [6]). Sometimes it compresses the template image to fit within some small region of the reference. Other times it may create visible folds in the image. One way to counteract this is by introducing a regularizer.

Two commonly used regularizers implemented in FAIR are elastic and hyperelastic [4]. Both of these support the same idea but vary in severity. An elastic regularizer will try to keep the image registration close to the original orientation by introducing linear strain to any change. A hyperelastic regularizer does this with a nonlinear strain equation. Because of this, the elastic regularizer is only accurate for smaller displacements. For anything larger, the elastic regularizer risks becoming inaccurate and causing folding in the transformation grid. The more realistic, yet computationally more expensive hyperelastic regularizer also prevents folding. Either of these regularizers would work for the Chiari data, but hyperelastic was used since seems to remove more of the edge cases and its computational costs for two-dimensional problems is feasible.

Regularization Parameter: When constructing the objective functional from Equation 4, there is a trade-off between whether to minimize the distance or follow the regularizer. The parameter $\alpha$ controls the balance between allowing for displacement between the original and transformed image, to better the distance measure, or to penalizing displacement, to preserve the general shape of the image throughout the registration [9]. Note that in the objective functional equation, $\alpha$ corresponds to the regularization parameter. This coefficient was chosen to be 500 in our final algorithm based on trial and error, but the parameter selection could be further refined for better results.

Multilevel Optimization: As common in FAIR, Gauss Newton optimization was applied to a discretization of the objective functional in Equation 4 to generate a transformation from template to reference. Because discretizing objective functional on a 256x256 grid can lead to many local minima, and the convergence of this Gauss Newton to an appropriate minimum is dependent on an initial guess, a multilevel approach was used. This involves creating coarser discretizations of the problem and leverages the continuous image model that we can find through linear interpolation of the image data as detailed below.

We model the template and reference images as continuous functions, ${\mathcal{T}}:\Omega\to\mathds{R}$ and ${\mathcal{R}}:\Omega\to\mathds{R}$ where $\Omega\subseteq\mathds{R}^{2}$ corresponds to coordinates of pixels in the image grid. Coarser representations are created by changing the grid for which our interpolated image function is called, that contain fewer pixels: in our case, three additional representations, with 32x32, 64x64, and 128x128 pixels each. The optimization is applied in a way that moves up through these representations, applying optimization on the coarsest first and using the found transformation as the initial guess for the image at the next level.

The coarsest representation, described above, is also used in an initial pre-registration step where a parametric transformation is estimated to better align the template and reference images. This step, much like the step to choose appropriate template images, helps simplify the registration. One natural instinct for this rough transformation may be to simply translate and rotate the template image to align the reference. This would be a rigid transformation. There are advantages to keeping the transformations rigid, especially in cases where we expect little difference in the shapes of two images, but more flexibility can be added by using an affine transformation. This allows for additional linear transformations, such as shearing and scaling. We use affine transformations because of this flexibility to further align our images of often widely varied brain shapes.

This subsection compares two examples of how the registration performed — one showing a successful registration and another showing how it can fail.

Using the same parameters described in Section 3.2, a single registration was run on both patients. Figure 6 compares how well the masks were aligned before and after the registration. The transformations found for each patient were reasonable and proved to be somewhat successful. However, there were still some regions that required more attention. Particularly, the lower right of the cerebellum overextends in both patients.

In an attempt to smooth some of the outlying regions, the average program was run on both patients as shown in Figure 7 and Table 1. Tuning the parameters of the average program (as seen in Section 5.1) resulted in it performing better on most patients. Averaging this with other patients lessened the severity of this by eliminating the outliers that result from running the registration once.

Looking at the particular patients from Table 1, the averaging improved the similarity of patient 1 by around 2-3% but decreased the similarity of patient 2 by around the same amount (particularly in the brain stem). This could represent the need for a more definitive metric or a larger dataset; either of these solutions would aim to increase the quality of the chosen template.

Given an input image $I$, our goal is to estimate a three-dimensional tensor $P$ in the set

$$\Delta=\left\{Q\in\mathds{R}^{256\times 256\times 3}:Q_{i,j,k}\geq 0\text{ }\forall\text{ }i,j,k\text{ and }\sum_{k=0}^{2}Q_{i,j,k}=1\right\},$$

(6)

where $P_{i,j,k}$ is the probability the pixel at $i,j$ belongs to class $k$. The relationship between $I$ and $P$ can then be represented by a parameterized function

$$f_{\theta}:\mathds{R}^{256\times 256\times 3}\to\Delta,$$

(7)

where $\theta$ represents the weights of the model. To then create a mask $M_{p}$, each pixel is assigned to the class with the highest probability; i.e., we maximize over the third dimension in $P$.

Our model is a convolutional neural network (CNN). CNNs are a type of neural network which process data that have a “grid like topology,” such as image data (a 2D grid of pixels), and utilizes a convolution operator. Since the data in this study consist of multidimensional arrays of pixels and the goal is to correctly predict which class each pixel belongs to, a CNN is a natural choice. Many deep networks require thousands of images, however, our data set, as is common in biomedical applications, is much smaller. A CNN architecture that has been successful in other biomedical segmentation tasks with limited amounts of training data is the U-Net [13]. A U-Net consists of two main sections: the contracting path (encoder) and the expansive path (decoder) shown on the left and right side of Figure 8, respectively.

In the first layer of the contracting path, a 3x3 convolution is applied to each channel increasing the number of total feature channels. This is followed by a max pooling operation to break down the image resolution by decreasing image size to better identify features. A similar convolution structure is applied at each layer doubling the number feature channels. This leads to the bottom most layer where 1x1 convolutions are used, maintaining number of channels. Increasing the size of the image data, the expansive path (decoder) replaces the max pooling with upsampling operators. In the same formation as the contracting path, 2x2 up convolutions are applied except these halve the number of feature channels. During the process, high resolution features found during encoding are ‘copied and cropped’ and combined with the decoded output increasing dimension size and providing information that will be used in the convolution layer to produce more accurate outputs.

As the input images were run through the U-Net, the network used information from each channel to assign each pixel a probability for each class. These values are used to calculate the loss and to find the final predicted segmentation, $M_{p}$.

In order to get a successful model, we want the predicted mask $M_{p}$ and the target mask $M$ to be as similar as possible. The first step in maximizing similarities is to create a satisfactory model. This was done through supervised training of the network using a training set, $\{\left(I_{t},M_{t}\right)\}_{t=1,2,\ldots 41}$, a subset of the 51 DENSE MR images and corresponding masks $\{(I,M)\}$. To test the accuracy of the model on ‘new’ data, a validation set, $\{\left(I_{v},M_{v}\right)\}_{v=1,2,\ldots,10}$, was defined using the remaining 10 DENSE images and masks not included in training.

The U-Net was setup based on code provided by Aman Arora [2]. The network was then trained by looping a sequence of inputting $I_{t}$ to the U-Net, computing the loss between $f_{\theta}(I_{t})$ and $M_{t}$, and updating the model with the optimizer for a set number of iterations. In each iteration, the updated model was also run on the validation data to gauge its generalization; the validation loss was used to tune hyperparameters but was neither used to calculate gradients nor to change model parameters.

Comparing the probability matrix $P=f_{\theta}(I)$ to the corresponding $M$, for both training and validation data, shows the success of the model and can be numerically calculated using a loss function.

By maximizing the probability at the class corresponding to that of the known mask in the same position, the model error is minimized. In other words, our goal is to maximize

$$[f_{\theta}(I)]_{i,j,[M]_{i,j}},\quad\text{for all}\quad i=1,\ldots,256,j=1,\ldots,256,$$

(8)

which is equivalent to minimizing

$$-\log\left([f_{\theta}(I)]_{i,j,[M]_{i,j}}\right),\quad\text{for all}\quad i,j\in[1,256],$$

(9)

where $i,j$ corresponds to the pixel position.

Averaging Equation 9 over all training image sand all pixels results in the objective function

$${\cal{J}}_{CNN}(\theta)=\frac{1}{41\cdot 256^{2}}\sum_{t=1}^{41}\sum_{i,j=1}^{256}-\log\left([f_{\theta}(I_{e})]_{i,j,[M_{e}]_{i,j}}\right).$$

(10)

It can be seen that ${\cal{J}}_{CNN}$ equals the cross-entropy loss between $[f_{\theta}(I)]_{i,j,:}$ and the probability distribution defined by the standard basis vector, $[M]_{i,j}$, which calculates how far the model prediction is from the expected output.

Inputting $f_{\theta}(I)$ and $M$ in the loss function is a forward pass; this is done for both training and validation data. By executing a backwards pass on the loss, gradients can be calculated and used by the optimizer to update model parameters. As emphasized before, a backward pass was only executed with respect to the training data.

A limited-memory Broyden–Fletcher–Goldfarb–Shanno (lBFGS) algorithm from the PyTorch library [12] was used to optimize the overall loss function, shown previously in Equation 10. lBFGS is a Quasi-Newton method which uses a line search to automatically find the optimal learning rate $\epsilon$ [7] which satisfies the strong Wolfe conditions set in the algorithm parameters. The linesearch ensures decrease of the objective function while keeping $\epsilon$ from becoming too small.

The lBFGS algorithm uses the gradient calculated from the loss function and the learning rate found in the line search to update the model weights. The images were then, once again, run through the U-Net in hopes of finding a better segmentation. The model with the lowest validation loss was saved.

When creating the averaging program, two main parameters were chosen arbitrarily: $n$ and threshold. The $n$ refers to the number of template images to average and the threshold refers to which values should be counted when converting the average image to binary. In order to decide on a more definitive value for these parameters, a program was made to iterate through different possible values and find which performed the best.

This program was run over two metrics for how well it performed: Dice similarity and biomarker difference. Both of these metrics compare the atlas-based generated mask (average program) to the manually segmented mask and then average together all of the patients. These produced approximately the same results since the biomarker difference depends on how similar the masks are. A plot showing averaged Dice indices and absolute biomarker error over changes of both $n$ and the threshold can be seen in Figure 9.

To choose these $n$ and threshold parameters, we performed a grid search, first using a wide range of $n$ and threshold values (ranging from 5 to 40 and 10% to 90%, respectively). We chose to complete all further registrations with $n=10$ and a threshold of 50%, as this maximizes the average Dice index. The maximizer of the Dice index was chosen over the minimizer of the biomarker error because of how sensitive this biomarker measure is to slightly different segmentations. We expect that the behavior of the Dice index average will generalize better to other data.

While training the model, cross-entropy loss and Dice similarities were used as a measure of model success. Our final model was saved at the lowest validation loss found, which occurred at iteration 309, and was used to produce the final results shown below. Cross-entropy loss and Dice similarities throughout model training are plotted in Figure 10 alongside a comparison of true to predicted masks for validation data using the final model. Final results of the model used can be found in Table 2.

The model was then used on the test data previously set aside. Biomarkers were calculated using test data segmentations which were compared to the corresponding biomarkers produced in the atlas-based method and those from the known masks $M$.

When analyzing atlas-based (AB) and neural-network-based (NN) in the end, we relied on the Dice and biomarker (BM) metrics. Table 3 shows the samples provided in the testing set and compares these methods to the manually segmented mask (true). The masks that were compared are shown in Figure 11.

The bars chart in Figure 12 can offer more visual insight into which method did better on each patient. Some other representations and related results of Table 3 are shown in Table 4, Figure 13 and Figure 14.