mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-08-02 03:42:50 +08:00
ee2b2feb98
Summary: Move the core of add_points_to_volumes to the new C++/CUDA implementation. Add new flag to let the user stop this happening. Avoids copies. About a 30% speedup on the larger cases, up to 50% on the smaller cases. New timings ``` Benchmark Avg Time(μs) Peak Time(μs) Iterations -------------------------------------------------------------------------------- ADD_POINTS_TO_VOLUMES_cpu_10_trilinear_[25, 25, 25]_1000 4575 12591 110 ADD_POINTS_TO_VOLUMES_cpu_10_trilinear_[25, 25, 25]_10000 25468 29186 20 ADD_POINTS_TO_VOLUMES_cpu_10_trilinear_[25, 25, 25]_100000 202085 209897 3 ADD_POINTS_TO_VOLUMES_cpu_10_trilinear_[101, 111, 121]_1000 46059 48188 11 ADD_POINTS_TO_VOLUMES_cpu_10_trilinear_[101, 111, 121]_10000 83759 95669 7 ADD_POINTS_TO_VOLUMES_cpu_10_trilinear_[101, 111, 121]_100000 326056 339393 2 ADD_POINTS_TO_VOLUMES_cpu_10_nearest_[25, 25, 25]_1000 2379 4738 211 ADD_POINTS_TO_VOLUMES_cpu_10_nearest_[25, 25, 25]_10000 12100 63099 42 ADD_POINTS_TO_VOLUMES_cpu_10_nearest_[25, 25, 25]_100000 63323 63737 8 ADD_POINTS_TO_VOLUMES_cpu_10_nearest_[101, 111, 121]_1000 45216 45479 12 ADD_POINTS_TO_VOLUMES_cpu_10_nearest_[101, 111, 121]_10000 57205 58524 9 ADD_POINTS_TO_VOLUMES_cpu_10_nearest_[101, 111, 121]_100000 139499 139926 4 ADD_POINTS_TO_VOLUMES_cpu_100_trilinear_[25, 25, 25]_1000 40129 40431 13 ADD_POINTS_TO_VOLUMES_cpu_100_trilinear_[25, 25, 25]_10000 204949 239293 3 ADD_POINTS_TO_VOLUMES_cpu_100_trilinear_[25, 25, 25]_100000 1664541 1664541 1 ADD_POINTS_TO_VOLUMES_cpu_100_trilinear_[101, 111, 121]_1000 391573 395108 2 ADD_POINTS_TO_VOLUMES_cpu_100_trilinear_[101, 111, 121]_10000 674869 674869 1 ADD_POINTS_TO_VOLUMES_cpu_100_trilinear_[101, 111, 121]_100000 2713632 2713632 1 ADD_POINTS_TO_VOLUMES_cpu_100_nearest_[25, 25, 25]_1000 12726 13506 40 ADD_POINTS_TO_VOLUMES_cpu_100_nearest_[25, 25, 25]_10000 73103 73299 7 ADD_POINTS_TO_VOLUMES_cpu_100_nearest_[25, 25, 25]_100000 598634 598634 1 ADD_POINTS_TO_VOLUMES_cpu_100_nearest_[101, 111, 121]_1000 398742 399256 2 ADD_POINTS_TO_VOLUMES_cpu_100_nearest_[101, 111, 121]_10000 543129 543129 1 ADD_POINTS_TO_VOLUMES_cpu_100_nearest_[101, 111, 121]_100000 1242956 1242956 1 ADD_POINTS_TO_VOLUMES_cuda:0_10_trilinear_[25, 25, 25]_1000 1814 8884 276 ADD_POINTS_TO_VOLUMES_cuda:0_10_trilinear_[25, 25, 25]_10000 1996 8851 251 ADD_POINTS_TO_VOLUMES_cuda:0_10_trilinear_[25, 25, 25]_100000 4608 11529 109 ADD_POINTS_TO_VOLUMES_cuda:0_10_trilinear_[101, 111, 121]_1000 5183 12508 97 ADD_POINTS_TO_VOLUMES_cuda:0_10_trilinear_[101, 111, 121]_10000 7106 14077 71 ADD_POINTS_TO_VOLUMES_cuda:0_10_trilinear_[101, 111, 121]_100000 25914 31818 20 ADD_POINTS_TO_VOLUMES_cuda:0_10_nearest_[25, 25, 25]_1000 1778 8823 282 ADD_POINTS_TO_VOLUMES_cuda:0_10_nearest_[25, 25, 25]_10000 1825 8613 274 ADD_POINTS_TO_VOLUMES_cuda:0_10_nearest_[25, 25, 25]_100000 3154 10161 159 ADD_POINTS_TO_VOLUMES_cuda:0_10_nearest_[101, 111, 121]_1000 4888 9404 103 ADD_POINTS_TO_VOLUMES_cuda:0_10_nearest_[101, 111, 121]_10000 5194 9963 97 ADD_POINTS_TO_VOLUMES_cuda:0_10_nearest_[101, 111, 121]_100000 8109 14933 62 ADD_POINTS_TO_VOLUMES_cuda:0_100_trilinear_[25, 25, 25]_1000 3320 10306 151 ADD_POINTS_TO_VOLUMES_cuda:0_100_trilinear_[25, 25, 25]_10000 7003 8595 72 ADD_POINTS_TO_VOLUMES_cuda:0_100_trilinear_[25, 25, 25]_100000 49140 52957 11 ADD_POINTS_TO_VOLUMES_cuda:0_100_trilinear_[101, 111, 121]_1000 35890 36918 14 ADD_POINTS_TO_VOLUMES_cuda:0_100_trilinear_[101, 111, 121]_10000 58890 59337 9 ADD_POINTS_TO_VOLUMES_cuda:0_100_trilinear_[101, 111, 121]_100000 286878 287600 2 ADD_POINTS_TO_VOLUMES_cuda:0_100_nearest_[25, 25, 25]_1000 2484 8805 202 ADD_POINTS_TO_VOLUMES_cuda:0_100_nearest_[25, 25, 25]_10000 3967 9090 127 ADD_POINTS_TO_VOLUMES_cuda:0_100_nearest_[25, 25, 25]_100000 19423 19799 26 ADD_POINTS_TO_VOLUMES_cuda:0_100_nearest_[101, 111, 121]_1000 33228 33329 16 ADD_POINTS_TO_VOLUMES_cuda:0_100_nearest_[101, 111, 121]_10000 37292 37370 14 ADD_POINTS_TO_VOLUMES_cuda:0_100_nearest_[101, 111, 121]_100000 73550 74017 7 -------------------------------------------------------------------------------- ``` Previous timings ``` Benchmark Avg Time(μs) Peak Time(μs) Iterations -------------------------------------------------------------------------------- ADD_POINTS_TO_VOLUMES_cpu_10_trilinear_[25, 25, 25]_1000 10100 46422 50 ADD_POINTS_TO_VOLUMES_cpu_10_trilinear_[25, 25, 25]_10000 28442 32100 18 ADD_POINTS_TO_VOLUMES_cpu_10_trilinear_[25, 25, 25]_100000 241127 254269 3 ADD_POINTS_TO_VOLUMES_cpu_10_trilinear_[101, 111, 121]_1000 54149 79480 10 ADD_POINTS_TO_VOLUMES_cpu_10_trilinear_[101, 111, 121]_10000 125459 212734 4 ADD_POINTS_TO_VOLUMES_cpu_10_trilinear_[101, 111, 121]_100000 512739 512739 1 ADD_POINTS_TO_VOLUMES_cpu_10_nearest_[25, 25, 25]_1000 2866 13365 175 ADD_POINTS_TO_VOLUMES_cpu_10_nearest_[25, 25, 25]_10000 7026 12604 72 ADD_POINTS_TO_VOLUMES_cpu_10_nearest_[25, 25, 25]_100000 48822 55607 11 ADD_POINTS_TO_VOLUMES_cpu_10_nearest_[101, 111, 121]_1000 38098 38576 14 ADD_POINTS_TO_VOLUMES_cpu_10_nearest_[101, 111, 121]_10000 48006 54120 11 ADD_POINTS_TO_VOLUMES_cpu_10_nearest_[101, 111, 121]_100000 131563 138536 4 ADD_POINTS_TO_VOLUMES_cpu_100_trilinear_[25, 25, 25]_1000 64615 91735 8 ADD_POINTS_TO_VOLUMES_cpu_100_trilinear_[25, 25, 25]_10000 228815 246095 3 ADD_POINTS_TO_VOLUMES_cpu_100_trilinear_[25, 25, 25]_100000 3086615 3086615 1 ADD_POINTS_TO_VOLUMES_cpu_100_trilinear_[101, 111, 121]_1000 464298 465292 2 ADD_POINTS_TO_VOLUMES_cpu_100_trilinear_[101, 111, 121]_10000 1053440 1053440 1 ADD_POINTS_TO_VOLUMES_cpu_100_trilinear_[101, 111, 121]_100000 6736236 6736236 1 ADD_POINTS_TO_VOLUMES_cpu_100_nearest_[25, 25, 25]_1000 11940 12440 42 ADD_POINTS_TO_VOLUMES_cpu_100_nearest_[25, 25, 25]_10000 56641 58051 9 ADD_POINTS_TO_VOLUMES_cpu_100_nearest_[25, 25, 25]_100000 711492 711492 1 ADD_POINTS_TO_VOLUMES_cpu_100_nearest_[101, 111, 121]_1000 326437 329846 2 ADD_POINTS_TO_VOLUMES_cpu_100_nearest_[101, 111, 121]_10000 418514 427911 2 ADD_POINTS_TO_VOLUMES_cpu_100_nearest_[101, 111, 121]_100000 1524285 1524285 1 ADD_POINTS_TO_VOLUMES_cuda:0_10_trilinear_[25, 25, 25]_1000 5949 13602 85 ADD_POINTS_TO_VOLUMES_cuda:0_10_trilinear_[25, 25, 25]_10000 5817 13001 86 ADD_POINTS_TO_VOLUMES_cuda:0_10_trilinear_[25, 25, 25]_100000 23833 25971 21 ADD_POINTS_TO_VOLUMES_cuda:0_10_trilinear_[101, 111, 121]_1000 9029 16178 56 ADD_POINTS_TO_VOLUMES_cuda:0_10_trilinear_[101, 111, 121]_10000 11595 18601 44 ADD_POINTS_TO_VOLUMES_cuda:0_10_trilinear_[101, 111, 121]_100000 46986 47344 11 ADD_POINTS_TO_VOLUMES_cuda:0_10_nearest_[25, 25, 25]_1000 2554 9747 196 ADD_POINTS_TO_VOLUMES_cuda:0_10_nearest_[25, 25, 25]_10000 2676 9537 187 ADD_POINTS_TO_VOLUMES_cuda:0_10_nearest_[25, 25, 25]_100000 6567 14179 77 ADD_POINTS_TO_VOLUMES_cuda:0_10_nearest_[101, 111, 121]_1000 5840 12811 86 ADD_POINTS_TO_VOLUMES_cuda:0_10_nearest_[101, 111, 121]_10000 6102 13128 82 ADD_POINTS_TO_VOLUMES_cuda:0_10_nearest_[101, 111, 121]_100000 11945 11995 42 ADD_POINTS_TO_VOLUMES_cuda:0_100_trilinear_[25, 25, 25]_1000 7642 13671 66 ADD_POINTS_TO_VOLUMES_cuda:0_100_trilinear_[25, 25, 25]_10000 25190 25260 20 ADD_POINTS_TO_VOLUMES_cuda:0_100_trilinear_[25, 25, 25]_100000 212018 212134 3 ADD_POINTS_TO_VOLUMES_cuda:0_100_trilinear_[101, 111, 121]_1000 40421 45692 13 ADD_POINTS_TO_VOLUMES_cuda:0_100_trilinear_[101, 111, 121]_10000 92078 92132 6 ADD_POINTS_TO_VOLUMES_cuda:0_100_trilinear_[101, 111, 121]_100000 457211 457229 2 ADD_POINTS_TO_VOLUMES_cuda:0_100_nearest_[25, 25, 25]_1000 3574 10377 140 ADD_POINTS_TO_VOLUMES_cuda:0_100_nearest_[25, 25, 25]_10000 7222 13023 70 ADD_POINTS_TO_VOLUMES_cuda:0_100_nearest_[25, 25, 25]_100000 48127 48165 11 ADD_POINTS_TO_VOLUMES_cuda:0_100_nearest_[101, 111, 121]_1000 34732 35295 15 ADD_POINTS_TO_VOLUMES_cuda:0_100_nearest_[101, 111, 121]_10000 43050 51064 12 ADD_POINTS_TO_VOLUMES_cuda:0_100_nearest_[101, 111, 121]_100000 106028 106058 5 -------------------------------------------------------------------------------- ``` Reviewed By: nikhilaravi Differential Revision: D29548609 fbshipit-source-id: 7026e832ea299145c3f6b55687f3c1601294f5c0
565 lines
20 KiB
Python
565 lines
20 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
|
|
from functools import partial
|
|
from itertools import product
|
|
from typing import Tuple
|
|
|
|
import numpy as np
|
|
import torch
|
|
from common_testing import TestCaseMixin
|
|
from pytorch3d.ops import add_pointclouds_to_volumes
|
|
from pytorch3d.ops.points_to_volumes import _points_to_volumes
|
|
from pytorch3d.ops.sample_points_from_meshes import sample_points_from_meshes
|
|
from pytorch3d.structures.meshes import Meshes
|
|
from pytorch3d.structures.pointclouds import Pointclouds
|
|
from pytorch3d.structures.volumes import Volumes
|
|
from pytorch3d.transforms.so3 import so3_exp_map
|
|
|
|
|
|
DEBUG = False
|
|
if DEBUG:
|
|
import os
|
|
import tempfile
|
|
|
|
from PIL import Image
|
|
|
|
|
|
def init_cube_point_cloud(batch_size: int, n_points: int, device: str, rotate_y: bool):
|
|
"""
|
|
Generate a random point cloud of `n_points` whose points
|
|
are sampled from faces of a 3D cube.
|
|
"""
|
|
|
|
# create the cube mesh batch_size times
|
|
meshes = TestPointsToVolumes.init_cube_mesh(batch_size=batch_size, device=device)
|
|
|
|
# generate point clouds by sampling points from the meshes
|
|
pcl = sample_points_from_meshes(meshes, num_samples=n_points, return_normals=False)
|
|
|
|
# colors of the cube sides
|
|
clrs = [
|
|
[1.0, 0.0, 0.0],
|
|
[1.0, 1.0, 0.0],
|
|
[0.0, 1.0, 0.0],
|
|
[0.0, 1.0, 1.0],
|
|
[1.0, 1.0, 1.0],
|
|
[1.0, 0.0, 1.0],
|
|
]
|
|
|
|
# init the color tensor "rgb"
|
|
rgb = torch.zeros_like(pcl)
|
|
|
|
# color each side of the cube with a constant color
|
|
clri = 0
|
|
for dim in (0, 1, 2):
|
|
for offs in (0.0, 1.0):
|
|
current_face_verts = (pcl[:, :, dim] - offs).abs() <= 1e-2
|
|
for bi in range(batch_size):
|
|
rgb[bi, current_face_verts[bi], :] = torch.tensor(clrs[clri]).type_as(
|
|
pcl
|
|
)
|
|
clri += 1
|
|
|
|
if rotate_y:
|
|
# uniformly spaced rotations around y axis
|
|
R = init_uniform_y_rotations(batch_size=batch_size, device=device)
|
|
# rotate the point clouds around y axis
|
|
pcl = torch.bmm(pcl - 0.5, R) + 0.5
|
|
|
|
return pcl, rgb
|
|
|
|
|
|
def init_volume_boundary_pointcloud(
|
|
batch_size: int,
|
|
volume_size: Tuple[int, int, int],
|
|
n_points: int,
|
|
interp_mode: str,
|
|
device: str,
|
|
require_grad: bool = False,
|
|
):
|
|
"""
|
|
Initialize a point cloud that closely follows a boundary of
|
|
a volume with a given size. The volume buffer is initialized as well.
|
|
"""
|
|
|
|
# generate a 3D point cloud sampled from sides of a [0,1] cube
|
|
xyz, rgb = init_cube_point_cloud(
|
|
batch_size, n_points=n_points, device=device, rotate_y=True
|
|
)
|
|
|
|
# make volume_size tensor
|
|
volume_size_t = torch.tensor(volume_size, dtype=xyz.dtype, device=xyz.device)
|
|
|
|
if interp_mode == "trilinear":
|
|
# make the xyz locations fall on the boundary of the
|
|
# first/last two voxels along each spatial dimension of the
|
|
# volume - this properly checks the correctness of the
|
|
# trilinear interpolation scheme
|
|
xyz = (xyz - 0.5) * ((volume_size_t - 2) / (volume_size_t - 1))[[2, 1, 0]] + 0.5
|
|
|
|
# rescale the cube pointcloud to overlap with the volume sides
|
|
# of the volume
|
|
rel_scale = volume_size_t / volume_size[0]
|
|
xyz = xyz * rel_scale[[2, 1, 0]][None, None]
|
|
|
|
# enable grad accumulation for the differentiability check
|
|
xyz.requires_grad = require_grad
|
|
rgb.requires_grad = require_grad
|
|
|
|
# create the pointclouds structure
|
|
pointclouds = Pointclouds(xyz, features=rgb)
|
|
|
|
# set the volume translation so that the point cloud is centered
|
|
# around 0
|
|
volume_translation = -0.5 * rel_scale[[2, 1, 0]]
|
|
|
|
# set the voxel size to 1 / (volume_size-1)
|
|
volume_voxel_size = 1 / (volume_size[0] - 1.0)
|
|
|
|
# instantiate the volumes
|
|
initial_volumes = Volumes(
|
|
features=xyz.new_zeros(batch_size, 3, *volume_size),
|
|
densities=xyz.new_zeros(batch_size, 1, *volume_size),
|
|
volume_translation=volume_translation,
|
|
voxel_size=volume_voxel_size,
|
|
)
|
|
|
|
return pointclouds, initial_volumes
|
|
|
|
|
|
def init_uniform_y_rotations(batch_size: int, device: torch.device):
|
|
"""
|
|
Generate a batch of `batch_size` 3x3 rotation matrices around y-axis
|
|
whose angles are uniformly distributed between 0 and 2 pi.
|
|
"""
|
|
axis = torch.tensor([0.0, 1.0, 0.0], device=device, dtype=torch.float32)
|
|
angles = torch.linspace(0, 2.0 * np.pi, batch_size + 1, device=device)
|
|
angles = angles[:batch_size]
|
|
log_rots = axis[None, :] * angles[:, None]
|
|
R = so3_exp_map(log_rots)
|
|
return R
|
|
|
|
|
|
class TestPointsToVolumes(TestCaseMixin, unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
np.random.seed(42)
|
|
torch.manual_seed(42)
|
|
|
|
@staticmethod
|
|
def add_points_to_volumes(
|
|
batch_size: int,
|
|
volume_size: Tuple[int, int, int],
|
|
n_points: int,
|
|
interp_mode: str,
|
|
device: str,
|
|
):
|
|
(pointclouds, initial_volumes) = init_volume_boundary_pointcloud(
|
|
batch_size=batch_size,
|
|
volume_size=volume_size,
|
|
n_points=n_points,
|
|
interp_mode=interp_mode,
|
|
require_grad=False,
|
|
device=device,
|
|
)
|
|
|
|
torch.cuda.synchronize()
|
|
|
|
def _add_points_to_volumes():
|
|
add_pointclouds_to_volumes(pointclouds, initial_volumes, mode=interp_mode)
|
|
torch.cuda.synchronize()
|
|
|
|
return _add_points_to_volumes
|
|
|
|
@staticmethod
|
|
def stack_4d_tensor_to_3d(arr):
|
|
n = arr.shape[0]
|
|
H = int(np.ceil(np.sqrt(n)))
|
|
W = int(np.ceil(n / H))
|
|
n_add = H * W - n
|
|
arr = torch.cat((arr, torch.zeros_like(arr[:1]).repeat(n_add, 1, 1, 1)))
|
|
rows = torch.chunk(arr, chunks=W, dim=0)
|
|
arr3d = torch.cat([torch.cat(list(row), dim=2) for row in rows], dim=1)
|
|
return arr3d
|
|
|
|
@staticmethod
|
|
def init_cube_mesh(batch_size: int, device: str):
|
|
"""
|
|
Generate a batch of `batch_size` cube meshes.
|
|
"""
|
|
|
|
device = torch.device(device)
|
|
|
|
verts, faces = [], []
|
|
|
|
for _ in range(batch_size):
|
|
v = torch.tensor(
|
|
[
|
|
[0.0, 0.0, 0.0],
|
|
[1.0, 0.0, 0.0],
|
|
[1.0, 1.0, 0.0],
|
|
[0.0, 1.0, 0.0],
|
|
[0.0, 1.0, 1.0],
|
|
[1.0, 1.0, 1.0],
|
|
[1.0, 0.0, 1.0],
|
|
[0.0, 0.0, 1.0],
|
|
],
|
|
dtype=torch.float32,
|
|
device=device,
|
|
)
|
|
verts.append(v)
|
|
faces.append(
|
|
torch.tensor(
|
|
[
|
|
[0, 2, 1],
|
|
[0, 3, 2],
|
|
[2, 3, 4],
|
|
[2, 4, 5],
|
|
[1, 2, 5],
|
|
[1, 5, 6],
|
|
[0, 7, 4],
|
|
[0, 4, 3],
|
|
[5, 4, 7],
|
|
[5, 7, 6],
|
|
[0, 6, 7],
|
|
[0, 1, 6],
|
|
],
|
|
dtype=torch.int64,
|
|
device=device,
|
|
)
|
|
)
|
|
|
|
faces = torch.stack(faces)
|
|
verts = torch.stack(verts)
|
|
|
|
simpleces = Meshes(verts=verts, faces=faces)
|
|
|
|
return simpleces
|
|
|
|
def test_from_point_cloud(self, interp_mode="trilinear"):
|
|
"""
|
|
Generates a volume from a random point cloud sampled from faces
|
|
of a 3D cube. Since each side of the cube is homogeneously colored with
|
|
a different color, this should result in a volume with a
|
|
predefined homogeneous color of the cells along its borders
|
|
and black interior. The test is run for both cube and non-cube shaped
|
|
volumes.
|
|
"""
|
|
|
|
# batch_size = 4 sides of the cube
|
|
batch_size = 4
|
|
|
|
for volume_size in ([25, 25, 25], [30, 25, 15]):
|
|
|
|
for python, interp_mode in product([True, False], ["trilinear", "nearest"]):
|
|
|
|
(pointclouds, initial_volumes) = init_volume_boundary_pointcloud(
|
|
volume_size=volume_size,
|
|
n_points=int(1e5),
|
|
interp_mode=interp_mode,
|
|
batch_size=batch_size,
|
|
require_grad=True,
|
|
device="cuda:0",
|
|
)
|
|
|
|
volumes = add_pointclouds_to_volumes(
|
|
pointclouds,
|
|
initial_volumes,
|
|
mode=interp_mode,
|
|
_python=python,
|
|
)
|
|
|
|
V_color, V_density = volumes.features(), volumes.densities()
|
|
|
|
# expected colors of different cube sides
|
|
clr_sides = torch.tensor(
|
|
[
|
|
[[1.0, 1.0, 1.0], [1.0, 0.0, 1.0]],
|
|
[[1.0, 0.0, 0.0], [1.0, 1.0, 0.0]],
|
|
[[1.0, 0.0, 1.0], [1.0, 1.0, 1.0]],
|
|
[[1.0, 1.0, 0.0], [1.0, 0.0, 0.0]],
|
|
],
|
|
dtype=V_color.dtype,
|
|
device=V_color.device,
|
|
)
|
|
clr_ambient = torch.tensor(
|
|
[0.0, 0.0, 0.0], dtype=V_color.dtype, device=V_color.device
|
|
)
|
|
clr_top_bot = torch.tensor(
|
|
[[0.0, 1.0, 0.0], [0.0, 1.0, 1.0]],
|
|
dtype=V_color.dtype,
|
|
device=V_color.device,
|
|
)
|
|
|
|
if DEBUG:
|
|
outdir = tempfile.gettempdir() + "/test_points_to_volumes"
|
|
os.makedirs(outdir, exist_ok=True)
|
|
|
|
for slice_dim in (1, 2):
|
|
for vidx in range(V_color.shape[0]):
|
|
vim = V_color.detach()[vidx].split(1, dim=slice_dim)
|
|
vim = torch.stack([v.squeeze() for v in vim])
|
|
vim = TestPointsToVolumes.stack_4d_tensor_to_3d(vim.cpu())
|
|
im = Image.fromarray(
|
|
(vim.numpy() * 255.0)
|
|
.astype(np.uint8)
|
|
.transpose(1, 2, 0)
|
|
)
|
|
outfile = (
|
|
outdir
|
|
+ f"/rgb_{interp_mode}"
|
|
+ f"_{str(volume_size).replace(' ','')}"
|
|
+ f"_{vidx:003d}_sldim{slice_dim}.png"
|
|
)
|
|
im.save(outfile)
|
|
print("exported %s" % outfile)
|
|
|
|
# check the density V_density
|
|
# first binarize the density
|
|
V_density_bin = (V_density > 1e-4).type_as(V_density)
|
|
d_one = V_density.new_ones(1)
|
|
d_zero = V_density.new_zeros(1)
|
|
for vidx in range(V_color.shape[0]):
|
|
# the first/last depth-wise slice has to be filled with 1.0
|
|
self._check_volume_slice_color_density(
|
|
V_density_bin[vidx], 1, interp_mode, d_one, "first"
|
|
)
|
|
self._check_volume_slice_color_density(
|
|
V_density_bin[vidx], 1, interp_mode, d_one, "last"
|
|
)
|
|
# the middle depth-wise slices have to be empty
|
|
self._check_volume_slice_color_density(
|
|
V_density_bin[vidx], 1, interp_mode, d_zero, "middle"
|
|
)
|
|
# the top/bottom slices have to be filled with 1.0
|
|
self._check_volume_slice_color_density(
|
|
V_density_bin[vidx], 2, interp_mode, d_one, "first"
|
|
)
|
|
self._check_volume_slice_color_density(
|
|
V_density_bin[vidx], 2, interp_mode, d_one, "last"
|
|
)
|
|
|
|
# check the colors
|
|
for vidx in range(V_color.shape[0]):
|
|
self._check_volume_slice_color_density(
|
|
V_color[vidx], 1, interp_mode, clr_sides[vidx][0], "first"
|
|
)
|
|
self._check_volume_slice_color_density(
|
|
V_color[vidx], 1, interp_mode, clr_sides[vidx][1], "last"
|
|
)
|
|
self._check_volume_slice_color_density(
|
|
V_color[vidx], 1, interp_mode, clr_ambient, "middle"
|
|
)
|
|
self._check_volume_slice_color_density(
|
|
V_color[vidx], 2, interp_mode, clr_top_bot[0], "first"
|
|
)
|
|
self._check_volume_slice_color_density(
|
|
V_color[vidx], 2, interp_mode, clr_top_bot[1], "last"
|
|
)
|
|
|
|
# check differentiability
|
|
loss = V_color.mean() + V_density.mean()
|
|
loss.backward()
|
|
rgb = pointclouds.features_padded()
|
|
xyz = pointclouds.points_padded()
|
|
for field in (xyz, rgb):
|
|
if interp_mode == "nearest" and (field is xyz):
|
|
# this does not produce grads w.r.t. xyz
|
|
self.assertIsNone(field.grad)
|
|
else:
|
|
self.assertTrue(torch.isfinite(field.grad.data).all())
|
|
|
|
def _check_volume_slice_color_density(
|
|
self, V, split_dim, interp_mode, clr_gt, slice_type, border=3
|
|
):
|
|
# decompose the volume to individual slices along split_dim
|
|
vim = V.detach().split(1, dim=split_dim)
|
|
vim = torch.stack([v.squeeze(split_dim) for v in vim])
|
|
|
|
# determine which slices should be compared to clr_gt based on
|
|
# the 'slice_type' input
|
|
if slice_type == "first":
|
|
slice_dims = (0, 1) if interp_mode == "trilinear" else (0,)
|
|
elif slice_type == "last":
|
|
slice_dims = (-1, -2) if interp_mode == "trilinear" else (-1,)
|
|
elif slice_type == "middle":
|
|
internal_border = 2 if interp_mode == "trilinear" else 1
|
|
slice_dims = torch.arange(internal_border, vim.shape[0] - internal_border)
|
|
else:
|
|
raise ValueError(slice_type)
|
|
|
|
# compute the average error within each slice
|
|
clr_diff = (
|
|
vim[slice_dims, :, border:-border, border:-border]
|
|
- clr_gt[None, :, None, None]
|
|
)
|
|
clr_diff = clr_diff.abs().mean(dim=(2, 3)).view(-1)
|
|
|
|
# check that all per-slice avg errors vanish
|
|
self.assertClose(clr_diff, torch.zeros_like(clr_diff), atol=1e-2)
|
|
|
|
|
|
class TestRawFunction(TestCaseMixin, unittest.TestCase):
|
|
"""
|
|
Testing the _C.points_to_volumes function through its wrapper
|
|
_points_to_volumes.
|
|
"""
|
|
|
|
def setUp(self) -> None:
|
|
torch.manual_seed(42)
|
|
|
|
def test_grad_corners_splat_cpu(self):
|
|
self.do_gradcheck(torch.device("cpu"), True, True)
|
|
|
|
def test_grad_corners_round_cpu(self):
|
|
self.do_gradcheck(torch.device("cpu"), False, True)
|
|
|
|
def test_grad_splat_cpu(self):
|
|
self.do_gradcheck(torch.device("cpu"), True, False)
|
|
|
|
def test_grad_round_cpu(self):
|
|
self.do_gradcheck(torch.device("cpu"), False, False)
|
|
|
|
def test_grad_corners_splat_cuda(self):
|
|
self.do_gradcheck(torch.device("cuda:0"), True, True)
|
|
|
|
def test_grad_corners_round_cuda(self):
|
|
self.do_gradcheck(torch.device("cuda:0"), False, True)
|
|
|
|
def test_grad_splat_cuda(self):
|
|
self.do_gradcheck(torch.device("cuda:0"), True, False)
|
|
|
|
def test_grad_round_cuda(self):
|
|
self.do_gradcheck(torch.device("cuda:0"), False, False)
|
|
|
|
def do_gradcheck(self, device, splat: bool, align_corners: bool):
|
|
"""
|
|
Use gradcheck to verify the gradient of _points_to_volumes
|
|
with random input.
|
|
"""
|
|
N, C, D, H, W, P = 2, 4, 5, 6, 7, 5
|
|
points_3d = (
|
|
torch.rand((N, P, 3), device=device, dtype=torch.float64) * 0.8 + 0.1
|
|
)
|
|
points_features = torch.rand((N, P, C), device=device, dtype=torch.float64)
|
|
volume_densities = torch.zeros((N, 1, D, H, W), device=device)
|
|
volume_features = torch.zeros((N, C, D, H, W), device=device)
|
|
volume_densities_scale = torch.rand_like(volume_densities)
|
|
volume_features_scale = torch.rand_like(volume_features)
|
|
grid_sizes = torch.tensor([D, H, W], dtype=torch.int64, device=device).expand(
|
|
N, 3
|
|
)
|
|
mask = torch.ones((N, P), device=device)
|
|
mask[:, 0] = 0
|
|
align_corners = False
|
|
|
|
def f(points_3d_, points_features_):
|
|
(volume_densities_, volume_features_) = _points_to_volumes(
|
|
points_3d_.to(torch.float32),
|
|
points_features_.to(torch.float32),
|
|
volume_densities.clone(),
|
|
volume_features.clone(),
|
|
grid_sizes,
|
|
2.0,
|
|
mask,
|
|
align_corners,
|
|
splat,
|
|
)
|
|
density = (volume_densities_ * volume_densities_scale).sum()
|
|
features = (volume_features_ * volume_features_scale).sum()
|
|
return density, features
|
|
|
|
base = f(points_3d.clone(), points_features.clone())
|
|
self.assertGreater(base[0], 0)
|
|
self.assertGreater(base[1], 0)
|
|
|
|
points_features.requires_grad = True
|
|
if splat:
|
|
points_3d.requires_grad = True
|
|
torch.autograd.gradcheck(
|
|
f,
|
|
(points_3d, points_features),
|
|
check_undefined_grad=False,
|
|
eps=2e-4,
|
|
atol=0.01,
|
|
)
|
|
else:
|
|
torch.autograd.gradcheck(
|
|
partial(f, points_3d),
|
|
points_features,
|
|
check_undefined_grad=False,
|
|
eps=2e-3,
|
|
atol=0.001,
|
|
)
|
|
|
|
def test_single_corners_round_cpu(self):
|
|
self.single_point(torch.device("cpu"), False, True)
|
|
|
|
def test_single_corners_splat_cpu(self):
|
|
self.single_point(torch.device("cpu"), True, True)
|
|
|
|
def test_single_round_cpu(self):
|
|
self.single_point(torch.device("cpu"), False, False)
|
|
|
|
def test_single_splat_cpu(self):
|
|
self.single_point(torch.device("cpu"), True, False)
|
|
|
|
def test_single_corners_round_cuda(self):
|
|
self.single_point(torch.device("cuda:0"), False, True)
|
|
|
|
def test_single_corners_splat_cuda(self):
|
|
self.single_point(torch.device("cuda:0"), True, True)
|
|
|
|
def test_single_round_cuda(self):
|
|
self.single_point(torch.device("cuda:0"), False, False)
|
|
|
|
def test_single_splat_cuda(self):
|
|
self.single_point(torch.device("cuda:0"), True, False)
|
|
|
|
def single_point(self, device, splat: bool, align_corners: bool):
|
|
"""
|
|
Check the outcome of _points_to_volumes where a single point
|
|
exists which lines up with a single voxel.
|
|
"""
|
|
D, H, W = (6, 6, 11) if align_corners else (5, 5, 10)
|
|
N, C, P = 1, 1, 1
|
|
if align_corners:
|
|
points_3d = torch.tensor([[[-0.2, 0.2, -0.2]]], device=device)
|
|
else:
|
|
points_3d = torch.tensor([[[-0.3, 0.4, -0.4]]], device=device)
|
|
points_features = torch.zeros((N, P, C), device=device)
|
|
volume_densities = torch.zeros((N, 1, D, H, W), device=device)
|
|
volume_densities_expected = torch.zeros((N, 1, D, H, W), device=device)
|
|
volume_features = torch.zeros((N, C, D, H, W), device=device)
|
|
grid_sizes = torch.tensor([D, H, W], dtype=torch.int64, device=device).expand(
|
|
N, 3
|
|
)
|
|
mask = torch.ones((N, P), device=device)
|
|
point_weight = 19.0
|
|
|
|
volume_densities_, volume_features_ = _points_to_volumes(
|
|
points_3d,
|
|
points_features,
|
|
volume_densities,
|
|
volume_features,
|
|
grid_sizes,
|
|
point_weight,
|
|
mask,
|
|
align_corners,
|
|
splat,
|
|
)
|
|
|
|
self.assertIs(volume_densities, volume_densities_)
|
|
self.assertIs(volume_features, volume_features_)
|
|
|
|
if align_corners:
|
|
volume_densities_expected[0, 0, 2, 3, 4] = point_weight
|
|
else:
|
|
volume_densities_expected[0, 0, 1, 3, 3] = point_weight
|
|
|
|
self.assertClose(volume_densities, volume_densities_expected)
|