# 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))