mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-08-02 03:42:50 +08:00
2293f1fed0
Summary: I have implemented an exact solution for 3D IoU of oriented 3D boxes. This file includes: * box3d_overlap: which computes the exact IoU of box1 and box2 * box3d_overlap_sampling: which computes an approximate IoU of box1 and box2 by sampling points within the boxes Note that both implementations currently do not support batching. Our exact IoU implementation is based on the fact that the intersecting shape of the two 3D boxes will be formed by segments of the surface of the boxes. Our algorithm computes these segments by reasoning whether triangles of one box are within the second box and vice versa. We deal with intersecting triangles by clipping them. Reviewed By: gkioxari Differential Revision: D30667497 fbshipit-source-id: 2f747f410f90b7f854eeaf3036794bc3ac982917
276 lines
11 KiB
Python
276 lines
11 KiB
Python
# Copyright (c) Facebook, Inc. and its 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 numpy as np
|
|
import torch
|
|
from common_testing import (
|
|
TestCaseMixin,
|
|
get_pytorch3d_dir,
|
|
get_random_cuda_device,
|
|
get_tests_dir,
|
|
)
|
|
from pytorch3d.io import load_obj
|
|
from pytorch3d.ops.sample_farthest_points import (
|
|
sample_farthest_points,
|
|
sample_farthest_points_naive,
|
|
)
|
|
from pytorch3d.ops.utils import masked_gather
|
|
|
|
|
|
DATA_DIR = get_tests_dir() / "data"
|
|
TUTORIAL_DATA_DIR = get_pytorch3d_dir() / "docs/tutorials/data"
|
|
DEBUG = False
|
|
|
|
|
|
class TestFPS(TestCaseMixin, unittest.TestCase):
|
|
def _test_simple(self, fps_func, device="cpu"):
|
|
# fmt: off
|
|
points = torch.tensor(
|
|
[
|
|
[
|
|
[-1.0, -1.0], # noqa: E241, E201
|
|
[-1.3, 1.1], # noqa: E241, E201
|
|
[ 0.2, -1.1], # noqa: E241, E201
|
|
[ 0.0, 0.0], # noqa: E241, E201
|
|
[ 1.3, 1.3], # noqa: E241, E201
|
|
[ 1.0, 0.5], # noqa: E241, E201
|
|
[-1.3, 0.2], # noqa: E241, E201
|
|
[ 1.5, -0.5], # noqa: E241, E201
|
|
],
|
|
[
|
|
[-2.2, -2.4], # noqa: E241, E201
|
|
[-2.1, 2.0], # noqa: E241, E201
|
|
[ 2.2, 2.1], # noqa: E241, E201
|
|
[ 2.1, -2.4], # noqa: E241, E201
|
|
[ 0.4, -1.0], # noqa: E241, E201
|
|
[ 0.3, 0.3], # noqa: E241, E201
|
|
[ 1.2, 0.5], # noqa: E241, E201
|
|
[ 4.5, 4.5], # noqa: E241, E201
|
|
],
|
|
],
|
|
dtype=torch.float32,
|
|
device=device,
|
|
)
|
|
# fmt: on
|
|
expected_inds = torch.tensor([[0, 4], [0, 7]], dtype=torch.int64, device=device)
|
|
out_points, out_inds = fps_func(points, K=2)
|
|
self.assertClose(out_inds, expected_inds)
|
|
|
|
# Gather the points
|
|
expected_inds = expected_inds[..., None].expand(-1, -1, points.shape[-1])
|
|
self.assertClose(out_points, points.gather(dim=1, index=expected_inds))
|
|
|
|
# Different number of points sampled for each pointcloud in the batch
|
|
expected_inds = torch.tensor(
|
|
[[0, 4, 1], [0, 7, -1]], dtype=torch.int64, device=device
|
|
)
|
|
out_points, out_inds = fps_func(points, K=[3, 2])
|
|
self.assertClose(out_inds, expected_inds)
|
|
|
|
# Gather the points
|
|
expected_points = masked_gather(points, expected_inds)
|
|
self.assertClose(out_points, expected_points)
|
|
|
|
def _test_compare_random_heterogeneous(self, device="cpu"):
|
|
N, P, D, K = 5, 20, 5, 8
|
|
points = torch.randn((N, P, D), device=device, dtype=torch.float32)
|
|
out_points_naive, out_idxs_naive = sample_farthest_points_naive(points, K=K)
|
|
out_points, out_idxs = sample_farthest_points(points, K=K)
|
|
self.assertTrue(out_idxs.min() >= 0)
|
|
self.assertClose(out_idxs, out_idxs_naive)
|
|
self.assertClose(out_points, out_points_naive)
|
|
for n in range(N):
|
|
self.assertEqual(out_idxs[n].ne(-1).sum(), K)
|
|
|
|
# Test case where K > P
|
|
K = 30
|
|
points1 = torch.randn((N, P, D), dtype=torch.float32, device=device)
|
|
points2 = points1.clone()
|
|
points1.requires_grad = True
|
|
points2.requires_grad = True
|
|
lengths = torch.randint(low=1, high=P, size=(N,), device=device)
|
|
out_points_naive, out_idxs_naive = sample_farthest_points_naive(
|
|
points1, lengths, K=K
|
|
)
|
|
out_points, out_idxs = sample_farthest_points(points2, lengths, K=K)
|
|
self.assertClose(out_idxs, out_idxs_naive)
|
|
self.assertClose(out_points, out_points_naive)
|
|
|
|
for n in range(N):
|
|
# Check that for heterogeneous batches, the max number of
|
|
# selected points is less than the length
|
|
self.assertTrue(out_idxs[n].ne(-1).sum() <= lengths[n])
|
|
self.assertTrue(out_idxs[n].max() <= lengths[n])
|
|
|
|
# Check there are no duplicate indices
|
|
val_mask = out_idxs[n].ne(-1)
|
|
vals, counts = torch.unique(out_idxs[n][val_mask], return_counts=True)
|
|
self.assertTrue(counts.le(1).all())
|
|
|
|
# Check gradients
|
|
grad_sampled_points = torch.ones((N, K, D), dtype=torch.float32, device=device)
|
|
loss1 = (out_points_naive * grad_sampled_points).sum()
|
|
loss1.backward()
|
|
loss2 = (out_points * grad_sampled_points).sum()
|
|
loss2.backward()
|
|
self.assertClose(points1.grad, points2.grad, atol=5e-6)
|
|
|
|
def _test_errors(self, fps_func, device="cpu"):
|
|
N, P, D, K = 5, 40, 5, 8
|
|
points = torch.randn((N, P, D), device=device)
|
|
wrong_batch_dim = torch.randint(low=1, high=K, size=(K,), device=device)
|
|
|
|
# K has diferent batch dimension to points
|
|
with self.assertRaisesRegex(ValueError, "K and points must have"):
|
|
sample_farthest_points_naive(points, K=wrong_batch_dim)
|
|
|
|
# lengths has diferent batch dimension to points
|
|
with self.assertRaisesRegex(ValueError, "points and lengths must have"):
|
|
sample_farthest_points_naive(points, lengths=wrong_batch_dim, K=K)
|
|
|
|
def _test_random_start(self, fps_func, device="cpu"):
|
|
N, P, D, K = 5, 40, 5, 8
|
|
points = torch.randn((N, P, D), dtype=torch.float32, device=device)
|
|
out_points, out_idxs = fps_func(points, K=K, random_start_point=True)
|
|
# Check the first index is not 0 or the same number for all batch elements
|
|
# when random_start_point = True
|
|
self.assertTrue(out_idxs[:, 0].sum() > 0)
|
|
self.assertFalse(out_idxs[:, 0].eq(out_idxs[0, 0]).all())
|
|
|
|
def _test_gradcheck(self, fps_func, device="cpu"):
|
|
N, P, D, K = 2, 10, 3, 2
|
|
points = torch.randn(
|
|
(N, P, D), dtype=torch.float32, device=device, requires_grad=True
|
|
)
|
|
lengths = torch.randint(low=1, high=P, size=(N,), device=device)
|
|
torch.autograd.gradcheck(
|
|
fps_func,
|
|
(points, lengths, K),
|
|
check_undefined_grad=False,
|
|
eps=2e-3,
|
|
atol=0.001,
|
|
)
|
|
|
|
def test_sample_farthest_points_naive(self):
|
|
device = get_random_cuda_device()
|
|
self._test_simple(sample_farthest_points_naive, device)
|
|
self._test_errors(sample_farthest_points_naive, device)
|
|
self._test_random_start(sample_farthest_points_naive, device)
|
|
self._test_gradcheck(sample_farthest_points_naive, device)
|
|
|
|
def test_sample_farthest_points_cpu(self):
|
|
self._test_simple(sample_farthest_points, "cpu")
|
|
self._test_errors(sample_farthest_points, "cpu")
|
|
self._test_compare_random_heterogeneous("cpu")
|
|
self._test_random_start(sample_farthest_points, "cpu")
|
|
self._test_gradcheck(sample_farthest_points, "cpu")
|
|
|
|
def test_sample_farthest_points_cuda(self):
|
|
device = get_random_cuda_device()
|
|
self._test_simple(sample_farthest_points, device)
|
|
self._test_errors(sample_farthest_points, device)
|
|
self._test_compare_random_heterogeneous(device)
|
|
self._test_random_start(sample_farthest_points, device)
|
|
self._test_gradcheck(sample_farthest_points, device)
|
|
|
|
def test_cuda_vs_cpu(self):
|
|
"""
|
|
Compare cuda vs cpu on a complex object
|
|
"""
|
|
obj_filename = TUTORIAL_DATA_DIR / "cow_mesh/cow.obj"
|
|
K = 250
|
|
|
|
# Run on CPU
|
|
device = "cpu"
|
|
points, _, _ = load_obj(obj_filename, device=device, load_textures=False)
|
|
points = points[None, ...]
|
|
out_points_cpu, out_idxs_cpu = sample_farthest_points(points, K=K)
|
|
|
|
# Run on GPU
|
|
device = get_random_cuda_device()
|
|
points_cuda = points.to(device)
|
|
out_points_cuda, out_idxs_cuda = sample_farthest_points(points_cuda, K=K)
|
|
|
|
# Check that the indices from CUDA and CPU match
|
|
self.assertClose(out_idxs_cpu, out_idxs_cuda.cpu())
|
|
|
|
# Check there are no duplicate indices
|
|
val_mask = out_idxs_cuda[0].ne(-1)
|
|
vals, counts = torch.unique(out_idxs_cuda[0][val_mask], return_counts=True)
|
|
self.assertTrue(counts.le(1).all())
|
|
|
|
# Plot all results
|
|
if DEBUG:
|
|
# mplot3d is required for 3d projection plots
|
|
import matplotlib.pyplot as plt
|
|
from mpl_toolkits import mplot3d # noqa: F401
|
|
|
|
# Move to cpu and convert to numpy for plotting
|
|
points = points.squeeze()
|
|
out_points_cpu = out_points_cpu.squeeze().numpy()
|
|
out_points_cuda = out_points_cuda.squeeze().cpu().numpy()
|
|
|
|
# Farthest point sampling CPU
|
|
fig = plt.figure(figsize=plt.figaspect(1.0 / 3))
|
|
ax1 = fig.add_subplot(1, 3, 1, projection="3d")
|
|
ax1.scatter(*points.T, alpha=0.1)
|
|
ax1.scatter(*out_points_cpu.T, color="black")
|
|
ax1.set_title("FPS CPU")
|
|
|
|
# Farthest point sampling CUDA
|
|
ax2 = fig.add_subplot(1, 3, 2, projection="3d")
|
|
ax2.scatter(*points.T, alpha=0.1)
|
|
ax2.scatter(*out_points_cuda.T, color="red")
|
|
ax2.set_title("FPS CUDA")
|
|
|
|
# Random Sampling
|
|
random_points = np.random.permutation(points)[:K]
|
|
ax3 = fig.add_subplot(1, 3, 3, projection="3d")
|
|
ax3.scatter(*points.T, alpha=0.1)
|
|
ax3.scatter(*random_points.T, color="green")
|
|
ax3.set_title("Random")
|
|
|
|
# Save image
|
|
filename = "DEBUG_fps.jpg"
|
|
filepath = DATA_DIR / filename
|
|
plt.savefig(filepath)
|
|
|
|
@staticmethod
|
|
def sample_farthest_points_naive(N: int, P: int, D: int, K: int, device: str):
|
|
device = torch.device(device)
|
|
pts = torch.randn(
|
|
N, P, D, dtype=torch.float32, device=device, requires_grad=True
|
|
)
|
|
grad_pts = torch.randn(N, K, D, dtype=torch.float32, device=device)
|
|
torch.cuda.synchronize()
|
|
|
|
def output():
|
|
out_points, _ = sample_farthest_points_naive(pts, K=K)
|
|
loss = (out_points * grad_pts).sum()
|
|
loss.backward()
|
|
torch.cuda.synchronize()
|
|
|
|
return output
|
|
|
|
@staticmethod
|
|
def sample_farthest_points(N: int, P: int, D: int, K: int, device: str):
|
|
device = torch.device(device)
|
|
pts = torch.randn(
|
|
N, P, D, dtype=torch.float32, device=device, requires_grad=True
|
|
)
|
|
grad_pts = torch.randn(N, K, D, dtype=torch.float32, device=device)
|
|
torch.cuda.synchronize()
|
|
|
|
def output():
|
|
out_points, _ = sample_farthest_points(pts, K=K)
|
|
loss = (out_points * grad_pts).sum()
|
|
loss.backward()
|
|
torch.cuda.synchronize()
|
|
|
|
return output
|