I’m training a 3D image classification network, and would like to modify the output of the network before the loss function is calculated. However, it’s not clear to me how best to do this in Tensorflow, so I’m looking for advice.
The output of the network (a U-net) is a one-hot-encoded 3D volume. I would like to dilate the predicted segmentation by an amount that depends on the location of each pixel within the volume. The image transforms I need to use are available in numpy (and/or scipy), but it’s not clear to me how best to incorporate that functionality between the network output and the loss function. I’ve written a custom model class with a train_step method that calls the loss function, but there doesn’t seem to be an clear way to move between tensors and numpy arrays. Or is there a different aproach I could be taking…?