OpenMotionLab/MotionGPT

1"""Meshes, conforming to the glTF 2.0 standards as specified in2https://github.com/KhronosGroup/glTF/tree/master/specification/2.0#reference-mesh3 4Author: Matthew Matl5"""6import copy7 8import numpy as np9import trimesh10 11from .primitive import Primitive12from .constants import GLTF13from .material import MetallicRoughnessMaterial14 15 16class Mesh(object):17    """A set of primitives to be rendered.18 19    Parameters20    ----------21    name : str22        The user-defined name of this object.23    primitives : list of :class:`Primitive`24        The primitives associated with this mesh.25    weights : (k,) float26        Array of weights to be applied to the Morph Targets.27    is_visible : bool28        If False, the mesh will not be rendered.29    """30 31    def __init__(self, primitives, name=None, weights=None, is_visible=True):32        self.primitives = primitives33        self.name = name34        self.weights = weights35        self.is_visible = is_visible36 37        self._bounds = None38 39    @property40    def name(self):41        """str : The user-defined name of this object.42        """43        return self._name44 45    @name.setter46    def name(self, value):47        if value is not None:48            value = str(value)49        self._name = value50 51    @property52    def primitives(self):53        """list of :class:`Primitive` : The primitives associated54        with this mesh.55        """56        return self._primitives57 58    @primitives.setter59    def primitives(self, value):60        self._primitives = value61 62    @property63    def weights(self):64        """(k,) float : Weights to be applied to morph targets.65        """66        return self._weights67 68    @weights.setter69    def weights(self, value):70        self._weights = value71 72    @property73    def is_visible(self):74        """bool : Whether the mesh is visible.75        """76        return self._is_visible77 78    @is_visible.setter79    def is_visible(self, value):80        self._is_visible = value81 82    @property83    def bounds(self):84        """(2,3) float : The axis-aligned bounds of the mesh.85        """86        if self._bounds is None:87            bounds = np.array([[np.infty, np.infty, np.infty],88                               [-np.infty, -np.infty, -np.infty]])89            for p in self.primitives:90                bounds[0] = np.minimum(bounds[0], p.bounds[0])91                bounds[1] = np.maximum(bounds[1], p.bounds[1])92            self._bounds = bounds93        return self._bounds94 95    @property96    def centroid(self):97        """(3,) float : The centroid of the mesh's axis-aligned bounding box98        (AABB).99        """100        return np.mean(self.bounds, axis=0)101 102    @property103    def extents(self):104        """(3,) float : The lengths of the axes of the mesh's AABB.105        """106        return np.diff(self.bounds, axis=0).reshape(-1)107 108    @property109    def scale(self):110        """(3,) float : The length of the diagonal of the mesh's AABB.111        """112        return np.linalg.norm(self.extents)113 114    @property115    def is_transparent(self):116        """bool : If True, the mesh is partially-transparent.117        """118        for p in self.primitives:119            if p.is_transparent:120                return True121        return False122 123    @staticmethod124    def from_points(points, colors=None, normals=None,125                    is_visible=True, poses=None):126        """Create a Mesh from a set of points.127 128        Parameters129        ----------130        points : (n,3) float131            The point positions.132        colors : (n,3) or (n,4) float, optional133            RGB or RGBA colors for each point.134        normals : (n,3) float, optionals135            The normal vectors for each point.136        is_visible : bool137            If False, the points will not be rendered.138        poses : (x,4,4)139            Array of 4x4 transformation matrices for instancing this object.140 141        Returns142        -------143        mesh : :class:`Mesh`144            The created mesh.145        """146        primitive = Primitive(147            positions=points,148            normals=normals,149            color_0=colors,150            mode=GLTF.POINTS,151            poses=poses152        )153        mesh = Mesh(primitives=[primitive], is_visible=is_visible)154        return mesh155 156    @staticmethod157    def from_trimesh(mesh, material=None, is_visible=True,158                     poses=None, wireframe=False, smooth=True):159        """Create a Mesh from a :class:`~trimesh.base.Trimesh`.160 161        Parameters162        ----------163        mesh : :class:`~trimesh.base.Trimesh` or list of them164            A triangular mesh or a list of meshes.165        material : :class:`Material`166            The material of the object. Overrides any mesh material.167            If not specified and the mesh has no material, a default material168            will be used.169        is_visible : bool170            If False, the mesh will not be rendered.171        poses : (n,4,4) float172            Array of 4x4 transformation matrices for instancing this object.173        wireframe : bool174            If `True`, the mesh will be rendered as a wireframe object175        smooth : bool176            If `True`, the mesh will be rendered with interpolated vertex177            normals. Otherwise, the mesh edges will stay sharp.178 179        Returns180        -------181        mesh : :class:`Mesh`182            The created mesh.183        """184 185        if isinstance(mesh, (list, tuple, set, np.ndarray)):186            meshes = list(mesh)187        elif isinstance(mesh, trimesh.Trimesh):188            meshes = [mesh]189        else:190            raise TypeError('Expected a Trimesh or a list, got a {}'191                            .format(type(mesh)))192 193        primitives = []194        for m in meshes:195            positions = None196            normals = None197            indices = None198 199            # Compute positions, normals, and indices200            if smooth:201                positions = m.vertices.copy()202                normals = m.vertex_normals.copy()203                indices = m.faces.copy()204            else:205                positions = m.vertices[m.faces].reshape((3 * len(m.faces), 3))206                normals = np.repeat(m.face_normals, 3, axis=0)207 208            # Compute colors, texture coords, and material properties209            color_0, texcoord_0, primitive_material = Mesh._get_trimesh_props(m, smooth=smooth, material=material)210 211            # Override if material is given.212            if material is not None:213                #primitive_material = copy.copy(material)214                primitive_material = copy.deepcopy(material)  # TODO215 216            if primitive_material is None:217                # Replace material with default if needed218                primitive_material = MetallicRoughnessMaterial(219                    alphaMode='BLEND',220                    baseColorFactor=[0.3, 0.3, 0.3, 1.0],221                    metallicFactor=0.2,222                    roughnessFactor=0.8223                )224 225            primitive_material.wireframe = wireframe226 227            # Create the primitive228            primitives.append(Primitive(229                positions=positions,230                normals=normals,231                texcoord_0=texcoord_0,232                color_0=color_0,233                indices=indices,234                material=primitive_material,235                mode=GLTF.TRIANGLES,236                poses=poses237            ))238 239        return Mesh(primitives=primitives, is_visible=is_visible)240 241    @staticmethod242    def _get_trimesh_props(mesh, smooth=False, material=None):243        """Gets the vertex colors, texture coordinates, and material properties244        from a :class:`~trimesh.base.Trimesh`.245        """246        colors = None247        texcoords = None248 249        # If the trimesh visual is undefined, return none for both250        if not mesh.visual.defined:251            return colors, texcoords, material252 253        # Process vertex colors254        if material is None:255            if mesh.visual.kind == 'vertex':256                vc = mesh.visual.vertex_colors.copy()257                if smooth:258                    colors = vc259                else:260                    colors = vc[mesh.faces].reshape(261                        (3 * len(mesh.faces), vc.shape[1])262                    )263                material = MetallicRoughnessMaterial(264                    alphaMode='BLEND',265                    baseColorFactor=[1.0, 1.0, 1.0, 1.0],266                    metallicFactor=0.2,267                    roughnessFactor=0.8268                )269            # Process face colors270            elif mesh.visual.kind == 'face':271                if smooth:272                    raise ValueError('Cannot use face colors with a smooth mesh')273                else:274                    colors = np.repeat(mesh.visual.face_colors, 3, axis=0)275 276                material = MetallicRoughnessMaterial(277                    alphaMode='BLEND',278                    baseColorFactor=[1.0, 1.0, 1.0, 1.0],279                    metallicFactor=0.2,280                    roughnessFactor=0.8281                )282 283        # Process texture colors284        if mesh.visual.kind == 'texture':285            # Configure UV coordinates286            if mesh.visual.uv is not None and len(mesh.visual.uv) != 0:287                uv = mesh.visual.uv.copy()288                if smooth:289                    texcoords = uv290                else:291                    texcoords = uv[mesh.faces].reshape(292                        (3 * len(mesh.faces), uv.shape[1])293                    )294 295            if material is None:296                # Configure mesh material297                mat = mesh.visual.material298 299                if isinstance(mat, trimesh.visual.texture.PBRMaterial):300                    material = MetallicRoughnessMaterial(301                        normalTexture=mat.normalTexture,302                        occlusionTexture=mat.occlusionTexture,303                        emissiveTexture=mat.emissiveTexture,304                        emissiveFactor=mat.emissiveFactor,305                        alphaMode='BLEND',306                        baseColorFactor=mat.baseColorFactor,307                        baseColorTexture=mat.baseColorTexture,308                        metallicFactor=mat.metallicFactor,309                        roughnessFactor=mat.roughnessFactor,310                        metallicRoughnessTexture=mat.metallicRoughnessTexture,311                        doubleSided=mat.doubleSided,312                        alphaCutoff=mat.alphaCutoff313                    )314                elif isinstance(mat, trimesh.visual.texture.SimpleMaterial):315                    glossiness = mat.kwargs.get('Ns', 1.0)316                    if isinstance(glossiness, list):317                        glossiness = float(glossiness[0])318                    roughness = (2 / (glossiness + 2)) ** (1.0 / 4.0)319                    material = MetallicRoughnessMaterial(320                        alphaMode='BLEND',321                        roughnessFactor=roughness,322                        baseColorFactor=mat.diffuse,323                        baseColorTexture=mat.image,324                    )325                elif isinstance(mat, MetallicRoughnessMaterial):326                    material = mat327 328        return colors, texcoords, material329