def spatial_gradient3d(input: torch.Tensor,
mode: str = 'diff',
order: int = 1,
normalized: bool = True) -> torch.Tensor:
r"""Computes the first and second order volume derivative in x, y and d using a diff
operator.
Args:
input (torch.Tensor): input features tensor with shape :math:`(B, C, D, H, W)`.
mode (str): derivatives modality, can be: `sobel` or `diff`. Default: `diff`.
order (int): the order of the derivatives. Default: 1.
Return:
torch.Tensor: the spatial gradients of the input feature map.
Shape:
- Input: :math:`(B, C, D, H, W)`. D, H, W are spatial dimensions, gradient is calculated w.r.t to them.
- Output: :math:`(B, C, 3, D, H, W)` or :math:`(B, C, 6, D, H, W)`
Examples:
>>> input = torch.rand(1, 4, 2, 4, 4)
>>> output = spatial_gradient3d(input)
>>> output.shape
torch.Size([1, 4, 3, 2, 4, 4])
"""
if not isinstance(input, torch.Tensor):
raise TypeError("Input type is not a torch.Tensor. Got {}".format(
type(input)))
if not len(input.shape) == 5:
raise ValueError(
"Invalid input shape, we expect BxCxDxHxW. Got: {}".format(
input.shape))
# allocate kernel
kernel: torch.Tensor = get_spatial_gradient_kernel3d(mode, order)
if normalized:
kernel = normalize_kernel3d(kernel)
# prepare kernel
b, c, d, h, w = input.shape
tmp_kernel: torch.Tensor = kernel.to(input).detach()
tmp_kernel = tmp_kernel.repeat(c, 1, 1, 1, 1)
# convolve input tensor with grad kernel
kernel_flip: torch.Tensor = tmp_kernel.flip(-3)
# Pad with "replicate for spatial dims, but with zeros for channel
spatial_pad = [
kernel.size(2) // 2,
kernel.size(2) // 2,
kernel.size(3) // 2,
kernel.size(3) // 2,
kernel.size(4) // 2,
kernel.size(4) // 2
]
out_ch: int = 6 if order == 2 else 3
return F.conv3d(F.pad(input, spatial_pad, 'replicate'),
kernel_flip,
padding=0,
groups=c).view(b, c, out_ch, d, h, w)
def __init__(self,
mode: str = 'diff',
order: int = 1) -> None:
super(SpatialGradient3d, self).__init__()
self.order: int = order
self.mode: str = mode
self.kernel = get_spatial_gradient_kernel3d(mode, order)
return