hyz317/StdGEN

1# modified from https://github.com/Profactor/continuous-remeshing2import spaces3import nvdiffrast.torch as dr4import torch5from typing import Tuple6import torch.nn.functional as tfunc7 8 9def _warmup(glctx, device=None):10    device = 'cuda' if device is None else device11    #windows workaround for https://github.com/NVlabs/nvdiffrast/issues/5912    def tensor(*args, **kwargs):13        return torch.tensor(*args, device=device, **kwargs)14    pos = tensor([[[-0.8, -0.8, 0, 1], [0.8, -0.8, 0, 1], [-0.8, 0.8, 0, 1]]], dtype=torch.float32)15    tri = tensor([[0, 1, 2]], dtype=torch.int32)16    dr.rasterize(glctx, pos, tri, resolution=[256, 256])17 18 19class NormalsRenderer:20    21    _glctx:dr.RasterizeCudaContext = None22    23    def __init__(24            self,25            mv: torch.Tensor, #C,4,426            proj: torch.Tensor, #C,4,427            image_size: Tuple[int,int],28            mvp = None,29            device=None,30            ):31        if mvp is None:32            self._mvp = proj @ mv #C,4,433        else:34            self._mvp = mvp35        self._image_size = image_size36        self._glctx = dr.RasterizeCudaContext(device=device)37        _warmup(self._glctx, device)38 39    def render(self,40            vertices: torch.Tensor, #V,3 float41            normals: torch.Tensor, #V,3 float   in [-1, 1]42            faces: torch.Tensor, #F,3 long43            ) ->torch.Tensor: #C,H,W,444 45        V = vertices.shape[0]46        faces = faces.type(torch.int32)47        vert_hom = torch.cat((vertices, torch.ones(V,1,device=vertices.device)),axis=-1) #V,3 -> V,448        vertices_clip = vert_hom @ self._mvp.transpose(-2,-1) #C,V,449        50        rast_out,_ = dr.rasterize(self._glctx, vertices_clip, faces, resolution=self._image_size, grad_db=False) #C,H,W,451        vert_col = (normals+1)/2 #V,352        col,_ = dr.interpolate(vert_col, rast_out, faces) #C,H,W,353        alpha = torch.clamp(rast_out[..., -1:], max=1) #C,H,W,154        col = torch.concat((col,alpha),dim=-1) #C,H,W,455        col = dr.antialias(col, rast_out, vertices_clip, faces) #C,H,W,456        return col #C,H,W,457 58 59 60from pytorch3d.structures import Meshes61from pytorch3d.renderer.mesh.shader import ShaderBase62from pytorch3d.renderer import (63    RasterizationSettings,64    MeshRendererWithFragments,65    TexturesVertex,66    MeshRasterizer,67    BlendParams,68    FoVOrthographicCameras,69    look_at_view_transform,70    hard_rgb_blend,71)72 73class VertexColorShader(ShaderBase):74    def forward(self, fragments, meshes, **kwargs) -> torch.Tensor:75        blend_params = kwargs.get("blend_params", self.blend_params)76        texels = meshes.sample_textures(fragments)77        return hard_rgb_blend(texels, fragments, blend_params)78 79def render_mesh_vertex_color(mesh, cameras, H, W, blur_radius=0.0, faces_per_pixel=1, bkgd=(0., 0., 0.), dtype=torch.float32, device="cuda"):80    if len(mesh) != len(cameras):81        if len(cameras) % len(mesh) == 0:82            mesh = mesh.extend(len(cameras))83        else:84            raise NotImplementedError()85    86    # render requires everything in float16 or float3287    input_dtype = dtype88    blend_params = BlendParams(1e-4, 1e-4, bkgd)89 90    # Define the settings for rasterization and shading91    raster_settings = RasterizationSettings(92        image_size=(H, W),93        blur_radius=blur_radius,94        faces_per_pixel=faces_per_pixel,95        clip_barycentric_coords=True,96        bin_size=None,97        max_faces_per_bin=None,98    )99 100    # Create a renderer by composing a rasterizer and a shader101    # We simply render vertex colors through the custom VertexColorShader (no lighting, materials are used)102    renderer = MeshRendererWithFragments(103        rasterizer=MeshRasterizer(104            cameras=cameras,105            raster_settings=raster_settings106        ),107        shader=VertexColorShader(108            device=device,109            cameras=cameras,110            blend_params=blend_params111        )112    )113 114    # render RGB and depth, get mask115    with torch.autocast(dtype=input_dtype, device_type=torch.device(device).type):116        images, _ = renderer(mesh)117    return images   # BHW4118 119class Pytorch3DNormalsRenderer: # 100 times slower!!!120    def __init__(self, cameras, image_size, device):121        self.cameras = cameras.to(device)122        self._image_size = image_size123        self.device = device124    125    def render(self,126            vertices: torch.Tensor, #V,3 float127            normals: torch.Tensor, #V,3 float   in [-1, 1]128            faces: torch.Tensor, #F,3 long129            ) ->torch.Tensor: #C,H,W,4130        mesh = Meshes(verts=[vertices], faces=[faces], textures=TexturesVertex(verts_features=[(normals + 1) / 2])).to(self.device)131        return render_mesh_vertex_color(mesh, self.cameras, self._image_size[0], self._image_size[1], device=self.device)132    133def save_tensor_to_img(tensor, save_dir):134    from PIL import Image135    import numpy as np136    for idx, img in enumerate(tensor):137        img = img[..., :3].cpu().numpy()138        img = (img * 255).astype(np.uint8)139        img = Image.fromarray(img)140        img.save(save_dir + f"{idx}.png")141 142if __name__ == "__main__":143    import sys144    import os145    sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))146    from mesh_reconstruction.func import make_star_cameras_orthographic, make_star_cameras_orthographic_py3d147    cameras = make_star_cameras_orthographic_py3d([0, 270, 180, 90], device="cuda", focal=1., dist=4.0)148    mv,proj = make_star_cameras_orthographic(4, 1)149    resolution = 1024150    renderer1 = NormalsRenderer(mv,proj, [resolution,resolution], device="cuda")151    renderer2 = Pytorch3DNormalsRenderer(cameras, [resolution,resolution], device="cuda")152    vertices = torch.tensor([[0,0,0],[0,0,1],[0,1,0],[1,0,0]], device="cuda", dtype=torch.float32)153    normals = torch.tensor([[-1,-1,-1],[1,-1,-1],[-1,-1,1],[-1,1,-1]], device="cuda", dtype=torch.float32)154    faces = torch.tensor([[0,1,2],[0,1,3],[0,2,3],[1,2,3]], device="cuda", dtype=torch.long)155    156    import time157    t0 = time.time()158    r1 = renderer1.render(vertices, normals, faces)159    print("time r1:", time.time() - t0)160    161    t0 = time.time()162    r2 = renderer2.render(vertices, normals, faces)163    print("time r2:", time.time() - t0)164    165    for i in range(4):166        print((r1[i]-r2[i]).abs().mean(), (r1[i]+r2[i]).abs().mean())167 168 169def calc_face_normals(170        vertices:torch.Tensor, #V,3 first vertex may be unreferenced171        faces:torch.Tensor, #F,3 long, first face may be all zero172        normalize:bool=False,173        )->torch.Tensor: #F,3174    """175         n176         |177         c0     corners ordered counterclockwise when178        / \     looking onto surface (in neg normal direction)179      c1---c2180    """181    full_vertices = vertices[faces] #F,C=3,3182    v0,v1,v2 = full_vertices.unbind(dim=1) #F,3183    face_normals = torch.cross(v1-v0,v2-v0, dim=1) #F,3184    if normalize:185        face_normals = tfunc.normalize(face_normals, eps=1e-6, dim=1) 186    return face_normals #F,3187 188 189def calc_vertex_normals(190        vertices:torch.Tensor, #V,3 first vertex may be unreferenced191        faces:torch.Tensor, #F,3 long, first face may be all zero192        face_normals:torch.Tensor=None, #F,3, not normalized193        )->torch.Tensor: #F,3194 195    F = faces.shape[0]196 197    if face_normals is None:198        face_normals = calc_face_normals(vertices,faces)199    200    vertex_normals = torch.zeros((vertices.shape[0],3,3),dtype=vertices.dtype,device=vertices.device) #V,C=3,3201    vertex_normals.scatter_add_(dim=0,index=faces[:,:,None].expand(F,3,3),src=face_normals[:,None,:].expand(F,3,3))202    vertex_normals = vertex_normals.sum(dim=1) #V,3203    return tfunc.normalize(vertex_normals, eps=1e-6, dim=1)204