File size: 1,876 Bytes
434b0b0 | 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 | import nvdiffrast.torch as dr
import torch
from matplotlib import image
def _warmup(glctx):
# windows workaround for https://github.com/NVlabs/nvdiffrast/issues/59
def tensor(*args, **kwargs):
return torch.tensor(*args, device="cuda", **kwargs)
pos = tensor(
[[[-0.8, -0.8, 0, 1], [0.8, -0.8, 0, 1], [-0.8, 0.8, 0, 1]]],
dtype=torch.float32,
)
tri = tensor([[0, 1, 2]], dtype=torch.int32)
dr.rasterize(glctx, pos, tri, resolution=[256, 256])
class NormalsRenderer:
def __init__(
self,
mv: torch.Tensor, # C,4,4
proj: torch.Tensor, # C,4,4
image_size: tuple[int, int],
):
self._mvp = proj @ mv # C,4,4
self._image_size = image_size
# self._glctx = dr.RasterizeGLContext()
self._glctx = dr.RasterizeCudaContext()
_warmup(self._glctx)
def render(
self,
vertices: torch.Tensor, # V,3 float
normals: torch.Tensor, # V,3 float
faces: torch.Tensor, # F,3 long
) -> torch.Tensor: # C,H,W,4
V = vertices.shape[0]
faces = faces.type(torch.int32)
vert_hom = torch.cat(
(vertices, torch.ones(V, 1, device=vertices.device)), axis=-1
) # V,3 -> V,4
vertices_clip = vert_hom @ self._mvp.transpose(-2, -1) # C,V,4
rast_out, _ = dr.rasterize(
self._glctx,
vertices_clip,
faces,
resolution=self._image_size,
grad_db=False,
) # C,H,W,4
vert_col = (normals + 1) / 2 # V,3
col, _ = dr.interpolate(vert_col, rast_out, faces) # C,H,W,3
alpha = torch.clamp(rast_out[..., -1:], max=1) # C,H,W,1
col = torch.concat((col, alpha), dim=-1) # C,H,W,4
col = dr.antialias(col, rast_out, vertices_clip, faces) # C,H,W,4
return col # C,H,W,4
|