use pytorch3d to render, instead of nvdiffrast

gradio_app/gradio_3dgen.py

@@ -10,13 +10,8 @@ from scripts.refine_lr_to_sr import run_sr_fast
 from scripts.utils import save_glb_and_video
 from scripts.multiview_inference import geo_reconstruct
 
-
-import nvdiffrast.torch as dr
-dr.RasterizeGLContext(output_db=False)
 @spaces.GPU
 def generate3dv2(preview_img, input_processing, seed, render_video=True, do_refine=True, expansion_weight=0.1, init_type="std"):
-    dr.RasterizeGLContext(output_db=False)  # BUG: cuda_runtime_api.h: No such file or directory
-
     if preview_img is None:
         raise gr.Error("preview_img is none")
     if isinstance(preview_img, str):

mesh_reconstruction/recon.py

@@ -6,14 +6,14 @@ from typing import List
 from mesh_reconstruction.remesh import calc_vertex_normals
 from mesh_reconstruction.opt import MeshOptimizer
 from mesh_reconstruction.func import make_star_cameras_orthographic
-from mesh_reconstruction.render import NormalsRenderer
+from mesh_reconstruction.render import NormalsRenderer, Pytorch3DNormalsRenderer
 from scripts.utils import to_py3d_mesh, init_target
 
 def reconstruct_stage1(pils: List[Image.Image], steps=100, vertices=None, faces=None, start_edge_len=0.15, end_edge_len=0.005, decay=0.995, return_mesh=True, loss_expansion_weight=0.1, gain=0.1):
     vertices, faces = vertices.to("cuda"), faces.to("cuda")
     assert len(pils) == 4
     mv,proj = make_star_cameras_orthographic(4, 1)
-    renderer = NormalsRenderer(mv,proj,list(pils[0].size))
+    renderer = Pytorch3DNormalsRenderer(mv,proj,list(pils[0].size))
     
     target_images = init_target(pils, new_bkgd=(0., 0., 0.)) # 4s
     # 1. no rotate

mesh_reconstruction/refine.py

@@ -5,7 +5,7 @@ from typing import List
 from mesh_reconstruction.remesh import calc_vertex_normals
 from mesh_reconstruction.opt import MeshOptimizer
 from mesh_reconstruction.func import make_star_cameras_orthographic
-from mesh_reconstruction.render import NormalsRenderer
+from mesh_reconstruction.render import NormalsRenderer, Pytorch3DNormalsRenderer
 from scripts.project_mesh import multiview_color_projection, get_cameras_list
 from scripts.utils import to_py3d_mesh, from_py3d_mesh, init_target
 
@@ -18,7 +18,7 @@ def run_mesh_refine(vertices, faces, pils: List[Image.Image], steps=100, start_e
 
     assert len(pils) == 4
     mv,proj = make_star_cameras_orthographic(4, 1)          
-    renderer = NormalsRenderer(mv,proj,list(pils[0].size))
+    renderer = Pytorch3DNormalsRenderer(mv,proj,list(pils[0].size))
 
     target_images = init_target(pils, new_bkgd=(0., 0., 0.)) # 4s
     # 1. no rotate

mesh_reconstruction/render.py

@@ -49,3 +49,121 @@ class NormalsRenderer:
         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
+
+from pytorch3d.structures import Meshes
+from pytorch3d.renderer.mesh.shader import ShaderBase
+from pytorch3d.renderer import (
+    RasterizationSettings,
+    MeshRendererWithFragments,
+    TexturesVertex,
+    MeshRasterizer,
+    BlendParams,
+    FoVOrthographicCameras,
+    look_at_view_transform,
+    hard_rgb_blend,
+)
+
+class VertexColorShader(ShaderBase):
+    def forward(self, fragments, meshes, **kwargs) -> torch.Tensor:
+        blend_params = kwargs.get("blend_params", self.blend_params)
+        texels = meshes.sample_textures(fragments)
+        return hard_rgb_blend(texels, fragments, blend_params)
+
+def 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"):
+    if len(mesh) != len(cameras):
+        if len(cameras) % len(mesh) == 0:
+            mesh = mesh.extend(len(cameras))
+        else:
+            raise NotImplementedError()
+    
+    # render requires everything in float16 or float32
+    input_dtype = dtype
+    blend_params = BlendParams(1e-4, 1e-4, bkgd)
+
+    # Define the settings for rasterization and shading
+    raster_settings = RasterizationSettings(
+        image_size=(H, W),
+        blur_radius=blur_radius,
+        faces_per_pixel=faces_per_pixel,
+        clip_barycentric_coords=True,
+        bin_size=None,
+        max_faces_per_bin=500000,
+    )
+
+    # Create a renderer by composing a rasterizer and a shader
+    # We simply render vertex colors through the custom VertexColorShader (no lighting, materials are used)
+    renderer = MeshRendererWithFragments(
+        rasterizer=MeshRasterizer(
+            cameras=cameras,
+            raster_settings=raster_settings
+        ),
+        shader=VertexColorShader(
+            device=device,
+            cameras=cameras,
+            blend_params=blend_params
+        )
+    )
+
+    # render RGB and depth, get mask
+    with torch.autocast(dtype=input_dtype, device_type=torch.device(device).type):
+        images, _ = renderer(mesh)
+    return images   # BHW4
+
+class Pytorch3DNormalsRenderer:
+    def __init__(self, cameras, image_size, device):
+        self.cameras = cameras.to(device)
+        self._image_size = image_size
+        self.device = device
+    
+    def render(self,
+            vertices: torch.Tensor, #V,3 float
+            normals: torch.Tensor, #V,3 float   in [-1, 1]
+            faces: torch.Tensor, #F,3 long
+            ) ->torch.Tensor: #C,H,W,4
+        mesh = Meshes(verts=[vertices], faces=[faces], textures=TexturesVertex(verts_features=[(normals + 1) / 2])).to(self.device)
+        return render_mesh_vertex_color(mesh, self.cameras, self._image_size[0], self._image_size[1], device=self.device)
+    
+def get_camera(R, T, focal_length=1 / (2**0.5)):
+    focal_length = 1 / focal_length
+    camera = FoVOrthographicCameras(device=R.device, R=R, T=T, min_x=-focal_length, max_x=focal_length, min_y=-focal_length, max_y=focal_length)
+    return camera
+
+def make_star_cameras_orthographic_py3d(azim_list, device, focal=2/1.35, dist=1.1):
+    R, T = look_at_view_transform(dist, 0, azim_list)
+    focal_length = 1 / focal
+    return FoVOrthographicCameras(device=R.device, R=R, T=T, min_x=-focal_length, max_x=focal_length, min_y=-focal_length, max_y=focal_length).to(device)
+
+def save_tensor_to_img(tensor, save_dir):
+    from PIL import Image
+    import numpy as np
+    for idx, img in enumerate(tensor):
+        img = img[..., :3].cpu().numpy()
+        img = (img * 255).astype(np.uint8)
+        img = Image.fromarray(img)
+        img.save(save_dir + f"{idx}.png")
+
+if __name__ == "__main__":
+    import sys
+    import os
+    sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+    from mesh_reconstruction.func import make_star_cameras_orthographic
+    cameras = make_star_cameras_orthographic_py3d([0, 270, 180, 90], device="cuda", focal=1., dist=4.0)
+    mv,proj = make_star_cameras_orthographic(4, 1)
+    resolution = 1024
+    renderer1 = NormalsRenderer(mv,proj, [resolution,resolution], device="cuda")
+    renderer2 = Pytorch3DNormalsRenderer(cameras, [resolution,resolution], device="cuda")
+    vertices = torch.tensor([[0,0,0],[0,0,1],[0,1,0],[1,0,0]], device="cuda", dtype=torch.float32)
+    normals = torch.tensor([[-1,-1,-1],[1,-1,-1],[-1,-1,1],[-1,1,-1]], device="cuda", dtype=torch.float32)
+    faces = torch.tensor([[0,1,2],[0,1,3],[0,2,3],[1,2,3]], device="cuda", dtype=torch.long)
+    
+    import time
+    t0 = time.time()
+    r1 = renderer1.render(vertices, normals, faces)
+    print("time r1:", time.time() - t0)
+    
+    t0 = time.time()
+    r2 = renderer2.render(vertices, normals, faces)
+    print("time r2:", time.time() - t0)
+    
+    for i in range(4):
+        print((r1[i]-r2[i]).abs().mean(), (r1[i]+r2[i]).abs().mean())

scripts/project_mesh.py

@@ -13,17 +13,6 @@ from pytorch3d.renderer import (
 )
 from pytorch3d.renderer import MeshRasterizer
 
-def get_camera(world_to_cam, fov_in_degrees=60, focal_length=1 / (2**0.5), cam_type='fov'):
-    # pytorch3d expects transforms as row-vectors, so flip rotation: https://github.com/facebookresearch/pytorch3d/issues/1183
-    R = world_to_cam[:3, :3].t()[None, ...]
-    T = world_to_cam[:3, 3][None, ...]
-    if cam_type == 'fov':
-        camera = FoVPerspectiveCameras(device=world_to_cam.device, R=R, T=T, fov=fov_in_degrees, degrees=True)
-    else:
-        focal_length = 1 / focal_length
-        camera = FoVOrthographicCameras(device=world_to_cam.device, R=R, T=T, min_x=-focal_length, max_x=focal_length, min_y=-focal_length, max_y=focal_length)
-    return camera
-
 def render_pix2faces_py3d(meshes, cameras, H=512, W=512, blur_radius=0.0, faces_per_pixel=1):
     """
     Renders pix2face of visible faces.
@@ -98,11 +87,11 @@ class Pix2FacesRenderer:
 pix2faces_renderer = None
 
 def get_visible_faces(meshes: Meshes, cameras: CamerasBase, resolution=1024):
-    global pix2faces_renderer
-    if pix2faces_renderer is None:
-        pix2faces_renderer = Pix2FacesRenderer()
-    # pix_to_face = render_pix2faces_py3d(meshes, cameras, H=resolution, W=resolution)['pix_to_face']
-    pix_to_face = pix2faces_renderer.render_pix2faces_nvdiff(meshes, cameras, H=resolution, W=resolution)
+    # global pix2faces_renderer
+    # if pix2faces_renderer is None:
+    #     pix2faces_renderer = Pix2FacesRenderer()
+    pix_to_face = render_pix2faces_py3d(meshes, cameras, H=resolution, W=resolution)['pix_to_face']
+    # pix_to_face = pix2faces_renderer.render_pix2faces_nvdiff(meshes, cameras, H=resolution, W=resolution)
 
     unique_faces = torch.unique(pix_to_face.flatten())
     unique_faces = unique_faces[unique_faces != -1]
@@ -313,12 +302,19 @@ def multiview_color_projection(meshes: Meshes, image_list: List[Image.Image], ca
     del meshes
     return ret_mesh
 
+def get_camera(R, T, fov_in_degrees=60, focal_length=1 / (2**0.5), cam_type='fov'):
+    if cam_type == 'fov':
+        camera = FoVPerspectiveCameras(device=R.device, R=R, T=T, fov=fov_in_degrees, degrees=True)
+    else:
+        focal_length = 1 / focal_length
+        camera = FoVOrthographicCameras(device=R.device, R=R, T=T, min_x=-focal_length, max_x=focal_length, min_y=-focal_length, max_y=focal_length)
+    return camera
+
 def get_cameras_list(azim_list, device, focal=2/1.35, dist=1.1):
     ret = []
     for azim in azim_list:
         R, T = look_at_view_transform(dist, 0, azim)
-        w2c = torch.cat([R[0].T, T[0, :, None]], dim=1)
-        cameras: OrthographicCameras = get_camera(w2c, focal_length=focal, cam_type='orthogonal').to(device)
+        cameras: OrthographicCameras = get_camera(R, T, focal_length=focal, cam_type='orthogonal').to(device)
         ret.append(cameras)
     return ret