trellis/representations/mesh/mc2mesh.py

import torch
from easydict import EasyDict as edict
from typing import Tuple, Optional
from diso import DiffDMC
from .cube2mesh import MeshExtractResult
from .utils_cube import *
from ...modules.sparse import SparseTensor

class EnhancedMarchingCubes:
    def __init__(self, device="cuda"):
        self.device = device
        self.diffdmc = DiffDMC(dtype=torch.float32)

    def __call__(self,
                 voxelgrid_vertices: torch.Tensor,
                 scalar_field: torch.Tensor,
                 voxelgrid_colors: Optional[torch.Tensor] = None,
                 training: bool = False
                 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
        """
        Enhanced Marching Cubes implementation using DiffDMC that handles deformations and colors
        """
        if scalar_field.dim() == 1:
            grid_size = int(round(scalar_field.shape[0] ** (1 / 3)))
            scalar_field = scalar_field.reshape(grid_size, grid_size, grid_size)
        elif scalar_field.dim() > 3:
            scalar_field = scalar_field.squeeze()

        if scalar_field.dim() != 3:
            raise ValueError(f"Expected 3D array, got shape {scalar_field.shape}")

        # Normalize coordinates for DiffDMC
        scalar_field = scalar_field.to(self.device)
        
        # Get deformation field if provided
        deform_field = None
        if voxelgrid_vertices is not None:
            if voxelgrid_vertices.dim() == 2:
                grid_size = int(round(voxelgrid_vertices.shape[0] ** (1 / 3)))
                voxelgrid_vertices = voxelgrid_vertices.reshape(grid_size, grid_size, grid_size, 3)
            deform_field = voxelgrid_vertices.to(self.device)

        # Run DiffDMC
        vertices, faces = self.diffdmc(
            scalar_field,
            deform_field,
            isovalue=0.0
        )

        # Handle colors if provided
        colors = None
        if voxelgrid_colors is not None:
            voxelgrid_colors = torch.sigmoid(voxelgrid_colors)
            if voxelgrid_colors.dim() == 2:
                grid_size = int(round(voxelgrid_colors.shape[0] ** (1/3)))
                voxelgrid_colors = voxelgrid_colors.reshape(grid_size, grid_size, grid_size, -1)

            grid_positions = vertices.clone() * grid_size
            grid_coords = grid_positions.long()
            local_coords = grid_positions - grid_coords.float()
            
            # Clamp coordinates to grid bounds
            grid_coords = torch.clamp(grid_coords, 0, voxelgrid_colors.shape[0] - 1)
            
            # Trilinear interpolation for colors
            colors = self._interpolate_color(grid_coords, local_coords, voxelgrid_colors)
            
        vertices = vertices * 2 - 1 # Normalize vertices to [-1, 1]
        vertices /= 2.0  # Normalize vertices to [-0.5, 0.5]

        # Compute deviation loss for training
        deviation_loss = torch.tensor(0.0, device=self.device)
        if training and deform_field is not None:
            # Compute deviation between original and deformed vertices
            deviation_loss = self._compute_deviation_loss(vertices, deform_field)

        # faces = faces.flip(dims=[1])  # Maintain consistent face orientation

        return vertices, faces, deviation_loss, colors

    def _interpolate_color(self, grid_coords: torch.Tensor,
                          local_coords: torch.Tensor,
                          color_field: torch.Tensor) -> torch.Tensor:
        """
        Interpolate colors using trilinear interpolation
        Args:
            grid_coords: (N, 3) integer grid coordinates
            local_coords: (N, 3) fractional positions within grid cells
            color_field: (res, res, res, C) color values
        """
        x, y, z = local_coords[:, 0], local_coords[:, 1], local_coords[:, 2]
        
        # Get corner values for each vertex
        c000 = color_field[grid_coords[:, 0], grid_coords[:, 1], grid_coords[:, 2]]
        c001 = color_field[grid_coords[:, 0], grid_coords[:, 1],
               torch.clamp(grid_coords[:, 2] + 1, 0, color_field.shape[2] - 1)]
        c010 = color_field[grid_coords[:, 0],
               torch.clamp(grid_coords[:, 1] + 1, 0, color_field.shape[1] - 1),
               grid_coords[:, 2]]
        c011 = color_field[grid_coords[:, 0],
               torch.clamp(grid_coords[:, 1] + 1, 0, color_field.shape[1] - 1),
               torch.clamp(grid_coords[:, 2] + 1, 0, color_field.shape[2] - 1)]
        c100 = color_field[torch.clamp(grid_coords[:, 0] + 1, 0, color_field.shape[0] - 1),
               grid_coords[:, 1], grid_coords[:, 2]]
        c101 = color_field[torch.clamp(grid_coords[:, 0] + 1, 0, color_field.shape[0] - 1),
               grid_coords[:, 1],
               torch.clamp(grid_coords[:, 2] + 1, 0, color_field.shape[2] - 1)]
        c110 = color_field[torch.clamp(grid_coords[:, 0] + 1, 0, color_field.shape[0] - 1),
               torch.clamp(grid_coords[:, 1] + 1, 0, color_field.shape[1] - 1),
               grid_coords[:, 2]]
        c111 = color_field[torch.clamp(grid_coords[:, 0] + 1, 0, color_field.shape[0] - 1),
               torch.clamp(grid_coords[:, 1] + 1, 0, color_field.shape[1] - 1),
               torch.clamp(grid_coords[:, 2] + 1, 0, color_field.shape[2] - 1)]
        
        # Interpolate along x axis
        c00 = c000 * (1 - x)[:, None] + c100 * x[:, None]
        c01 = c001 * (1 - x)[:, None] + c101 * x[:, None]
        c10 = c010 * (1 - x)[:, None] + c110 * x[:, None]
        c11 = c011 * (1 - x)[:, None] + c111 * x[:, None]
        
        # Interpolate along y axis
        c0 = c00 * (1 - y)[:, None] + c10 * y[:, None]
        c1 = c01 * (1 - y)[:, None] + c11 * y[:, None]
        
        # Interpolate along z axis
        colors = c0 * (1 - z)[:, None] + c1 * z[:, None]
        
        return colors

    def _compute_deviation_loss(self, vertices: torch.Tensor,
                               deform_field: torch.Tensor) -> torch.Tensor:
        """Compute deviation loss for training"""
        # Since DiffDMC already handles the deformation, we compute the loss
        # based on the magnitude of the deformation field
        return torch.mean(deform_field ** 2)

class SparseFeatures2MCMesh:
    def __init__(self, device="cuda", res=128, use_color=True):
        super().__init__()
        self.device = device

        self.res = res

        self.mesh_extractor = EnhancedMarchingCubes(device=device)
        self.sdf_bias = -1.0 / res
        verts, cube = construct_dense_grid(self.res, self.device)
        self.reg_c = cube.to(self.device)
        self.reg_v = verts.to(self.device)
        self.use_color = use_color
        self._calc_layout()

    def _calc_layout(self):
        LAYOUTS = {
            'sdf': {'shape': (8, 1), 'size': 8},
            'deform': {'shape': (8, 3), 'size': 8 * 3},
            'weights': {'shape': (21,), 'size': 21}
        }
        if self.use_color:
            '''
            6 channel color including normal map
            '''
            LAYOUTS['color'] = {'shape': (8, 6,), 'size': 8 * 6}
        self.layouts = edict(LAYOUTS)
        start = 0
        for k, v in self.layouts.items():
            v['range'] = (start, start + v['size'])
            start += v['size']
        self.feats_channels = start

    def get_layout(self, feats: torch.Tensor, name: str):
        if name not in self.layouts:
            return None
        return feats[:, self.layouts[name]['range'][0]:self.layouts[name]['range'][1]].reshape(-1, *self.layouts[name][
            'shape'])

    def __call__(self, cubefeats: SparseTensor, training=False):
        coords = cubefeats.coords[:, 1:]
        feats = cubefeats.feats

        sdf, deform, color, weights = [self.get_layout(feats, name)
                                       for name in ['sdf', 'deform', 'color', 'weights']]
        sdf += self.sdf_bias
        v_attrs = [sdf, deform, color] if self.use_color else [sdf, deform]
        v_pos, v_attrs, reg_loss = sparse_cube2verts(coords, torch.cat(v_attrs, dim=-1),
                                                     training=training)

        v_attrs_d = get_dense_attrs(v_pos, v_attrs, res=self.res + 1, sdf_init=True)

        if self.use_color:
            sdf_d, deform_d, colors_d = (v_attrs_d[..., 0], v_attrs_d[..., 1:4],
                                         v_attrs_d[..., 4:])
        else:
            sdf_d, deform_d = v_attrs_d[..., 0], v_attrs_d[..., 1:4]
            colors_d = None

        x_nx3 = get_defomed_verts(self.reg_v, deform_d, self.res)

        vertices, faces, L_dev, colors = self.mesh_extractor(
            voxelgrid_vertices=x_nx3,
            scalar_field=sdf_d,
            voxelgrid_colors=colors_d,
            training=training
        )

        mesh = MeshExtractResult(vertices=vertices, faces=faces,
                                 vertex_attrs=colors, res=self.res)

        if training:
            if mesh.success:
                reg_loss += L_dev.mean() * 0.5
            reg_loss += (weights[:, :20]).abs().mean() * 0.2
            mesh.reg_loss = reg_loss
            mesh.tsdf_v = get_defomed_verts(v_pos, v_attrs[:, 1:4], self.res)
            mesh.tsdf_s = v_attrs[:, 0]

        return mesh