mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-07-31 10:52:50 +08:00
055ab3a2e3
Summary:
Converts the directory specified to use the Ruff formatter in pyfmt
ruff_dog
If this diff causes merge conflicts when rebasing, please run
`hg status -n -0 --change . -I '**/*.{py,pyi}' | xargs -0 arc pyfmt`
on your diff, and amend any changes before rebasing onto latest.
That should help reduce or eliminate any merge conflicts.
allow-large-files
Reviewed By: bottler
Differential Revision: D66472063
fbshipit-source-id: 35841cb397e4f8e066e2159550d2f56b403b1bef
275 lines
9.1 KiB
Python
275 lines
9.1 KiB
Python
# 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
|
|
|
|
import numpy as np
|
|
import torch
|
|
from pytorch3d.common.compat import meshgrid_ij
|
|
from pytorch3d.ops import eyes
|
|
from pytorch3d.renderer import (
|
|
AlphaCompositor,
|
|
PerspectiveCameras,
|
|
PointsRasterizationSettings,
|
|
PointsRasterizer,
|
|
PointsRenderer,
|
|
)
|
|
from pytorch3d.renderer.utils import (
|
|
ndc_grid_sample,
|
|
ndc_to_grid_sample_coords,
|
|
TensorProperties,
|
|
)
|
|
from pytorch3d.structures import Pointclouds
|
|
|
|
from .common_testing import TestCaseMixin
|
|
|
|
|
|
# Example class for testing
|
|
class TensorPropertiesTestClass(TensorProperties):
|
|
def __init__(self, x=None, y=None, device="cpu"):
|
|
super().__init__(device=device, x=x, y=y)
|
|
|
|
def clone(self):
|
|
other = TensorPropertiesTestClass()
|
|
return super().clone(other)
|
|
|
|
|
|
class TestTensorProperties(TestCaseMixin, unittest.TestCase):
|
|
def test_init(self):
|
|
example = TensorPropertiesTestClass(x=10.0, y=(100.0, 200.0))
|
|
# Check kwargs set as attributes + converted to tensors
|
|
self.assertTrue(torch.is_tensor(example.x))
|
|
self.assertTrue(torch.is_tensor(example.y))
|
|
# Check broadcasting
|
|
self.assertTrue(example.x.shape == (2,))
|
|
self.assertTrue(example.y.shape == (2,))
|
|
self.assertTrue(len(example) == 2)
|
|
|
|
def test_to(self):
|
|
# Check to method
|
|
example = TensorPropertiesTestClass(x=10.0, y=(100.0, 200.0))
|
|
device = torch.device("cuda:0")
|
|
new_example = example.to(device=device)
|
|
self.assertEqual(new_example.device, device)
|
|
|
|
example_cpu = example.cpu()
|
|
self.assertEqual(example_cpu.device, torch.device("cpu"))
|
|
|
|
example_gpu = example.cuda()
|
|
self.assertEqual(example_gpu.device.type, "cuda")
|
|
self.assertIsNotNone(example_gpu.device.index)
|
|
|
|
example_gpu1 = example.cuda(1)
|
|
self.assertEqual(example_gpu1.device, torch.device("cuda:1"))
|
|
|
|
def test_clone(self):
|
|
# Check clone method
|
|
example = TensorPropertiesTestClass(x=10.0, y=(100.0, 200.0))
|
|
new_example = example.clone()
|
|
self.assertSeparate(example.x, new_example.x)
|
|
self.assertSeparate(example.y, new_example.y)
|
|
|
|
def test_get_set(self):
|
|
# Test getitem returns an accessor which can be used to modify
|
|
# attributes at a particular index
|
|
example = TensorPropertiesTestClass(x=10.0, y=(100.0, 200.0, 300.0))
|
|
|
|
# update y1
|
|
example[1].y = 5.0
|
|
self.assertTrue(example.y[1] == 5.0)
|
|
|
|
# Get item and get value
|
|
ex0 = example[0]
|
|
self.assertTrue(ex0.y == 100.0)
|
|
|
|
def test_empty_input(self):
|
|
example = TensorPropertiesTestClass(x=(), y=())
|
|
self.assertTrue(len(example) == 0)
|
|
self.assertTrue(example.isempty())
|
|
|
|
def test_gather_props(self):
|
|
N = 4
|
|
x = torch.randn((N, 3, 4))
|
|
y = torch.randn((N, 5))
|
|
test_class = TensorPropertiesTestClass(x=x, y=y)
|
|
|
|
S = 15
|
|
idx = torch.tensor(np.random.choice(N, S))
|
|
test_class_gathered = test_class.gather_props(idx)
|
|
|
|
self.assertTrue(test_class_gathered.x.shape == (S, 3, 4))
|
|
self.assertTrue(test_class_gathered.y.shape == (S, 5))
|
|
|
|
for i in range(N):
|
|
inds = idx == i
|
|
if inds.sum() > 0:
|
|
# Check the gathered points in the output have the same value from
|
|
# the input.
|
|
self.assertClose(test_class_gathered.x[inds].mean(dim=0), x[i, ...])
|
|
self.assertClose(test_class_gathered.y[inds].mean(dim=0), y[i, ...])
|
|
|
|
def test_ndc_grid_sample_rendering(self):
|
|
"""
|
|
Use PyTorch3D point renderer to render a colored point cloud, then
|
|
sample the image at the locations of the point projections with
|
|
`ndc_grid_sample`. Finally, assert that the sampled colors are equal to the
|
|
original point cloud colors.
|
|
|
|
Note that, in order to ensure correctness, we use a nearest-neighbor
|
|
assignment point renderer (i.e. no soft splatting).
|
|
"""
|
|
|
|
# generate a bunch of 3D points on a regular grid lying in the z-plane
|
|
n_grid_pts = 10
|
|
grid_scale = 0.9
|
|
z_plane = 2.0
|
|
image_size = [128, 128]
|
|
point_radius = 0.015
|
|
n_pts = n_grid_pts * n_grid_pts
|
|
pts = torch.stack(
|
|
meshgrid_ij(
|
|
[torch.linspace(-grid_scale, grid_scale, n_grid_pts)] * 2,
|
|
),
|
|
dim=-1,
|
|
)
|
|
pts = torch.cat([pts, z_plane * torch.ones_like(pts[..., :1])], dim=-1)
|
|
pts = pts.reshape(1, n_pts, 3)
|
|
|
|
# color the points randomly
|
|
pts_colors = torch.rand(1, n_pts, 3)
|
|
|
|
# make trivial rendering cameras
|
|
cameras = PerspectiveCameras(
|
|
R=eyes(dim=3, N=1),
|
|
device=pts.device,
|
|
T=torch.zeros(1, 3, dtype=torch.float32, device=pts.device),
|
|
)
|
|
|
|
# render the point cloud
|
|
pcl = Pointclouds(points=pts, features=pts_colors)
|
|
renderer = NearestNeighborPointsRenderer(
|
|
rasterizer=PointsRasterizer(
|
|
cameras=cameras,
|
|
raster_settings=PointsRasterizationSettings(
|
|
image_size=image_size,
|
|
radius=point_radius,
|
|
points_per_pixel=1,
|
|
),
|
|
),
|
|
compositor=AlphaCompositor(),
|
|
)
|
|
im_render = renderer(pcl)
|
|
|
|
# sample the render at projected pts
|
|
pts_proj = cameras.transform_points(pcl.points_padded())[..., :2]
|
|
pts_colors_sampled = ndc_grid_sample(
|
|
im_render,
|
|
pts_proj,
|
|
mode="nearest",
|
|
align_corners=False,
|
|
).permute(0, 2, 1)
|
|
|
|
# assert that the samples are the same as original points
|
|
self.assertClose(pts_colors, pts_colors_sampled, atol=1e-4)
|
|
|
|
def test_ndc_to_grid_sample_coords(self):
|
|
"""
|
|
Test the conversion from ndc to grid_sample coords by comparing
|
|
to known conversion results.
|
|
"""
|
|
|
|
# square image tests
|
|
image_size_square = [100, 100]
|
|
xy_ndc_gs_square = torch.FloatTensor(
|
|
[
|
|
# 4 corners
|
|
[[-1.0, -1.0], [1.0, 1.0]],
|
|
[[1.0, 1.0], [-1.0, -1.0]],
|
|
[[1.0, -1.0], [-1.0, 1.0]],
|
|
[[1.0, 1.0], [-1.0, -1.0]],
|
|
# center
|
|
[[0.0, 0.0], [0.0, 0.0]],
|
|
]
|
|
)
|
|
|
|
# non-batched version
|
|
for xy_ndc, xy_gs in xy_ndc_gs_square:
|
|
xy_gs_predicted = ndc_to_grid_sample_coords(
|
|
xy_ndc,
|
|
image_size_square,
|
|
)
|
|
self.assertClose(xy_gs_predicted, xy_gs)
|
|
|
|
# batched version
|
|
xy_ndc, xy_gs = xy_ndc_gs_square[:, 0], xy_ndc_gs_square[:, 1]
|
|
xy_gs_predicted = ndc_to_grid_sample_coords(
|
|
xy_ndc,
|
|
image_size_square,
|
|
)
|
|
self.assertClose(xy_gs_predicted, xy_gs)
|
|
|
|
# non-square image tests
|
|
image_size = [100, 200]
|
|
xy_ndc_gs = torch.FloatTensor(
|
|
[
|
|
# 4 corners
|
|
[[-2.0, -1.0], [1.0, 1.0]],
|
|
[[2.0, -1.0], [-1.0, 1.0]],
|
|
[[-2.0, 1.0], [1.0, -1.0]],
|
|
[[2.0, 1.0], [-1.0, -1.0]],
|
|
# center
|
|
[[0.0, 0.0], [0.0, 0.0]],
|
|
# non-corner points
|
|
[[4.0, 0.5], [-2.0, -0.5]],
|
|
[[1.0, -0.5], [-0.5, 0.5]],
|
|
]
|
|
)
|
|
|
|
# check both H > W and W > H
|
|
for flip_axes in [False, True]:
|
|
# non-batched version
|
|
for xy_ndc, xy_gs in xy_ndc_gs:
|
|
xy_gs_predicted = ndc_to_grid_sample_coords(
|
|
xy_ndc.flip(dims=(-1,)) if flip_axes else xy_ndc,
|
|
list(reversed(image_size)) if flip_axes else image_size,
|
|
)
|
|
self.assertClose(
|
|
xy_gs_predicted,
|
|
xy_gs.flip(dims=(-1,)) if flip_axes else xy_gs,
|
|
)
|
|
|
|
# batched version
|
|
xy_ndc, xy_gs = xy_ndc_gs[:, 0], xy_ndc_gs[:, 1]
|
|
xy_gs_predicted = ndc_to_grid_sample_coords(
|
|
xy_ndc.flip(dims=(-1,)) if flip_axes else xy_ndc,
|
|
list(reversed(image_size)) if flip_axes else image_size,
|
|
)
|
|
self.assertClose(
|
|
xy_gs_predicted,
|
|
xy_gs.flip(dims=(-1,)) if flip_axes else xy_gs,
|
|
)
|
|
|
|
|
|
class NearestNeighborPointsRenderer(PointsRenderer):
|
|
"""
|
|
A class for rendering a batch of points by a trivial nearest
|
|
neighbor assignment.
|
|
"""
|
|
|
|
def forward(self, point_clouds, **kwargs) -> torch.Tensor:
|
|
fragments = self.rasterizer(point_clouds, **kwargs)
|
|
# set all weights trivially to one
|
|
dists2 = fragments.dists.permute(0, 3, 1, 2)
|
|
weights = torch.ones_like(dists2)
|
|
images = self.compositor(
|
|
fragments.idx.long().permute(0, 3, 1, 2),
|
|
weights,
|
|
point_clouds.features_packed().permute(1, 0),
|
|
**kwargs,
|
|
)
|
|
return images
|