Segmentation 3D Unet checkerboard artifacts in slices above and below segmentation voxels

Question

I suppose an image is worth too many words, so here is the image:

As you can see, in the middle where there are voxels to be segmented, no artifacts are present. Whereas on the top and bottom I get a nice checkerboard artifact. This is the prediction. The model was trained on three 3D images, and the validation set consisted of the same three images. This training was an overfitting test. It overfits well, but the artifacts at the top and bottom are present. The slices at the top and bottom with the artifacts have no segmented pixels, which leads me to believe that the neural network comes up with fake data for the slices where there is nothing to come up with. I just don't understand why and how to get rid of it. Any ideas?

Edit: I believe it has less to do with the checkerboard artifact than with the fact that the NN inserts data into slices that have no information, i.e., the voxels in the slices on the top and bottom are all zeros (0). I got rid of the checkerboard pattern using this method: https://arxiv.org/pdf/2002.02117.pdf (TL;DR: add a fixed convolutional layer right after the transposed convolution). The change reduced the checkerboard pattern but I'm still getting the filled in top and bottom parts:

Answer 1

I've found the solution.

In order to train the model and thus segment the lesions, the label has to have two classes: (1) every voxel that denotes the lesion, and (2) every voxel that is not a lesion. The image label provided only contained voxels with lesions. To obtain two classes, I used tensorflow.keras.utils.to_categorical. The model trains just fine and the output is a softmax resulting in values from 0 to 1.

The output shows the segmentation without reverting from the two categories. Each slice is plotted individually and the intensity is scaled according to its maximum value. The slices with the artifacts have really low values (e.g., 0.000001834), while those with no artifacts (those towards the middle) have much higher values (e.g., 0.3421). Thus, really low intensity artifacts are shown as being the highest values.

The solution is reverting from tensorflow.keras.utils.to_categorical. We would need to get all relevant voxels that have very high values and turn these into ones while turning those irrelevant voxels (i.e., artifacts) into zeros.

This is how to do it:

predicted = np.argmax(model.predict(data_to_predict), axis=-1)

The numpy function argmax() returns the indices of the maximum values along an axis. We return them along the last axis which is the axis of the two channels (no lesion, lesion).