You could try an earthmovers distance (https://en.m.wikipedia.org/wiki/Earth_mover%27s_distance) in 2d or 3d over the image? For example you could do this, but call sequentially (https://discuss.pytorch.org/t/implementation-of-squared-earth-movers-distance-loss-function-for-ordinal-scale/107927/2) The idea would be something like this (untested and written on my cell phone):
def cumsum_3d(a):
a = torch.cumsum(a, -1)
a = torch.cumsum(a, -2)
a = torch.cumsum(a, -3)
return a
def norm_3d(a):
return a / torch.sum(a, dim=(-1,-2,-3), keepdim=True)
def emd_3d(a, b):
a = norm_3d(a)
b = norm_3d(b)
return torch.mean(torch.square(cumsum_3d(a) - cumsum_3d(b)), dim=(-1,-2,-3))
This should also work with batched data. I would also try normalizing the images first (so they each sum to 1) unless you want to account for changes in intensity.