mirror of
https://github.com/deepseek-ai/DreamCraft3D
synced 2024-12-05 02:25:45 +00:00
79 lines
2.6 KiB
Python
79 lines
2.6 KiB
Python
import nvdiffrast.torch as dr
|
|
import torch
|
|
|
|
from threestudio.utils.typing import *
|
|
|
|
|
|
class NVDiffRasterizerContext:
|
|
def __init__(self, context_type: str, device: torch.device) -> None:
|
|
self.device = device
|
|
self.ctx = self.initialize_context(context_type, device)
|
|
|
|
def initialize_context(
|
|
self, context_type: str, device: torch.device
|
|
) -> Union[dr.RasterizeGLContext, dr.RasterizeCudaContext]:
|
|
if context_type == "gl":
|
|
return dr.RasterizeGLContext(device=device)
|
|
elif context_type == "cuda":
|
|
return dr.RasterizeCudaContext(device=device)
|
|
else:
|
|
raise ValueError(f"Unknown rasterizer context type: {context_type}")
|
|
|
|
def vertex_transform(
|
|
self, verts: Float[Tensor, "Nv 3"], mvp_mtx: Float[Tensor, "B 4 4"]
|
|
) -> Float[Tensor, "B Nv 4"]:
|
|
verts_homo = torch.cat(
|
|
[verts, torch.ones([verts.shape[0], 1]).to(verts)], dim=-1
|
|
)
|
|
return torch.matmul(verts_homo, mvp_mtx.permute(0, 2, 1))
|
|
|
|
def rasterize(
|
|
self,
|
|
pos: Float[Tensor, "B Nv 4"],
|
|
tri: Integer[Tensor, "Nf 3"],
|
|
resolution: Union[int, Tuple[int, int]],
|
|
):
|
|
# rasterize in instance mode (single topology)
|
|
return dr.rasterize(self.ctx, pos.float(), tri.int(), resolution, grad_db=True)
|
|
|
|
def rasterize_one(
|
|
self,
|
|
pos: Float[Tensor, "Nv 4"],
|
|
tri: Integer[Tensor, "Nf 3"],
|
|
resolution: Union[int, Tuple[int, int]],
|
|
):
|
|
# rasterize one single mesh under a single viewpoint
|
|
rast, rast_db = self.rasterize(pos[None, ...], tri, resolution)
|
|
return rast[0], rast_db[0]
|
|
|
|
def antialias(
|
|
self,
|
|
color: Float[Tensor, "B H W C"],
|
|
rast: Float[Tensor, "B H W 4"],
|
|
pos: Float[Tensor, "B Nv 4"],
|
|
tri: Integer[Tensor, "Nf 3"],
|
|
) -> Float[Tensor, "B H W C"]:
|
|
return dr.antialias(color.float(), rast, pos.float(), tri.int())
|
|
|
|
def interpolate(
|
|
self,
|
|
attr: Float[Tensor, "B Nv C"],
|
|
rast: Float[Tensor, "B H W 4"],
|
|
tri: Integer[Tensor, "Nf 3"],
|
|
rast_db=None,
|
|
diff_attrs=None,
|
|
) -> Float[Tensor, "B H W C"]:
|
|
return dr.interpolate(
|
|
attr.float(), rast, tri.int(), rast_db=rast_db, diff_attrs=diff_attrs
|
|
)
|
|
|
|
def interpolate_one(
|
|
self,
|
|
attr: Float[Tensor, "Nv C"],
|
|
rast: Float[Tensor, "B H W 4"],
|
|
tri: Integer[Tensor, "Nf 3"],
|
|
rast_db=None,
|
|
diff_attrs=None,
|
|
) -> Float[Tensor, "B H W C"]:
|
|
return self.interpolate(attr[None, ...], rast, tri, rast_db, diff_attrs)
|