mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-12-19 05:40:34 +08:00
Convert from Pytorch3D NDC coordinates to grid_sample coordinates.
Summary: Implements a utility function to convert from 2D coordinates in Pytorch3D NDC space to the coordinates in grid_sample. Reviewed By: shapovalov Differential Revision: D33741394 fbshipit-source-id: 88981653356588fe646e6dea48fe7f7298738437
This commit is contained in:
David Novotny
committed by
Facebook GitHub Bot
Facebook GitHub Bot
parent
47c0997227
commit
12f20d799e
@@ -10,7 +10,20 @@ import unittest
|
||||
import numpy as np
|
||||
import torch
|
||||
from common_testing import TestCaseMixin
|
||||
from pytorch3d.renderer.utils import TensorProperties
|
||||
from pytorch3d.ops import eyes
|
||||
from pytorch3d.renderer import (
|
||||
PerspectiveCameras,
|
||||
AlphaCompositor,
|
||||
PointsRenderer,
|
||||
PointsRasterizationSettings,
|
||||
PointsRasterizer,
|
||||
)
|
||||
from pytorch3d.renderer.utils import (
|
||||
TensorProperties,
|
||||
ndc_to_grid_sample_coords,
|
||||
ndc_grid_sample,
|
||||
)
|
||||
from pytorch3d.structures import Pointclouds
|
||||
|
||||
|
||||
# Example class for testing
|
||||
@@ -96,3 +109,165 @@ class TestTensorProperties(TestCaseMixin, unittest.TestCase):
|
||||
# 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(
|
||||
torch.meshgrid(
|
||||
[torch.linspace(-grid_scale, grid_scale, n_grid_pts)] * 2, indexing="ij"
|
||||
),
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user