pytorch3d/tests/test_perspective_n_points.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

import unittest

import torch
from pytorch3d.ops import perspective_n_points
from pytorch3d.transforms import rotation_conversions

from .common_testing import TestCaseMixin


def reproj_error(x_world, y, R, T, weight=None):
    # applies the affine transform, projects, and computes the reprojection error
    y_hat = torch.matmul(x_world, R) + T[:, None, :]
    y_hat = y_hat / y_hat[..., 2:]
    if weight is None:
        weight = y.new_ones((1, 1))
    return (((weight[:, :, None] * (y - y_hat[..., :2])) ** 2).sum(dim=-1) ** 0.5).mean(
        dim=-1
    )


class TestPerspectiveNPoints(TestCaseMixin, unittest.TestCase):
    def setUp(self) -> None:
        super().setUp()
        torch.manual_seed(42)

    @classmethod
    def _generate_epnp_test_from_2d(cls, y):
        """
        Instantiate random x_world, x_cam, R, T given a set of input
        2D projections y.
        """
        batch_size = y.shape[0]
        x_cam = torch.cat((y, torch.rand_like(y[:, :, :1]) * 2.0 + 3.5), dim=2)
        x_cam[:, :, :2] *= x_cam[:, :, 2:]  # unproject
        R = rotation_conversions.random_rotations(batch_size).to(y)
        T = torch.randn_like(R[:, :1, :])
        T[:, :, 2] = (T[:, :, 2] + 3.0).clamp(2.0)
        x_world = torch.matmul(x_cam - T, R.transpose(1, 2))
        return x_cam, x_world, R, T

    def _run_and_print(self, x_world, y, R, T, print_stats, skip_q, check_output=False):
        sol = perspective_n_points.efficient_pnp(
            x_world, y.expand_as(x_world[:, :, :2]), skip_quadratic_eq=skip_q
        )

        err_2d = reproj_error(x_world, y, sol.R, sol.T)
        R_est_quat = rotation_conversions.matrix_to_quaternion(sol.R)
        R_quat = rotation_conversions.matrix_to_quaternion(R)

        num_pts = x_world.shape[-2]
        if check_output:
            assert_msg = (
                f"test_perspective_n_points assertion failure for "
                f"n_points={num_pts}, "
                f"skip_quadratic={skip_q}, "
                f"no noise."
            )

            self.assertClose(err_2d, sol.err_2d, msg=assert_msg)
            self.assertTrue((err_2d < 1e-3).all(), msg=assert_msg)

            def norm_fn(t):
                return t.norm(dim=-1)

            self.assertNormsClose(
                T, sol.T[:, None, :], rtol=4e-3, norm_fn=norm_fn, msg=assert_msg
            )
            self.assertNormsClose(
                R_quat, R_est_quat, rtol=3e-3, norm_fn=norm_fn, msg=assert_msg
            )

        if print_stats:
            torch.set_printoptions(precision=5, sci_mode=False)
            for err_2d, err_3d, R_gt, T_gt in zip(
                sol.err_2d,
                sol.err_3d,
                torch.cat((sol.R, R), dim=-1),
                torch.stack((sol.T, T[:, 0, :]), dim=-1),
            ):
                print("2D Error: %1.4f" % err_2d.item())
                print("3D Error: %1.4f" % err_3d.item())
                print("R_hat | R_gt\n", R_gt)
                print("T_hat | T_gt\n", T_gt)

    def _testcase_from_2d(
        self, y, print_stats, benchmark, skip_q=False, skip_check_thresh=5
    ):
        """
        In case num_pts < 6, EPnP gets unstable, so we check it doesn't crash
        """
        x_cam, x_world, R, T = TestPerspectiveNPoints._generate_epnp_test_from_2d(
            y[None].repeat(16, 1, 1)
        )

        if print_stats:
            print("Run without noise")

        if benchmark:  # return curried call
            torch.cuda.synchronize()

            def result():
                self._run_and_print(x_world, y, R, T, False, skip_q)
                torch.cuda.synchronize()

            return result

        self._run_and_print(
            x_world,
            y,
            R,
            T,
            print_stats,
            skip_q,
            check_output=True if y.shape[1] > skip_check_thresh else False,
        )

        # in the noisy case, there are no guarantees, so we check it doesn't crash
        if print_stats:
            print("Run with noise")
        x_world += torch.randn_like(x_world) * 0.1
        self._run_and_print(x_world, y, R, T, print_stats, skip_q)

    def case_with_gaussian_points(
        self, batch_size=10, num_pts=20, print_stats=False, benchmark=True, skip_q=False
    ):
        return self._testcase_from_2d(
            torch.randn((num_pts, 2)).cuda() / 3.0,
            print_stats=print_stats,
            benchmark=benchmark,
            skip_q=skip_q,
        )

    def test_perspective_n_points(self, print_stats=False):
        if print_stats:
            print("RUN ON A DENSE GRID")
        u = torch.linspace(-1.0, 1.0, 20)
        v = torch.linspace(-1.0, 1.0, 15)
        for skip_q in [False, True]:
            self._testcase_from_2d(
                torch.cartesian_prod(u, v).cuda(), print_stats, False, skip_q
            )

        for num_pts in range(6, 3, -1):
            for skip_q in [False, True]:
                if print_stats:
                    print(f"RUN ON {num_pts} points; skip_quadratic: {skip_q}")

                self.case_with_gaussian_points(
                    num_pts=num_pts,
                    print_stats=print_stats,
                    benchmark=False,
                    skip_q=skip_q,
                )

    def test_weighted_perspective_n_points(self, batch_size=16, num_pts=200):
        # instantiate random x_world and y
        y = torch.randn((batch_size, num_pts, 2)).cuda() / 3.0
        x_cam, x_world, R, T = TestPerspectiveNPoints._generate_epnp_test_from_2d(y)

        # randomly drop 50% of the rows
        weights = (torch.rand_like(x_world[:, :, 0]) > 0.5).float()

        # make sure we retain at least 6 points for each case
        weights[:, :6] = 1.0

        # fill ignored y with trash to ensure that we get different
        # solution in case the weighting is wrong
        y = y + (1 - weights[:, :, None]) * 100.0

        def norm_fn(t):
            return t.norm(dim=-1)

        for skip_quadratic_eq in (True, False):
            # get the solution for the 0/1 weighted case
            sol = perspective_n_points.efficient_pnp(
                x_world, y, skip_quadratic_eq=skip_quadratic_eq, weights=weights
            )
            sol_R_quat = rotation_conversions.matrix_to_quaternion(sol.R)
            sol_T = sol.T

            # check that running only on points with non-zero weights ends in the
            # same place as running the 0/1 weighted version
            for i in range(batch_size):
                ok = weights[i] > 0
                x_world_ok = x_world[i, ok][None]
                y_ok = y[i, ok][None]
                sol_ok = perspective_n_points.efficient_pnp(
                    x_world_ok, y_ok, skip_quadratic_eq=False
                )
                R_est_quat_ok = rotation_conversions.matrix_to_quaternion(sol_ok.R)

                self.assertNormsClose(sol_T[i], sol_ok.T[0], rtol=3e-3, norm_fn=norm_fn)
                self.assertNormsClose(
                    sol_R_quat[i], R_est_quat_ok[0], rtol=3e-4, norm_fn=norm_fn
                )