How to implement random cropping during training?

user_15860992 0
Viewed 265

I'm developing a U-net like model which segments the damaged tissue of the brain between two time-points in Multiple Sclerosis patients. The model is given the baseline and follow-up images as x and the segmentation mask as y. The images are 3D (192, 218, 192) and the model input size is (128, 128, 128) and is being trained by the dice loss. The model is based on this paper: https://www.sciencedirect.com/science/article/abs/pii/S0895611120300732

What I'm currently doing is a centre crop for each image before the training, and the training dice score seems to be learning properly, but in the validation data it is not.

I have read that random cropping 10 times or so helps to reduce overfitting, but I don't exactly know how to implement it. I have written this function to do the random cropping:

def random_crop(img_bl, img_fu, mask, width=128, height=128, depth=128):
    x = random.randint(0, img_bl.shape[1] - width)
    y = random.randint(0, img_bl.shape[0] - height)
    z = random.randint(0, img_bl.shape[2] - depth)
    img_bl = img_bl[y:y + height, x:x + width, z:z + depth]
    img_fu = img_fu[y:y + height, x:x + width, z:z + depth]
    mask = mask[y:y + height, x:x + width, z:z + depth]

    return img_bl, img_fu, mask

Do I just apply this function 10 times before the training for each image? Or is there a way to include the random cropping inside the model and overlap the predictions of these 10 subvolumes?

data-augmentation conv-neural-network tensorflow deep-learning python
1 Answers

You can crop images in Tensorflow by the tf.image.random_crop function.

cropped=tf.image.random_crop(image, (PATCH_SIZE,PATCH_SIZE,3))

To crop 3D models in your case, you can set the channel parameter also

cropped=tf.image.random_crop(model, (PATCH_SIZE_1,PATCH_SIZE_2,PATCH_SIZE_3))

To beautifully apply this cropping to your input pipeline, you can map this function using tf.keras.layers.Lambda,

def crop_layer():
  def _crop_layer(x):
    return cropped=tf.image.random_crop(x, (PATCH_SIZE_1,PATCH_SIZE_2,PATCH_SIZE_3))
  return lambda x:_crop_layer(x)

layer=tf.keras.layers.Lambda(crop_layer)(input)
...
user_112454750
Related