NVlabs / nvdiffrast

Nvdiffrast - Modular Primitives for High-Performance Differentiable Rendering
Other
1.28k stars 139 forks source link

extract mesh with NeuS mesh #184

Closed ansj11 closed 1 day ago

ansj11 commented 1 month ago

Dear author, I have trained a NeuS mesh and I want to use nvdiffrast to extract texture, but I can not figure out the pose transform from Neus to nvdiffrast. Can you help me to find the bug? Here is my code. ` image_paths = sorted(glob(os.path.join(source_path, 'image/*png'))) cameras = np.load(os.path.join(source_path, 'cameras_sphere.npz')) training_cameras = {'camera_mat':[], 'world_mat':[]} for i in range(len(image_paths)): key0 = 'worldmat%d' % i # inv right key1 = 'scalemat%d' % i world_mat = cameras[key0] scale_mat = cameras[key1] P = world_mat #@ scale_mat # no scale right P = P[:3, :4] intrinsics, pose = load_K_Rt_from_P(None, P) training_cameras['camera_mat'].append(intrinsics) training_cameras['world_mat'].append(pose)

    for key in training_cameras.keys():
        training_cameras[key] = torch.from_numpy(np.stack(training_cameras[key], axis=0)).to(device)
    # set_trace()
    image = cv2.imread(image_paths[0], -1)
    # filterPointCloud(verts, verts, training_cameras, image)
    image_height, image_width = image.shape[:2]
    p3d_cameras = convert_camera_from_gs_to_pytorch3d(training_cameras, image_height, image_width, device)

    mesh = Mesh(v=verts, f=faces, device='cuda')

    n_triangles = len(faces)
    n_squares = n_triangles // 2 + 1
    n_square_per_axis = int(np.sqrt(n_squares) + 1) # 385
    texture_size = square_size * (n_square_per_axis)
    h = w = texture_size
    mesh.auto_uv()
    mesh.auto_normal()

    albedo = torch.zeros((h, w, 3), device=device, dtype=torch.float32)
    cnt = torch.zeros((h, w, 1), device=device, dtype=torch.float32)

    render_resolution = (1280, 720) # wh

    import nvdiffrast.torch as dr
    if device != torch.device("cuda"):
        glctx = dr.RasterizeGLContext()
    else:
        glctx = dr.RasterizeCudaContext()

    # for ver, hor in zip(vers, hors):
    for i, image_path in enumerate(image_paths):
        w2c = training_cameras['world_mat'][i]
        c2w = w2c.inverse()
        c2w[:3, 1:3] *= -1  # r,u,f->r,u,b; 2 wrong
        # c2w[:3,[0,2]] *= 1
        pose = c2w  # 1
        intr = training_cameras['camera_mat'][i]

        image = cv2.imread(image_path, -1)
        height, width = image.shape[:2]
        if height > width:
            render_resolution = render_resolution[::-1]
        image = cv2.resize(image, render_resolution, cv2.INTER_AREA)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        rgbs = torch.from_numpy(image/255.).permute(2, 0, 1).unsqueeze(0).to(device)
        # get coordinate in texture image
        proj = intr.float()
        if not isinstance(pose, torch.Tensor):
            pose = torch.from_numpy(pose.astype(np.float32)).to(device)
            proj = torch.from_numpy(proj.astype(np.float32)).to(device)
        proj[:2,:2] *= -1
        mesh.v[...,1:3] *= -1

        v_cam = torch.matmul(F.pad(mesh.v, pad=(0, 1), mode='constant', value=1.0), torch.inverse(pose).T).float().unsqueeze(0)
        v_clip = v_cam @ proj.T
        rast, rast_db = dr.rasterize(glctx, v_clip, mesh.f, (render_resolution[1], render_resolution[0]))
        uv = v_clip[...,:-1] / v_clip[...,2:3]
        depth, _ = dr.interpolate(-v_cam[..., [2]], rast, mesh.f) # [1, H, W, 1]
        depth = depth.squeeze(0) # [H, W, 1]

        alpha = (rast[0, ..., 3:] > 0).float()

`

s-laine commented 1 month ago

I don't know the conventions used in NeuS, but for nvdiffrast all you need are standard OpenGL (or Direct3D) style modelview and projection matrices. Perhaps you can find more information from the NeuS side? They have at least one GitHub issue related to cameras and transformations here, perhaps that will help you forward.