How to use kornia - 10 common examples

def test_jit(self, device):
        @torch.jit.script
        def op_script(input: torch.Tensor, height: int,
                      width: int) -> torch.Tensor:
            return kornia.denormalize_pixel_coordinates(input, height, width)
        height, width = 3, 4
        grid = kornia.utils.create_meshgrid(
            height, width, normalized_coordinates=True).to(device)

        actual = op_script(grid, height, width)
        expected = kornia.denormalize_pixel_coordinates(
            grid, height, width)

        assert_allclose(actual, expected)

    @pytest.mark.skip(reason="turn off all jit for a while")
    def test_jit(self, device):
        op = kornia.denormalize_pixel_coordinates
        op_script = torch.jit.script(op)

        height, width = 3, 4
        grid = kornia.utils.create_meshgrid(
            height, width, normalized_coordinates=True).to(device)

        actual = op_script(grid, height, width)
        expected = op(grid, height, width)

        assert_allclose(actual, expected)

def test_jit(self):
        batch_size, channels, height, width = 2, 3, 64, 64
        img = torch.ones(batch_size, channels, height, width)
        gray = kornia.color.RgbToGrayscale()
        gray_traced = torch.jit.trace(kornia.color.RgbToGrayscale(), img)
        assert_allclose(gray(img), gray_traced(img))

def test_jit(self):
        batch_size, channels, height, width = 2, 3, 64, 64
        img = torch.ones(batch_size, channels, height, width)
        gray = kornia.color.RgbToGrayscale()
        gray_traced = torch.jit.trace(kornia.color.RgbToGrayscale(), img)
        assert_allclose(gray(img), gray_traced(img))

def test_jit(self, device):
        @torch.jit.script
        def op_script(input, k):
            return kornia.feature.harris_response(input, k)
        k = torch.tensor(0.04)
        img = torch.rand(2, 3, 4, 5, device=device)
        actual = op_script(img, k)
        expected = kornia.feature.harris_response(img, k)
        assert_allclose(actual, expected)

    @pytest.mark.skip(reason="turn off all jit for a while")
    def test_jit(self, device):
        @torch.jit.script
        def op_script(input):
            return kornia.quaternion_to_rotation_matrix(input)

        quaternion = torch.tensor([0., 0., 1., 0.]).to(device)
        actual = op_script(quaternion)
        expected = kornia.quaternion_to_rotation_matrix(quaternion)
        assert_allclose(actual, expected)

def test_jit(self, device):
        @torch.jit.script
        def op_script(input):
            return kornia.quaternion_to_rotation_matrix(input)

        quaternion = torch.tensor([0., 0., 1., 0.]).to(device)
        actual = op_script(quaternion)
        expected = kornia.quaternion_to_rotation_matrix(quaternion)
        assert_allclose(actual, expected)

def test_jit(self, device):
        @torch.jit.script
        def op_script(input):
            return kornia.filters.sobel(input)
        img = torch.rand(2, 3, 4, 5).to(device)
        actual = op_script(img)
        expected = kornia.filters.sobel(img)
        assert_allclose(actual, expected)

def test_jit(self, device):
        @torch.jit.script
        def op_script(
            input: torch.Tensor,
            ksize: int,
            angle: float,
            direction: float
        ) -> torch.Tensor:
            return kornia.filters.motion_blur(input, ksize, angle, direction)

        img = torch.rand(2, 3, 4, 5).to(device)
        ksize = 5
        angle = 65.
        direction = .1
        actual = op_script(img, ksize, angle, direction)
        expected = kornia.filters.motion_blur(img, ksize, angle, direction)
        assert_allclose(actual, expected)

def test_jit(self):
        @torch.jit.script
        def op_script(transform, points):
            return kornia.transform_points(transform, points)

        points = torch.ones(1, 2, 2)
        transform = torch.eye(3)[None]
        actual = op_script(transform, points)
        expected = kornia.transform_points(transform, points)

        assert_allclose(actual, expected)