Lets say the output to my model is a stack of images of size [B, C, M, N].
I wish to binarize each image channel and apply a skimage function such as remove_small_objects to the images, and then apply a loss function to those images.
What would be the best way for me to use one of these function to the tensors images and still be able to keep the gradients and use loss functions?
As an extension to this, how could one use a function skimage.measure library to that produces an array of centroids and then apply a loss function to the prediction and target centroids that propagates back into the model parameters?