You could try an earth mover distance in 2d or 3d over the image? For example you could do this, but call sequentially 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.
