Source code for miop.surface_reconstruction

# Copyright (c) 2025, Maxime Paschoud.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
#
# (http://opensource.org/licenses/BSD-3-Clause)
#
# __author__ = "Maxime Paschoud, ETHZ: CMBM"
#

from pipeline import DAGNode

import numpy as np
import open3d as o3d

import os




class PoissonSurfaceReconstruction(DAGNode):

    """
    Class for performing surface reconstruction from 3D point clouds using the Poisson Surface Reconstruction method.

    This class wraps the Open3D implementation of Poisson reconstruction and includes methods for visualizing and saving the resulting mesh.

    Attributes
    ----------
    depth : int
        Octree depth controlling reconstruction resolution. Higher values give finer meshes at the cost of more computation. Default is 5.
    color : np.ndarray, shape (3, 1)
        RGB color used to uniformly paint the reconstructed mesh. Default is gray ([[0.5], [0.5], [0.5]]).
    mesh : o3d.geometry.TriangleMesh or None
        Reconstructed mesh, set after `eval()` is called.
    """

    def __init__(self, depth: int = 5, color: np.ndarray = np.array([[0.5], [0.5], [0.5]])):
        """
        Initializes the PoissonSurfaceReconstruction module.

        Parameters
        ----------
        depth : int, optional
            Depth of the octree used in the Poisson Surface Reconstruction algorithm. Controls mesh resolution. Default is 5.
        color : np.ndarray, shape (3, 1), optional
            RGB color to paint the output mesh. Default is medium gray.
        """

        super().__init__()
        self.mesh = None
        self.depth = depth
        self.color = color



    def eval(self, pcds):

        """
        Estimates normals and performs Poisson Surface Reconstruction on the first point cloud in the list.

        Parameters
        ----------
        pcds : list of o3d.geometry.PointCloud
            A list of Open3D point clouds. Only the first point cloud (`pcds[0]`) is used for reconstruction.

        Returns
        -------
        list of o3d.geometry.TriangleMesh
            A list containing one element: the reconstructed mesh.
        
        Raises
        ------
        ValueError
            If the input list is empty or the input is not a valid point cloud.
        """

        print("Start of Poisson Reconstruction...")

        o3d.utility.set_verbosity_level(o3d.utility.VerbosityLevel.Error)

        pcd = pcds[0]

        # estimate normals
        pcd.estimate_normals()
        pcd.orient_normals_consistent_tangent_plane(500)

        # check if normals point in +z direction
        sign = np.asarray(pcd.normals)[:, 2].mean()

        if sign < 0:
            for i, n in enumerate(pcd.normals):
                pcd.normals[i] = -n
        
        with o3d.utility.VerbosityContextManager(
                o3d.utility.VerbosityLevel.Error) as cm:
            mesh, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(
                pcd, depth=self.depth)
        
        mesh.compute_vertex_normals()
        mesh.paint_uniform_color(self.color)

        self.mesh = mesh

        print("End of Poisson Reconstruction.")

        return [mesh]




    def save(self, path):
        """
        Saves the reconstructed mesh to a PLY file.

        Parameters
        ----------
        path : str
            Path to save the mesh file. Must end with '.ply'.

        Raises
        ------
        ValueError
            If the file extension is not '.ply' or the mesh has not been computed yet (i.e., `eval()` hasn't been called).
        """
        
        if not path.endswith('.ply'):
            raise ValueError(f"path should end with .ply. path is {path}")
        if self.mesh is None:
            raise ValueError(f"self.mesh is currently None. self.mesh is typically set by running self.eval().")
        
        o3d.io.write_triangle_mesh(path, self.mesh)




    def show(self):
        """
        Displays the reconstructed mesh in an Open3D visualization window.

        Raises
        ------
        ValueError
            If the mesh has not been computed yet (i.e., `eval()` hasn't been called).
        """
        
        if self.mesh is None:
            raise ValueError("Run eval() before calling show().")

        o3d.visualization.draw_geometries([self.mesh])