423A35C7/pytorch3d
1
0
Fork 0
mirror of https://github.com/facebookresearch/pytorch3d.git synced 2025-08-02 20:02:49 +08:00
Code Issues Packages Projects Releases Wiki Activity

Example and test updates.

Browse Source 
Summary: This commit performs pulsar example and test refinements. The examples are fully adjusted to adhere to PEP style guide and additional comments are added.

Reviewed By: nikhilaravi

Differential Revision: D24723391

fbshipit-source-id: 6d289006f080140159731e7f3a8c98b582164f1a
This commit is contained in:
Christoph Lassner 2020-11-04 09:53:19 -08:00 committed by Facebook GitHub Bot
parent e9a26f263a
commit b6be3b95fb
9 changed files with 569 additions and 448 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

94
docs/examples/pulsar_basic.py
View File

@ -7,49 +7,65 @@ Output: basic.png.
""" """
import math import math
from os import path from os import path
import logging
import imageio import imageio
import torch import torch
from pytorch3d.renderer.points.pulsar import Renderer from pytorch3d.renderer.points.pulsar import Renderer
torch.manual_seed(1) LOGGER = logging.getLogger(__name__)
n_points = 10
width = 1_000 def cli():
height = 1_000 """
device = torch.device("cuda") Basic example for the pulsar sphere renderer.
# The PyTorch3D system is right handed; in pulsar you can choose the handedness.
# For easy reproducibility we use a right handed coordinate system here. Writes to `basic.png`.
renderer = Renderer(width, height, n_points, right_handed_system=True).to(device) """
# Generate sample data. LOGGER.info("Rendering on GPU...")
vert_pos = torch.rand(n_points, 3, dtype=torch.float32, device=device) * 10.0 torch.manual_seed(1)
vert_pos[:, 2] += 25.0 n_points = 10
vert_pos[:, :2] -= 5.0 width = 1_000
vert_col = torch.rand(n_points, 3, dtype=torch.float32, device=device) height = 1_000
vert_rad = torch.rand(n_points, dtype=torch.float32, device=device) device = torch.device("cuda")
cam_params = torch.tensor( # The PyTorch3D system is right handed; in pulsar you can choose the handedness.
[ # For easy reproducibility we use a right handed coordinate system here.
0.0, renderer = Renderer(width, height, n_points, right_handed_system=True).to(device)
0.0, # Generate sample data.
0.0, # Position 0, 0, 0 (x, y, z). vert_pos = torch.rand(n_points, 3, dtype=torch.float32, device=device) * 10.0
0.0, vert_pos[:, 2] += 25.0
math.pi, # Because of the right handed system, the camera must look 'back'. vert_pos[:, :2] -= 5.0
0.0, # Rotation 0, 0, 0 (in axis-angle format). vert_col = torch.rand(n_points, 3, dtype=torch.float32, device=device)
5.0, # Focal length in world size. vert_rad = torch.rand(n_points, dtype=torch.float32, device=device)
2.0, # Sensor size in world size. cam_params = torch.tensor(
], [
dtype=torch.float32, 0.0,
device=device, 0.0,
) 0.0, # Position 0, 0, 0 (x, y, z).
# Render. 0.0,
image = renderer( math.pi, # Because of the right handed system, the camera must look 'back'.
vert_pos, 0.0, # Rotation 0, 0, 0 (in axis-angle format).
vert_col, 5.0, # Focal length in world size.
vert_rad, 2.0, # Sensor size in world size.
cam_params, ],
1.0e-1, # Renderer blending parameter gamma, in [1., 1e-5]. dtype=torch.float32,
45.0, # Maximum depth. device=device,
) )
print("Writing image to `%s`." % (path.abspath("basic.png"))) # Render.
imageio.imsave("basic.png", (image.cpu().detach() * 255.0).to(torch.uint8).numpy()) image = renderer(
vert_pos,
vert_col,
vert_rad,
cam_params,
1.0e-1, # Renderer blending parameter gamma, in [1., 1e-5].
45.0, # Maximum depth.
)
LOGGER.info("Writing image to `%s`.", path.abspath("basic.png"))
imageio.imsave("basic.png", (image.cpu().detach() * 255.0).to(torch.uint8).numpy())
LOGGER.info("Done.")
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO)
cli()

108
docs/examples/pulsar_basic_unified.py
View File

@ -6,10 +6,14 @@ interface for sphere renderering. It renders and saves an image with
10 random spheres. 10 random spheres.
Output: basic-pt3d.png. Output: basic-pt3d.png.
""" """
import logging
from os import path from os import path
import imageio import imageio
import torch import torch
# Import `look_at_view_transform` as needed in the suggestion later in the
# example.
from pytorch3d.renderer import PerspectiveCameras # , look_at_view_transform from pytorch3d.renderer import PerspectiveCameras # , look_at_view_transform
from pytorch3d.renderer import ( from pytorch3d.renderer import (
PointsRasterizationSettings, PointsRasterizationSettings,
@ -19,49 +23,65 @@ from pytorch3d.renderer import (
from pytorch3d.structures import Pointclouds from pytorch3d.structures import Pointclouds
torch.manual_seed(1) LOGGER = logging.getLogger(__name__)
n_points = 10
width = 1_000
height = 1_000
device = torch.device("cuda")
# Generate sample data. def cli():
vert_pos = torch.rand(n_points, 3, dtype=torch.float32, device=device) * 10.0 """
vert_pos[:, 2] += 25.0 Basic example for the pulsar sphere renderer using the PyTorch3D interface.
vert_pos[:, :2] -= 5.0
vert_col = torch.rand(n_points, 3, dtype=torch.float32, device=device) Writes to `basic-pt3d.png`.
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: LOGGER.info("Rendering on GPU...")
# R, T = look_at_view_transform( torch.manual_seed(1)
# dist=30.0, elev=0.0, azim=180.0, at=((0.0, 0.0, 30.0),), up=((0, 1, 0),), n_points = 10
# ) width = 1_000
cameras = PerspectiveCameras( height = 1_000
# The focal length must be double the size for PyTorch3D because of the NDC device = torch.device("cuda")
# coordinates spanning a range of two - and they must be normalized by the # Generate sample data.
# sensor width (see the pulsar example). This means we need here vert_pos = torch.rand(n_points, 3, dtype=torch.float32, device=device) * 10.0
# 5.0 * 2.0 / 2.0 to get the equivalent results as in pulsar. vert_pos[:, 2] += 25.0
focal_length=(5.0 * 2.0 / 2.0,), vert_pos[:, :2] -= 5.0
R=torch.eye(3, dtype=torch.float32, device=device)[None, ...], vert_col = torch.rand(n_points, 3, dtype=torch.float32, device=device)
T=torch.zeros((1, 3), dtype=torch.float32, device=device), pcl = Pointclouds(points=vert_pos[None, ...], features=vert_col[None, ...])
image_size=((width, height),), # Alternatively, you can also use the look_at_view_transform to get R and T:
device=device, # 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),),
vert_rad = torch.rand(n_points, dtype=torch.float32, device=device) # )
raster_settings = PointsRasterizationSettings( cameras = PerspectiveCameras(
image_size=(width, height), # The focal length must be double the size for PyTorch3D because of the NDC
radius=vert_rad, # coordinates spanning a range of two - and they must be normalized by the
) # sensor width (see the pulsar example). This means we need here
rasterizer = PointsRasterizer(cameras=cameras, raster_settings=raster_settings) # 5.0 * 2.0 / 2.0 to get the equivalent results as in pulsar.
renderer = PulsarPointsRenderer(rasterizer=rasterizer).to(device) focal_length=(5.0 * 2.0 / 2.0,),
# Render. R=torch.eye(3, dtype=torch.float32, device=device)[None, ...],
image = renderer( T=torch.zeros((1, 3), dtype=torch.float32, device=device),
pcl, image_size=((width, height),),
gamma=(1.0e-1,), # Renderer blending parameter gamma, in [1., 1e-5]. device=device,
znear=(1.0,), )
zfar=(45.0,), vert_rad = torch.rand(n_points, dtype=torch.float32, device=device)
radius_world=True, raster_settings = PointsRasterizationSettings(
bg_col=torch.ones((3,), dtype=torch.float32, device=device), image_size=(width, height),
)[0] radius=vert_rad,
print("Writing image to `%s`." % (path.abspath("basic-pt3d.png"))) )
imageio.imsave("basic-pt3d.png", (image.cpu().detach() * 255.0).to(torch.uint8).numpy()) 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]
LOGGER.info("Writing image to `%s`.", path.abspath("basic-pt3d.png"))
imageio.imsave(
"basic-pt3d.png", (image.cpu().detach() * 255.0).to(torch.uint8).numpy()
)
LOGGER.info("Done.")
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO)
cli()

135
docs/examples/pulsar_cam.py
View File

@ -9,6 +9,7 @@ distorted. Gradient-based optimization is used to converge towards the
original camera parameters. original camera parameters.
Output: cam.gif. Output: cam.gif.
""" """
import logging
import math import math
from os import path from os import path
@ -21,10 +22,11 @@ from pytorch3d.transforms import axis_angle_to_matrix, matrix_to_rotation_6d
from torch import nn, optim from torch import nn, optim
n_points = 20 LOGGER = logging.getLogger(__name__)
width = 1_000 N_POINTS = 20
height = 1_000 WIDTH = 1_000
device = torch.device("cuda") HEIGHT = 1_000
DEVICE = torch.device("cuda")
class SceneModel(nn.Module): class SceneModel(nn.Module):
@ -45,20 +47,20 @@ class SceneModel(nn.Module):
self.gamma = 0.1 self.gamma = 0.1
# Points. # Points.
torch.manual_seed(1) torch.manual_seed(1)
vert_pos = torch.rand(n_points, 3, dtype=torch.float32) * 10.0 vert_pos = torch.rand(N_POINTS, 3, dtype=torch.float32) * 10.0
vert_pos[:, 2] += 25.0 vert_pos[:, 2] += 25.0
vert_pos[:, :2] -= 5.0 vert_pos[:, :2] -= 5.0
self.register_parameter("vert_pos", nn.Parameter(vert_pos, requires_grad=False)) self.register_parameter("vert_pos", nn.Parameter(vert_pos, requires_grad=False))
self.register_parameter( self.register_parameter(
"vert_col", "vert_col",
nn.Parameter( nn.Parameter(
torch.rand(n_points, 3, dtype=torch.float32), requires_grad=False torch.rand(N_POINTS, 3, dtype=torch.float32), requires_grad=False
), ),
) )
self.register_parameter( self.register_parameter(
"vert_rad", "vert_rad",
nn.Parameter( nn.Parameter(
torch.rand(n_points, dtype=torch.float32), requires_grad=False torch.rand(N_POINTS, dtype=torch.float32), requires_grad=False
), ),
) )
self.register_parameter( self.register_parameter(
@ -90,7 +92,7 @@ class SceneModel(nn.Module):
torch.tensor([4.8, 1.8], dtype=torch.float32), requires_grad=True torch.tensor([4.8, 1.8], dtype=torch.float32), requires_grad=True
), ),
) )
self.renderer = Renderer(width, height, n_points, right_handed_system=True) self.renderer = Renderer(WIDTH, HEIGHT, N_POINTS, right_handed_system=True)
def forward(self): def forward(self):
return self.renderer.forward( return self.renderer.forward(
@ -103,58 +105,71 @@ class SceneModel(nn.Module):
) )
# Load reference. def cli():
ref = ( """
torch.from_numpy( Camera optimization example using pulsar.
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}, # 1e-3
{"params": [model.cam_rot], "lr": 5e-6},
{"params": [model.cam_sensor], "lr": 1e-4},
]
)
print("Writing video to `%s`." % (path.abspath("cam.gif"))) Writes to `cam.gif`.
writer = imageio.get_writer("cam.gif", format="gif", fps=25) """
LOGGER.info("Loading reference...")
# Optimize. # Load reference.
for i in range(300): ref = (
optimizer.zero_grad() torch.from_numpy(
result = model() imageio.imread(
# Visualize. "../../tests/pulsar/reference/examples_TestRenderer_test_cam.png"
result_im = (result.cpu().detach().numpy() * 255).astype(np.uint8) )[:, ::-1, :].copy()
cv2.imshow("opt", result_im[:, :, ::-1]) ).to(torch.float32)
writer.append_data(result_im) / 255.0
overlay_img = np.ascontiguousarray( ).to(DEVICE)
((result * 0.5 + ref * 0.5).cpu().detach().numpy() * 255).astype(np.uint8)[ # Set up model.
:, :, ::-1 model = SceneModel().to(DEVICE)
# Optimizer.
optimizer = optim.SGD(
[
{"params": [model.cam_pos], "lr": 1e-4}, # 1e-3
{"params": [model.cam_rot], "lr": 5e-6},
{"params": [model.cam_sensor], "lr": 1e-4},
] ]
) )
overlay_img = cv2.putText(
overlay_img, LOGGER.info("Writing video to `%s`.", path.abspath("cam.gif"))
"Step %d" % (i), writer = imageio.get_writer("cam.gif", format="gif", fps=25)
(10, 40),
cv2.FONT_HERSHEY_SIMPLEX, # Optimize.
1, for i in range(300):
(0, 0, 0), optimizer.zero_grad()
2, result = model()
cv2.LINE_AA, # Visualize.
False, result_im = (result.cpu().detach().numpy() * 255).astype(np.uint8)
) cv2.imshow("opt", result_im[:, :, ::-1])
cv2.imshow("overlay", overlay_img) writer.append_data(result_im)
cv2.waitKey(1) overlay_img = np.ascontiguousarray(
# Update. ((result * 0.5 + ref * 0.5).cpu().detach().numpy() * 255).astype(np.uint8)[
loss = ((result - ref) ** 2).sum() :, :, ::-1
print("loss {}: {}".format(i, loss.item())) ]
loss.backward() )
optimizer.step() overlay_img = cv2.putText(
writer.close() 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()
LOGGER.info("loss %d: %f", i, loss.item())
loss.backward()
optimizer.step()
writer.close()
LOGGER.info("Done.")
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO)
cli()

148
docs/examples/pulsar_cam_unified.py
View File

@ -10,11 +10,15 @@ original camera parameters.
Output: cam-pt3d.gif Output: cam-pt3d.gif
""" """
from os import path from os import path
import logging
import cv2 import cv2
import imageio import imageio
import numpy as np import numpy as np
import torch import torch
# Import `look_at_view_transform` as needed in the suggestion later in the
# example.
from pytorch3d.renderer.cameras import PerspectiveCameras # , look_at_view_transform from pytorch3d.renderer.cameras import PerspectiveCameras # , look_at_view_transform
from pytorch3d.renderer.points import ( from pytorch3d.renderer.points import (
PointsRasterizationSettings, PointsRasterizationSettings,
@ -26,10 +30,11 @@ from pytorch3d.transforms import axis_angle_to_matrix
from torch import nn, optim from torch import nn, optim
n_points = 20 LOGGER = logging.getLogger(__name__)
width = 1_000 N_POINTS = 20
height = 1_000 WIDTH = 1_000
device = torch.device("cuda") HEIGHT = 1_000
DEVICE = torch.device("cuda")
class SceneModel(nn.Module): class SceneModel(nn.Module):
@ -50,21 +55,21 @@ class SceneModel(nn.Module):
self.gamma = 0.1 self.gamma = 0.1
# Points. # Points.
torch.manual_seed(1) torch.manual_seed(1)
vert_pos = torch.rand(n_points, 3, dtype=torch.float32) * 10.0 vert_pos = torch.rand(N_POINTS, 3, dtype=torch.float32) * 10.0
vert_pos[:, 2] += 25.0 vert_pos[:, 2] += 25.0
vert_pos[:, :2] -= 5.0 vert_pos[:, :2] -= 5.0
self.register_parameter("vert_pos", nn.Parameter(vert_pos, requires_grad=False)) self.register_parameter("vert_pos", nn.Parameter(vert_pos, requires_grad=False))
self.register_parameter( self.register_parameter(
"vert_col", "vert_col",
nn.Parameter( nn.Parameter(
torch.rand(n_points, 3, dtype=torch.float32), torch.rand(N_POINTS, 3, dtype=torch.float32),
requires_grad=False, requires_grad=False,
), ),
) )
self.register_parameter( self.register_parameter(
"vert_rad", "vert_rad",
nn.Parameter( nn.Parameter(
torch.rand(n_points, dtype=torch.float32), torch.rand(N_POINTS, dtype=torch.float32),
requires_grad=False, requires_grad=False,
), ),
) )
@ -118,11 +123,11 @@ class SceneModel(nn.Module):
focal_length=self.focal_length, focal_length=self.focal_length,
R=self.cam_rot[None, ...], R=self.cam_rot[None, ...],
T=self.cam_pos[None, ...], T=self.cam_pos[None, ...],
image_size=((width, height),), image_size=((WIDTH, HEIGHT),),
device=device, device=DEVICE,
) )
raster_settings = PointsRasterizationSettings( raster_settings = PointsRasterizationSettings(
image_size=(width, height), image_size=(WIDTH, HEIGHT),
radius=self.vert_rad, radius=self.vert_rad,
) )
rasterizer = PointsRasterizer( rasterizer = PointsRasterizer(
@ -142,7 +147,7 @@ class SceneModel(nn.Module):
zfar=(45.0,), zfar=(45.0,),
znear=(1.0,), znear=(1.0,),
radius_world=True, radius_world=True,
bg_col=torch.ones((3,), dtype=torch.float32, device=device), bg_col=torch.ones((3,), dtype=torch.float32, device=DEVICE),
# As mentioned above: workaround for device placement of gradients for # As mentioned above: workaround for device placement of gradients for
# camera parameters. # camera parameters.
focal_length=self.focal_length, focal_length=self.focal_length,
@ -151,60 +156,73 @@ class SceneModel(nn.Module):
)[0] )[0]
# Load reference. def cli():
ref = ( """
torch.from_numpy( Camera optimization example using pulsar.
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"))) Writes to `cam.gif`.
writer = imageio.get_writer("cam-pt3d.gif", format="gif", fps=25) """
LOGGER.info("Loading reference...")
# Optimize. # Load reference.
for i in range(300): ref = (
optimizer.zero_grad() torch.from_numpy(
result = model() imageio.imread(
# Visualize. "../../tests/pulsar/reference/examples_TestRenderer_test_cam.png"
result_im = (result.cpu().detach().numpy() * 255).astype(np.uint8) )[:, ::-1, :].copy()
cv2.imshow("opt", result_im[:, :, ::-1]) ).to(torch.float32)
writer.append_data(result_im) / 255.0
overlay_img = np.ascontiguousarray( ).to(DEVICE)
((result * 0.5 + ref * 0.5).cpu().detach().numpy() * 255).astype(np.uint8)[ # Set up model.
:, :, ::-1 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},
] ]
) )
overlay_img = cv2.putText(
overlay_img, LOGGER.info("Writing video to `%s`.", path.abspath("cam-pt3d.gif"))
"Step %d" % (i), writer = imageio.get_writer("cam-pt3d.gif", format="gif", fps=25)
(10, 40),
cv2.FONT_HERSHEY_SIMPLEX, # Optimize.
1, for i in range(300):
(0, 0, 0), optimizer.zero_grad()
2, result = model()
cv2.LINE_AA, # Visualize.
False, result_im = (result.cpu().detach().numpy() * 255).astype(np.uint8)
) cv2.imshow("opt", result_im[:, :, ::-1])
cv2.imshow("overlay", overlay_img) writer.append_data(result_im)
cv2.waitKey(1) overlay_img = np.ascontiguousarray(
# Update. ((result * 0.5 + ref * 0.5).cpu().detach().numpy() * 255).astype(np.uint8)[
loss = ((result - ref) ** 2).sum() :, :, ::-1
print("loss {}: {}".format(i, loss.item())) ]
loss.backward() )
optimizer.step() overlay_img = cv2.putText(
writer.close() 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()
LOGGER.info("loss %d: %f", i, loss.item())
loss.backward()
optimizer.step()
writer.close()
LOGGER.info("Done.")
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO)
cli()

219
docs/examples/pulsar_multiview.py
View File

@ -3,7 +3,8 @@
""" """
This example demonstrates multiview 3D reconstruction using the plain This example demonstrates multiview 3D reconstruction using the plain
pulsar interface. For this, reference images have been pre-generated pulsar interface. For this, reference images have been pre-generated
(you can find them at `../../tests/pulsar/reference/examples_TestRenderer_test_multiview_%d.png`). (you can find them at
`../../tests/pulsar/reference/examples_TestRenderer_test_multiview_%d.png`).
The camera parameters are assumed given. The scene is initialized with The camera parameters are assumed given. The scene is initialized with
random spheres. Gradient-based optimization is used to optimize sphere random spheres. Gradient-based optimization is used to optimize sphere
parameters and prune spheres to converge to a 3D representation. parameters and prune spheres to converge to a 3D representation.
@ -14,6 +15,7 @@ structures yet.
""" """
import math import math
from os import path from os import path
import logging
import cv2 import cv2
import imageio import imageio
@ -23,11 +25,12 @@ from pytorch3d.renderer.points.pulsar import Renderer
from torch import nn, optim from torch import nn, optim
n_points = 400_000 LOGGER = logging.getLogger(__name__)
width = 1_000 N_POINTS = 400_000
height = 1_000 WIDTH = 1_000
visualize_ids = [0, 1] HEIGHT = 1_000
device = torch.device("cuda") VISUALIZE_IDS = [0, 1]
DEVICE = torch.device("cuda")
class SceneModel(nn.Module): class SceneModel(nn.Module):
@ -50,27 +53,27 @@ class SceneModel(nn.Module):
self.gamma = 1.0 self.gamma = 1.0
# Points. # Points.
torch.manual_seed(1) torch.manual_seed(1)
vert_pos = torch.rand((1, n_points, 3), dtype=torch.float32) * 10.0 vert_pos = torch.rand((1, N_POINTS, 3), dtype=torch.float32) * 10.0
vert_pos[:, :, 2] += 25.0 vert_pos[:, :, 2] += 25.0
vert_pos[:, :, :2] -= 5.0 vert_pos[:, :, :2] -= 5.0
self.register_parameter("vert_pos", nn.Parameter(vert_pos, requires_grad=True)) self.register_parameter("vert_pos", nn.Parameter(vert_pos, requires_grad=True))
self.register_parameter( self.register_parameter(
"vert_col", "vert_col",
nn.Parameter( nn.Parameter(
torch.ones(1, n_points, 3, dtype=torch.float32) * 0.5, torch.ones(1, N_POINTS, 3, dtype=torch.float32) * 0.5,
requires_grad=True, requires_grad=True,
), ),
) )
self.register_parameter( self.register_parameter(
"vert_rad", "vert_rad",
nn.Parameter( nn.Parameter(
torch.ones(1, n_points, dtype=torch.float32) * 0.05, requires_grad=True torch.ones(1, N_POINTS, dtype=torch.float32) * 0.05, requires_grad=True
), ),
) )
self.register_parameter( self.register_parameter(
"vert_opy", "vert_opy",
nn.Parameter( nn.Parameter(
torch.ones(1, n_points, dtype=torch.float32), requires_grad=True torch.ones(1, N_POINTS, dtype=torch.float32), requires_grad=True
), ),
) )
self.register_buffer( self.register_buffer(
@ -92,7 +95,7 @@ class SceneModel(nn.Module):
dtype=torch.float32, dtype=torch.float32,
), ),
) )
self.renderer = Renderer(width, height, n_points, right_handed_system=True) self.renderer = Renderer(WIDTH, HEIGHT, N_POINTS, right_handed_system=True)
def forward(self, cam=None): def forward(self, cam=None):
if cam is None: if cam is None:
@ -110,97 +113,113 @@ class SceneModel(nn.Module):
) )
# Load reference. def cli():
ref = torch.stack( """
[ Simple demonstration for a multi-view 3D reconstruction using pulsar.
torch.from_numpy(
imageio.imread(
"../../tests/pulsar/reference/examples_TestRenderer_test_multiview_%d.png"
% idx
)
).to(torch.float32)
/ 255.0
for idx in range(8)
]
).to(device)
# Set up model.
model = SceneModel().to(device)
# Optimizer.
optimizer = optim.SGD(
[
{"params": [model.vert_col], "lr": 1e-1},
{"params": [model.vert_rad], "lr": 1e-3},
{"params": [model.vert_pos], "lr": 1e-3},
]
)
# For visualization. This example makes use of opacity, which is not yet supported through
angle = 0.0 the unified PyTorch3D interface.
print("Writing video to `%s`." % (path.abspath("multiview.avi")))
writer = imageio.get_writer("multiview.gif", format="gif", fps=25)
# Optimize. Writes to `multiview.gif`.
for i in range(300): """
optimizer.zero_grad() LOGGER.info("Loading reference...")
result = model() # Load reference.
# Visualize. ref = torch.stack(
result_im = (result.cpu().detach().numpy() * 255).astype(np.uint8) [
cv2.imshow("opt", result_im[0, :, :, ::-1]) torch.from_numpy(
overlay_img = np.ascontiguousarray( imageio.imread(
((result * 0.5 + ref * 0.5).cpu().detach().numpy() * 255).astype(np.uint8)[ "../../tests/pulsar/reference/examples_TestRenderer_test_multiview_%d.png"
0, :, :, ::-1 % idx
)
).to(torch.float32)
/ 255.0
for idx in range(8)
]
).to(DEVICE)
# Set up model.
model = SceneModel().to(DEVICE)
# Optimizer.
optimizer = optim.SGD(
[
{"params": [model.vert_col], "lr": 1e-1},
{"params": [model.vert_rad], "lr": 1e-3},
{"params": [model.vert_pos], "lr": 1e-3},
] ]
) )
overlay_img = cv2.putText(
overlay_img, # For visualization.
"Step %d" % (i), angle = 0.0
(10, 40), LOGGER.info("Writing video to `%s`.", path.abspath("multiview.avi"))
cv2.FONT_HERSHEY_SIMPLEX, writer = imageio.get_writer("multiview.gif", format="gif", fps=25)
1,
(0, 0, 0), # Optimize.
2, for i in range(300):
cv2.LINE_AA, optimizer.zero_grad()
False, result = model()
) # Visualize.
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(1, 1, 3, dtype=torch.float32).to(device))
.abs()
.sum(dim=2)
)
model.vert_pos.data[vd <= 0.2] = -1000.0
# Rotating visualization.
cam_control = torch.tensor(
[
[
np.sin(angle) * 35.0,
0.0,
30.0 - np.cos(angle) * 35.0,
0.0,
-angle + math.pi,
0.0,
5.0,
2.0,
]
],
dtype=torch.float32,
).to(device)
with torch.no_grad():
result = model.forward(cam=cam_control)[0]
result_im = (result.cpu().detach().numpy() * 255).astype(np.uint8) result_im = (result.cpu().detach().numpy() * 255).astype(np.uint8)
cv2.imshow("vis", result_im[:, :, ::-1]) cv2.imshow("opt", result_im[0, :, :, ::-1])
writer.append_data(result_im) overlay_img = np.ascontiguousarray(
angle += 0.05 ((result * 0.5 + ref * 0.5).cpu().detach().numpy() * 255).astype(np.uint8)[
writer.close() 0, :, :, ::-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()
LOGGER.info("loss %d: %f", 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(1, 1, 3, dtype=torch.float32).to(DEVICE))
.abs()
.sum(dim=2)
)
model.vert_pos.data[vd <= 0.2] = -1000.0
# Rotating visualization.
cam_control = torch.tensor(
[
[
np.sin(angle) * 35.0,
0.0,
30.0 - np.cos(angle) * 35.0,
0.0,
-angle + math.pi,
0.0,
5.0,
2.0,
]
],
dtype=torch.float32,
).to(DEVICE)
with torch.no_grad():
result = model.forward(cam=cam_control)[0]
result_im = (result.cpu().detach().numpy() * 255).astype(np.uint8)
cv2.imshow("vis", result_im[:, :, ::-1])
writer.append_data(result_im)
angle += 0.05
writer.close()
LOGGER.info("Done.")
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO)
cli()

147
docs/examples/pulsar_optimization.py
View File

@ -9,6 +9,7 @@ optimization is used to converge towards a faithful
scene representation. scene representation.
""" """
import math import math
import logging
import cv2 import cv2
import imageio import imageio
@ -18,10 +19,11 @@ from pytorch3d.renderer.points.pulsar import Renderer
from torch import nn, optim from torch import nn, optim
n_points = 10_000 LOGGER = logging.getLogger(__name__)
width = 1_000 N_POINTS = 10_000
height = 1_000 WIDTH = 1_000
device = torch.device("cuda") HEIGHT = 1_000
DEVICE = torch.device("cuda")
class SceneModel(nn.Module): class SceneModel(nn.Module):
@ -42,20 +44,20 @@ class SceneModel(nn.Module):
self.gamma = 1.0 self.gamma = 1.0
# Points. # Points.
torch.manual_seed(1) torch.manual_seed(1)
vert_pos = torch.rand(n_points, 3, dtype=torch.float32) * 10.0 vert_pos = torch.rand(N_POINTS, 3, dtype=torch.float32) * 10.0
vert_pos[:, 2] += 25.0 vert_pos[:, 2] += 25.0
vert_pos[:, :2] -= 5.0 vert_pos[:, :2] -= 5.0
self.register_parameter("vert_pos", nn.Parameter(vert_pos, requires_grad=True)) self.register_parameter("vert_pos", nn.Parameter(vert_pos, requires_grad=True))
self.register_parameter( self.register_parameter(
"vert_col", "vert_col",
nn.Parameter( nn.Parameter(
torch.ones(n_points, 3, dtype=torch.float32) * 0.5, requires_grad=True torch.ones(N_POINTS, 3, dtype=torch.float32) * 0.5, requires_grad=True
), ),
) )
self.register_parameter( self.register_parameter(
"vert_rad", "vert_rad",
nn.Parameter( nn.Parameter(
torch.ones(n_points, dtype=torch.float32) * 0.3, requires_grad=True torch.ones(N_POINTS, dtype=torch.float32) * 0.3, requires_grad=True
), ),
) )
self.register_buffer( self.register_buffer(
@ -67,7 +69,7 @@ class SceneModel(nn.Module):
# The volumetric optimization works better with a higher number of tracked # The volumetric optimization works better with a higher number of tracked
# intersections per ray. # intersections per ray.
self.renderer = Renderer( self.renderer = Renderer(
width, height, n_points, n_track=32, right_handed_system=True WIDTH, HEIGHT, N_POINTS, n_track=32, right_handed_system=True
) )
def forward(self): def forward(self):
@ -82,65 +84,76 @@ class SceneModel(nn.Module):
) )
# Load reference. def cli():
ref = ( """
torch.from_numpy( Scene optimization example using pulsar.
imageio.imread( """
"../../tests/pulsar/reference/examples_TestRenderer_test_smallopt.png" LOGGER.info("Loading reference...")
)[:, ::-1, :].copy() # Load reference.
).to(torch.float32) ref = (
/ 255.0 torch.from_numpy(
).to(device) imageio.imread(
# Set up model. "../../tests/pulsar/reference/examples_TestRenderer_test_smallopt.png"
model = SceneModel().to(device) )[:, ::-1, :].copy()
# Optimizer. ).to(torch.float32)
optimizer = optim.SGD( / 255.0
[ ).to(DEVICE)
{"params": [model.vert_col], "lr": 1e0}, # Set up model.
{"params": [model.vert_rad], "lr": 5e-3}, model = SceneModel().to(DEVICE)
{"params": [model.vert_pos], "lr": 1e-2}, # Optimizer.
] optimizer = optim.SGD(
) [
{"params": [model.vert_col], "lr": 1e0},
# Optimize. {"params": [model.vert_rad], "lr": 5e-3},
for i in range(500): {"params": [model.vert_pos], "lr": 1e-2},
optimizer.zero_grad()
result, result_info = 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( LOGGER.info("Optimizing...")
overlay_img, # Optimize.
"Step %d" % (i), for i in range(500):
(10, 40), optimizer.zero_grad()
cv2.FONT_HERSHEY_SIMPLEX, result, result_info = model()
1, # Visualize.
(0, 0, 0), result_im = (result.cpu().detach().numpy() * 255).astype(np.uint8)
2, cv2.imshow("opt", result_im[:, :, ::-1])
cv2.LINE_AA, overlay_img = np.ascontiguousarray(
False, ((result * 0.5 + ref * 0.5).cpu().detach().numpy() * 255).astype(np.uint8)[
) :, :, ::-1
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 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()
LOGGER.info("loss %d: %f", 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
LOGGER.info("Done.")
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO)
cli()

160
docs/examples/pulsar_optimization_unified.py
View File

@ -9,11 +9,15 @@ optimization is used to converge towards a faithful
scene representation. scene representation.
""" """
import math import math
import logging
import cv2 import cv2
import imageio import imageio
import numpy as np import numpy as np
import torch import torch
# Import `look_at_view_transform` as needed in the suggestion later in the
# example.
from pytorch3d.renderer.cameras import PerspectiveCameras # , look_at_view_transform from pytorch3d.renderer.cameras import PerspectiveCameras # , look_at_view_transform
from pytorch3d.renderer.points import ( from pytorch3d.renderer.points import (
PointsRasterizationSettings, PointsRasterizationSettings,
@ -24,10 +28,11 @@ from pytorch3d.structures.pointclouds import Pointclouds
from torch import nn, optim from torch import nn, optim
n_points = 10_000 LOGGER = logging.getLogger(__name__)
width = 1_000 N_POINTS = 10_000
height = 1_000 WIDTH = 1_000
device = torch.device("cuda") HEIGHT = 1_000
DEVICE = torch.device("cuda")
class SceneModel(nn.Module): class SceneModel(nn.Module):
@ -48,21 +53,21 @@ class SceneModel(nn.Module):
self.gamma = 1.0 self.gamma = 1.0
# Points. # Points.
torch.manual_seed(1) torch.manual_seed(1)
vert_pos = torch.rand(n_points, 3, dtype=torch.float32, device=device) * 10.0 vert_pos = torch.rand(N_POINTS, 3, dtype=torch.float32, device=DEVICE) * 10.0
vert_pos[:, 2] += 25.0 vert_pos[:, 2] += 25.0
vert_pos[:, :2] -= 5.0 vert_pos[:, :2] -= 5.0
self.register_parameter("vert_pos", nn.Parameter(vert_pos, requires_grad=True)) self.register_parameter("vert_pos", nn.Parameter(vert_pos, requires_grad=True))
self.register_parameter( self.register_parameter(
"vert_col", "vert_col",
nn.Parameter( nn.Parameter(
torch.ones(n_points, 3, dtype=torch.float32, device=device) * 0.5, torch.ones(N_POINTS, 3, dtype=torch.float32, device=DEVICE) * 0.5,
requires_grad=True, requires_grad=True,
), ),
) )
self.register_parameter( self.register_parameter(
"vert_rad", "vert_rad",
nn.Parameter( nn.Parameter(
torch.ones(n_points, dtype=torch.float32) * 0.3, requires_grad=True torch.ones(N_POINTS, dtype=torch.float32) * 0.3, requires_grad=True
), ),
) )
self.register_buffer( self.register_buffer(
@ -77,13 +82,13 @@ class SceneModel(nn.Module):
# sensor width (see the pulsar example). This means we need here # 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. # 5.0 * 2.0 / 2.0 to get the equivalent results as in pulsar.
focal_length=5.0, focal_length=5.0,
R=torch.eye(3, dtype=torch.float32, device=device)[None, ...], R=torch.eye(3, dtype=torch.float32, device=DEVICE)[None, ...],
T=torch.zeros((1, 3), dtype=torch.float32, device=device), T=torch.zeros((1, 3), dtype=torch.float32, device=DEVICE),
image_size=((width, height),), image_size=((WIDTH, HEIGHT),),
device=device, device=DEVICE,
) )
raster_settings = PointsRasterizationSettings( raster_settings = PointsRasterizationSettings(
image_size=(width, height), image_size=(WIDTH, HEIGHT),
radius=self.vert_rad, radius=self.vert_rad,
) )
rasterizer = PointsRasterizer( rasterizer = PointsRasterizer(
@ -103,69 +108,80 @@ class SceneModel(nn.Module):
zfar=(45.0,), zfar=(45.0,),
znear=(1.0,), znear=(1.0,),
radius_world=True, radius_world=True,
bg_col=torch.ones((3,), dtype=torch.float32, device=device), bg_col=torch.ones((3,), dtype=torch.float32, device=DEVICE),
)[0] )[0]
# Load reference. def cli():
ref = ( """
torch.from_numpy( Scene optimization example using pulsar and the unified PyTorch3D interface.
imageio.imread( """
"../../tests/pulsar/reference/examples_TestRenderer_test_smallopt.png" LOGGER.info("Loading reference...")
)[:, ::-1, :].copy() # Load reference.
).to(torch.float32) ref = (
/ 255.0 torch.from_numpy(
).to(device) imageio.imread(
# Set up model. "../../tests/pulsar/reference/examples_TestRenderer_test_smallopt.png"
model = SceneModel().to(device) )[:, ::-1, :].copy()
# Optimizer. ).to(torch.float32)
optimizer = optim.SGD( / 255.0
[ ).to(DEVICE)
{"params": [model.vert_col], "lr": 1e0}, # Set up model.
{"params": [model.vert_rad], "lr": 5e-3}, model = SceneModel().to(DEVICE)
{"params": [model.vert_pos], "lr": 1e-2}, # Optimizer.
] optimizer = optim.SGD(
) [
{"params": [model.vert_col], "lr": 1e0},
# Optimize. {"params": [model.vert_rad], "lr": 5e-3},
for i in range(500): {"params": [model.vert_pos], "lr": 1e-2},
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( LOGGER.info("Optimizing...")
overlay_img, # Optimize.
"Step %d" % (i), for i in range(500):
(10, 40), optimizer.zero_grad()
cv2.FONT_HERSHEY_SIMPLEX, result = model()
1, # Visualize.
(0, 0, 0), result_im = (result.cpu().detach().numpy() * 255).astype(np.uint8)
2, cv2.imshow("opt", result_im[:, :, ::-1])
cv2.LINE_AA, overlay_img = np.ascontiguousarray(
False, ((result * 0.5 + ref * 0.5).cpu().detach().numpy() * 255).astype(np.uint8)[
) :, :, ::-1
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 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()
LOGGER.info("loss %d: %f", 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
LOGGER.info("Done.")
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO)
cli()

4
tests/pulsar/test_depth.py
View File

@ -44,6 +44,8 @@ class TestDepth(TestCaseMixin, unittest.TestCase):
n_channels=1, n_channels=1,
).to(device) ).to(device)
data = torch.load(IN_REF_FP, map_location="cpu") data = torch.load(IN_REF_FP, map_location="cpu")
# For creating the reference files.
# Use in case of updates.
# data["pos"] = torch.rand_like(data["pos"]) # data["pos"] = torch.rand_like(data["pos"])
# data["pos"][:, 0] = data["pos"][:, 0] * 2. - 1. # data["pos"][:, 0] = data["pos"][:, 0] * 2. - 1.
# data["pos"][:, 1] = data["pos"][:, 1] * 2. - 1. # data["pos"][:, 1] = data["pos"][:, 1] * 2. - 1.
@ -74,6 +76,8 @@ class TestDepth(TestCaseMixin, unittest.TestCase):
), ),
depth_vis.cpu().numpy().astype(np.uint8), depth_vis.cpu().numpy().astype(np.uint8),
) )
# For creating the reference files.
# Use in case of updates.
# torch.save( # torch.save(
# data, path.join(path.dirname(__file__), "reference", "nr0000-in.pth") # data, path.join(path.dirname(__file__), "reference", "nr0000-in.pth")
# ) # )

2
tests/pulsar/test_small_spheres.py
View File

@ -123,7 +123,7 @@ class TestSmallSpheres(unittest.TestCase):
self.assertTrue( self.assertTrue(
(sphere_ids == idx).sum() > 0, "Sphere ID %d missing!" % (idx) (sphere_ids == idx).sum() > 0, "Sphere ID %d missing!" % (idx)
) )
# Visualize. # Visualization code. Activate for debugging.
# result_im = (result.cpu().detach().numpy() * 255).astype(np.uint8) # result_im = (result.cpu().detach().numpy() * 255).astype(np.uint8)
# cv2.imshow("res", result_im[0, :, :, ::-1]) # cv2.imshow("res", result_im[0, :, :, ::-1])
# cv2.waitKey(0) # cv2.waitKey(0)