Spaces:
Runtime error
Runtime error
File size: 4,111 Bytes
2fa4776 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 |
from dataclasses import dataclass
import nerfacc
import torch
import torch.nn.functional as F
import threestudio
from threestudio.models.background.base import BaseBackground
from threestudio.models.geometry.base import BaseImplicitGeometry
from threestudio.models.materials.base import BaseMaterial
from threestudio.models.renderers.base import Rasterizer, VolumeRenderer
from threestudio.utils.misc import get_device
from threestudio.utils.rasterize import NVDiffRasterizerContext
from threestudio.utils.typing import *
@threestudio.register("nvdiff-rasterizer")
class NVDiffRasterizer(Rasterizer):
@dataclass
class Config(VolumeRenderer.Config):
context_type: str = "gl"
cfg: Config
def configure(
self,
geometry: BaseImplicitGeometry,
material: BaseMaterial,
background: BaseBackground,
) -> None:
super().configure(geometry, material, background)
self.ctx = NVDiffRasterizerContext(self.cfg.context_type, get_device())
def forward(
self,
mvp_mtx: Float[Tensor, "B 4 4"],
camera_positions: Float[Tensor, "B 3"],
light_positions: Float[Tensor, "B 3"],
height: int,
width: int,
render_rgb: bool = True,
**kwargs
) -> Dict[str, Any]:
batch_size = mvp_mtx.shape[0]
mesh = self.geometry.isosurface()
v_pos_clip: Float[Tensor, "B Nv 4"] = self.ctx.vertex_transform(
mesh.v_pos, mvp_mtx
)
rast, _ = self.ctx.rasterize(v_pos_clip, mesh.t_pos_idx, (height, width))
mask = rast[..., 3:] > 0
mask_aa = self.ctx.antialias(mask.float(), rast, v_pos_clip, mesh.t_pos_idx)
out = {"opacity": mask_aa, "mesh": mesh}
gb_normal, _ = self.ctx.interpolate_one(mesh.v_nrm, rast, mesh.t_pos_idx)
gb_normal = F.normalize(gb_normal, dim=-1)
gb_normal_aa = torch.lerp(
torch.zeros_like(gb_normal), (gb_normal + 1.0) / 2.0, mask.float()
)
gb_normal_aa = self.ctx.antialias(
gb_normal_aa, rast, v_pos_clip, mesh.t_pos_idx
)
out.update({"comp_normal": gb_normal_aa}) # in [0, 1]
# TODO: make it clear whether to compute the normal, now we compute it in all cases
# consider using: require_normal_computation = render_normal or (render_rgb and material.requires_normal)
# or
# render_normal = render_normal or (render_rgb and material.requires_normal)
if render_rgb:
selector = mask[..., 0]
gb_pos, _ = self.ctx.interpolate_one(mesh.v_pos, rast, mesh.t_pos_idx)
gb_viewdirs = F.normalize(
gb_pos - camera_positions[:, None, None, :], dim=-1
)
gb_light_positions = light_positions[:, None, None, :].expand(
-1, height, width, -1
)
positions = gb_pos[selector]
geo_out = self.geometry(positions, output_normal=False)
extra_geo_info = {}
if self.material.requires_normal:
extra_geo_info["shading_normal"] = gb_normal[selector]
if self.material.requires_tangent:
gb_tangent, _ = self.ctx.interpolate_one(
mesh.v_tng, rast, mesh.t_pos_idx
)
gb_tangent = F.normalize(gb_tangent, dim=-1)
extra_geo_info["tangent"] = gb_tangent[selector]
rgb_fg = self.material(
viewdirs=gb_viewdirs[selector],
positions=positions,
light_positions=gb_light_positions[selector],
**extra_geo_info,
**geo_out
)
gb_rgb_fg = torch.zeros(batch_size, height, width, 3).to(rgb_fg)
gb_rgb_fg[selector] = rgb_fg
gb_rgb_bg = self.background(dirs=gb_viewdirs)
gb_rgb = torch.lerp(gb_rgb_bg, gb_rgb_fg, mask.float())
gb_rgb_aa = self.ctx.antialias(gb_rgb, rast, v_pos_clip, mesh.t_pos_idx)
out.update({"comp_rgb": gb_rgb_aa, "comp_rgb_bg": gb_rgb_bg})
return out
|