def test_bottom_right(self, device, dtype):
input = torch.zeros(1, 1, 2, 3, device=device, dtype=dtype)
input[..., -1, -1] = 1e16
coord = kornia.spatial_soft_argmax2d(input, normalized_coordinates=False)
assert_allclose(coord[..., 0].item(), 2.0, atol=1e-4, rtol=1e-4)
assert_allclose(coord[..., 1].item(), 1.0, atol=1e-4, rtol=1e-4)
def spatial_soft_argmax2d(input,
temperature,
normalized_coordinates,
eps=1e-08):
""" Docs of kornia.spatial_soft_argmax2d:
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, 1, False)
tensor([[[1.0000, 1.0000]]])
"""
coors = kornia.spatial_soft_argmax2d(input, temperature,
normalized_coordinates, eps)
# kornia 基于 PyTorch 提供可微分操作,所以要 request_grad = False
coors = coors.detach()
return coors
def test_bottom_right_normalized(self):
input = torch.zeros(1, 1, 2, 3)
input[..., -1, 1] = 10.
coord = kornia.spatial_soft_argmax2d(input, False)
assert pytest.approx(coord[..., 0].item(), 2.0)
assert pytest.approx(coord[..., 1].item(), 1.0)
def test_top_left_normalized(self, device, dtype):
input = torch.zeros(1, 1, 2, 3, device=device, dtype=dtype)
input[..., 0, 0] = 1e16
coord = kornia.spatial_soft_argmax2d(input, normalized_coordinates=True)
assert_allclose(coord[..., 0].item(), -1.0, atol=1e-4, rtol=1e-4)
assert_allclose(coord[..., 1].item(), -1.0, atol=1e-4, rtol=1e-4)
def test_top_left_normalized(self):
input = torch.zeros(1, 1, 2, 3)
input[..., 0, 0] = 10.
coord = kornia.spatial_soft_argmax2d(input, False)
assert pytest.approx(coord[..., 0].item(), 0.0)
assert pytest.approx(coord[..., 1].item(), 0.0)
def test_bottom_right_normalized(self, device):
input = torch.zeros(1, 1, 2, 3).to(device)
input[..., -1, 1] = 10.
coord = kornia.spatial_soft_argmax2d(input, False)
assert_allclose(coord[..., 0].item(), 2.0)
assert_allclose(coord[..., 1].item(), 1.0)
def test_top_left(self, device):
input = torch.zeros(1, 1, 2, 3).to(device)
input[..., 0, 0] = 10.
coord = kornia.spatial_soft_argmax2d(input, True)
assert_allclose(coord[..., 0].item(), -1.0)
assert_allclose(coord[..., 1].item(), -1.0)
def test_top_left(self, device):
input = torch.zeros(1, 1, 2, 3).to(device)
input[..., 0, 0] = 1e16
coord = kornia.spatial_soft_argmax2d(input,
normalized_coordinates=False)
assert_allclose(coord[..., 0].item(), 0.0)
assert_allclose(coord[..., 1].item(), 0.0)
def test_jit(self):
@torch.jit.script
def op_script(input: torch.Tensor, temperature: torch.Tensor,
normalize_coords: bool, eps: float) -> torch.Tensor:
return kornia.spatial_soft_argmax2d(input, temperature,
normalize_coords, eps)
input = torch.rand(1, 2, 3, 4)
actual = op_script(input, torch.tensor(1.0), True, 1e-8)
expected = kornia.spatial_soft_argmax2d(input)
assert_allclose(actual, expected)
def test_batch2_n2(self, device):
input = torch.zeros(2, 2, 2, 3).to(device)
input[0, 0, 0, 0] = 1e16 # top-left
input[0, 1, 0, -1] = 1e16 # top-right
input[1, 0, -1, 0] = 1e16 # bottom-left
input[1, 1, -1, -1] = 1e16 # bottom-right
coord = kornia.spatial_soft_argmax2d(input)
assert_allclose(coord[0, 0, 0].item(), -1.0) # top-left
assert_allclose(coord[0, 0, 1].item(), -1.0)
assert_allclose(coord[0, 1, 0].item(), 1.0) # top-right
assert_allclose(coord[0, 1, 1].item(), -1.0)
assert_allclose(coord[1, 0, 0].item(), -1.0) # bottom-left
assert_allclose(coord[1, 0, 1].item(), 1.0)
assert_allclose(coord[1, 1, 0].item(), 1.0) # bottom-right
assert_allclose(coord[1, 1, 1].item(), 1.0)
def test_batch2_n2(self):
input = torch.zeros(2, 2, 2, 3)
input[0, 0, 0, 0] = 10. # top-left
input[0, 1, 0, -1] = 10. # top-right
input[1, 0, -1, 0] = 10. # bottom-left
input[1, 1, -1, -1] = 10. # bottom-right
coord = kornia.spatial_soft_argmax2d(input)
assert pytest.approx(coord[0, 0, 0].item(), -1.0) # top-left
assert pytest.approx(coord[0, 0, 1].item(), -1.0)
assert pytest.approx(coord[0, 1, 0].item(), 1.0) # top-right
assert pytest.approx(coord[0, 1, 1].item(), -1.0)
assert pytest.approx(coord[1, 0, 0].item(), -1.0) # bottom-left
assert pytest.approx(coord[1, 0, 1].item(), 1.0)
assert pytest.approx(coord[1, 1, 0].item(), 1.0) # bottom-right
assert pytest.approx(coord[1, 1, 1].item(), 1.0)
def test_batch2_n2(self, device, dtype):
input = torch.zeros(2, 2, 2, 3, device=device, dtype=dtype)
input[0, 0, 0, 0] = 1e16 # top-left
input[0, 1, 0, -1] = 1e16 # top-right
input[1, 0, -1, 0] = 1e16 # bottom-left
input[1, 1, -1, -1] = 1e16 # bottom-right
coord = kornia.spatial_soft_argmax2d(input)
assert_close(coord[0, 0, 0].item(), -1.0, atol=1e-4,
rtol=1e-4) # top-left
assert_close(coord[0, 0, 1].item(), -1.0, atol=1e-4, rtol=1e-4)
assert_close(coord[0, 1, 0].item(), 1.0, atol=1e-4,
rtol=1e-4) # top-right
assert_close(coord[0, 1, 1].item(), -1.0, atol=1e-4, rtol=1e-4)
assert_close(coord[1, 0, 0].item(), -1.0, atol=1e-4,
rtol=1e-4) # bottom-left
assert_close(coord[1, 0, 1].item(), 1.0, atol=1e-4, rtol=1e-4)
assert_close(coord[1, 1, 0].item(), 1.0, atol=1e-4,
rtol=1e-4) # bottom-right
assert_close(coord[1, 1, 1].item(), 1.0, atol=1e-4, rtol=1e-4)
def op_script(input: torch.Tensor, temperature: torch.Tensor,
normalize_coords: bool, eps: float) -> torch.Tensor:
return kornia.spatial_soft_argmax2d(input, temperature,
normalize_coords, eps)
