TSDF笔记

def transform_dtu_mesh(scene, mesh):
    # Taking the biggest connected component
    print("Taking the biggest connected component")
    cleaned_mesh = mesh
    triangle_clusters, cluster_n_triangles, cluster_area = cleaned_mesh.cluster_connected_triangles()
    cluster_n_triangles = np.asarray(cluster_n_triangles)
    largest_cluster_idx = cluster_n_triangles.argmax()
    triangles_to_remove = np.array(triangle_clusters) != largest_cluster_idx
    cleaned_mesh.remove_triangles_by_mask(triangles_to_remove)
    cleaned_mesh.remove_unreferenced_vertices()
    
    # transform to world
    cam_file = f"{scene.source_path}/cameras_sphere.npz"
    scale_mat = np.identity(4)
    if os.path.exists(cam_file):
        camera_param = dict(np.load(cam_file))
        scale_mat = camera_param['scale_mat_0']

    vertices = np.asarray(cleaned_mesh.vertices)
    vertices = vertices * scale_mat[0,0] + scale_mat[:3,3][None]
    cleaned_mesh.vertices = o3d.utility.Vector3dVector(vertices)

    return cleaned_mesh

def render_set(model_path, name, iteration, views, scene, gaussians, pipeline, background, 
               app_model=None, max_depth=5.0, volume=None, use_depth_filter=False):
    gts_path = os.path.join(model_path, name, "ours_{}".format(iteration), "gt")
    render_path = os.path.join(model_path, name, "ours_{}".format(iteration), "renders")
    render_depth_path = os.path.join(model_path, name, "ours_{}".format(iteration), "renders_depth")
    render_normal_path = os.path.join(model_path, name, "ours_{}".format(iteration), "renders_normal")

    makedirs(gts_path, exist_ok=True)
    makedirs(render_path, exist_ok=True)
    makedirs(render_depth_path, exist_ok=True)
    makedirs(render_normal_path, exist_ok=True)
    print(len(views))
    depths_tsdf_fusion = []
    #初始化
    for idx, view in enumerate(tqdm(views, desc="Rendering progress")):
        gt, _ = view.get_image()
        out = render(view, gaussians, pipeline, background, app_model=app_model)
        #按照训练流程对每一张视图输入splatting出对应的渲染结果，其中包括颜色，深度，法线等
        rendering = out["render"].clamp(0.0, 1.0)
        #渲染的RGB限制在【0，1】
        _, H, W = rendering.shape

        depth = out["plane_depth"].squeeze()
        depth_tsdf = depth.clone()
        depths_tsdf_fusion.append(depth_tsdf.squeeze())
        #
    if volume is not None:
        depths_tsdf_fusion = torch.stack(depths_tsdf_fusion, dim=0)
        for idx, view in enumerate(tqdm(views, desc="TSDF Fusion progress")):
            ref_depth = depths_tsdf_fusion[idx]
            #pts = gaussians.get_points_from_depth(view, ref_depth)
            ref_depth[ref_depth>max_depth] = 0
            ref_depth = ref_depth.detach().cpu().numpy()
            #把所有视图的深度整合在一起得到depths_tsdf_fusion，对于每一个视图去掉一些深度大于阈值的设置为0。
            pose = np.identity(4)#创建单位矩阵
            pose[:3,:3] = view.R.transpose(-1,-2)
            pose[:3, 3] = view.T
            color = o3d.io.read_image(os.path.join(render_path, view.image_name + ".jpg"))
            depth = o3d.geometry.Image((ref_depth*1000).astype(np.uint16))
            rgbd = o3d.geometry.RGBDImage.create_from_color_and_depth(
                color, depth, depth_scale=1000.0, depth_trunc=max_depth, convert_rgb_to_intensity=False)
            #彩色图像和深度图像合并成一个 RGBD 图像
            volume.integrate(
                rgbd,
                o3d.camera.PinholeCameraIntrinsic(W, H, view.Fx, view.Fy, view.Cx, view.Cy),
                pose)
            #将RGBD图像，相机内参，相机姿态输入到integrate方法里对volume体素网格进行更新，每个视图都重复以上步骤

def render_sets(dataset : ModelParams, iteration : int, pipeline : PipelineParams, skip_train : bool, skip_test : bool,
                 max_depth : float, voxel_size : float, use_depth_filter : bool):
    with torch.no_grad():
        gaussians = GaussianModel(dataset.sh_degree)
        scene = Scene(dataset, gaussians, load_iteration=iteration, shuffle=False)
        bg_color = [1,1,1] if dataset.white_background else [0, 0, 0]
        background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")
        print(f"TSDF voxel_size {voxel_size}")
        #首先在GaussianModel中初始化高斯的参数，在Scene中载入训练后的结果，默认设置背景为白色
        volume = o3d.pipelines.integration.ScalableTSDFVolume(
        voxel_length=voxel_size,
        sdf_trunc=4.0*voxel_size,
        color_type=o3d.pipelines.integration.TSDFVolumeColorType.RGB8)
        #初始化TSDF中的体素网格，设置voxel_size，越小结果越精细。sdf_trunc设置了TSDF过程中SDF的上下限。
        if not skip_train:
            render_set(dataset.model_path, "train", scene.loaded_iter, scene.getTrainCameras(), scene, gaussians, pipeline, background, 
                       max_depth=max_depth, volume=volume, use_depth_filter=use_depth_filter)
            print(f"extract_triangle_mesh")
            mesh = volume.extract_triangle_mesh()
            mesh = transform_dtu_mesh(scene, mesh)
            #网格变换操作，得到mesh后转到世界坐标系

            path = os.path.join(dataset.model_path, "mesh")
            os.makedirs(path, exist_ok=True)
            
            o3d.io.write_triangle_mesh(os.path.join(path, "tsdf_fusion.ply"), mesh, 
                                       write_triangle_uvs=True, write_vertex_colors=True, write_vertex_normals=True)
        # if not skip_test:
        #     print(222)
        #     input()
        #     render_set(dataset.model_path, "test", scene.loaded_iter, scene.getTestCameras(), scene, gaussians, pipeline, background)

TSDF中o3d.piplines.integration.TSDF的计算流程

TSDF的输入是RGBD（颜色+深度），相机内外参。具体计算方式如下图。给定一个相机视角，我们首先通过一系列深度图估计出其隐式表面（surface）。其中ds是相机到隐式表面的深度（即深度图的深度值）。我们计算体素到隐式表面的距离（dx）作为SDF。因为隐式表面是由深度图估计的因此位于像素坐标系下。因此我们需要利用相机内外参将体素从世界坐标系转换为图像坐标系。此时就可以计算相机到voxel的距离dv。此时就知道了SDF（dx）=ds-dv。T指的是对大于或小于阈值的SDF值设定了一个边界值。TSDF>0说明体素在表面前，反之说明在表面后