pytorch3d/tests/implicitron/test_circle_fitting.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
from math import pi

import torch
from pytorch3d.implicitron.tools.circle_fitting import (
    _signed_area,
    fit_circle_in_2d,
    fit_circle_in_3d,
    get_rotation_to_best_fit_xy,
)
from pytorch3d.transforms import random_rotation, random_rotations
from tests.common_testing import TestCaseMixin


class TestCircleFitting(TestCaseMixin, unittest.TestCase):
    def setUp(self):
        torch.manual_seed(42)

    def _assertParallel(self, a, b, **kwargs):
        """
        Given a and b of shape (..., 3) each containing 3D vectors,
        assert that correspnding vectors are parallel. Changed sign is ok.
        """
        self.assertClose(torch.cross(a, b, dim=-1), torch.zeros_like(a), **kwargs)

    def test_plane_levelling(self):
        device = torch.device("cuda:0")
        B = 16
        N = 1024
        random = torch.randn((B, N, 3), device=device)

        # first, check that we always return a vaild rotation
        rot = get_rotation_to_best_fit_xy(random)
        self.assertClose(rot.det(), torch.ones_like(rot[:, 0, 0]))
        self.assertClose(rot.norm(dim=-1), torch.ones_like(rot[:, 0]))

        # then, check the result is what we expect
        z_squeeze = 0.1
        random[..., -1] *= z_squeeze
        rot_gt = random_rotations(B, device=device)
        rotated = random @ rot_gt.transpose(-1, -2)
        rot_hat = get_rotation_to_best_fit_xy(rotated)
        self.assertClose(rot.det(), torch.ones_like(rot[:, 0, 0]))
        self.assertClose(rot.norm(dim=-1), torch.ones_like(rot[:, 0]))
        # covariance matrix of the levelled points is by design diag(1, 1, z_squeeze²)
        self.assertClose(
            (rotated @ rot_hat)[..., -1].std(dim=-1),
            torch.ones_like(rot_hat[:, 0, 0]) * z_squeeze,
            rtol=0.1,
        )

    def test_simple_3d(self):
        device = torch.device("cuda:0")
        for _ in range(7):
            radius = 10 * torch.rand(1, device=device)[0]
            center = 10 * torch.rand(3, device=device)
            rot = random_rotation(device=device)
            offset = torch.rand(3, device=device)
            up = torch.rand(3, device=device)
            self._simple_3d_test(radius, center, rot, offset, up)

    def _simple_3d_test(self, radius, center, rot, offset, up):
        # angles are increasing so the points move in a well defined direction.
        angles = torch.cumsum(torch.rand(17, device=rot.device), dim=0)
        many = torch.stack(
            [torch.cos(angles), torch.sin(angles), torch.zeros_like(angles)], dim=1
        )
        source_points = (many * radius) @ rot + center[None]

        # case with no generation
        result = fit_circle_in_3d(source_points)
        self.assertClose(result.radius, radius)
        self.assertClose(result.center, center)
        self._assertParallel(result.normal, rot[2], atol=1e-5)
        self.assertEqual(result.generated_points.shape, (0, 3))

        # Generate 5 points around the circle
        n_new_points = 5
        result2 = fit_circle_in_3d(source_points, n_points=n_new_points)
        self.assertClose(result2.radius, radius)
        self.assertClose(result2.center, center)
        self.assertClose(result2.normal, result.normal)
        self.assertEqual(result2.generated_points.shape, (5, 3))

        observed_points = result2.generated_points
        self.assertClose(observed_points[0], observed_points[4], atol=1e-4)
        self.assertClose(observed_points[0], source_points[0], atol=1e-5)
        observed_normal = torch.cross(
            observed_points[0] - observed_points[2],
            observed_points[1] - observed_points[3],
            dim=-1,
        )
        self._assertParallel(observed_normal, result.normal, atol=1e-4)
        diameters = observed_points[:2] - observed_points[2:4]
        self.assertClose(
            torch.norm(diameters, dim=1), diameters.new_full((2,), 2 * radius)
        )

        # Regenerate the input points
        result3 = fit_circle_in_3d(source_points, angles=angles - angles[0])
        self.assertClose(result3.radius, radius)
        self.assertClose(result3.center, center)
        self.assertClose(result3.normal, result.normal)
        self.assertClose(result3.generated_points, source_points, atol=1e-5)

        # Test with offset
        result4 = fit_circle_in_3d(
            source_points, angles=angles - angles[0], offset=offset, up=up
        )
        self.assertClose(result4.radius, radius)
        self.assertClose(result4.center, center)
        self.assertClose(result4.normal, result.normal)
        observed_offsets = result4.generated_points - source_points

        # observed_offset is constant
        self.assertClose(
            observed_offsets.min(0).values, observed_offsets.max(0).values, atol=1e-5
        )
        # observed_offset has the right length
        self.assertClose(observed_offsets[0].norm(), offset.norm())

        self.assertClose(result.normal.norm(), torch.ones(()))
        # component of observed_offset along normal
        component = torch.dot(observed_offsets[0], result.normal)
        self.assertClose(component.abs(), offset[2].abs(), atol=1e-5)
        agree_normal = torch.dot(result.normal, up) > 0
        agree_signs = component * offset[2] > 0
        self.assertEqual(agree_normal, agree_signs)

    def test_simple_2d(self):
        radius = 7.0
        center = torch.tensor([9, 2.5])
        angles = torch.cumsum(torch.rand(17), dim=0)
        many = torch.stack([torch.cos(angles), torch.sin(angles)], dim=1)
        source_points = (many * radius) + center[None]

        result = fit_circle_in_2d(source_points)
        self.assertClose(result.radius, torch.tensor(radius))
        self.assertClose(result.center, center)
        self.assertEqual(result.generated_points.shape, (0, 2))

        # Generate 5 points around the circle
        n_new_points = 5
        result2 = fit_circle_in_2d(source_points, n_points=n_new_points)
        self.assertClose(result2.radius, torch.tensor(radius))
        self.assertClose(result2.center, center)
        self.assertEqual(result2.generated_points.shape, (5, 2))

        observed_points = result2.generated_points
        self.assertClose(observed_points[0], observed_points[4])
        self.assertClose(observed_points[0], source_points[0], atol=1e-5)
        diameters = observed_points[:2] - observed_points[2:4]
        self.assertClose(torch.norm(diameters, dim=1), torch.full((2,), 2 * radius))

        # Regenerate the input points
        result3 = fit_circle_in_2d(source_points, angles=angles - angles[0])
        self.assertClose(result3.radius, torch.tensor(radius))
        self.assertClose(result3.center, center)
        self.assertClose(result3.generated_points, source_points, atol=1e-5)

    def test_minimum_inputs(self):
        fit_circle_in_3d(torch.rand(3, 3), n_points=10)

        with self.assertRaisesRegex(
            ValueError, "2 points are not enough to determine a circle"
        ):
            fit_circle_in_3d(torch.rand(2, 3))

    def test_signed_area(self):
        n_points = 1001
        angles = torch.linspace(0, 2 * pi, n_points)
        radius = 0.85
        center = torch.rand(2)
        circle = center + radius * torch.stack(
            [torch.cos(angles), torch.sin(angles)], dim=1
        )
        circle_area = torch.tensor(pi * radius * radius)
        self.assertClose(_signed_area(circle), circle_area)
        # clockwise is negative
        self.assertClose(_signed_area(circle.flip(0)), -circle_area)

        # Semicircles
        self.assertClose(_signed_area(circle[: (n_points + 1) // 2]), circle_area / 2)
        self.assertClose(_signed_area(circle[n_points // 2 :]), circle_area / 2)

        # A straight line bounds no area
        self.assertClose(_signed_area(torch.rand(2, 2)), torch.tensor(0.0))

        # Letter 'L' written anticlockwise.
        L_shape = [[0, 1], [0, 0], [1, 0]]
        # Triangle area is 0.5 * b * h.
        self.assertClose(_signed_area(torch.tensor(L_shape)), torch.tensor(0.5))