def spatial_soft_argmax2d(
input: torch.Tensor,
temperature: torch.Tensor = torch.tensor(1.0),
normalized_coordinates: bool = True,
eps: float = 1e-8) -> torch.Tensor:
r"""Function that computes the Spatial Soft-Argmax 2D
of a given input heatmap.
Returns the index of the maximum 2d coordinates of the give map.
The output order is x-coord and y-coord.
Arguments:
temperature (torch.Tensor): factor to apply to input. Default is 1.
normalized_coordinates (bool): whether to return the
coordinates normalized in the range of [-1, 1]. Otherwise,
it will return the coordinates in the range of the input shape.
Default is True.
eps (float): small value to avoid zero division. Default is 1e-8.
Shape:
- Input: :math:`(B, N, H, W)`
- Output: :math:`(B, N, 2)`
Examples:
>>> input = torch.tensor([[[
[0., 0., 0.],
[0., 10., 0.],
[0., 0., 0.]]]])
>>> coords = kornia.spatial_soft_argmax2d(input, False)
tensor([[[1.0000, 1.0000]]])
"""
input_soft: torch.Tensor = dsnt.spatial_softmax2d(input, temperature)
output: torch.Tensor = dsnt.spatial_expectation2d(input_soft, normalized_coordinates)
return output
def spatial_soft_argmax2d(
input: torch.Tensor,
temperature: torch.Tensor = torch.tensor(1.0),
normalized_coordinates: bool = True,
eps: float = 1e-8,
) -> torch.Tensor:
r"""Function that computes the Spatial Soft-Argmax 2D of a given input heatmap.
Args:
input: the given heatmap with shape :math:`(B, N, H, W)`.
temperature: factor to apply to input.
normalized_coordinates: whether to return the coordinates normalized in the range of :math:`[-1, 1]`.
Otherwise, it will return the coordinates in the range of the input shape.
eps: small value to avoid zero division.
Returns:
the index of the maximum 2d coordinates of the give map :math:`(B, N, 2)`.
The output order is x-coord and y-coord.
Examples:
>>> input = torch.tensor([[[
... [0., 0., 0.],
... [0., 10., 0.],
... [0., 0., 0.]]]])
>>> spatial_soft_argmax2d(input, normalized_coordinates=False)
tensor([[[1.0000, 1.0000]]])
"""
input_soft: torch.Tensor = dsnt.spatial_softmax2d(input, temperature)
output: torch.Tensor = dsnt.spatial_expectation2d(input_soft, normalized_coordinates)
return output
def forward(self, inputs: th.Tensor) -> Tuple[th.Tensor, th.Tensor]:
"""NCHW -> NC2."""
# Convert logits to probabilities.
# TODO(ycho): Consider if this is preferable to elementwise sigmoid.
prob = spatial_softmax2d(inputs, temperature=self.temperature)
kpts = spatial_expectation2d(prob, normalized_coordinates=True)
return (prob, kpts)