Blame for diffusionGS/utils/structure.py

from typing import Optional
import torch
import torch.nn.functional as F
import numpy as np
import trimesh
from trimesh.visual import create_visual

def dot(x, y, dim=-1):
    return torch.sum(x * y, dim=dim, keepdim=True)

class Mesh:
    def __init__(self, v_pos:torch.Tensor, t_pos_idx:torch.Tensor, **kwargs):
        self.v_pos = v_pos
        self.t_pos_idx = t_pos_idx
        self._v_nrm = None
        self._v_tng = None
        self._v_tex = None
        self._t_tex_idx = None
        self._v_tex_map = None
        self._edges = None

        assert v_pos.device == t_pos_idx.device, \
            f'v_pos is on {v_pos.device} but t_pos_idx is on {t_pos_idx.device}'
        self.device = v_pos.device
    
    def to(self, device):
        for key in [
            'v_pos',
            't_pos_idx',
            '_v_nrm',
            '_v_tng',
            '_v_tex',
            '_t_tex_idx',
            '_v_tex_map',
            '_edges',
        ]:
            value = getattr(self, key)
            if isinstance(value, torch.Tensor) and value.device != torch.device(device):
                setattr(self, key, value.to(device))
        self.device = torch.device(device)
        return self
    
    @classmethod
    def from_trimesh(cls, mesh: trimesh.Trimesh):
        v_pos = torch.as_tensor(mesh.vertices, dtype=torch.float32)
        t_pos_idx = torch.as_tensor(mesh.faces, dtype=torch.float32)
        return Mesh(v_pos=v_pos, t_pos_idx=t_pos_idx)
    
    def to_trimesh(self) -> trimesh.Trimesh:
        mesh = trimesh.Trimesh(
            vertices=self.v_pos.detach().cpu().numpy(), 
            faces=self.t_pos_idx.cpu().numpy(), 
            process=True, validate=True,
        )
        return mesh

    def remove_outlier(self, outlier_n_faces_threshold):
        if self.requires_grad:
            print("Mesh is differentiable, not removing outliers")
            return self

        # use trimesh to first split the mesh into connected components
        # then remove the components with less than n_face_threshold faces
        import trimesh

        # construct a trimesh object
        mesh = trimesh.Trimesh(
            vertices=self.v_pos.detach().cpu().numpy(),
            faces=self.t_pos_idx.detach().cpu().numpy(),
        )

        # split the mesh into connected components
        components = mesh.split(only_watertight=False)
        # log the number of faces in each component
        print(
            "Mesh has {} components, with faces: {}".format(
                len(components), [c.faces.shape[0] for c in components]
            )
        )

        n_faces_threshold: int
        if isinstance(outlier_n_faces_threshold, float):
            # set the threshold to the number of faces in the largest component multiplied by outlier_n_faces_threshold
            n_faces_threshold = int(
                max([c.faces.shape[0] for c in components]) * outlier_n_faces_threshold
            )
        else:
            # set the threshold directly to outlier_n_faces_threshold
            n_faces_threshold = outlier_n_faces_threshold

        # log the threshold
        print(
            "Removing components with less than {} faces".format(n_faces_threshold)
        )

        # remove the components with less than n_face_threshold faces
        components = [c for c in components if c.faces.shape[0] >= n_faces_threshold]

        # log the number of faces in each component after removing outliers
        print(
            "Mesh has {} components after removing outliers, with faces: {}".format(
                len(components), [c.faces.shape[0] for c in components]
            )
        )
        # merge the components
        mesh = trimesh.util.concatenate(components)

        # convert back to our mesh format
        v_pos = torch.from_numpy(mesh.vertices).to(self.v_pos)
        t_pos_idx = torch.from_numpy(mesh.faces).to(self.t_pos_idx)

        clean_mesh = Mesh(v_pos, t_pos_idx)
        # keep the extras unchanged

        if len(self.extras) > 0:
            clean_mesh.extras = self.extras
            print(
                f"The following extra attributes are inherited from the original mesh unchanged: {list(self.extras.keys())}"
            )
        return clean_mesh

    def blender_to_threejs(self):
        v_pos = self.v_pos[:, [0, 2, 1]]
        v_pos[:, 2] = -v_pos[:, 2]
        self.v_pos = v_pos

    @property
    def v_nrm(self):
        if self._v_nrm is None:
            self._v_nrm = self._compute_vertex_normal()
        return self._v_nrm

    @property
    def v_tng(self):
        if self._v_tng is None:
            self._v_tng = self._compute_vertex_tangent()
        return self._v_tng

    # @property
    # def v_tex(self):
    #     if self._v_tex is None:
    #         self._v_tex_map, self._t_tex_idx, self._v_tex = self._unwrap_uv()
    #     return self._v_tex

    # @property
    # def t_tex_idx(self):
    #     if self._t_tex_idx is None:
    #         self._v_tex_map, self._t_tex_idx, self._v_tex = self._unwrap_uv()
    #     return self._t_tex_idx
    
    # @property
    # def v_tex_map(self):
    #     if self._v_tex_map is None:
    #         self._v_tex_map, self._t_tex_idx, self._v_tex = self._unwrap_uv()
    #     return self._v_tex_map

    @property
    def edges(self):
        if self._edges is None:
            self._edges = self._compute_edges()
        return self._edges

    def _compute_vertex_normal(self):
        i0 = self.t_pos_idx[:, 0].long()
        i1 = self.t_pos_idx[:, 1].long()
        i2 = self.t_pos_idx[:, 2].long()

        v0 = self.v_pos[i0, :]
        v1 = self.v_pos[i1, :]
        v2 = self.v_pos[i2, :]

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

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

        # Normalize, replace zero (degenerated) normals with some default value
        v_nrm = torch.where(
            dot(v_nrm, v_nrm) > 1e-20, v_nrm, torch.as_tensor([0.0, 0.0, 1.0]).to(v_nrm)
        )
        v_nrm = F.normalize(v_nrm, dim=1)

        if torch.is_anomaly_enabled():
            assert torch.all(torch.isfinite(v_nrm))

        return v_nrm

    def _compute_vertex_tangent(self):
        vn_idx = [None] * 3
        pos = [None] * 3
        tex = [None] * 3
        for i in range(0, 3):
            pos[i] = self.v_pos[self.t_pos_idx[:, i]]
            tex[i] = self.v_tex[self.t_tex_idx[:, i]]
            # t_nrm_idx is always the same as t_pos_idx
            vn_idx[i] = self.t_pos_idx[:, i]

        tangents = torch.zeros_like(self.v_nrm)
        tansum = torch.zeros_like(self.v_nrm)

        # Compute tangent space for each triangle
        uve1 = tex[1] - tex[0]
        uve2 = tex[2] - tex[0]
        pe1 = pos[1] - pos[0]
        pe2 = pos[2] - pos[0]

        nom = pe1 * uve2[..., 1:2] - pe2 * uve1[..., 1:2]
        denom = uve1[..., 0:1] * uve2[..., 1:2] - uve1[..., 1:2] * uve2[..., 0:1]

        # Avoid division by zero for degenerated texture coordinates
        tang = nom / torch.where(
            denom > 0.0, torch.clamp(denom, min=1e-6), torch.clamp(denom, max=-1e-6)
        )

        # Update all 3 vertices
        for i in range(0, 3):
            idx = vn_idx[i][:, None].repeat(1, 3)
            tangents.scatter_add_(0, idx, tang)  # tangents[n_i] = tangents[n_i] + tang
            tansum.scatter_add_(
                0, idx, torch.ones_like(tang)
            )  # tansum[n_i] = tansum[n_i] + 1
        tangents = tangents / tansum

        # Normalize and make sure tangent is perpendicular to normal
        tangents = F.normalize(tangents, dim=1)
        tangents = F.normalize(tangents - dot(tangents, self.v_nrm) * self.v_nrm)

        if torch.is_anomaly_enabled():
            assert torch.all(torch.isfinite(tangents))

        return tangents

    def _unwrap_uv(
        self, xatlas_chart_options: dict = {}, xatlas_pack_options: dict = {}
    ):
        xatlas_chart_options = {
            'max_chart_area': 0.0,
            'max_boundary_length': 0.0,
            'normal_deviation_weight': 2.0,
            'roundness_weight': 0.01,
            'straightness_weight': 6.0,
            'normal_seam_weight': 4.0,
            'texture_seam_weight': 0.5,
            'max_cost': 16.0,  # avoid small charts
            'max_iterations': 1,
            'use_input_mesh_uvs': False,
            'fix_winding': False,
        }
        xatlas_pack_options = {
            'max_chart_size': 0,
            'padding': 4,  # avoid adjoint
            'texels_per_unit': 0.0,
            'resolution': 2048,
            'bilinear': True,
            'blockAlign': False,
            'bruteForce': False,
            'create_image': False,
            'rotate_charts_to_axis': True,
            'rotate_charts': True,
        }
        print("Using xatlas to perform UV unwrapping, may take a while ...")

        import xatlas

        atlas = xatlas.Atlas()
        atlas.add_mesh(
            self.v_pos.detach().cpu().numpy(),
            self.t_pos_idx.cpu().numpy(),
        )
        co = xatlas.ChartOptions()
        po = xatlas.PackOptions()
        for k, v in xatlas_chart_options.items():
            setattr(co, k, v)
        for k, v in xatlas_pack_options.items():
            setattr(po, k, v)
        
        atlas.generate(co, po, verbose=True)
        vmapping, indices, uvs = atlas.get_mesh(0)
        vmapping = (
            torch.from_numpy(
                vmapping.astype(np.uint64, casting="same_kind").view(np.int64)
            )
            .to(self.v_pos.device)
            .long()
        )
        uvs = torch.from_numpy(uvs).to(self.v_pos.device).float()
        indices = (
            torch.from_numpy(
                indices.astype(np.uint64, casting="same_kind").view(np.int64)
            )
            .to(self.v_pos.device)
            .long()
        )
        return vmapping, indices, uvs

    def unwrap_uv(
        self, xatlas_chart_options: dict = {}, xatlas_pack_options: dict = {}
    ):
        self._v_tex_map, self._t_tex_idx, self._v_tex = self._unwrap_uv(
            xatlas_chart_options, xatlas_pack_options
        )

    def _compute_edges(self):
        # Compute edges
        edges = torch.cat(
            [
                self.t_pos_idx[:, [0, 1]],
                self.t_pos_idx[:, [1, 2]],
                self.t_pos_idx[:, [2, 0]],
            ],
            dim=0,
        )
        edges = edges.sort()[0]
        edges = torch.unique(edges, dim=0)
        return edges

    def normal_consistency(self):
        edge_nrm = self.v_nrm[self.edges]
        nc = (
            1.0 - torch.cosine_similarity(edge_nrm[:, 0], edge_nrm[:, 1], dim=-1)
        ).mean()
        return nc

    def _laplacian_uniform(self):
        # from stable-dreamfusion
        # https://github.com/ashawkey/stable-dreamfusion/blob/8fb3613e9e4cd1ded1066b46e80ca801dfb9fd06/nerf/renderer.py#L224
        verts, faces = self.v_pos, self.t_pos_idx

        V = verts.shape[0]
        F = faces.shape[0]

        # Neighbor indices
        ii = faces[:, [1, 2, 0]].flatten()
        jj = faces[:, [2, 0, 1]].flatten()
        adj = torch.stack([torch.cat([ii, jj]), torch.cat([jj, ii])], dim=0).unique(
            dim=1
        )
        adj_values = torch.ones(adj.shape[1]).to(verts)

        # Diagonal indices
        diag_idx = adj[0]

        # Build the sparse matrix
        idx = torch.cat((adj, torch.stack((diag_idx, diag_idx), dim=0)), dim=1)
        values = torch.cat((-adj_values, adj_values))

        # The coalesce operation sums the duplicate indices, resulting in the
        # correct diagonal
        return torch.sparse_coo_tensor(idx, values, (V, V)).coalesce()

    def laplacian(self):
        with torch.no_grad():
            L = self._laplacian_uniform()
        loss = L.mm(self.v_pos)
        loss = loss.norm(dim=1)
        loss = loss.mean()
        return loss


class Texture:
    def __init__(
        self, 
        mesh:Mesh, 
        v_rgb:Optional[torch.Tensor]=None, 
        map_Kd:Optional[torch.Tensor]=None,
        map_Ks:Optional[torch.Tensor]=None,
        **kwargs,
    ) -> None:
        self.mesh: Mesh = mesh
        self.v_rgb = v_rgb
        self.map_Kd = map_Kd
        self.map_Ks = map_Ks

        assert v_rgb is None or mesh.device == v_rgb.device, \
            f'mesh is on {mesh.device} but v_rgb is on {v_rgb.device}'
        assert map_Kd is None or mesh.device == map_Kd.device, \
            f'mesh is on {mesh.device} but map_Kd is on {map_Kd.device}'
        assert map_Ks is None or mesh.device == map_Ks.device, \
            f'mesh is on {mesh.device} but map_Ks is on {map_Ks.device}'
        self.device = mesh.device
    
    def to(self, device):
        for key in [
            'mesh',
        ]:
            value = getattr(self, key)
            if value is not None and value.device != torch.device(device):
                setattr(self, key, value.to(device))
        for key in [
            'v_rgb',
            'map_Kd',
            'map_Ks',
        ]:
            value = getattr(self, key)
            if isinstance(value, torch.Tensor) and value.device != torch.device(device):
                setattr(self, key, value.to(device))
        self.device = torch.device(device)
        return self

    def to_trimesh(self) -> trimesh.Trimesh:
        mesh = self.mesh.to_trimesh()
        if self.v_rgb is not None:
            mesh.visual = create_visual(
                face_colors = None,
                vertex_colors = self.v_rgb.detach().cpu().numpy(),
                mesh = mesh,
            )
        return mesh
    
    @classmethod
    def from_trimesh(cls, mesh: trimesh.Trimesh):
        if mesh.visual is not None:
            v_rgb = torch.as_tensor(mesh.visual.vertex_colors, dtype=torch.float32)
        else:
            v_rgb = None
        mesh = Mesh.from_trimesh(mesh)
        return Texture(mesh=mesh, v_rgb=v_rgb, map_Kd=None, map_Ks=None)