Spaces:
Running
Running
File size: 8,282 Bytes
7088d16 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 |
# 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))
|