Max Pooling across MRI Slices

Question

I am trying to implement a Machine Learning Model for MRI scan diagnosis. I have Inputs of shape (x, 256, 256, 3), where we have 3 color channels and where x is the number of slices in a sequence. I read the MRNet paper and I want to implement a similar architecture in TensorFlow Keras. Instead of using the AlexNet feature extractor, I'd like to use VGG16.

The model architecture in the paper:

The primary building block of our prediction system is MRNet, a convolutional neural network (CNN) mapping a 3-dimensional MRI series to a probability [15] (Fig 2). The input to MRNet has dimensions s × 3 × 256 × 256, where s is the number of images in the MRI series (3 is the number of color channels). First, each 2-dimensional MRI image slice was passed through a feature extractor based on AlexNet to obtain a s × 256 × 7 × 7 tensor containing features for each slice. A global average pooling layer was then applied to reduce these features to s × 256. We then applied max pooling across slices to obtain a 256-dimensional vector, which was passed to a fully connected layer and sigmoid activation function to obtain a prediction in the 0 to 1 range.

So far so good. I have a sequential model, added the feature extractor as the first step, then apply a GlobalAveragePooling2D() to reduce features to shape (x, 512). Then I must MaxPool across the slices but I have no approach for this problem.

feature_extractor = VGG16(weights='imagenet', include_top=False, input_shape=(256, 256, 3))
model = Sequential()
model.add(feature_extractor)         #output shape: (x, 8, 8, 512)
model.add(GlobalAveragePooling2D())  #output shape: (x, 512)
# Here i have to add a Layer witch Pools over the slices.
model.add(                         )  #output shape(1, 512)

model.add(Dense(1, activation='sigmoid'))
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

The example Scan has a shape of (44, 256, 256, 3). When it runs through the VGG16 its features have the Dimension of (44, 8, 8, 512). After GlobalAverage Pooling I got (44, 512). This 2-D Array must then somehow be transformed into the shape of (1, 512). I mean if I do the operation on a simple 2-D NumPy array I need a function like np.max over the 0-axis

np.max(x, axis=0)

Maybe you can give me a hint or have an approach for this. Thanks a lot for your help :)

################################################################################ Edit: 01.05.2021

I played around with your approach @Aaron Keesing, but fitting the model does not train it at all somehow. Ater 25 epochs i still have the same accuracy. And the accuracy is the distribution over my 2 classes (i was just training on coronal plane and abnormal)

In this case for example i have 500 cases, and 80% of the cases do have an abnormality and 20% doesn't.

# Dataset train, overall 500 cases
Absolute:
 abnormal  acl  meniscus
1         0    0           184
               1           118
0         0    0           100
1         1    1            63
               0            35
dtype: int64
Relative:
 abnormal  acl  meniscus
1         0    0           0.368
               1           0.236
0         0    0           0.200
1         1    1           0.126
               0           0.070

###########################################################
# Dataset valid, overall 100 cases
Absolute:
 abnormal  acl  meniscus
1         1    1           27
0         0    0           25
1         1    0           23
          0    0           20
               1            5
dtype: int64
Relative:
 abnormal  acl  meniscus
1         1    1           0.27
0         0    0           0.25
1         1    0           0.23
          0    0           0.20
               1           0.05

user_10044861 · Answer

You should be able to use a GlobalAveragePooling1D layer. But note that it expects a batch dimension. Since you're inputting a sequence of images, your input should be 5-dimensional with the first dimension the batch_size (which can be 1).

I don't think the image CNNs work with 5D input, so you might be able to use a TimeDistributed layer to apply to the image sequence, which will give you your sequence of features of shape (x, 512), then apply GlobalAveragePooling1D to get a final feature vector.

So perhaps something like this might work. Note that you must specify the number of images in the sequence, x (which could be None):

vgg16 = VGG16(weights='imagenet', include_top=False, input_shape=(256, 256, 3))
feature_extractor = Sequential()
feature_extractor.add(vgg16)         #output shape: (bs, 8, 8, 512)
feature_extractor.add(GlobalAveragePooling2D())  #output shape: (bs, 512)

model = Sequential()
model.add(TimeDistributed(feature_extractor, input_shape=(x, 256, 256, 3)))  #output shape(bs, x, 512)
# Here i have to add a Layer witch Pools over the slices.
model.add(GlobalAveragePooling1D())   #output shape(bs, 512)
model.add(Dense(1, activation='sigmoid'))   #output shape(bs, 1)

You can just put a batch of one MRI sequence at a time so that bs = 1.

This yields the following model structure with x = None:

Model: "sequential_1"
_________________________________________________________________
Layer (type)                 Output Shape              Param #
=================================================================
time_distributed (TimeDistri (None, None, 512)         14714688
_________________________________________________________________
global_average_pooling1d (Gl (None, 512)               0
_________________________________________________________________
dense (Dense)                (None, 1)                 513
=================================================================
Total params: 14,715,201
Trainable params: 14,715,201
Non-trainable params: 0
_________________________________________________________________

user_8013686 · Answer

After thinking about this problem I have found one solution which could work

    vgg16 = VGG16(weights='imagenet', include_top=False, input_shape=(256, 256, 3))  #

    average_pool = Sequential(name='AveragePool')
    average_pool.add(layers.AveragePooling2D(input_shape=(8, 8, 512)))
    average_pool.add(layers.Flatten())
    self.average_pool = average_pool
    
    self.model = Sequential([
        vgg16,
        average_pool], name='MyModel')
    self.model.summary()
    
    # Max-pooling
    self.model.add(Dense(256, activation='relu', kernel_constraint=constraints.MaxNorm(max_value=2, axis=0)))
    self.model.add(Dense(1, activation='sigmoid'))

    self.model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

This leads to the following summary

Model: "AveragePool"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
average_pooling2d (AveragePo (None, 4, 4, 512)         0         
_________________________________________________________________
flatten (Flatten)            (None, 8192)              0         
=================================================================
Total params: 0
Trainable params: 0
Non-trainable params: 0
_________________________________________________________________
Model: "MyModel"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
vgg16 (Model)                (None, 8, 8, 512)         14714688  
_________________________________________________________________
AveragePool (Sequential)     (None, 8192)              0         
_________________________________________________________________
dense (Dense)                (None, 256)               2097408   
_________________________________________________________________
dense_1 (Dense)              (None, 1)                 257       
=================================================================
Total params: 16,812,353
Trainable params: 16,812,353
Non-trainable params: 0

If you got any other ideas or ideas for improvement, let me know!

Max Pooling across MRI Slices

2 Answers