HorizonRobotics/EmbodiedGen-Image-to-3D

1# Project EmbodiedGen2#3# Copyright (c) 2025 Horizon Robotics. All Rights Reserved.4#5# Licensed under the Apache License, Version 2.0 (the "License");6# you may not use this file except in compliance with the License.7# You may obtain a copy of the License at8#9#       http://www.apache.org/licenses/LICENSE-2.010#11# Unless required by applicable law or agreed to in writing, software12# distributed under the License is distributed on an "AS IS" BASIS,13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or14# implied. See the License for the specific language governing15# permissions and limitations under the License.16 17 18import logging19from typing import Tuple, Union20 21import igraph22import numpy as np23import pyvista as pv24import spaces25import torch26import utils3d27from pymeshfix import _meshfix28from tqdm import tqdm29 30logging.basicConfig(31    format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO32)33logger = logging.getLogger(__name__)34 35 36__all__ = [37    "MeshFixer",38]39 40 41def _radical_inverse(base, n):42    val = 043    inv_base = 1.0 / base44    inv_base_n = inv_base45    while n > 0:46        digit = n % base47        val += digit * inv_base_n48        n //= base49        inv_base_n *= inv_base50    return val51 52 53def _halton_sequence(dim, n):54    PRIMES = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53]55    return [_radical_inverse(PRIMES[dim], n) for dim in range(dim)]56 57 58def _hammersley_sequence(dim, n, num_samples):59    return [n / num_samples] + _halton_sequence(dim - 1, n)60 61 62def _sphere_hammersley_seq(n, num_samples, offset=(0, 0), remap=False):63    """Generate a point on a unit sphere using the Hammersley sequence.64 65    Args:66        n (int): The index of the sample.67        num_samples (int): The total number of samples.68        offset (tuple, optional): Offset for the u and v coordinates.69        remap (bool, optional): Whether to remap the u coordinate.70 71    Returns:72        list: A list containing the spherical coordinates [phi, theta].73    """74    u, v = _hammersley_sequence(2, n, num_samples)75    u += offset[0] / num_samples76    v += offset[1]77 78    if remap:79        u = 2 * u if u < 0.25 else 2 / 3 * u + 1 / 380 81    theta = np.arccos(1 - 2 * u) - np.pi / 282    phi = v * 2 * np.pi83    return [phi, theta]84 85 86class MeshFixer(object):87    """MeshFixer simplifies and repairs 3D triangle meshes by TSDF.88 89    Attributes:90        vertices (torch.Tensor): A tensor of shape (V, 3) representing vertex positions.91        faces (torch.Tensor): A tensor of shape (F, 3) representing face indices.92        device (str): Device to run computations on, typically "cuda" or "cpu".93 94    Main logic reference: https://github.com/microsoft/TRELLIS/blob/main/trellis/utils/postprocessing_utils.py#L2295    """96 97    def __init__(98        self,99        vertices: Union[torch.Tensor, np.ndarray],100        faces: Union[torch.Tensor, np.ndarray],101        device: str = "cuda",102    ) -> None:103        self.device = device104        if isinstance(vertices, np.ndarray):105            vertices = torch.tensor(vertices)106        self.vertices = vertices107 108        if isinstance(faces, np.ndarray):109            faces = torch.tensor(faces)110        self.faces = faces111 112    @staticmethod113    def log_mesh_changes(method):114        def wrapper(self, *args, **kwargs):115            logger.info(116                f"Before {method.__name__}: {self.vertices.shape[0]} vertices, {self.faces.shape[0]} faces"  # noqa117            )118            result = method(self, *args, **kwargs)119            logger.info(120                f"After {method.__name__}: {self.vertices.shape[0]} vertices, {self.faces.shape[0]} faces"  # noqa121            )122            return result123 124        return wrapper125 126    @log_mesh_changes127    def fill_holes(128        self,129        max_hole_size: float,130        max_hole_nbe: int,131        resolution: int,132        num_views: int,133        norm_mesh_ratio: float = 1.0,134    ) -> None:135        self.vertices = self.vertices * norm_mesh_ratio136        vertices, self.faces = self._fill_holes(137            self.vertices,138            self.faces,139            max_hole_size,140            max_hole_nbe,141            resolution,142            num_views,143        )144        self.vertices = vertices / norm_mesh_ratio145 146    @staticmethod147    @torch.no_grad()148    def _fill_holes(149        vertices: torch.Tensor,150        faces: torch.Tensor,151        max_hole_size: float,152        max_hole_nbe: int,153        resolution: int,154        num_views: int,155    ) -> Union[torch.Tensor, torch.Tensor]:156        yaws, pitchs = [], []157        for i in range(num_views):158            y, p = _sphere_hammersley_seq(i, num_views)159            yaws.append(y)160            pitchs.append(p)161 162        yaws, pitchs = (163            torch.tensor(yaws).to(vertices),164            torch.tensor(pitchs).to(vertices),165        )166        radius, fov = 2.0, torch.deg2rad(torch.tensor(40)).to(vertices)167        projection = utils3d.torch.perspective_from_fov_xy(fov, fov, 1, 3)168 169        views = []170        for yaw, pitch in zip(yaws, pitchs):171            orig = (172                torch.tensor(173                    [174                        torch.sin(yaw) * torch.cos(pitch),175                        torch.cos(yaw) * torch.cos(pitch),176                        torch.sin(pitch),177                    ]178                ).to(vertices)179                * radius180            )181            view = utils3d.torch.view_look_at(182                orig,183                torch.tensor([0, 0, 0]).to(vertices),184                torch.tensor([0, 0, 1]).to(vertices),185            )186            views.append(view)187        views = torch.stack(views, dim=0)188 189        # Rasterize the mesh190        visibility = torch.zeros(191            faces.shape[0], dtype=torch.int32, device=faces.device192        )193        rastctx = utils3d.torch.RastContext(backend="cuda")194 195        for i in tqdm(196            range(views.shape[0]), total=views.shape[0], desc="Rasterizing"197        ):198            view = views[i]199            buffers = utils3d.torch.rasterize_triangle_faces(200                rastctx,201                vertices[None],202                faces,203                resolution,204                resolution,205                view=view,206                projection=projection,207            )208            face_id = buffers["face_id"][0][buffers["mask"][0] > 0.95] - 1209            face_id = torch.unique(face_id).long()210            visibility[face_id] += 1211 212        # Normalize visibility by the number of views213        visibility = visibility.float() / num_views214 215        # Mincut: Identify outer and inner faces216        edges, face2edge, edge_degrees = utils3d.torch.compute_edges(faces)217        boundary_edge_indices = torch.nonzero(edge_degrees == 1).reshape(-1)218        connected_components = utils3d.torch.compute_connected_components(219            faces, edges, face2edge220        )221 222        outer_face_indices = torch.zeros(223            faces.shape[0], dtype=torch.bool, device=faces.device224        )225        for i in range(len(connected_components)):226            outer_face_indices[connected_components[i]] = visibility[227                connected_components[i]228            ] > min(229                max(230                    visibility[connected_components[i]].quantile(0.75).item(),231                    0.25,232                ),233                0.5,234            )235 236        outer_face_indices = outer_face_indices.nonzero().reshape(-1)237        inner_face_indices = torch.nonzero(visibility == 0).reshape(-1)238 239        if inner_face_indices.shape[0] == 0:240            return vertices, faces241 242        # Construct dual graph (faces as nodes, edges as edges)243        dual_edges, dual_edge2edge = utils3d.torch.compute_dual_graph(244            face2edge245        )246        dual_edge2edge = edges[dual_edge2edge]247        dual_edges_weights = torch.norm(248            vertices[dual_edge2edge[:, 0]] - vertices[dual_edge2edge[:, 1]],249            dim=1,250        )251 252        # Mincut: Construct main graph and solve the mincut problem253        g = igraph.Graph()254        g.add_vertices(faces.shape[0])255        g.add_edges(dual_edges.cpu().numpy())256        g.es["weight"] = dual_edges_weights.cpu().numpy()257 258        g.add_vertex("s")  # source259        g.add_vertex("t")  # target260 261        g.add_edges(262            [(f, "s") for f in inner_face_indices],263            attributes={264                "weight": torch.ones(265                    inner_face_indices.shape[0], dtype=torch.float32266                )267                .cpu()268                .numpy()269            },270        )271        g.add_edges(272            [(f, "t") for f in outer_face_indices],273            attributes={274                "weight": torch.ones(275                    outer_face_indices.shape[0], dtype=torch.float32276                )277                .cpu()278                .numpy()279            },280        )281 282        cut = g.mincut("s", "t", (np.array(g.es["weight"]) * 1000).tolist())283        remove_face_indices = torch.tensor(284            [v for v in cut.partition[0] if v < faces.shape[0]],285            dtype=torch.long,286            device=faces.device,287        )288 289        # Check if the cut is valid with each connected component290        to_remove_cc = utils3d.torch.compute_connected_components(291            faces[remove_face_indices]292        )293        valid_remove_cc = []294        cutting_edges = []295        for cc in to_remove_cc:296            # Check visibility median for connected component297            visibility_median = visibility[remove_face_indices[cc]].median()298            if visibility_median > 0.25:299                continue300 301            # Check if the cutting loop is small enough302            cc_edge_indices, cc_edges_degree = torch.unique(303                face2edge[remove_face_indices[cc]], return_counts=True304            )305            cc_boundary_edge_indices = cc_edge_indices[cc_edges_degree == 1]306            cc_new_boundary_edge_indices = cc_boundary_edge_indices[307                ~torch.isin(cc_boundary_edge_indices, boundary_edge_indices)308            ]309            if len(cc_new_boundary_edge_indices) > 0:310                cc_new_boundary_edge_cc = (311                    utils3d.torch.compute_edge_connected_components(312                        edges[cc_new_boundary_edge_indices]313                    )314                )315                cc_new_boundary_edges_cc_center = [316                    vertices[edges[cc_new_boundary_edge_indices[edge_cc]]]317                    .mean(dim=1)318                    .mean(dim=0)319                    for edge_cc in cc_new_boundary_edge_cc320                ]321                cc_new_boundary_edges_cc_area = []322                for i, edge_cc in enumerate(cc_new_boundary_edge_cc):323                    _e1 = (324                        vertices[325                            edges[cc_new_boundary_edge_indices[edge_cc]][:, 0]326                        ]327                        - cc_new_boundary_edges_cc_center[i]328                    )329                    _e2 = (330                        vertices[331                            edges[cc_new_boundary_edge_indices[edge_cc]][:, 1]332                        ]333                        - cc_new_boundary_edges_cc_center[i]334                    )335                    cc_new_boundary_edges_cc_area.append(336                        torch.norm(torch.cross(_e1, _e2, dim=-1), dim=1).sum()337                        * 0.5338                    )339                cutting_edges.append(cc_new_boundary_edge_indices)340                if any(341                    [342                        _l > max_hole_size343                        for _l in cc_new_boundary_edges_cc_area344                    ]345                ):346                    continue347 348            valid_remove_cc.append(cc)349 350        if len(valid_remove_cc) > 0:351            remove_face_indices = remove_face_indices[352                torch.cat(valid_remove_cc)353            ]354            mask = torch.ones(355                faces.shape[0], dtype=torch.bool, device=faces.device356            )357            mask[remove_face_indices] = 0358            faces = faces[mask]359            faces, vertices = utils3d.torch.remove_unreferenced_vertices(360                faces, vertices361            )362 363            tqdm.write(f"Removed {(~mask).sum()} faces by mincut")364        else:365            tqdm.write("Removed 0 faces by mincut")366 367        # Fill small boundaries (holes)368        mesh = _meshfix.PyTMesh()369        mesh.load_array(vertices.cpu().numpy(), faces.cpu().numpy())370        mesh.fill_small_boundaries(nbe=max_hole_nbe, refine=True)371 372        _vertices, _faces = mesh.return_arrays()373        vertices = torch.tensor(_vertices).to(vertices)374        faces = torch.tensor(_faces).to(faces)375 376        return vertices, faces377 378    @property379    def vertices_np(self) -> np.ndarray:380        return self.vertices.cpu().numpy()381 382    @property383    def faces_np(self) -> np.ndarray:384        return self.faces.cpu().numpy()385 386    @log_mesh_changes387    def simplify(self, ratio: float) -> None:388        """Simplify the mesh using quadric edge collapse decimation.389 390        Args:391            ratio (float): Ratio of faces to filter out.392        """393        if ratio <= 0 or ratio >= 1:394            raise ValueError("Simplify ratio must be between 0 and 1.")395 396        # Convert to PyVista format for simplification397        mesh = pv.PolyData(398            self.vertices_np,399            np.hstack([np.full((self.faces.shape[0], 1), 3), self.faces_np]),400        )401        mesh.clean(inplace=True)402        mesh.clear_data()403        mesh = mesh.triangulate()404        mesh = mesh.decimate(ratio, progress_bar=True)405 406        # Update vertices and faces407        self.vertices = torch.tensor(408            mesh.points, device=self.device, dtype=torch.float32409        )410        self.faces = torch.tensor(411            mesh.faces.reshape(-1, 4)[:, 1:],412            device=self.device,413            dtype=torch.int32,414        )415 416    @spaces.GPU417    def __call__(418        self,419        filter_ratio: float,420        max_hole_size: float,421        resolution: int,422        num_views: int,423        norm_mesh_ratio: float = 1.0,424    ) -> Tuple[np.ndarray, np.ndarray]:425        """Post-process the mesh by simplifying and filling holes.426 427        This method performs a two-step process:428        1. Simplifies mesh by reducing faces using quadric edge decimation.429        2. Fills holes by removing invisible faces, repairing small boundaries.430 431        Args:432            filter_ratio (float): Ratio of faces to simplify out.433                Must be in the range (0, 1).434            max_hole_size (float): Maximum area of a hole to fill. Connected435                components of holes larger than this size will not be repaired.436            resolution (int): Resolution of the rasterization buffer.437            num_views (int): Number of viewpoints to sample for rasterization.438            norm_mesh_ratio (float, optional): A scaling factor applied to the439                vertices of the mesh during processing.440 441        Returns:442            Tuple[np.ndarray, np.ndarray]:443                - vertices: Simplified and repaired vertex array of (V, 3).444                - faces: Simplified and repaired face array of (F, 3).445        """446        self.vertices = self.vertices.to(self.device)447        self.faces = self.faces.to(self.device)448 449        self.simplify(ratio=filter_ratio)450        self.fill_holes(451            max_hole_size=max_hole_size,452            max_hole_nbe=int(250 * np.sqrt(1 - filter_ratio)),453            resolution=resolution,454            num_views=num_views,455            norm_mesh_ratio=norm_mesh_ratio,456        )457 458        return self.vertices_np, self.faces_np459