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)