examples and docs.

Summary: This diff updates the documentation and tutorials with information about the new pulsar backend. For more information about the pulsar backend, see the release notes and the paper (https://arxiv.org/abs/2004.07484). For information on how to use the backend, see the point cloud rendering notebook and the examples in the folder docs/examples.

Reviewed By: nikhilaravi

Differential Revision: D24498129

fbshipit-source-id: e312b0169a72b13590df6e4db36bfe6190d742f9
This commit is contained in:
Christoph Lassner
2020-11-03 13:05:02 -08:00
committed by Facebook GitHub Bot
parent 960fd6d8b6
commit 039e02601d
21 changed files with 759 additions and 60 deletions

View File

@@ -5,6 +5,7 @@ This example demonstrates the most trivial, direct interface of the pulsar
sphere renderer. It renders and saves an image with 10 random spheres.
Output: basic.png.
"""
import math
from os import path
import imageio
@@ -12,11 +13,15 @@ import torch
from pytorch3d.renderer.points.pulsar import Renderer
torch.manual_seed(1)
n_points = 10
width = 1_000
height = 1_000
device = torch.device("cuda")
renderer = Renderer(width, height, n_points).to(device)
# The PyTorch3D system is right handed; in pulsar you can choose the handedness.
# For easy reproducibility we use a right handed coordinate system here.
renderer = Renderer(width, height, n_points, right_handed_system=True).to(device)
# Generate sample data.
vert_pos = torch.rand(n_points, 3, dtype=torch.float32, device=device) * 10.0
vert_pos[:, 2] += 25.0
@@ -29,7 +34,7 @@ cam_params = torch.tensor(
0.0,
0.0, # Position 0, 0, 0 (x, y, z).
0.0,
0.0,
math.pi, # Because of the right handed system, the camera must look 'back'.
0.0, # Rotation 0, 0, 0 (in axis-angle format).
5.0, # Focal length in world size.
2.0, # Sensor size in world size.

View File

@@ -0,0 +1,67 @@
#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
"""
This example demonstrates the most trivial use of the pulsar PyTorch3D
interface for sphere renderering. It renders and saves an image with
10 random spheres.
Output: basic-pt3d.png.
"""
from os import path
import imageio
import torch
from pytorch3d.renderer import PerspectiveCameras # , look_at_view_transform
from pytorch3d.renderer import (
PointsRasterizationSettings,
PointsRasterizer,
PulsarPointsRenderer,
)
from pytorch3d.structures import Pointclouds
torch.manual_seed(1)
n_points = 10
width = 1_000
height = 1_000
device = torch.device("cuda")
# Generate sample data.
vert_pos = torch.rand(n_points, 3, dtype=torch.float32, device=device) * 10.0
vert_pos[:, 2] += 25.0
vert_pos[:, :2] -= 5.0
vert_col = torch.rand(n_points, 3, dtype=torch.float32, device=device)
pcl = Pointclouds(points=vert_pos[None, ...], features=vert_col[None, ...])
# Alternatively, you can also use the look_at_view_transform to get R and T:
# R, T = look_at_view_transform(
# dist=30.0, elev=0.0, azim=180.0, at=((0.0, 0.0, 30.0),), up=((0, 1, 0),),
# )
cameras = PerspectiveCameras(
# The focal length must be double the size for PyTorch3D because of the NDC
# coordinates spanning a range of two - and they must be normalized by the
# sensor width (see the pulsar example). This means we need here
# 5.0 * 2.0 / 2.0 to get the equivalent results as in pulsar.
focal_length=(5.0 * 2.0 / 2.0,),
R=torch.eye(3, dtype=torch.float32, device=device)[None, ...],
T=torch.zeros((1, 3), dtype=torch.float32, device=device),
image_size=((width, height),),
device=device,
)
vert_rad = torch.rand(n_points, dtype=torch.float32, device=device)
raster_settings = PointsRasterizationSettings(
image_size=(width, height),
radius=vert_rad,
)
rasterizer = PointsRasterizer(cameras=cameras, raster_settings=raster_settings)
renderer = PulsarPointsRenderer(rasterizer=rasterizer).to(device)
# Render.
image = renderer(
pcl,
gamma=(1.0e-1,), # Renderer blending parameter gamma, in [1., 1e-5].
znear=(1.0,),
zfar=(45.0,),
radius_world=True,
bg_col=torch.ones((3,), dtype=torch.float32, device=device),
)[0]
print("Writing image to `%s`." % (path.abspath("basic-pt3d.png")))
imageio.imsave("basic-pt3d.png", (image.cpu().detach() * 255.0).to(torch.uint8).numpy())

View File

@@ -7,7 +7,9 @@ pulsar interface. For this, a reference image has been pre-generated
The same scene parameterization is loaded and the camera parameters
distorted. Gradient-based optimization is used to converge towards the
original camera parameters.
Output: cam.gif.
"""
import math
from os import path
import cv2
@@ -15,6 +17,7 @@ import imageio
import numpy as np
import torch
from pytorch3d.renderer.points.pulsar import Renderer
from pytorch3d.transforms import axis_angle_to_matrix, matrix_to_rotation_6d
from torch import nn, optim
@@ -66,19 +69,18 @@ class SceneModel(nn.Module):
)
self.register_parameter(
"cam_rot",
# We're using the 6D rot. representation for better gradients.
nn.Parameter(
torch.tensor(
[
# We're using the 6D rot. representation for better gradients.
0.9995,
0.0300445,
-0.0098482,
-0.0299445,
0.9995,
0.0101482,
],
dtype=torch.float32,
),
matrix_to_rotation_6d(
axis_angle_to_matrix(
torch.tensor(
[
[0.02, math.pi + 0.02, 0.01],
],
dtype=torch.float32,
)
)
)[0],
requires_grad=True,
),
)
@@ -88,7 +90,7 @@ class SceneModel(nn.Module):
torch.tensor([4.8, 1.8], dtype=torch.float32), requires_grad=True
),
)
self.renderer = Renderer(width, height, n_points)
self.renderer = Renderer(width, height, n_points, right_handed_system=True)
def forward(self):
return self.renderer.forward(
@@ -106,7 +108,7 @@ ref = (
torch.from_numpy(
imageio.imread(
"../../tests/pulsar/reference/examples_TestRenderer_test_cam.png"
)
)[:, ::-1, :].copy()
).to(torch.float32)
/ 255.0
).to(device)

View File

@@ -0,0 +1,210 @@
#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
"""
This example demonstrates camera parameter optimization with the pulsar
PyTorch3D interface. For this, a reference image has been pre-generated
(you can find it at `../../tests/pulsar/reference/examples_TestRenderer_test_cam.png`).
The same scene parameterization is loaded and the camera parameters
distorted. Gradient-based optimization is used to converge towards the
original camera parameters.
Output: cam-pt3d.gif
"""
from os import path
import cv2
import imageio
import numpy as np
import torch
from pytorch3d.renderer.cameras import PerspectiveCameras # , look_at_view_transform
from pytorch3d.renderer.points import (
PointsRasterizationSettings,
PointsRasterizer,
PulsarPointsRenderer,
)
from pytorch3d.structures.pointclouds import Pointclouds
from pytorch3d.transforms import axis_angle_to_matrix
from torch import nn, optim
n_points = 20
width = 1_000
height = 1_000
device = torch.device("cuda")
class SceneModel(nn.Module):
"""
A simple scene model to demonstrate use of pulsar in PyTorch modules.
The scene model is parameterized with sphere locations (vert_pos),
channel content (vert_col), radiuses (vert_rad), camera position (cam_pos),
camera rotation (cam_rot) and sensor focal length and width (cam_sensor).
The forward method of the model renders this scene description. Any
of these parameters could instead be passed as inputs to the forward
method and come from a different model.
"""
def __init__(self):
super(SceneModel, self).__init__()
self.gamma = 0.1
# Points.
torch.manual_seed(1)
vert_pos = torch.rand(n_points, 3, dtype=torch.float32) * 10.0
vert_pos[:, 2] += 25.0
vert_pos[:, :2] -= 5.0
self.register_parameter("vert_pos", nn.Parameter(vert_pos, requires_grad=False))
self.register_parameter(
"vert_col",
nn.Parameter(
torch.rand(n_points, 3, dtype=torch.float32),
requires_grad=False,
),
)
self.register_parameter(
"vert_rad",
nn.Parameter(
torch.rand(n_points, dtype=torch.float32),
requires_grad=False,
),
)
self.register_parameter(
"cam_pos",
nn.Parameter(
torch.tensor([0.1, 0.1, 0.0], dtype=torch.float32),
requires_grad=True,
),
)
self.register_parameter(
"cam_rot",
# We're using the 6D rot. representation for better gradients.
nn.Parameter(
axis_angle_to_matrix(
torch.tensor(
[
[0.02, 0.02, 0.01],
],
dtype=torch.float32,
)
)[0],
requires_grad=True,
),
)
self.register_parameter(
"focal_length",
nn.Parameter(
torch.tensor(
[
4.8 * 2.0 / 2.0,
],
dtype=torch.float32,
),
requires_grad=True,
),
)
self.cameras = PerspectiveCameras(
# The focal length must be double the size for PyTorch3D because of the NDC
# coordinates spanning a range of two - and they must be normalized by the
# sensor width (see the pulsar example). This means we need here
# 5.0 * 2.0 / 2.0 to get the equivalent results as in pulsar.
#
# R, T and f are provided here, but will be provided again
# at every call to the forward method. The reason are problems
# with PyTorch which makes device placement for gradients problematic
# for tensors which are themselves on a 'gradient path' but not
# leafs in the calculation tree. This will be addressed by an architectural
# change in PyTorch3D in the future. Until then, this workaround is
# recommended.
focal_length=self.focal_length,
R=self.cam_rot[None, ...],
T=self.cam_pos[None, ...],
image_size=((width, height),),
device=device,
)
raster_settings = PointsRasterizationSettings(
image_size=(width, height),
radius=self.vert_rad,
)
rasterizer = PointsRasterizer(
cameras=self.cameras, raster_settings=raster_settings
)
self.renderer = PulsarPointsRenderer(rasterizer=rasterizer)
def forward(self):
# The Pointclouds object creates copies of it's arguments - that's why
# we have to create a new object in every forward step.
pcl = Pointclouds(
points=self.vert_pos[None, ...], features=self.vert_col[None, ...]
)
return self.renderer(
pcl,
gamma=(self.gamma,),
zfar=(45.0,),
znear=(1.0,),
radius_world=True,
bg_col=torch.ones((3,), dtype=torch.float32, device=device),
# As mentioned above: workaround for device placement of gradients for
# camera parameters.
focal_length=self.focal_length,
R=self.cam_rot[None, ...],
T=self.cam_pos[None, ...],
)[0]
# Load reference.
ref = (
torch.from_numpy(
imageio.imread(
"../../tests/pulsar/reference/examples_TestRenderer_test_cam.png"
)[:, ::-1, :].copy()
).to(torch.float32)
/ 255.0
).to(device)
# Set up model.
model = SceneModel().to(device)
# Optimizer.
optimizer = optim.SGD(
[
{"params": [model.cam_pos], "lr": 1e-4},
{"params": [model.cam_rot], "lr": 5e-6},
# Using a higher lr for the focal length here, because
# the sensor width can not be optimized directly.
{"params": [model.focal_length], "lr": 1e-3},
]
)
print("Writing video to `%s`." % (path.abspath("cam-pt3d.gif")))
writer = imageio.get_writer("cam-pt3d.gif", format="gif", fps=25)
# Optimize.
for i in range(300):
optimizer.zero_grad()
result = model()
# Visualize.
result_im = (result.cpu().detach().numpy() * 255).astype(np.uint8)
cv2.imshow("opt", result_im[:, :, ::-1])
writer.append_data(result_im)
overlay_img = np.ascontiguousarray(
((result * 0.5 + ref * 0.5).cpu().detach().numpy() * 255).astype(np.uint8)[
:, :, ::-1
]
)
overlay_img = cv2.putText(
overlay_img,
"Step %d" % (i),
(10, 40),
cv2.FONT_HERSHEY_SIMPLEX,
1,
(0, 0, 0),
2,
cv2.LINE_AA,
False,
)
cv2.imshow("overlay", overlay_img)
cv2.waitKey(1)
# Update.
loss = ((result - ref) ** 2).sum()
print("loss {}: {}".format(i, loss.item()))
loss.backward()
optimizer.step()
writer.close()

View File

@@ -7,7 +7,12 @@ pulsar interface. For this, reference images have been pre-generated
The camera parameters are assumed given. The scene is initialized with
random spheres. Gradient-based optimization is used to optimize sphere
parameters and prune spheres to converge to a 3D representation.
This example is not available yet through the 'unified' interface,
because opacity support has not landed in PyTorch3D for general data
structures yet.
"""
import math
from os import path
import cv2
@@ -77,7 +82,7 @@ class SceneModel(nn.Module):
0.0,
30.0 - np.cos(angle) * 35.0,
0.0,
-angle,
-angle + math.pi,
0.0,
5.0,
2.0,
@@ -87,7 +92,7 @@ class SceneModel(nn.Module):
dtype=torch.float32,
),
)
self.renderer = Renderer(width, height, n_points)
self.renderer = Renderer(width, height, n_points, right_handed_system=True)
def forward(self, cam=None):
if cam is None:
@@ -184,7 +189,7 @@ for i in range(300):
0.0,
30.0 - np.cos(angle) * 35.0,
0.0,
-angle,
-angle + math.pi,
0.0,
5.0,
2.0,

View File

@@ -8,6 +8,8 @@ The scene is initialized with random spheres. Gradient-based
optimization is used to converge towards a faithful
scene representation.
"""
import math
import cv2
import imageio
import numpy as np
@@ -58,11 +60,15 @@ class SceneModel(nn.Module):
)
self.register_buffer(
"cam_params",
torch.tensor([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 5.0, 2.0], dtype=torch.float32),
torch.tensor(
[0.0, 0.0, 0.0, 0.0, math.pi, 0.0, 5.0, 2.0], dtype=torch.float32
),
)
# The volumetric optimization works better with a higher number of tracked
# intersections per ray.
self.renderer = Renderer(width, height, n_points, n_track=32)
self.renderer = Renderer(
width, height, n_points, n_track=32, right_handed_system=True
)
def forward(self):
return self.renderer.forward(
@@ -81,7 +87,7 @@ ref = (
torch.from_numpy(
imageio.imread(
"../../tests/pulsar/reference/examples_TestRenderer_test_smallopt.png"
)
)[:, ::-1, :].copy()
).to(torch.float32)
/ 255.0
).to(device)

View File

@@ -0,0 +1,171 @@
#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
"""
This example demonstrates scene optimization with the PyTorch3D
pulsar interface. For this, a reference image has been pre-generated
(you can find it at `../../tests/pulsar/reference/examples_TestRenderer_test_smallopt.png`).
The scene is initialized with random spheres. Gradient-based
optimization is used to converge towards a faithful
scene representation.
"""
import math
import cv2
import imageio
import numpy as np
import torch
from pytorch3d.renderer.cameras import PerspectiveCameras # , look_at_view_transform
from pytorch3d.renderer.points import (
PointsRasterizationSettings,
PointsRasterizer,
PulsarPointsRenderer,
)
from pytorch3d.structures.pointclouds import Pointclouds
from torch import nn, optim
n_points = 10_000
width = 1_000
height = 1_000
device = torch.device("cuda")
class SceneModel(nn.Module):
"""
A simple scene model to demonstrate use of pulsar in PyTorch modules.
The scene model is parameterized with sphere locations (vert_pos),
channel content (vert_col), radiuses (vert_rad), camera position (cam_pos),
camera rotation (cam_rot) and sensor focal length and width (cam_sensor).
The forward method of the model renders this scene description. Any
of these parameters could instead be passed as inputs to the forward
method and come from a different model.
"""
def __init__(self):
super(SceneModel, self).__init__()
self.gamma = 1.0
# Points.
torch.manual_seed(1)
vert_pos = torch.rand(n_points, 3, dtype=torch.float32, device=device) * 10.0
vert_pos[:, 2] += 25.0
vert_pos[:, :2] -= 5.0
self.register_parameter("vert_pos", nn.Parameter(vert_pos, requires_grad=True))
self.register_parameter(
"vert_col",
nn.Parameter(
torch.ones(n_points, 3, dtype=torch.float32, device=device) * 0.5,
requires_grad=True,
),
)
self.register_parameter(
"vert_rad",
nn.Parameter(
torch.ones(n_points, dtype=torch.float32) * 0.3, requires_grad=True
),
)
self.register_buffer(
"cam_params",
torch.tensor(
[0.0, 0.0, 0.0, 0.0, math.pi, 0.0, 5.0, 2.0], dtype=torch.float32
),
)
self.cameras = PerspectiveCameras(
# The focal length must be double the size for PyTorch3D because of the NDC
# coordinates spanning a range of two - and they must be normalized by the
# sensor width (see the pulsar example). This means we need here
# 5.0 * 2.0 / 2.0 to get the equivalent results as in pulsar.
focal_length=5.0,
R=torch.eye(3, dtype=torch.float32, device=device)[None, ...],
T=torch.zeros((1, 3), dtype=torch.float32, device=device),
image_size=((width, height),),
device=device,
)
raster_settings = PointsRasterizationSettings(
image_size=(width, height),
radius=self.vert_rad,
)
rasterizer = PointsRasterizer(
cameras=self.cameras, raster_settings=raster_settings
)
self.renderer = PulsarPointsRenderer(rasterizer=rasterizer, n_track=32)
def forward(self):
# The Pointclouds object creates copies of it's arguments - that's why
# we have to create a new object in every forward step.
pcl = Pointclouds(
points=self.vert_pos[None, ...], features=self.vert_col[None, ...]
)
return self.renderer(
pcl,
gamma=(self.gamma,),
zfar=(45.0,),
znear=(1.0,),
radius_world=True,
bg_col=torch.ones((3,), dtype=torch.float32, device=device),
)[0]
# Load reference.
ref = (
torch.from_numpy(
imageio.imread(
"../../tests/pulsar/reference/examples_TestRenderer_test_smallopt.png"
)[:, ::-1, :].copy()
).to(torch.float32)
/ 255.0
).to(device)
# Set up model.
model = SceneModel().to(device)
# Optimizer.
optimizer = optim.SGD(
[
{"params": [model.vert_col], "lr": 1e0},
{"params": [model.vert_rad], "lr": 5e-3},
{"params": [model.vert_pos], "lr": 1e-2},
]
)
# Optimize.
for i in range(500):
optimizer.zero_grad()
result = model()
# Visualize.
result_im = (result.cpu().detach().numpy() * 255).astype(np.uint8)
cv2.imshow("opt", result_im[:, :, ::-1])
overlay_img = np.ascontiguousarray(
((result * 0.5 + ref * 0.5).cpu().detach().numpy() * 255).astype(np.uint8)[
:, :, ::-1
]
)
overlay_img = cv2.putText(
overlay_img,
"Step %d" % (i),
(10, 40),
cv2.FONT_HERSHEY_SIMPLEX,
1,
(0, 0, 0),
2,
cv2.LINE_AA,
False,
)
cv2.imshow("overlay", overlay_img)
cv2.waitKey(1)
# Update.
loss = ((result - ref) ** 2).sum()
print("loss {}: {}".format(i, loss.item()))
loss.backward()
optimizer.step()
# Cleanup.
with torch.no_grad():
model.vert_col.data = torch.clamp(model.vert_col.data, 0.0, 1.0)
# Remove points.
model.vert_pos.data[model.vert_rad < 0.001, :] = -1000.0
model.vert_rad.data[model.vert_rad < 0.001] = 0.0001
vd = (
(model.vert_col - torch.ones(3, dtype=torch.float32).to(device))
.abs()
.sum(dim=1)
)
model.vert_pos.data[vd <= 0.2] = -1000.0