core/models/rendering/mesh_utils.py

"""Mesh Define
"""

import os
import pdb
from typing import Optional, Union

import cv2
import numpy as np
import torch
import trimesh
from numpy import ndarray
from torch import Tensor


def length(x: Union[Tensor, ndarray], eps=1e-20) -> Union[Tensor, ndarray]:
    """length of an array (along the last dim).

    Args:
        x (Union[Tensor, ndarray]): x, [..., C]
        eps (float, optional): eps. Defaults to 1e-20.

    Returns:
        Union[Tensor, ndarray]: length, [..., 1]
    """
    if isinstance(x, np.ndarray):
        return np.sqrt(np.maximum(np.sum(x * x, axis=-1, keepdims=True), eps))
    else:
        return torch.sqrt(torch.clamp(dot(x, x), min=eps))


def safe_normalize(x: Union[Tensor, ndarray], eps=1e-20) -> Union[Tensor, ndarray]:
    """normalize an array (along the last dim).

    Args:
        x (Union[Tensor, ndarray]): x, [..., C]
        eps (float, optional): eps. Defaults to 1e-20.

    Returns:
        Union[Tensor, ndarray]: normalized x, [..., C]
    """

    return x / length(x, eps)


def dot(x: Union[Tensor, ndarray], y: Union[Tensor, ndarray]) -> Union[Tensor, ndarray]:
    """dot product (along the last dim).

    Args:
        x (Union[Tensor, ndarray]): x, [..., C]
        y (Union[Tensor, ndarray]): y, [..., C]

    Returns:
        Union[Tensor, ndarray]: x dot y, [..., 1]
    """
    if isinstance(x, np.ndarray):
        return np.sum(x * y, -1, keepdims=True)
    else:
        return torch.sum(x * y, -1, keepdim=True)


class Mesh:
    """
    A torch-native trimesh class, with support for ``ply/obj/glb`` formats.

    Note:
        This class only supports one mesh with a single texture image (an albedo texture and a metallic-roughness texture).
    """

    def __init__(
        self,
        v: Optional[Tensor] = None,
        f: Optional[Tensor] = None,
        vn: Optional[Tensor] = None,
        fn: Optional[Tensor] = None,
        vt: Optional[Tensor] = None,
        ft: Optional[Tensor] = None,
        vc: Optional[Tensor] = None,  # vertex color
        albedo: Optional[Tensor] = None,
        metallicRoughness: Optional[Tensor] = None,
        device: Optional[torch.device] = None,
    ):
        """Init a mesh directly using all attributes.

        Args:
            v (Optional[Tensor]): vertices, float [N, 3]. Defaults to None.
            f (Optional[Tensor]): faces, int [M, 3]. Defaults to None.
            vn (Optional[Tensor]): vertex normals, float [N, 3]. Defaults to None.
            fn (Optional[Tensor]): faces for normals, int [M, 3]. Defaults to None.
            vt (Optional[Tensor]): vertex uv coordinates, float [N, 2]. Defaults to None.
            ft (Optional[Tensor]): faces for uvs, int [M, 3]. Defaults to None.
            vc (Optional[Tensor]): vertex colors, float [N, 3]. Defaults to None.
            albedo (Optional[Tensor]): albedo texture, float [H, W, 3], RGB format. Defaults to None.
            metallicRoughness (Optional[Tensor]): metallic-roughness texture, float [H, W, 3], metallic(Blue) = metallicRoughness[..., 2], roughness(Green) = metallicRoughness[..., 1]. Defaults to None.
            device (Optional[torch.device]): torch device. Defaults to None.
        """
        self.device = device
        self.v = v
        self.vn = vn
        self.vt = vt
        self.f = f
        self.fn = fn
        self.ft = ft
        # will first see if there is vertex color to use
        self.vc = vc
        # only support a single albedo image
        self.albedo = albedo
        # pbr extension, metallic(Blue) = metallicRoughness[..., 2], roughness(Green) = metallicRoughness[..., 1]
        # ref: https://registry.khronos.org/glTF/specs/2.0/glTF-2.0.html
        self.metallicRoughness = metallicRoughness

        self.ori_center = 0
        self.ori_scale = 1

    @classmethod
    def load(
        cls,
        path,
        resize=True,
        clean=False,
        renormal=True,
        retex=False,
        bound=0.9,
        front_dir="+z",
        **kwargs,
    ):
        """load mesh from path.

        Args:
            path (str): path to mesh file, supports ply, obj, glb.
            clean (bool, optional): perform mesh cleaning at load (e.g., merge close vertices). Defaults to False.
            resize (bool, optional): auto resize the mesh using ``bound`` into [-bound, bound]^3. Defaults to True.
            renormal (bool, optional): re-calc the vertex normals. Defaults to True.
            retex (bool, optional): re-calc the uv coordinates, will overwrite the existing uv coordinates. Defaults to False.
            bound (float, optional): bound to resize. Defaults to 0.9.
            front_dir (str, optional): front-view direction of the mesh, should be [+-][xyz][ 123]. Defaults to '+z'.
            device (torch.device, optional): torch device. Defaults to None.

        Note:
            a ``device`` keyword argument can be provided to specify the torch device.
            If it's not provided, we will try to use ``'cuda'`` as the device if it's available.

        Returns:
            Mesh: the loaded Mesh object.
        """
        # obj supports face uv
        if path.endswith(".obj"):
            mesh = cls.load_obj(path, **kwargs)
        # trimesh only supports vertex uv, but can load more formats
        else:
            try:
                kwargs.pop("albedo_path")
            except:
                pass
            mesh = cls.load_trimesh(path, **kwargs)

        # clean
        if clean:
            from kiui.mesh_utils import clean_mesh

            vertices = mesh.v.detach().cpu().numpy()
            triangles = mesh.f.detach().cpu().numpy()
            vertices, triangles = clean_mesh(vertices, triangles, remesh=False)
            mesh.v = torch.from_numpy(vertices).contiguous().float().to(mesh.device)
            mesh.f = torch.from_numpy(triangles).contiguous().int().to(mesh.device)

        print(f"[INFO] load mesh, v: {mesh.v.shape}, f: {mesh.f.shape}")
        # auto-normalize
        if resize:
            mesh.auto_size(bound=bound)
        # auto-fix normal
        if renormal or mesh.vn is None:
            mesh.auto_normal()
            print(f"[INFO] load mesh, vn: {mesh.vn.shape}, fn: {mesh.fn.shape}")
        # auto-fix texcoords
        if retex or (mesh.albedo is not None and mesh.vt is None):
            mesh.auto_uv(cache_path=path)
            print(f"[INFO] load mesh, vt: {mesh.vt.shape}, ft: {mesh.ft.shape}")

        # rotate front dir to +z
        if front_dir != "+z":
            # axis switch
            if "-z" in front_dir:
                T = torch.tensor(
                    [[1, 0, 0], [0, 1, 0], [0, 0, -1]],
                    device=mesh.device,
                    dtype=torch.float32,
                )
            elif "+x" in front_dir:
                T = torch.tensor(
                    [[0, 0, 1], [0, 1, 0], [1, 0, 0]],
                    device=mesh.device,
                    dtype=torch.float32,
                )
            elif "-x" in front_dir:
                T = torch.tensor(
                    [[0, 0, -1], [0, 1, 0], [1, 0, 0]],
                    device=mesh.device,
                    dtype=torch.float32,
                )
            elif "+y" in front_dir:
                T = torch.tensor(
                    [[1, 0, 0], [0, 0, 1], [0, 1, 0]],
                    device=mesh.device,
                    dtype=torch.float32,
                )
            elif "-y" in front_dir:
                T = torch.tensor(
                    [[1, 0, 0], [0, 0, -1], [0, 1, 0]],
                    device=mesh.device,
                    dtype=torch.float32,
                )
            else:
                T = torch.tensor(
                    [[1, 0, 0], [0, 1, 0], [0, 0, 1]],
                    device=mesh.device,
                    dtype=torch.float32,
                )
            # rotation (how many 90 degrees)
            if "1" in front_dir:
                T @= torch.tensor(
                    [[0, -1, 0], [1, 0, 0], [0, 0, 1]],
                    device=mesh.device,
                    dtype=torch.float32,
                )
            elif "2" in front_dir:
                T @= torch.tensor(
                    [[1, 0, 0], [0, -1, 0], [0, 0, 1]],
                    device=mesh.device,
                    dtype=torch.float32,
                )
            elif "3" in front_dir:
                T @= torch.tensor(
                    [[0, 1, 0], [-1, 0, 0], [0, 0, 1]],
                    device=mesh.device,
                    dtype=torch.float32,
                )
            mesh.v @= T
            mesh.vn @= T

        return mesh

    @classmethod
    def load_processed(
        cls,
        path,
        resize=True,
        clean=False,
        renormal=True,
        retex=False,
        bound=0.9,
        front_dir="+z",
        scale=None,
        center=None,
        **kwargs,
    ):
        """load mesh from path.

        Args:
            path (str): path to mesh file, supports ply, obj, glb.
            clean (bool, optional): perform mesh cleaning at load (e.g., merge close vertices). Defaults to False.
            resize (bool, optional): auto resize the mesh using ``bound`` into [-bound, bound]^3. Defaults to True.
            renormal (bool, optional): re-calc the vertex normals. Defaults to True.
            retex (bool, optional): re-calc the uv coordinates, will overwrite the existing uv coordinates. Defaults to False.
            bound (float, optional): bound to resize. Defaults to 0.9.
            front_dir (str, optional): front-view direction of the mesh, should be [+-][xyz][ 123]. Defaults to '+z'.
            device (torch.device, optional): torch device. Defaults to None.

        Note:
            a ``device`` keyword argument can be provided to specify the torch device.
            If it's not provided, we will try to use ``'cuda'`` as the device if it's available.

        Returns:
            Mesh: the loaded Mesh object.
        """
        # obj supports face uv
        if path.endswith(".obj"):
            mesh = cls.load_obj(path, **kwargs)
        # trimesh only supports vertex uv, but can load more formats
        else:
            try:
                kwargs.pop("albedo_path")
            except:
                pass
            mesh = cls.load_trimesh(path, **kwargs)

        # clean
        if clean:
            from kiui.mesh_utils import clean_mesh

            vertices = mesh.v.detach().cpu().numpy()
            triangles = mesh.f.detach().cpu().numpy()
            vertices, triangles = clean_mesh(vertices, triangles, remesh=False)
            mesh.v = torch.from_numpy(vertices).contiguous().float().to(mesh.device)
            mesh.f = torch.from_numpy(triangles).contiguous().int().to(mesh.device)

        print(f"[INFO] load mesh, v: {mesh.v.shape}, f: {mesh.f.shape}")
        # auto-normalize
        if resize:
            mesh.ori_center = torch.tensor(center).to(mesh.v)
            mesh.ori_scale = scale
            mesh.v = (mesh.v - mesh.ori_center) * mesh.ori_scale
        # auto-fix normal
        if renormal or mesh.vn is None:
            mesh.auto_normal()
            print(f"[INFO] load mesh, vn: {mesh.vn.shape}, fn: {mesh.fn.shape}")
        # auto-fix texcoords
        if retex or (mesh.albedo is not None and mesh.vt is None):
            mesh.auto_uv(cache_path=path)
            print(f"[INFO] load mesh, vt: {mesh.vt.shape}, ft: {mesh.ft.shape}")

        # rotate front dir to +z
        if front_dir != "+z":
            # axis switch
            if "-z" in front_dir:
                T = torch.tensor(
                    [[1, 0, 0], [0, 1, 0], [0, 0, -1]],
                    device=mesh.device,
                    dtype=torch.float32,
                )
            elif "+x" in front_dir:
                T = torch.tensor(
                    [[0, 0, 1], [0, 1, 0], [1, 0, 0]],
                    device=mesh.device,
                    dtype=torch.float32,
                )
            elif "-x" in front_dir:
                T = torch.tensor(
                    [[0, 0, -1], [0, 1, 0], [1, 0, 0]],
                    device=mesh.device,
                    dtype=torch.float32,
                )
            elif "+y" in front_dir:
                T = torch.tensor(
                    [[1, 0, 0], [0, 0, 1], [0, 1, 0]],
                    device=mesh.device,
                    dtype=torch.float32,
                )
            elif "-y" in front_dir:
                T = torch.tensor(
                    [[1, 0, 0], [0, 0, -1], [0, 1, 0]],
                    device=mesh.device,
                    dtype=torch.float32,
                )
            else:
                T = torch.tensor(
                    [[1, 0, 0], [0, 1, 0], [0, 0, 1]],
                    device=mesh.device,
                    dtype=torch.float32,
                )
            # rotation (how many 90 degrees)
            if "1" in front_dir:
                T @= torch.tensor(
                    [[0, -1, 0], [1, 0, 0], [0, 0, 1]],
                    device=mesh.device,
                    dtype=torch.float32,
                )
            elif "2" in front_dir:
                T @= torch.tensor(
                    [[1, 0, 0], [0, -1, 0], [0, 0, 1]],
                    device=mesh.device,
                    dtype=torch.float32,
                )
            elif "3" in front_dir:
                T @= torch.tensor(
                    [[0, 1, 0], [-1, 0, 0], [0, 0, 1]],
                    device=mesh.device,
                    dtype=torch.float32,
                )
            mesh.v @= T
            mesh.vn @= T

        return mesh

    # load from obj file
    @classmethod
    def load_obj(cls, path, albedo_path=None, device=None):
        """load an ``obj`` mesh.

        Args:
            path (str): path to mesh.
            albedo_path (str, optional): path to the albedo texture image, will overwrite the existing texture path if specified in mtl. Defaults to None.
            device (torch.device, optional): torch device. Defaults to None.

        Note:
            We will try to read `mtl` path from `obj`, else we assume the file name is the same as `obj` but with `mtl` extension.
            The `usemtl` statement is ignored, and we only use the last material path in `mtl` file.

        Returns:
            Mesh: the loaded Mesh object.
        """
        assert os.path.splitext(path)[-1] == ".obj"

        mesh = cls()

        # device
        if device is None:
            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        mesh.device = device

        # load obj
        with open(path, "r") as f:
            lines = f.readlines()

        def parse_f_v(fv):
            # pass in a vertex term of a face, return {v, vt, vn} (-1 if not provided)
            # supported forms:
            # f v1 v2 v3
            # f v1/vt1 v2/vt2 v3/vt3
            # f v1/vt1/vn1 v2/vt2/vn2 v3/vt3/vn3
            # f v1//vn1 v2//vn2 v3//vn3
            xs = [int(x) - 1 if x != "" else -1 for x in fv.split("/")]
            xs.extend([-1] * (3 - len(xs)))
            return xs[0], xs[1], xs[2]

        vertices, texcoords, normals = [], [], []
        faces, tfaces, nfaces = [], [], []
        mtl_path = None

        for line in lines:
            split_line = line.split()
            # empty line
            if len(split_line) == 0:
                continue
            prefix = split_line[0].lower()
            # mtllib
            if prefix == "mtllib":
                mtl_path = split_line[1]
            # usemtl
            elif prefix == "usemtl":
                pass  # ignored
            # v/vn/vt
            elif prefix == "v":
                vertices.append([float(v) for v in split_line[1:]])
            elif prefix == "vn":
                normals.append([float(v) for v in split_line[1:]])
            elif prefix == "vt":
                val = [float(v) for v in split_line[1:]]
                texcoords.append([val[0], 1.0 - val[1]])
            elif prefix == "f":
                vs = split_line[1:]
                nv = len(vs)
                v0, t0, n0 = parse_f_v(vs[0])
                for i in range(nv - 2):  # triangulate (assume vertices are ordered)
                    v1, t1, n1 = parse_f_v(vs[i + 1])
                    v2, t2, n2 = parse_f_v(vs[i + 2])
                    faces.append([v0, v1, v2])
                    tfaces.append([t0, t1, t2])
                    nfaces.append([n0, n1, n2])

        mesh.v = torch.tensor(vertices, dtype=torch.float32, device=device)
        mesh.vt = (
            torch.tensor(texcoords, dtype=torch.float32, device=device)
            if len(texcoords) > 0
            else None
        )
        mesh.vn = (
            torch.tensor(normals, dtype=torch.float32, device=device)
            if len(normals) > 0
            else None
        )

        mesh.f = torch.tensor(faces, dtype=torch.int32, device=device)
        mesh.ft = (
            torch.tensor(tfaces, dtype=torch.int32, device=device)
            if len(texcoords) > 0
            else None
        )
        mesh.fn = (
            torch.tensor(nfaces, dtype=torch.int32, device=device)
            if len(normals) > 0
            else None
        )

        # see if there is vertex color
        use_vertex_color = False
        if mesh.v.shape[1] == 6:
            use_vertex_color = True
            mesh.vc = mesh.v[:, 3:]
            mesh.v = mesh.v[:, :3]
            print(f"[INFO] load obj mesh: use vertex color: {mesh.vc.shape}")

        return mesh

    @classmethod
    def load_trimesh(cls, path, device=None):
        """load a mesh using ``trimesh.load()``.

        Can load various formats like ``glb`` and serves as a fallback.

        Note:
            We will try to merge all meshes if the glb contains more than one,
            but **this may cause the texture to lose**, since we only support one texture image!

        Args:
            path (str): path to the mesh file.
            device (torch.device, optional): torch device. Defaults to None.

        Returns:
            Mesh: the loaded Mesh object.
        """
        mesh = cls()

        # device
        if device is None:
            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        mesh.device = device

        # use trimesh to load ply/glb
        _data = trimesh.load(path)
        # always convert scene to mesh, and apply all transforms...
        if isinstance(_data, trimesh.Scene):
            print(f"[INFO] load trimesh: concatenating {len(_data.geometry)} meshes.")
            _concat = []
            # loop the scene graph and apply transform to each mesh
            scene_graph = (
                _data.graph.to_flattened()
            )  # dict {name: {transform: 4x4 mat, geometry: str}}
            for k, v in scene_graph.items():
                name = v["geometry"]
                if name in _data.geometry and isinstance(
                    _data.geometry[name], trimesh.Trimesh
                ):
                    transform = v["transform"]
                    _concat.append(_data.geometry[name].apply_transform(transform))
            _mesh = trimesh.util.concatenate(_concat)
        else:
            _mesh = _data

        if _mesh.visual.kind == "vertex":
            vertex_colors = _mesh.visual.vertex_colors
            vertex_colors = np.array(vertex_colors[..., :3]).astype(np.float32) / 255
            mesh.vc = torch.tensor(vertex_colors, dtype=torch.float32, device=device)
            print(f"[INFO] load trimesh: use vertex color: {mesh.vc.shape}")
        elif _mesh.visual.kind == "texture":
            _material = _mesh.visual.material
            if isinstance(_material, trimesh.visual.material.PBRMaterial):
                texture = np.array(_material.baseColorTexture).astype(np.float32) / 255
                # load metallicRoughness if present
                if _material.metallicRoughnessTexture is not None:
                    metallicRoughness = (
                        np.array(_material.metallicRoughnessTexture).astype(np.float32)
                        / 255
                    )
                    mesh.metallicRoughness = torch.tensor(
                        metallicRoughness, dtype=torch.float32, device=device
                    ).contiguous()
            elif isinstance(_material, trimesh.visual.material.SimpleMaterial):
                texture = (
                    np.array(_material.to_pbr().baseColorTexture).astype(np.float32)
                    / 255
                )
            else:
                raise NotImplementedError(
                    f"material type {type(_material)} not supported!"
                )
            mesh.albedo = torch.tensor(
                texture[..., :3], dtype=torch.float32, device=device
            ).contiguous()
            print(f"[INFO] load trimesh: load texture: {texture.shape}")
        else:
            mesh.albedo = None
            print(f"[INFO] load trimesh: failed to load texture.")

        vertices = _mesh.vertices

        try:
            texcoords = _mesh.visual.uv
            texcoords[:, 1] = 1 - texcoords[:, 1]
        except Exception as e:
            texcoords = None

        try:
            normals = _mesh.vertex_normals
        except Exception as e:
            normals = None

        # trimesh only support vertex uv...
        faces = tfaces = nfaces = _mesh.faces

        mesh.v = torch.tensor(vertices, dtype=torch.float32, device=device)
        mesh.vt = (
            torch.tensor(texcoords, dtype=torch.float32, device=device)
            if texcoords is not None
            else None
        )
        mesh.vn = (
            torch.tensor(normals, dtype=torch.float32, device=device)
            if normals is not None
            else None
        )

        mesh.f = torch.tensor(faces, dtype=torch.int32, device=device)
        mesh.ft = (
            torch.tensor(tfaces, dtype=torch.int32, device=device)
            if texcoords is not None
            else None
        )
        mesh.fn = (
            torch.tensor(nfaces, dtype=torch.int32, device=device)
            if normals is not None
            else None
        )

        return mesh

    # sample surface (using trimesh)
    def sample_surface(self, count: int):
        """sample points on the surface of the mesh.

        Args:
            count (int): number of points to sample.

        Returns:
            torch.Tensor: the sampled points, float [count, 3].
        """
        _mesh = trimesh.Trimesh(
            vertices=self.v.detach().cpu().numpy(), faces=self.f.detach().cpu().numpy()
        )
        points, face_idx = trimesh.sample.sample_surface(_mesh, count)
        points = torch.from_numpy(points).float().to(self.device)
        return points

    # aabb
    def aabb(self):
        """get the axis-aligned bounding box of the mesh.

        Returns:
            Tuple[torch.Tensor]: the min xyz and max xyz of the mesh.
        """
        return torch.min(self.v, dim=0).values, torch.max(self.v, dim=0).values

    # unit size
    @torch.no_grad()
    def auto_size(self, bound=0.9):
        """auto resize the mesh.

        Args:
            bound (float, optional): resizing into ``[-bound, bound]^3``. Defaults to 0.9.
        """
        vmin, vmax = self.aabb()
        self.ori_center = (vmax + vmin) / 2
        self.ori_scale = 2 * bound / torch.max(vmax - vmin).item()
        self.v = (self.v - self.ori_center) * self.ori_scale

    def auto_normal(self):
        """auto calculate the vertex normals."""
        i0, i1, i2 = self.f[:, 0].long(), self.f[:, 1].long(), self.f[:, 2].long()
        v0, v1, v2 = self.v[i0, :], self.v[i1, :], self.v[i2, :]

        face_normals = torch.cross(v1 - v0, v2 - v0)

        # Splat face normals to vertices
        vn = torch.zeros_like(self.v)
        vn.scatter_add_(0, i0[:, None].repeat(1, 3), face_normals)
        vn.scatter_add_(0, i1[:, None].repeat(1, 3), face_normals)
        vn.scatter_add_(0, i2[:, None].repeat(1, 3), face_normals)

        # Normalize, replace zero (degenerated) normals with some default value
        vn = torch.where(
            dot(vn, vn) > 1e-20,
            vn,
            torch.tensor([0.0, 0.0, 1.0], dtype=torch.float32, device=vn.device),
        )
        vn = safe_normalize(vn)

        self.vn = vn

        # Normalize, replace zero (degenerated) normals with some default value
        face_normals = torch.where(
            dot(face_normals, face_normals) > 1e-20,
            face_normals,
            torch.tensor([0.0, 0.0, 1.0], dtype=torch.float32, device=vn.device),
        )
        face_normals = safe_normalize(face_normals)
        self.fn = face_normals

    def auto_uv(self, cache_path=None, vmap=True):
        """auto calculate the uv coordinates.

        Args:
            cache_path (str, optional): path to save/load the uv cache as a npz file, this can avoid calculating uv every time when loading the same mesh, which is time-consuming. Defaults to None.
            vmap (bool, optional): remap vertices based on uv coordinates, so each v correspond to a unique vt (necessary for formats like gltf).
                Usually this will duplicate the vertices on the edge of uv atlas. Defaults to True.
        """
        # try to load cache
        if cache_path is not None:
            cache_path = os.path.splitext(cache_path)[0] + "_uv.npz"
        if cache_path is not None and os.path.exists(cache_path):
            data = np.load(cache_path)
            vt_np, ft_np, vmapping = data["vt"], data["ft"], data["vmapping"]
        else:
            import xatlas

            v_np = self.v.detach().cpu().numpy()
            f_np = self.f.detach().int().cpu().numpy()
            atlas = xatlas.Atlas()
            atlas.add_mesh(v_np, f_np)
            chart_options = xatlas.ChartOptions()
            # chart_options.max_iterations = 4
            atlas.generate(chart_options=chart_options)
            vmapping, ft_np, vt_np = atlas[0]  # [N], [M, 3], [N, 2]

            # save to cache
            if cache_path is not None:
                np.savez(cache_path, vt=vt_np, ft=ft_np, vmapping=vmapping)

        vt = torch.from_numpy(vt_np.astype(np.float32)).to(self.device)
        ft = torch.from_numpy(ft_np.astype(np.int32)).to(self.device)
        self.vt = vt
        self.ft = ft

        if vmap:
            vmapping = (
                torch.from_numpy(vmapping.astype(np.int64)).long().to(self.device)
            )
            self.align_v_to_vt(vmapping)

    def align_v_to_vt(self, vmapping=None):
        """remap v/f and vn/fn to vt/ft.

        Args:
            vmapping (np.ndarray, optional): the mapping relationship from f to ft. Defaults to None.
        """
        if vmapping is None:
            ft = self.ft.view(-1).long()
            f = self.f.view(-1).long()
            vmapping = torch.zeros(
                self.vt.shape[0], dtype=torch.long, device=self.device
            )
            vmapping[ft] = f  # scatter, randomly choose one if index is not unique

        self.v = self.v[vmapping]
        self.f = self.ft

        if self.vn is not None:
            self.vn = self.vn[vmapping]
            self.fn = self.ft

    def to(self, device):
        """move all tensor attributes to device.

        Args:
            device (torch.device): target device.

        Returns:
            Mesh: self.
        """
        self.device = device
        for name in [
            "v",
            "f",
            "vn",
            "fn",
            "vt",
            "ft",
            "albedo",
            "vc",
            "metallicRoughness",
        ]:
            tensor = getattr(self, name)
            if tensor is not None:
                setattr(self, name, tensor.to(device))
        return self

    def write(self, path):
        """write the mesh to a path.

        Args:
            path (str): path to write, supports ply, obj and glb.
        """
        if path.endswith(".ply"):
            self.write_ply(path)
        elif path.endswith(".obj"):
            self.write_obj(path)
        elif path.endswith(".glb") or path.endswith(".gltf"):
            self.write_glb(path)
        else:
            raise NotImplementedError(f"format {path} not supported!")

    def write_ply(self, path):
        """write the mesh in ply format. Only for geometry!

        Args:
            path (str): path to write.
        """

        if self.albedo is not None:
            print(f"[WARN] ply format does not support exporting texture, will ignore!")

        v_np = self.v.detach().cpu().numpy()
        f_np = self.f.detach().cpu().numpy()

        _mesh = trimesh.Trimesh(vertices=v_np, faces=f_np)
        _mesh.export(path)

    def write_glb(self, path):
        """write the mesh in glb/gltf format.
          This will create a scene with a single mesh.

        Args:
            path (str): path to write.
        """

        # assert self.v.shape[0] == self.vn.shape[0] and self.v.shape[0] == self.vt.shape[0]
        if self.vt is not None and self.v.shape[0] != self.vt.shape[0]:
            self.align_v_to_vt()

        import pygltflib

        f_np = self.f.detach().cpu().numpy().astype(np.uint32)
        f_np_blob = f_np.flatten().tobytes()

        v_np = self.v.detach().cpu().numpy().astype(np.float32)
        v_np_blob = v_np.tobytes()

        blob = f_np_blob + v_np_blob
        byteOffset = len(blob)

        # base mesh
        gltf = pygltflib.GLTF2(
            scene=0,
            scenes=[pygltflib.Scene(nodes=[0])],
            nodes=[pygltflib.Node(mesh=0)],
            meshes=[
                pygltflib.Mesh(
                    primitives=[
                        pygltflib.Primitive(
                            # indices to accessors (0 is triangles)
                            attributes=pygltflib.Attributes(
                                POSITION=1,
                            ),
                            indices=0,
                        )
                    ]
                )
            ],
            buffers=[pygltflib.Buffer(byteLength=len(f_np_blob) + len(v_np_blob))],
            # buffer view (based on dtype)
            bufferViews=[
                # triangles; as flatten (element) array
                pygltflib.BufferView(
                    buffer=0,
                    byteLength=len(f_np_blob),
                    target=pygltflib.ELEMENT_ARRAY_BUFFER,  # GL_ELEMENT_ARRAY_BUFFER (34963)
                ),
                # positions; as vec3 array
                pygltflib.BufferView(
                    buffer=0,
                    byteOffset=len(f_np_blob),
                    byteLength=len(v_np_blob),
                    byteStride=12,  # vec3
                    target=pygltflib.ARRAY_BUFFER,  # GL_ARRAY_BUFFER (34962)
                ),
            ],
            accessors=[
                # 0 = triangles
                pygltflib.Accessor(
                    bufferView=0,
                    componentType=pygltflib.UNSIGNED_INT,  # GL_UNSIGNED_INT (5125)
                    count=f_np.size,
                    type=pygltflib.SCALAR,
                    max=[int(f_np.max())],
                    min=[int(f_np.min())],
                ),
                # 1 = positions
                pygltflib.Accessor(
                    bufferView=1,
                    componentType=pygltflib.FLOAT,  # GL_FLOAT (5126)
                    count=len(v_np),
                    type=pygltflib.VEC3,
                    max=v_np.max(axis=0).tolist(),
                    min=v_np.min(axis=0).tolist(),
                ),
            ],
        )

        # append texture info
        if self.vt is not None:

            vt_np = self.vt.detach().cpu().numpy().astype(np.float32)
            vt_np_blob = vt_np.tobytes()

            albedo = self.albedo.detach().cpu().numpy()
            albedo = (albedo * 255).astype(np.uint8)
            albedo = cv2.cvtColor(albedo, cv2.COLOR_RGB2BGR)
            albedo_blob = cv2.imencode(".png", albedo)[1].tobytes()

            # update primitive
            gltf.meshes[0].primitives[0].attributes.TEXCOORD_0 = 2
            gltf.meshes[0].primitives[0].material = 0

            # update materials
            gltf.materials.append(
                pygltflib.Material(
                    pbrMetallicRoughness=pygltflib.PbrMetallicRoughness(
                        baseColorTexture=pygltflib.TextureInfo(index=0, texCoord=0),
                        metallicFactor=0.0,
                        roughnessFactor=1.0,
                    ),
                    alphaMode=pygltflib.OPAQUE,
                    alphaCutoff=None,
                    doubleSided=True,
                )
            )

            gltf.textures.append(pygltflib.Texture(sampler=0, source=0))
            gltf.samplers.append(
                pygltflib.Sampler(
                    magFilter=pygltflib.LINEAR,
                    minFilter=pygltflib.LINEAR_MIPMAP_LINEAR,
                    wrapS=pygltflib.REPEAT,
                    wrapT=pygltflib.REPEAT,
                )
            )
            gltf.images.append(pygltflib.Image(bufferView=3, mimeType="image/png"))

            # update buffers
            gltf.bufferViews.append(
                # index = 2, texcoords; as vec2 array
                pygltflib.BufferView(
                    buffer=0,
                    byteOffset=byteOffset,
                    byteLength=len(vt_np_blob),
                    byteStride=8,  # vec2
                    target=pygltflib.ARRAY_BUFFER,
                )
            )

            gltf.accessors.append(
                # 2 = texcoords
                pygltflib.Accessor(
                    bufferView=2,
                    componentType=pygltflib.FLOAT,
                    count=len(vt_np),
                    type=pygltflib.VEC2,
                    max=vt_np.max(axis=0).tolist(),
                    min=vt_np.min(axis=0).tolist(),
                )
            )

            blob += vt_np_blob
            byteOffset += len(vt_np_blob)

            gltf.bufferViews.append(
                # index = 3, albedo texture; as none target
                pygltflib.BufferView(
                    buffer=0,
                    byteOffset=byteOffset,
                    byteLength=len(albedo_blob),
                )
            )

            blob += albedo_blob
            byteOffset += len(albedo_blob)

            gltf.buffers[0].byteLength = byteOffset

            # append metllic roughness
            if self.metallicRoughness is not None:
                metallicRoughness = self.metallicRoughness.detach().cpu().numpy()
                metallicRoughness = (metallicRoughness * 255).astype(np.uint8)
                metallicRoughness = cv2.cvtColor(metallicRoughness, cv2.COLOR_RGB2BGR)
                metallicRoughness_blob = cv2.imencode(".png", metallicRoughness)[
                    1
                ].tobytes()

                # update texture definition
                gltf.materials[0].pbrMetallicRoughness.metallicFactor = 1.0
                gltf.materials[0].pbrMetallicRoughness.roughnessFactor = 1.0
                gltf.materials[0].pbrMetallicRoughness.metallicRoughnessTexture = (
                    pygltflib.TextureInfo(index=1, texCoord=0)
                )

                gltf.textures.append(pygltflib.Texture(sampler=1, source=1))
                gltf.samplers.append(
                    pygltflib.Sampler(
                        magFilter=pygltflib.LINEAR,
                        minFilter=pygltflib.LINEAR_MIPMAP_LINEAR,
                        wrapS=pygltflib.REPEAT,
                        wrapT=pygltflib.REPEAT,
                    )
                )
                gltf.images.append(pygltflib.Image(bufferView=4, mimeType="image/png"))

                # update buffers
                gltf.bufferViews.append(
                    # index = 4, metallicRoughness texture; as none target
                    pygltflib.BufferView(
                        buffer=0,
                        byteOffset=byteOffset,
                        byteLength=len(metallicRoughness_blob),
                    )
                )

                blob += metallicRoughness_blob
                byteOffset += len(metallicRoughness_blob)

                gltf.buffers[0].byteLength = byteOffset

        # set actual data
        gltf.set_binary_blob(blob)

        # glb = b"".join(gltf.save_to_bytes())
        gltf.save(path)

    def write_obj(self, path):
        """write the mesh in obj format. Will also write the texture and mtl files.

        Args:
            path (str): path to write.
        """

        v_np = self.v.detach().cpu().numpy()
        vt_np = self.vt.detach().cpu().numpy() if self.vt is not None else None
        vn_np = self.vn.detach().cpu().numpy() if self.vn is not None else None
        f_np = self.f.detach().cpu().numpy()
        ft_np = self.ft.detach().cpu().numpy() if self.ft is not None else None
        fn_np = self.fn.detach().cpu().numpy() if self.fn is not None else None

        with open(path, "w") as fp:

            for v in v_np:
                fp.write(f"v {v[0]} {v[1]} {v[2]} \n")

            for i in range(len(f_np)):
                fp.write(
                    f'f {f_np[i, 0] + 1}/{ft_np[i, 0] + 1 if ft_np is not None else ""}/{fn_np[i, 0] + 1 if fn_np is not None else ""} \
                             {f_np[i, 1] + 1}/{ft_np[i, 1] + 1 if ft_np is not None else ""}/{fn_np[i, 1] + 1 if fn_np is not None else ""} \
                             {f_np[i, 2] + 1}/{ft_np[i, 2] + 1 if ft_np is not None else ""}/{fn_np[i, 2] + 1 if fn_np is not None else ""} \n'
                )


if __name__ == "__main__":
    import os
    import pdb
    import sys

    import tqlt

    sys.path.append("./")
    obj_mesh = Mesh().load(path="./test.glb")

    pdb.set_trace()