Apply a skimage (or any) function to output before loss

vision
1

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?

2

Since skimage uses numpy under the hood, if I’m not mistaken, you would have to use a custom autograd.Function and implement the backward method manually.
Autograd won’t be able to create the backward pass automatically.

I’m not sure which operations are applied for the mentioned operations, but you might be able to implement them using PyTorch operations. If that’s the case, you wouldn’t need to write the backward manually.