mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-12-14 19:36:23 +08:00
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:
committed by
Facebook GitHub Bot
parent
960fd6d8b6
commit
039e02601d
@@ -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.
|
||||
|
||||
67
docs/examples/pulsar_basic_unified.py
Executable file
67
docs/examples/pulsar_basic_unified.py
Executable 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())
|
||||
@@ -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)
|
||||
|
||||
210
docs/examples/pulsar_cam_unified.py
Executable file
210
docs/examples/pulsar_cam_unified.py
Executable 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()
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
171
docs/examples/pulsar_optimization_unified.py
Executable file
171
docs/examples/pulsar_optimization_unified.py
Executable 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
|
||||
Reference in New Issue
Block a user