mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-12-24 08:10:34 +08:00
Example and test updates.
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
committed by
Facebook GitHub Bot
Facebook GitHub Bot
parent
e9a26f263a
commit
b6be3b95fb
@@ -3,7 +3,8 @@
|
||||
"""
|
||||
This example demonstrates multiview 3D reconstruction using the plain
|
||||
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
|
||||
random spheres. Gradient-based optimization is used to optimize sphere
|
||||
parameters and prune spheres to converge to a 3D representation.
|
||||
@@ -14,6 +15,7 @@ structures yet.
|
||||
"""
|
||||
import math
|
||||
from os import path
|
||||
import logging
|
||||
|
||||
import cv2
|
||||
import imageio
|
||||
@@ -23,11 +25,12 @@ from pytorch3d.renderer.points.pulsar import Renderer
|
||||
from torch import nn, optim
|
||||
|
||||
|
||||
n_points = 400_000
|
||||
width = 1_000
|
||||
height = 1_000
|
||||
visualize_ids = [0, 1]
|
||||
device = torch.device("cuda")
|
||||
LOGGER = logging.getLogger(__name__)
|
||||
N_POINTS = 400_000
|
||||
WIDTH = 1_000
|
||||
HEIGHT = 1_000
|
||||
VISUALIZE_IDS = [0, 1]
|
||||
DEVICE = torch.device("cuda")
|
||||
|
||||
|
||||
class SceneModel(nn.Module):
|
||||
@@ -50,27 +53,27 @@ class SceneModel(nn.Module):
|
||||
self.gamma = 1.0
|
||||
# Points.
|
||||
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] -= 5.0
|
||||
self.register_parameter("vert_pos", nn.Parameter(vert_pos, requires_grad=True))
|
||||
self.register_parameter(
|
||||
"vert_col",
|
||||
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,
|
||||
),
|
||||
)
|
||||
self.register_parameter(
|
||||
"vert_rad",
|
||||
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(
|
||||
"vert_opy",
|
||||
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(
|
||||
@@ -92,7 +95,7 @@ class SceneModel(nn.Module):
|
||||
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):
|
||||
if cam is None:
|
||||
@@ -110,97 +113,113 @@ class SceneModel(nn.Module):
|
||||
)
|
||||
|
||||
|
||||
# Load reference.
|
||||
ref = torch.stack(
|
||||
[
|
||||
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},
|
||||
]
|
||||
)
|
||||
def cli():
|
||||
"""
|
||||
Simple demonstration for a multi-view 3D reconstruction using pulsar.
|
||||
|
||||
# For visualization.
|
||||
angle = 0.0
|
||||
print("Writing video to `%s`." % (path.abspath("multiview.avi")))
|
||||
writer = imageio.get_writer("multiview.gif", format="gif", fps=25)
|
||||
This example makes use of opacity, which is not yet supported through
|
||||
the unified PyTorch3D interface.
|
||||
|
||||
# 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[0, :, :, ::-1])
|
||||
overlay_img = np.ascontiguousarray(
|
||||
((result * 0.5 + ref * 0.5).cpu().detach().numpy() * 255).astype(np.uint8)[
|
||||
0, :, :, ::-1
|
||||
Writes to `multiview.gif`.
|
||||
"""
|
||||
LOGGER.info("Loading reference...")
|
||||
# Load reference.
|
||||
ref = torch.stack(
|
||||
[
|
||||
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},
|
||||
]
|
||||
)
|
||||
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(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]
|
||||
|
||||
# For visualization.
|
||||
angle = 0.0
|
||||
LOGGER.info("Writing video to `%s`.", path.abspath("multiview.avi"))
|
||||
writer = imageio.get_writer("multiview.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("vis", result_im[:, :, ::-1])
|
||||
writer.append_data(result_im)
|
||||
angle += 0.05
|
||||
writer.close()
|
||||
cv2.imshow("opt", result_im[0, :, :, ::-1])
|
||||
overlay_img = np.ascontiguousarray(
|
||||
((result * 0.5 + ref * 0.5).cpu().detach().numpy() * 255).astype(np.uint8)[
|
||||
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()
|
||||
|
||||
Reference in New Issue
Block a user