mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-08-02 03:42:50 +08:00
960fd6d8b6
Summary: This diff builds on top of the `pulsar integration` diff to provide a unified interface for the existing PyTorch3D point renderer and Pulsar. 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. The unified interfaces are completely consistent. Switching the render backend is as easy as using `renderer = PulsarPointsRenderer(rasterizer=rasterizer).to(device)` instead of `renderer = PointsRenderer(rasterizer=rasterizer, compositor=compositor)` and adding the `gamma` parameter to the forward function. All PyTorch3D camera types are supported as far as possible; keyword arguments are properly forwarded to the camera. The `PerspectiveCamera` and `OrthographicCamera` require znear and zfar as additional parameters for the forward pass. Reviewed By: nikhilaravi Differential Revision: D21421443 fbshipit-source-id: 4aa0a83a419592d9a0bb5d62486a1cdea9d73ce6
56 lines
1.7 KiB
Python
56 lines
1.7 KiB
Python
#!/usr/bin/env python3
|
|
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
# A renderer class should be initialized with a
|
|
# function for rasterization and a function for compositing.
|
|
# The rasterizer should:
|
|
# - transform inputs from world -> screen space
|
|
# - rasterize inputs
|
|
# - return fragments
|
|
# The compositor can take fragments as input along with any other properties of
|
|
# the scene and generate images.
|
|
|
|
# E.g. rasterize inputs and then shade
|
|
#
|
|
# fragments = self.rasterize(point_clouds)
|
|
# images = self.compositor(fragments, point_clouds)
|
|
# return images
|
|
|
|
|
|
class PointsRenderer(nn.Module):
|
|
"""
|
|
A class for rendering a batch of points. The class should
|
|
be initialized with a rasterizer and compositor class which each have a forward
|
|
function.
|
|
"""
|
|
|
|
def __init__(self, rasterizer, compositor):
|
|
super().__init__()
|
|
self.rasterizer = rasterizer
|
|
self.compositor = compositor
|
|
|
|
def forward(self, point_clouds, **kwargs) -> torch.Tensor:
|
|
fragments = self.rasterizer(point_clouds, **kwargs)
|
|
|
|
# Construct weights based on the distance of a point to the true point.
|
|
# However, this could be done differently: e.g. predicted as opposed
|
|
# to a function of the weights.
|
|
r = self.rasterizer.raster_settings.radius
|
|
|
|
dists2 = fragments.dists.permute(0, 3, 1, 2)
|
|
weights = 1 - dists2 / (r * r)
|
|
images = self.compositor(
|
|
fragments.idx.long().permute(0, 3, 1, 2),
|
|
weights,
|
|
point_clouds.features_packed().permute(1, 0),
|
|
**kwargs,
|
|
)
|
|
|
|
# permute so image comes at the end
|
|
images = images.permute(0, 2, 3, 1)
|
|
|
|
return images
|