[논문리뷰] SuGaR: Surface-Aligned Gaussian Splatting for Efficient 3D Mesh Reconstruction and High-Quality Mesh Rendering

SuGaR: Surface-Aligned Gaussian Splatting for Efficient 3D Mesh Reconstruction and High-Quality Mesh Rendering

CVPR 2024. [Paper] [Page] [Github][Youtube] Antoine Guédon Vincent Lepetit LIGM, Ecole des Ponts, Univ Gustave Eiffel, CNRS 21 Nov 2023

Gaussian Splatting 원리

• Novel View Synthesis (NVS)를 하기 위해 3DGS는 millions of tiny 3D Gaussain Primitives를 제공하여 scene의 geometry를 represent하도록 합니다.
• 그리고 3D Gaussian에 tailored된 new rasterization strategy로 image를 훨씬 빠르게 render합니다. (volumetric rendering보다 빠릅니다)

• 비록 3DGS이 photorealistic rendering을 가능하게 하지만, 3D 아티스트는 surfaces의 explicit representation이 필요합니다.
• 일반적으로 explicit representation은 triangle mesh라고 부릅니다.
• 실제로, mesh는 컴퓨터 그래픽스에서 default representation입니다.

• Standard practice로 radiance fields에서 meshes를 extract하는 방법은 Marching Cubes를 적용하는 것입니다.
• 하지만 3dgs는 real scene을 reconstruct할 때, millions of tiny and unorganized Gaussians을 생성합니다.
• Marching Cubes는 그렇게 tiny Primitives에 대해서는 does not scale well 합니다.
• 그리고 이와 같이 작은 tiny Primitives로 meshes를 output하면 많은 holes과 artifacts가 생깁니다.

• 게다가, 일부 방법이 NeRFs에서 meshes을 추출하도록 제안되었지만, 이러한 접근 방식은 매우 느리거나 부정확할 수 있습니다.

• SuGaR를 통해 우리는 3D 가우시안 스플래팅 표현에서 정확하고 편집 가능한 메쉬를 단일 GPU에서 몇 분 만에 추출하는 방법을 소개합니다. 이 메쉬는 편집, 스컬핑, 애니메이션화, 합성 및 매우 현실적인 가우시안 스플래팅 렌더링이 가능합니다.

• 우리의 접근 방식은 다음 세 가지 기여를 포함합니다.

• Poisson reconstruction 수행시, points와 points에 해당하는 normals이 주어져야 합니다.
• Poisson reconstruction을 수행하면, triangle mesh를 얻을 수 있습니다.
• Poisson surface reconstruction creates watertight surfaces from oriented point sets. (watertight는 닫힌 표면이라는 의미입니다.)

• 메쉬를 추출한 후 새로운 가우시안을 결합하여 매우 현실적인 가우시안 스플래팅 렌더링을 통해 메쉬를 렌더링할 수 있습니다.
• 메쉬의 삼각형을 매우 얇은 가우시안으로 텍스처링합니다. 이러한 가우시안의 매개변수는 삼각형의 좌표 공간에서 refined되어 메쉬를 편집하거나 애니메이션화할 때 자동으로 조정됩니다.
• We simply texture the triangles of the mesh with very thin gaussians.
• The parameters of these gaussians are refined in the coordinate space of the triangle so that they are automatically adjusted when editing or animating the mesh.

• SuGaR의 train.py 코드를 보면 n_gaussian_per_surface_triangle이라는 변수로 triangle에 bind한 new gaussian의 수를 결정합니다.

# train.py

# ----- Extract mesh and texture from refined SuGaR -----
    if args.export_uv_textured_mesh:
        refined_mesh_args = AttrDict({
            'scene_path': args.scene_path,
            'iteration_to_load': args.iteration_to_load,
            'checkpoint_path': args.checkpoint_path,
            'refined_model_path': refined_sugar_path,
            'mesh_output_dir': None,
            'n_gaussians_per_surface_triangle': args.gaussians_per_triangle,
            'square_size': args.square_size,
            'eval': args.eval,
            'gpu': args.gpu,
            'postprocess_mesh': args.postprocess_mesh,
            'postprocess_density_threshold': args.postprocess_density_threshold,
            'postprocess_iterations': args.postprocess_iterations,
        })
        refined_mesh_path = extract_mesh_and_texture_from_refined_sugar(refined_mesh_args)

• 그리고 sugar_scene/sugar_model.py에서 n_gaussian_per_surface_triangle의 수만큼 triangle 내부의 barycentric coordinates를 manual하게 정의해줍니다.
• 이후 barycentric coordinates 위치들에 new gaussian들을 생성해줍니다.
• barycentric coordinates는 triangle의 3개의 vertices에서 triangle 내부의 한점 P까지의 거리를 구할 때, 각 정점으로부터의 거리의 weights를 의미합니다.
• 자세한 내용은 글쓴이 본인이 작성한 barycentric coordinates 관련 포스트를 참고하세요.
• 아래 코드에서와 같이 barycentric coordinates를 정하면, 각 가우시안이 삼각형 내부에서 고르게 분포되도록 하여, 텍스처링 및 렌더링 시 메쉬의 표면을 더욱 자연스럽고 부드럽게 표현할 수 있게 합니다.

# sugar_scene/sugar_model.py

class SuGaR(nn.Module):
    """Main class for SuGaR models.
    Because SuGaR optimization starts with first optimizing a vanilla Gaussian Splatting model for 7k iterations,
    we built this class as a wrapper of a vanilla Gaussian Splatting model.
    Consequently, a corresponding Gaussian Splatting model trained for 7k iterations must be provided.
    However, this wrapper implementation may not be the most optimal one for memory usage, so we might change it in the future.
    """
    def __init__(
        self, 
        nerfmodel: GaussianSplattingWrapper,
        points: torch.Tensor,
        colors: torch.Tensor,
        initialize:bool=True,
        sh_levels:int=4,
        learnable_positions:bool=True,
        triangle_scale:float=2.,
        keep_track_of_knn:bool=False,
        knn_to_track:int=16,
        learn_color_only=False,
        beta_mode='average',  # 'learnable', 'average', 'weighted_average'
        freeze_gaussians=False,
        primitive_types='diamond',  # 'diamond', 'square'
        surface_mesh_to_bind=None,  # Open3D mesh
        surface_mesh_thickness=None,
        learn_surface_mesh_positions=True,
        learn_surface_mesh_opacity=True,
        learn_surface_mesh_scales=True,
        n_gaussians_per_surface_triangle=6,  # 1, 3, 4 or 6
        editable=False,  # If True, allows for automatically rescaling Gaussians in real time if triangles are deformed from their original shape. 
        # We wrote about this functionality in the paper, and it was previously part of sugar_compositor.py, which we haven't finished cleaning yet for this repo.
        # We now moved it to this script as it is more related to the SuGaR model than to the compositor.
        *args, **kwargs) -> None:
        """
        Args:
            nerfmodel (GaussianSplattingWrapper): A vanilla Gaussian Splatting model trained for 7k iterations.
            points (torch.Tensor): Initial positions of the Gaussians (not used when wrapping).
            colors (torch.Tensor): Initial colors of the Gaussians (not used when wrapping).
            initialize (bool, optional): Whether to initialize the radiuses. Defaults to True.
            sh_levels (int, optional): Number of spherical harmonics levels to use for the color features. Defaults to 4.
            learnable_positions (bool, optional): Whether to learn the positions of the Gaussians. Defaults to True.
            triangle_scale (float, optional): Scale of the triangles used to replace the Gaussians. Defaults to 2.
            keep_track_of_knn (bool, optional): Whether to keep track of the KNN information for training regularization. Defaults to False.
            knn_to_track (int, optional): Number of KNN to track. Defaults to 16.
            learn_color_only (bool, optional): Whether to learn only the color features. Defaults to False.
            beta_mode (str, optional): Whether to use a learnable beta, or to average the beta values. Defaults to 'average'.
            freeze_gaussians (bool, optional): Whether to freeze the Gaussians. Defaults to False.
            primitive_types (str, optional): Type of primitive to use to replace the Gaussians. Defaults to 'diamond'.
            surface_mesh_to_bind (None, optional): Surface mesh to bind the Gaussians to. Defaults to None.
            surface_mesh_thickness (None, optional): Thickness of the bound Gaussians. Defaults to None.
            learn_surface_mesh_positions (bool, optional): Whether to learn the positions of the bound Gaussians. Defaults to True.
            learn_surface_mesh_opacity (bool, optional): Whether to learn the opacity of the bound Gaussians. Defaults to True.
            learn_surface_mesh_scales (bool, optional): Whether to learn the scales of the bound Gaussians. Defaults to True.
            n_gaussians_per_surface_triangle (int, optional): Number of bound Gaussians per surface triangle. Defaults to 6.
        """
        
        super(SuGaR, self).__init__()
        
        self.nerfmodel = nerfmodel
        self.freeze_gaussians = freeze_gaussians
        
        self.learn_positions = ((not learn_color_only) and learnable_positions) and (not freeze_gaussians)
        self.learn_opacities = (not learn_color_only) and (not freeze_gaussians)
        self.learn_scales = (not learn_color_only) and (not freeze_gaussians)
        self.learn_quaternions = (not learn_color_only) and (not freeze_gaussians)
        self.learnable_positions = learnable_positions
        
        if surface_mesh_to_bind is not None:
            self.learn_surface_mesh_positions = learn_surface_mesh_positions
            self.binded_to_surface_mesh = True
            self.learn_surface_mesh_opacity = learn_surface_mesh_opacity
            self.learn_surface_mesh_scales = learn_surface_mesh_scales
            self.n_gaussians_per_surface_triangle = n_gaussians_per_surface_triangle
            self.editable = editable
            
            self.learn_positions = self.learn_surface_mesh_positions
            self.learn_scales = self.learn_surface_mesh_scales
            self.learn_quaternions = self.learn_surface_mesh_scales
            self.learn_opacities = self.learn_surface_mesh_opacity
            
            self._surface_mesh_faces = torch.nn.Parameter(
                torch.tensor(np.array(surface_mesh_to_bind.triangles)).to(nerfmodel.device), 
                requires_grad=False).to(nerfmodel.device)
            if surface_mesh_thickness is None:
                surface_mesh_thickness = nerfmodel.training_cameras.get_spatial_extent() / 1_000_000
            self.surface_mesh_thickness = torch.nn.Parameter(
                torch.tensor(surface_mesh_thickness).to(nerfmodel.device), 
                requires_grad=False).to(nerfmodel.device)
            
            print("Binding radiance cloud to surface mesh...")
            if n_gaussians_per_surface_triangle == 1:
                self.surface_triangle_circle_radius = 1. / 2. / np.sqrt(3.)
                self.surface_triangle_bary_coords = torch.tensor(
                    [[1/3, 1/3, 1/3]],
                    dtype=torch.float32,
                    device=nerfmodel.device,
                )[..., None]
            
            if n_gaussians_per_surface_triangle == 3:
                self.surface_triangle_circle_radius = 1. / 2. / (np.sqrt(3.) + 1.)
                self.surface_triangle_bary_coords = torch.tensor(
                    [[1/2, 1/4, 1/4],
                    [1/4, 1/2, 1/4],
                    [1/4, 1/4, 1/2]],
                    dtype=torch.float32,
                    device=nerfmodel.device,
                )[..., None]
            
            if n_gaussians_per_surface_triangle == 4:
                self.surface_triangle_circle_radius = 1 / (4. * np.sqrt(3.))
                self.surface_triangle_bary_coords = torch.tensor(
                    [[1/3, 1/3, 1/3],
                    [2/3, 1/6, 1/6],
                    [1/6, 2/3, 1/6],
                    [1/6, 1/6, 2/3]],
                    dtype=torch.float32,
                    device=nerfmodel.device,
                )[..., None]  # n_gaussians_per_face, 3, 1
                
            if n_gaussians_per_surface_triangle == 6:
                self.surface_triangle_circle_radius = 1 / (4. + 2.*np.sqrt(3.))
                self.surface_triangle_bary_coords = torch.tensor(
                    [[2/3, 1/6, 1/6],
                    [1/6, 2/3, 1/6],
                    [1/6, 1/6, 2/3],
                    [1/6, 5/12, 5/12],
                    [5/12, 1/6, 5/12],
                    [5/12, 5/12, 1/6]],
                    dtype=torch.float32,
                    device=nerfmodel.device,
                )[..., None]
                
            points = torch.tensor(np.array(surface_mesh_to_bind.vertices)).float().to(nerfmodel.device)
            # verts_normals = torch.tensor(np.array(surface_mesh_to_bind.vertex_normals)).float().to(nerfmodel.device)
            self._vertex_colors = torch.tensor(np.array(surface_mesh_to_bind.vertex_colors)).float().to(nerfmodel.device)
            faces_colors = self._vertex_colors[self._surface_mesh_faces]  # n_faces, 3, n_coords
            colors = faces_colors[:, None] * self.surface_triangle_bary_coords[None]  # n_faces, n_gaussians_per_face, 3, n_colors
            colors = colors.sum(dim=-2)  # n_faces, n_gaussians_per_face, n_colors
            colors = colors.reshape(-1, 3)  # n_faces * n_gaussians_per_face, n_colors
                
            self._points = nn.Parameter(points, requires_grad=self.learn_positions).to(nerfmodel.device)
            n_points = len(np.array(surface_mesh_to_bind.triangles)) * n_gaussians_per_surface_triangle
            self._n_points = n_points
            
        else:
            self.binded_to_surface_mesh = False
            self._points = nn.Parameter(points, requires_grad=self.learn_positions).to(nerfmodel.device)
            n_points = len(self._points)