mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-08-02 20:02:49 +08:00
956d3a010c
Summary: Support for moving all the tensors of the renderer to another device by calling `renderer.to(new_device)` Currently the `MeshRenderer`, `MeshRasterizer` and `SoftPhongShader` (and other shaders) are all of type `nn.Module` which already supports easily moving tensors of submodules (defined as class attributes) to a different device. However the class attributes of the rasterizer and shader (e.g. cameras, lights, materials), are of type `TensorProperties`, not nn.Module so we need to explicity create a `to` method to move these tensors to device. Note that the `TensorProperties` class already has a `to` method so we only need to call `cameras.to(device)` and don't need to worry about the internal tensors. The other option is of course making these other classes (cameras, lights etc) also of type nn.Module. Reviewed By: gkioxari Differential Revision: D23885107 fbshipit-source-id: d71565c442181f739de4d797076ed5d00fb67f8e
58 lines
1.9 KiB
Python
58 lines
1.9 KiB
Python
# 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 shading.
|
|
# The rasterizer should:
|
|
# - transform inputs from world -> screen space
|
|
# - rasterize inputs
|
|
# - return fragments
|
|
# The shader 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(meshes)
|
|
# images = self.shader(fragments, meshes)
|
|
# return images
|
|
|
|
|
|
class MeshRenderer(nn.Module):
|
|
"""
|
|
A class for rendering a batch of heterogeneous meshes. The class should
|
|
be initialized with a rasterizer and shader class which each have a forward
|
|
function.
|
|
"""
|
|
|
|
def __init__(self, rasterizer, shader):
|
|
super().__init__()
|
|
self.rasterizer = rasterizer
|
|
self.shader = shader
|
|
|
|
def to(self, device):
|
|
# Rasterizer and shader have submodules which are not of type nn.Module
|
|
self.rasterizer.to(device)
|
|
self.shader.to(device)
|
|
|
|
def forward(self, meshes_world, **kwargs) -> torch.Tensor:
|
|
"""
|
|
Render a batch of images from a batch of meshes by rasterizing and then
|
|
shading.
|
|
|
|
NOTE: If the blur radius for rasterization is > 0.0, some pixels can
|
|
have one or more barycentric coordinates lying outside the range [0, 1].
|
|
For a pixel with out of bounds barycentric coordinates with respect to a
|
|
face f, clipping is required before interpolating the texture uv
|
|
coordinates and z buffer so that the colors and depths are limited to
|
|
the range for the corresponding face.
|
|
For this set rasterizer.raster_settings.clip_barycentric_coords=True
|
|
"""
|
|
fragments = self.rasterizer(meshes_world, **kwargs)
|
|
images = self.shader(fragments, meshes_world, **kwargs)
|
|
|
|
return images
|