A neural network for elastography (part 1)
Let's train a neural network to solve a typical problem encountered in elastography.
Elastography is a medical imaging modality that informs us about the stiffness of soft tissue. This information is relevant because many diseases form hard tissues. For instance, tumors, thrombus, and diseased liver tissue are usually stiffer than their surrounding tissues.
The basic idea is to:
- Apply some deformation in the tissue. It can be done manually, with a probe, or we could even use our heartbeat.
- Measure how much the tissue displaces. This is the displacement field, which can be measured with ultrasound, magnetic resonance imaging (MRI), etc.
- Calculate the stiffness of the tissue or some other elastic property.
The third step is called Inverse Problem. It receives this name because the usual problem is determining the displacement field from a known stiffness (this would be the Direct Problem), and the inverse problem does the opposite. There are mainly two ways to solve the inverse problem:
- Directly: we solve the governing equations of the problem. This approach is relatively fast and simple, but it suffers from the noise present in the displacement field measured.
- Iteratively: we keep updating the stiffness until we find one such as when we solve the direct problem, we obtain a displacement field very similar to the measured one. This approach is, in theory, more robust to noise, but it is much slower and more complicated to use in practice because it involves some parameter tweaking.
If you want to know more about elastography and these two approaches, check out Fovargue, Nordsletten and Sinkus (2018).
We will solve this inverse problem differently. We will use a neural network! Recently, some authors seem to be giving some attention to this approach. See, for instance, Kibria and Rivaz (2018) and Solamen, Shi and Amoh (2018). The main advantage is that this approach is fast (after the neural network is trained, of course) and robust to noise. A major limitation is the necessity of training data, but transfer learning could be a way to overcome it. Models could be pre-trained with synthetic data and then fine-tuned with a small number of real clinical data samples.
Creating synthetic data takes some time, but lucky for us, Lejeune (2020) did it. She built the Mechanical MNIST, which is a collection of datasets of finite element simulations, where heterogeneous blocks composed of Neo-Hookean material are subjected to some kind of loading. It can be a uniaxial extension, confined compression, shear, etc. The blocks are essentially a soft material with a hard inclusion, whose shape is a handwritten digit from the famous MNIST Dataset. The value of each pixel was converted to a Young modulus, which was used in the finite element simulations to obtain the displacement field. Young modulus is a measure of the stiffness of solid material, so let's call it just "stiffness".
Enough talk, let's code!
We will make our neural network with PyTorch and fastai. PyTorch is a widely used Python library developed by Facebook for machine and deep learning. Fastai is a deep learning library built on Pytorch. It provides high-level and low-level components that simplifies many aspects of a deep learning project.
It's very easy to run the code we show. Just click on the "Open in Colab" button on top of this page. It will open this page as a Colab notebook. PyTorch and fastai are not installed in Colab by default, so we need to install them manually.
!pip install torch
!pip install fastai
!pip install fastai --upgrade -q
Now we download and unzip the data that we will use for the training. We will use only the fifth step of the uniaxial extension from Mechanical MNIST.
The data contain the displacement field and handwritten digits from MNIST Dataset, which represents the heterogeneous material (hard inclusion in a soft background).
import os.path
if not os.path.exists('MNIST_input_files.zip'):
!wget https://open.bu.edu/bitstream/handle/2144/38693/MNIST_input_files.zip
if not os.path.exists('FEA_displacement_results_step5.zip'):
!wget https://open.bu.edu/bitstream/handle/2144/38693/FEA_displacement_results_step5.zip
!unzip -n "MNIST_input_files.zip" -d "." # replace "-n" with "-o" if you wish to overwrite existing files
!unzip -n "FEA_displacement_results_step5.zip" -d "uniaxial/"
Let's import some libraries.
from pathlib import Path
import pandas as pd
import torch
from fastai.vision.all import *
We load the data from the text files to PyTorch tensors. Don't worry, the data is small.
df = pd.read_csv('MNIST_input_files/mnist_img_train.txt', sep=' ', header=None)
imgs_train = torch.tensor(df.values)
df = pd.read_csv('MNIST_input_files/mnist_img_test.txt', sep=' ', header=None)
imgs_valid = torch.tensor(df.values)
imgs_train.shape, imgs_valid.shape
df = pd.read_csv('uniaxial/FEA_displacement_results_step5/summary_dispx_train_step5.txt', sep=' ', header=None)
ux5s_train = torch.tensor(df.values)
df = pd.read_csv('uniaxial/FEA_displacement_results_step5/summary_dispy_train_step5.txt', sep=' ', header=None)
uy5s_train = torch.tensor(df.values)
df = pd.read_csv('uniaxial/FEA_displacement_results_step5/summary_dispx_test_step5.txt', sep=' ', header=None)
ux5s_valid = torch.tensor(df.values)
df = pd.read_csv('uniaxial/FEA_displacement_results_step5/summary_dispy_test_step5.txt', sep=' ', header=None)
uy5s_valid = torch.tensor(df.values)
ux5s_train.shape, uy5s_train.shape, ux5s_valid.shape, uy5s_valid.shape
We reshape the tensors to show them as regular images with one channel.
ux5s_train = torch.reshape(ux5s_train, (-1,1,28,28))
uy5s_train = torch.reshape(uy5s_train, (-1,1,28,28))
ux5s_valid = torch.reshape(ux5s_valid, (-1,1,28,28))
uy5s_valid = torch.reshape(uy5s_valid, (-1,1,28,28))
imgs_train = torch.reshape(imgs_train, (-1,1,28,28))
imgs_valid = torch.reshape(imgs_valid, (-1,1,28,28))
ux5s_train.shape, uy5s_train.shape, ux5s_valid.shape, uy5s_valid.shape, imgs_train.shape, imgs_valid.shape
Let's take a look at our data.
show_image(imgs_train[2], cmap='Greys')
show_image(uy5s_train[2], cmap='Greys')
The displacement field seems inverted. Let's fix it.
ux5s_train = torch.flip(ux5s_train, dims=[1,2])
uy5s_train = torch.flip(uy5s_train, dims=[1,2])
ux5s_valid = torch.flip(ux5s_valid, dims=[1,2])
uy5s_valid = torch.flip(uy5s_valid, dims=[1,2])
show_image(imgs_train[2], cmap='Greys')
show_image(uy5s_train[2], cmap='Greys')
Ok, much better.
Now let's convert the pixel values (0-255) to stiffness values (0.01-1). This step is not really necessary. Our network could directly predict pixel values instead of stiffness.
The original formula from Lejeune (2020) is $$ \mathrm{stiffness} = \frac{\mathrm{pixel}}{255}\, 99 + 1, $$ but we include a division by 100 because it is convenient for visualization purposes. Again, we could skip this conversion.
imgs_train = (imgs_train / 255.) * 99. + 1.
imgs_valid = (imgs_valid / 255.) * 99. + 1.
imgs_train = imgs_train / 100.
imgs_valid = imgs_valid / 100.
imgs_train.min(), imgs_valid.min(), imgs_train.max(), imgs_valid.max()
Let's merge the horizontal and vertical displacements into one tensor, so each component of the displacement field is in a separate channel.
us_train = torch.cat([ux5s_train, uy5s_train], dim=1)
us_valid = torch.cat([ux5s_valid, uy5s_valid], dim=1)
us_train.shape, us_valid.shape
We do some cleaning.
del ux5s_train; del uy5s_train; del ux5s_valid; del uy5s_valid
Now, we create fastai DataLoaders from a low-level API.
The first step is to create our Transform. Fastai Transforms are classes that modify the dataset when the batch is being prepared.
Our Transform is simple. It receives an integer i and returns a tuple (x,y), where x and y are the i-th element of us_train and imgs_train, respectively. Our Transform also normalizes these elements.
class GetNormalizedData(Transform):
def __init__(self, us, imgs, mean, std):
self.us, self.imgs = us, imgs
self.mean, self.std = mean, std
def encodes(self, i):
us_normalized = torch.true_divide((self.us[i] - self.mean.view(2,1,1)), self.std.view(2,1,1))
return (us_normalized.float(), self.imgs[i].float())
We calculate the mean and the standard deviation of our data.
us_mean = torch.mean(us_train, dim=[0,2,3])
us_std = torch.std(us_train, dim=[0,2,3])
We create a TfmdLists (Transformed Lists), which is just an object that will apply some Transforms in a list when the time is right. It is lazy.
So it is like the range(...) is our dataset, but the Transform will receive the elements of this range(...) and return a tuple of the actual data.
train_tl= TfmdLists(range(len(us_train)),
GetNormalizedData(us_train, imgs_train, us_mean, us_std)
)
valid_tl= TfmdLists(range(len(us_valid)),
GetNormalizedData(us_valid, imgs_valid, us_mean, us_std)
)
Once we have the TfmdLists, we can easily create the Dataloaders. If we have access to a GPU, it will use it.
dls = DataLoaders.from_dsets(train_tl, valid_tl, bs=64)
if torch.cuda.is_available(): dls = dls.cuda()
Let's take a look at a batch. The mean is next to zero and the standard deviation to one, it seems right.
x,y = dls.one_batch()
x.shape, y.shape, x.mean(), x.std()
From the images we have seen above, our task does not seem difficult.
Let's use a simple architecture. It is just a CNN with an input layer and an output layer. We know that the targets range from 0.01 to 1, so we also apply a scaled sigmoid to return values between 0 and 1.1. The target values are numbers between 0 and 1, so why we multiply the sigmoid layer by 1.1? We do it because, otherwise, the preceding layer would have to return a very high value to the output of the sigmoid layer be equal to 1.
We choose the parameters of the convolutions so that the activations have shape a (n,28,28), where n is free to be any number.
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.conv1 = nn.Conv2d(2, 4, 5, 1, 2, padding_mode='reflect', bias=True)
self.conv2 = nn.Conv2d(4, 1, 5, 1, 2, padding_mode='reflect', bias=True)
self.sigmoid = nn.Sigmoid()
self.relu = nn.ReLU()
def forward(self, x):
h = self.relu(self.conv1(x))
return self.sigmoid(self.conv2(h)) * 1.1
We create a fastai Learner. This object groups the data (dataloaders), the architecture, the loss function, and the optimizer (ADAM by default). It has everything that we need to start the training.
l = Learner(dls, CNN(), loss_func=F.mse_loss, model_dir='')
We find a good learning rate. lr_find trains the model a bit with exponentially growing learning rates and report the loss. This technique was introduced by Smith (2017). We choose 0.01 because it is close to the point where the slope is the steepest.
l.lr_find()
Now we train!
fit_one_cycle uses the 1cycle policy proposed by Smith and Topin (2017). It's a training strategy where the learning rate increases until it reaches lr_max, and then it decreases. The momentum does the opposite. Check out the documentation for more details.
l.fit_one_cycle(10,lr_max=1e-2)
We can save the model.
l.path = Path('')
l.save('models/01_cnn_ep10')
Let's check the predictions of the model.
preds, targets, losses = l.get_preds(with_loss=True)
On the left, we have the predictions, and on the right, the targets.
for i in range(0,10):
show_image(torch.cat([preds, torch.ones(preds.shape[0],1,28,1), targets], dim=3)[i], cmap='Greys_r')
It looks good. Let's now check the predictions that had the top losses.
_, indices = torch.sort(losses, descending=True)
for i in indices[:10]:
show_image(torch.cat([preds, torch.ones(preds.shape[0],1,28,1), targets], dim=3)[i], cmap='Greys_r')
It looks good, as well. Now, we check the lowest losses.
for i in indices[-10:]:
show_image(torch.cat([preds, torch.ones(preds.shape[0],1,28,1), targets], dim=3)[i], cmap='Greys_r')
Okay, it seems like our model is doing a good job.
Lejeune (2020) also made a model for the same task. She evaluated her model using the mean and the standard deviation of the absolute error of the test set (which is our validation set). She constructed a simple FNN with scikit-learn and obtained a mean equals 13 and a standard deviation equals 79, in pixel values.
Let's calculate these metrics for our model.
px_preds = ((preds*100)-1)*(255/99)
px_targets = ((targets*100)-1)*(255/99)
px_error = torch.abs(px_preds - px_targets)
px_mean = torch.mean(px_error)
px_std = torch.std(px_error)
px_mean, px_std