Upload 11 files

trellis/representations/__init__.py

@@ -0,0 +1,4 @@
+from .radiance_field import Strivec
+from .octree import DfsOctree as Octree
+from .gaussian import Gaussian
+from .mesh import MeshExtractResult

trellis/representations/gaussian/__init__.py

@@ -0,0 +1 @@
+from .gaussian_model import Gaussian

trellis/representations/gaussian/gaussian_model.py

@@ -0,0 +1,209 @@
+import torch
+import numpy as np
+from plyfile import PlyData, PlyElement
+from .general_utils import inverse_sigmoid, strip_symmetric, build_scaling_rotation
+import utils3d
+
+
+class Gaussian:
+    def __init__(
+            self, 
+            aabb : list,
+            sh_degree : int = 0,
+            mininum_kernel_size : float = 0.0,
+            scaling_bias : float = 0.01,
+            opacity_bias : float = 0.1,
+            scaling_activation : str = "exp",
+            device='cuda'
+        ):
+        self.init_params = {
+            'aabb': aabb,
+            'sh_degree': sh_degree,
+            'mininum_kernel_size': mininum_kernel_size,
+            'scaling_bias': scaling_bias,
+            'opacity_bias': opacity_bias,
+            'scaling_activation': scaling_activation,
+        }
+        
+        self.sh_degree = sh_degree
+        self.active_sh_degree = sh_degree
+        self.mininum_kernel_size = mininum_kernel_size 
+        self.scaling_bias = scaling_bias
+        self.opacity_bias = opacity_bias
+        self.scaling_activation_type = scaling_activation
+        self.device = device
+        self.aabb = torch.tensor(aabb, dtype=torch.float32, device=device)
+        self.setup_functions()
+
+        self._xyz = None
+        self._features_dc = None
+        self._features_rest = None
+        self._scaling = None
+        self._rotation = None
+        self._opacity = None
+
+    def setup_functions(self):
+        def build_covariance_from_scaling_rotation(scaling, scaling_modifier, rotation):
+            L = build_scaling_rotation(scaling_modifier * scaling, rotation)
+            actual_covariance = L @ L.transpose(1, 2)
+            symm = strip_symmetric(actual_covariance)
+            return symm
+        
+        if self.scaling_activation_type == "exp":
+            self.scaling_activation = torch.exp
+            self.inverse_scaling_activation = torch.log
+        elif self.scaling_activation_type == "softplus":
+            self.scaling_activation = torch.nn.functional.softplus
+            self.inverse_scaling_activation = lambda x: x + torch.log(-torch.expm1(-x))
+
+        self.covariance_activation = build_covariance_from_scaling_rotation
+
+        self.opacity_activation = torch.sigmoid
+        self.inverse_opacity_activation = inverse_sigmoid
+
+        self.rotation_activation = torch.nn.functional.normalize
+        
+        self.scale_bias = self.inverse_scaling_activation(torch.tensor(self.scaling_bias)).cuda()
+        self.rots_bias = torch.zeros((4)).cuda()
+        self.rots_bias[0] = 1
+        self.opacity_bias = self.inverse_opacity_activation(torch.tensor(self.opacity_bias)).cuda()
+
+    @property
+    def get_scaling(self):
+        scales = self.scaling_activation(self._scaling + self.scale_bias)
+        scales = torch.square(scales) + self.mininum_kernel_size ** 2
+        scales = torch.sqrt(scales)
+        return scales
+    
+    @property
+    def get_rotation(self):
+        return self.rotation_activation(self._rotation + self.rots_bias[None, :])
+    
+    @property
+    def get_xyz(self):
+        return self._xyz * self.aabb[None, 3:] + self.aabb[None, :3]
+    
+    @property
+    def get_features(self):
+        return torch.cat((self._features_dc, self._features_rest), dim=2) if self._features_rest is not None else self._features_dc
+    
+    @property
+    def get_opacity(self):
+        return self.opacity_activation(self._opacity + self.opacity_bias)
+    
+    def get_covariance(self, scaling_modifier = 1):
+        return self.covariance_activation(self.get_scaling, scaling_modifier, self._rotation + self.rots_bias[None, :])
+    
+    def from_scaling(self, scales):
+        scales = torch.sqrt(torch.square(scales) - self.mininum_kernel_size ** 2)
+        self._scaling = self.inverse_scaling_activation(scales) - self.scale_bias
+        
+    def from_rotation(self, rots):
+        self._rotation = rots - self.rots_bias[None, :]
+    
+    def from_xyz(self, xyz):
+        self._xyz = (xyz - self.aabb[None, :3]) / self.aabb[None, 3:]
+        
+    def from_features(self, features):
+        self._features_dc = features
+        
+    def from_opacity(self, opacities):
+        self._opacity = self.inverse_opacity_activation(opacities) - self.opacity_bias
+
+    def construct_list_of_attributes(self):
+        l = ['x', 'y', 'z', 'nx', 'ny', 'nz']
+        # All channels except the 3 DC
+        for i in range(self._features_dc.shape[1]*self._features_dc.shape[2]):
+            l.append('f_dc_{}'.format(i))
+        l.append('opacity')
+        for i in range(self._scaling.shape[1]):
+            l.append('scale_{}'.format(i))
+        for i in range(self._rotation.shape[1]):
+            l.append('rot_{}'.format(i))
+        return l
+        
+    def save_ply(self, path, transform=[[1, 0, 0], [0, 0, -1], [0, 1, 0]]):
+        xyz = self.get_xyz.detach().cpu().numpy()
+        normals = np.zeros_like(xyz)
+        f_dc = self._features_dc.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
+        opacities = inverse_sigmoid(self.get_opacity).detach().cpu().numpy()
+        scale = torch.log(self.get_scaling).detach().cpu().numpy()
+        rotation = (self._rotation + self.rots_bias[None, :]).detach().cpu().numpy()
+        
+        if transform is not None:
+            transform = np.array(transform)
+            xyz = np.matmul(xyz, transform.T)
+            rotation = utils3d.numpy.quaternion_to_matrix(rotation)
+            rotation = np.matmul(transform, rotation)
+            rotation = utils3d.numpy.matrix_to_quaternion(rotation)
+
+        dtype_full = [(attribute, 'f4') for attribute in self.construct_list_of_attributes()]
+
+        elements = np.empty(xyz.shape[0], dtype=dtype_full)
+        attributes = np.concatenate((xyz, normals, f_dc, opacities, scale, rotation), axis=1)
+        elements[:] = list(map(tuple, attributes))
+        el = PlyElement.describe(elements, 'vertex')
+        PlyData([el]).write(path)
+
+    def load_ply(self, path, transform=[[1, 0, 0], [0, 0, -1], [0, 1, 0]]):
+        plydata = PlyData.read(path)
+
+        xyz = np.stack((np.asarray(plydata.elements[0]["x"]),
+                        np.asarray(plydata.elements[0]["y"]),
+                        np.asarray(plydata.elements[0]["z"])),  axis=1)
+        opacities = np.asarray(plydata.elements[0]["opacity"])[..., np.newaxis]
+
+        features_dc = np.zeros((xyz.shape[0], 3, 1))
+        features_dc[:, 0, 0] = np.asarray(plydata.elements[0]["f_dc_0"])
+        features_dc[:, 1, 0] = np.asarray(plydata.elements[0]["f_dc_1"])
+        features_dc[:, 2, 0] = np.asarray(plydata.elements[0]["f_dc_2"])
+
+        if self.sh_degree > 0:
+            extra_f_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("f_rest_")]
+            extra_f_names = sorted(extra_f_names, key = lambda x: int(x.split('_')[-1]))
+            assert len(extra_f_names)==3*(self.sh_degree + 1) ** 2 - 3
+            features_extra = np.zeros((xyz.shape[0], len(extra_f_names)))
+            for idx, attr_name in enumerate(extra_f_names):
+                features_extra[:, idx] = np.asarray(plydata.elements[0][attr_name])
+            # Reshape (P,F*SH_coeffs) to (P, F, SH_coeffs except DC)
+            features_extra = features_extra.reshape((features_extra.shape[0], 3, (self.max_sh_degree + 1) ** 2 - 1))
+
+        scale_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("scale_")]
+        scale_names = sorted(scale_names, key = lambda x: int(x.split('_')[-1]))
+        scales = np.zeros((xyz.shape[0], len(scale_names)))
+        for idx, attr_name in enumerate(scale_names):
+            scales[:, idx] = np.asarray(plydata.elements[0][attr_name])
+
+        rot_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("rot")]
+        rot_names = sorted(rot_names, key = lambda x: int(x.split('_')[-1]))
+        rots = np.zeros((xyz.shape[0], len(rot_names)))
+        for idx, attr_name in enumerate(rot_names):
+            rots[:, idx] = np.asarray(plydata.elements[0][attr_name])
+            
+        if transform is not None:
+            transform = np.array(transform)
+            xyz = np.matmul(xyz, transform)
+            rotation = utils3d.numpy.quaternion_to_matrix(rotation)
+            rotation = np.matmul(rotation, transform)
+            rotation = utils3d.numpy.matrix_to_quaternion(rotation)
+            
+        # convert to actual gaussian attributes
+        xyz = torch.tensor(xyz, dtype=torch.float, device=self.device)
+        features_dc = torch.tensor(features_dc, dtype=torch.float, device=self.device).transpose(1, 2).contiguous()
+        if self.sh_degree > 0:
+            features_extra = torch.tensor(features_extra, dtype=torch.float, device=self.device).transpose(1, 2).contiguous()
+        opacities = torch.sigmoid(torch.tensor(opacities, dtype=torch.float, device=self.device))
+        scales = torch.exp(torch.tensor(scales, dtype=torch.float, device=self.device))
+        rots = torch.tensor(rots, dtype=torch.float, device=self.device)
+        
+        # convert to _hidden attributes
+        self._xyz = (xyz - self.aabb[None, :3]) / self.aabb[None, 3:]
+        self._features_dc = features_dc
+        if self.sh_degree > 0:
+            self._features_rest = features_extra
+        else:
+            self._features_rest = None
+        self._opacity = self.inverse_opacity_activation(opacities) - self.opacity_bias
+        self._scaling = self.inverse_scaling_activation(torch.sqrt(torch.square(scales) - self.mininum_kernel_size ** 2)) - self.scale_bias
+        self._rotation = rots - self.rots_bias[None, :]
+        

trellis/representations/gaussian/general_utils.py

@@ -0,0 +1,133 @@
+#
+# Copyright (C) 2023, Inria
+# GRAPHDECO research group, https://team.inria.fr/graphdeco
+# All rights reserved.
+#
+# This software is free for non-commercial, research and evaluation use 
+# under the terms of the LICENSE.md file.
+#
+# For inquiries contact  george.drettakis@inria.fr
+#
+
+import torch
+import sys
+from datetime import datetime
+import numpy as np
+import random
+
+def inverse_sigmoid(x):
+    return torch.log(x/(1-x))
+
+def PILtoTorch(pil_image, resolution):
+    resized_image_PIL = pil_image.resize(resolution)
+    resized_image = torch.from_numpy(np.array(resized_image_PIL)) / 255.0
+    if len(resized_image.shape) == 3:
+        return resized_image.permute(2, 0, 1)
+    else:
+        return resized_image.unsqueeze(dim=-1).permute(2, 0, 1)
+
+def get_expon_lr_func(
+    lr_init, lr_final, lr_delay_steps=0, lr_delay_mult=1.0, max_steps=1000000
+):
+    """
+    Copied from Plenoxels
+
+    Continuous learning rate decay function. Adapted from JaxNeRF
+    The returned rate is lr_init when step=0 and lr_final when step=max_steps, and
+    is log-linearly interpolated elsewhere (equivalent to exponential decay).
+    If lr_delay_steps>0 then the learning rate will be scaled by some smooth
+    function of lr_delay_mult, such that the initial learning rate is
+    lr_init*lr_delay_mult at the beginning of optimization but will be eased back
+    to the normal learning rate when steps>lr_delay_steps.
+    :param conf: config subtree 'lr' or similar
+    :param max_steps: int, the number of steps during optimization.
+    :return HoF which takes step as input
+    """
+
+    def helper(step):
+        if step < 0 or (lr_init == 0.0 and lr_final == 0.0):
+            # Disable this parameter
+            return 0.0
+        if lr_delay_steps > 0:
+            # A kind of reverse cosine decay.
+            delay_rate = lr_delay_mult + (1 - lr_delay_mult) * np.sin(
+                0.5 * np.pi * np.clip(step / lr_delay_steps, 0, 1)
+            )
+        else:
+            delay_rate = 1.0
+        t = np.clip(step / max_steps, 0, 1)
+        log_lerp = np.exp(np.log(lr_init) * (1 - t) + np.log(lr_final) * t)
+        return delay_rate * log_lerp
+
+    return helper
+
+def strip_lowerdiag(L):
+    uncertainty = torch.zeros((L.shape[0], 6), dtype=torch.float, device="cuda")
+
+    uncertainty[:, 0] = L[:, 0, 0]
+    uncertainty[:, 1] = L[:, 0, 1]
+    uncertainty[:, 2] = L[:, 0, 2]
+    uncertainty[:, 3] = L[:, 1, 1]
+    uncertainty[:, 4] = L[:, 1, 2]
+    uncertainty[:, 5] = L[:, 2, 2]
+    return uncertainty
+
+def strip_symmetric(sym):
+    return strip_lowerdiag(sym)
+
+def build_rotation(r):
+    norm = torch.sqrt(r[:,0]*r[:,0] + r[:,1]*r[:,1] + r[:,2]*r[:,2] + r[:,3]*r[:,3])
+
+    q = r / norm[:, None]
+
+    R = torch.zeros((q.size(0), 3, 3), device='cuda')
+
+    r = q[:, 0]
+    x = q[:, 1]
+    y = q[:, 2]
+    z = q[:, 3]
+
+    R[:, 0, 0] = 1 - 2 * (y*y + z*z)
+    R[:, 0, 1] = 2 * (x*y - r*z)
+    R[:, 0, 2] = 2 * (x*z + r*y)
+    R[:, 1, 0] = 2 * (x*y + r*z)
+    R[:, 1, 1] = 1 - 2 * (x*x + z*z)
+    R[:, 1, 2] = 2 * (y*z - r*x)
+    R[:, 2, 0] = 2 * (x*z - r*y)
+    R[:, 2, 1] = 2 * (y*z + r*x)
+    R[:, 2, 2] = 1 - 2 * (x*x + y*y)
+    return R
+
+def build_scaling_rotation(s, r):
+    L = torch.zeros((s.shape[0], 3, 3), dtype=torch.float, device="cuda")
+    R = build_rotation(r)
+
+    L[:,0,0] = s[:,0]
+    L[:,1,1] = s[:,1]
+    L[:,2,2] = s[:,2]
+
+    L = R @ L
+    return L
+
+def safe_state(silent):
+    old_f = sys.stdout
+    class F:
+        def __init__(self, silent):
+            self.silent = silent
+
+        def write(self, x):
+            if not self.silent:
+                if x.endswith("\n"):
+                    old_f.write(x.replace("\n", " [{}]\n".format(str(datetime.now().strftime("%d/%m %H:%M:%S")))))
+                else:
+                    old_f.write(x)
+
+        def flush(self):
+            old_f.flush()
+
+    sys.stdout = F(silent)
+
+    random.seed(0)
+    np.random.seed(0)
+    torch.manual_seed(0)
+    torch.cuda.set_device(torch.device("cuda:0"))

trellis/representations/mesh/__init__.py

@@ -0,0 +1 @@
+from .cube2mesh import SparseFeatures2Mesh, MeshExtractResult

trellis/representations/mesh/cube2mesh.py

@@ -0,0 +1,143 @@
+import torch
+from ...modules.sparse import SparseTensor
+from easydict import EasyDict as edict
+from .utils_cube import *
+from .flexicubes.flexicubes import FlexiCubes
+
+
+class MeshExtractResult:
+    def __init__(self,
+        vertices,
+        faces,
+        vertex_attrs=None,
+        res=64
+    ):
+        self.vertices = vertices
+        self.faces = faces.long()
+        self.vertex_attrs = vertex_attrs
+        self.face_normal = self.comput_face_normals(vertices, faces)
+        self.res = res
+        self.success = (vertices.shape[0] != 0 and faces.shape[0] != 0)
+
+        # training only
+        self.tsdf_v = None
+        self.tsdf_s = None
+        self.reg_loss = None
+        
+    def comput_face_normals(self, verts, faces):
+        i0 = faces[..., 0].long()
+        i1 = faces[..., 1].long()
+        i2 = faces[..., 2].long()
+
+        v0 = verts[i0, :]
+        v1 = verts[i1, :]
+        v2 = verts[i2, :]
+        face_normals = torch.cross(v1 - v0, v2 - v0, dim=-1)
+        face_normals = torch.nn.functional.normalize(face_normals, dim=1)
+        # print(face_normals.min(), face_normals.max(), face_normals.shape)
+        return face_normals[:, None, :].repeat(1, 3, 1)
+                
+    def comput_v_normals(self, verts, faces):
+        i0 = faces[..., 0].long()
+        i1 = faces[..., 1].long()
+        i2 = faces[..., 2].long()
+
+        v0 = verts[i0, :]
+        v1 = verts[i1, :]
+        v2 = verts[i2, :]
+        face_normals = torch.cross(v1 - v0, v2 - v0, dim=-1)
+        v_normals = torch.zeros_like(verts)
+        v_normals.scatter_add_(0, i0[..., None].repeat(1, 3), face_normals)
+        v_normals.scatter_add_(0, i1[..., None].repeat(1, 3), face_normals)
+        v_normals.scatter_add_(0, i2[..., None].repeat(1, 3), face_normals)
+
+        v_normals = torch.nn.functional.normalize(v_normals, dim=1)
+        return v_normals   
+
+
+class SparseFeatures2Mesh:
+    def __init__(self, device="cuda", res=64, use_color=True):
+        '''
+        a model to generate a mesh from sparse features structures using flexicube
+        '''
+        super().__init__()
+        self.device=device
+        self.res = res
+        self.mesh_extractor = FlexiCubes(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):
+        """
+        Generates a mesh based on the specified sparse voxel structures.
+        Args:
+            cube_attrs [Nx21] : Sparse Tensor attrs about cube weights
+            verts_attrs [Nx10] : [0:1] SDF [1:4] deform [4:7] color [7:10] normal 
+        Returns:
+            return the success tag and ni you loss, 
+        """
+        # add sdf bias to verts_attrs
+        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)
+        weights_d = get_dense_attrs(coords, weights, res=self.res, sdf_init=False)
+        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,
+            cube_idx=self.reg_c,
+            resolution=self.res,
+            beta=weights_d[:, :12],
+            alpha=weights_d[:, 12:20],
+            gamma_f=weights_d[:, 20],
+            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

trellis/representations/mesh/utils_cube.py

@@ -0,0 +1,61 @@
+import torch
+cube_corners = torch.tensor([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], [0, 0, 1], [
+        1, 0, 1], [0, 1, 1], [1, 1, 1]], dtype=torch.int)
+cube_neighbor = torch.tensor([[1, 0, 0], [-1, 0, 0], [0, 1, 0], [0, -1, 0], [0, 0, 1], [0, 0, -1]])
+cube_edges = torch.tensor([0, 1, 1, 5, 4, 5, 0, 4, 2, 3, 3, 7, 6, 7, 2, 6,
+                2, 0, 3, 1, 7, 5, 6, 4], dtype=torch.long, requires_grad=False)
+     
+def construct_dense_grid(res, device='cuda'):
+    '''construct a dense grid based on resolution'''
+    res_v = res + 1
+    vertsid = torch.arange(res_v ** 3, device=device)
+    coordsid = vertsid.reshape(res_v, res_v, res_v)[:res, :res, :res].flatten()
+    cube_corners_bias = (cube_corners[:, 0] * res_v + cube_corners[:, 1]) * res_v + cube_corners[:, 2]
+    cube_fx8 = (coordsid.unsqueeze(1) + cube_corners_bias.unsqueeze(0).to(device))
+    verts = torch.stack([vertsid // (res_v ** 2), (vertsid // res_v) % res_v, vertsid % res_v], dim=1)
+    return verts, cube_fx8
+
+
+def construct_voxel_grid(coords):
+    verts = (cube_corners.unsqueeze(0).to(coords) + coords.unsqueeze(1)).reshape(-1, 3)
+    verts_unique, inverse_indices = torch.unique(verts, dim=0, return_inverse=True)
+    cubes = inverse_indices.reshape(-1, 8)
+    return verts_unique, cubes
+
+
+def cubes_to_verts(num_verts, cubes, value, reduce='mean'):
+    """
+    Args:
+        cubes [Vx8] verts index for each cube
+        value [Vx8xM] value to be scattered
+    Operation:
+        reduced[cubes[i][j]][k] += value[i][k]
+    """
+    M = value.shape[2] # number of channels
+    reduced = torch.zeros(num_verts, M, device=cubes.device)
+    return torch.scatter_reduce(reduced, 0, 
+        cubes.unsqueeze(-1).expand(-1, -1, M).flatten(0, 1), 
+        value.flatten(0, 1), reduce=reduce, include_self=False)
+    
+def sparse_cube2verts(coords, feats, training=True):
+    new_coords, cubes = construct_voxel_grid(coords)
+    new_feats = cubes_to_verts(new_coords.shape[0], cubes, feats)
+    if training:
+        con_loss = torch.mean((feats - new_feats[cubes]) ** 2)
+    else:
+        con_loss = 0.0
+    return new_coords, new_feats, con_loss
+    
+
+def get_dense_attrs(coords : torch.Tensor, feats : torch.Tensor, res : int, sdf_init=True):
+    F = feats.shape[-1]
+    dense_attrs = torch.zeros([res] * 3 + [F], device=feats.device)
+    if sdf_init:
+        dense_attrs[..., 0] = 1 # initial outside sdf value
+    dense_attrs[coords[:, 0], coords[:, 1], coords[:, 2], :] = feats
+    return dense_attrs.reshape(-1, F)
+
+
+def get_defomed_verts(v_pos : torch.Tensor, deform : torch.Tensor, res):
+    return v_pos / res - 0.5 + (1 - 1e-8) / (res * 2) * torch.tanh(deform)
+        

trellis/representations/octree/__init__.py

@@ -0,0 +1 @@
+from .octree_dfs import DfsOctree

trellis/representations/octree/octree_dfs.py

@@ -0,0 +1,347 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class DfsOctree:
+    """
+    Sparse Voxel Octree (SVO) implementation for PyTorch.
+    Using Depth-First Search (DFS) order to store the octree.
+    DFS order suits rendering and ray tracing.
+
+    The structure and data are separatedly stored.
+    Structure is stored as a continuous array, each element is a 3*32 bits descriptor.
+    |-----------------------------------------|
+    |      0:3 bits      |      4:31 bits     |
+    |      leaf num      |       unused       |
+    |-----------------------------------------|
+    |               0:31  bits                |
+    |                child ptr                |
+    |-----------------------------------------|
+    |               0:31  bits                |
+    |                data ptr                 |
+    |-----------------------------------------|
+    Each element represents a non-leaf node in the octree.
+    The valid mask is used to indicate whether the children are valid.
+    The leaf mask is used to indicate whether the children are leaf nodes.
+    The child ptr is used to point to the first non-leaf child. Non-leaf children descriptors are stored continuously from the child ptr.
+    The data ptr is used to point to the data of leaf children. Leaf children data are stored continuously from the data ptr.
+
+    There are also auxiliary arrays to store the additional structural information to facilitate parallel processing.
+      - Position: the position of the octree nodes.
+      - Depth: the depth of the octree nodes.
+
+    Args:
+        depth (int): the depth of the octree.
+    """
+
+    def __init__(
+            self,
+            depth,
+            aabb=[0,0,0,1,1,1],
+            sh_degree=2,
+            primitive='voxel',
+            primitive_config={},
+            device='cuda',
+        ):
+        self.max_depth = depth
+        self.aabb = torch.tensor(aabb, dtype=torch.float32, device=device)
+        self.device = device
+        self.sh_degree = sh_degree
+        self.active_sh_degree = sh_degree
+        self.primitive = primitive
+        self.primitive_config = primitive_config
+
+        self.structure = torch.tensor([[8, 1, 0]], dtype=torch.int32, device=self.device)
+        self.position = torch.zeros((8, 3), dtype=torch.float32, device=self.device)
+        self.depth = torch.zeros((8, 1), dtype=torch.uint8, device=self.device)
+        self.position[:, 0] = torch.tensor([0.25, 0.75, 0.25, 0.75, 0.25, 0.75, 0.25, 0.75], device=self.device)
+        self.position[:, 1] = torch.tensor([0.25, 0.25, 0.75, 0.75, 0.25, 0.25, 0.75, 0.75], device=self.device)
+        self.position[:, 2] = torch.tensor([0.25, 0.25, 0.25, 0.25, 0.75, 0.75, 0.75, 0.75], device=self.device)
+        self.depth[:, 0] = 1
+
+        self.data = ['position', 'depth']
+        self.param_names = []
+
+        if primitive == 'voxel':
+            self.features_dc = torch.zeros((8, 1, 3), dtype=torch.float32, device=self.device)
+            self.features_ac = torch.zeros((8, (sh_degree+1)**2-1, 3), dtype=torch.float32, device=self.device)
+            self.data += ['features_dc', 'features_ac']
+            self.param_names += ['features_dc', 'features_ac']
+            if not primitive_config.get('solid', False):
+                self.density = torch.zeros((8, 1), dtype=torch.float32, device=self.device)
+                self.data.append('density')
+                self.param_names.append('density')
+        elif primitive == 'gaussian':
+            self.features_dc = torch.zeros((8, 1, 3), dtype=torch.float32, device=self.device)
+            self.features_ac = torch.zeros((8, (sh_degree+1)**2-1, 3), dtype=torch.float32, device=self.device)
+            self.opacity = torch.zeros((8, 1), dtype=torch.float32, device=self.device)
+            self.data += ['features_dc', 'features_ac', 'opacity']
+            self.param_names += ['features_dc', 'features_ac', 'opacity']
+        elif primitive == 'trivec':
+            self.trivec = torch.zeros((8, primitive_config['rank'], 3, primitive_config['dim']), dtype=torch.float32, device=self.device)
+            self.density = torch.zeros((8, primitive_config['rank']), dtype=torch.float32, device=self.device)
+            self.features_dc = torch.zeros((8, primitive_config['rank'], 1, 3), dtype=torch.float32, device=self.device)
+            self.features_ac = torch.zeros((8, primitive_config['rank'], (sh_degree+1)**2-1, 3), dtype=torch.float32, device=self.device)
+            self.density_shift = 0
+            self.data += ['trivec', 'density', 'features_dc', 'features_ac']
+            self.param_names += ['trivec', 'density', 'features_dc', 'features_ac']
+        elif primitive == 'decoupoly':
+            self.decoupoly_V = torch.zeros((8, primitive_config['rank'], 3), dtype=torch.float32, device=self.device)
+            self.decoupoly_g = torch.zeros((8, primitive_config['rank'], primitive_config['degree']), dtype=torch.float32, device=self.device)
+            self.density = torch.zeros((8, primitive_config['rank']), dtype=torch.float32, device=self.device)
+            self.features_dc = torch.zeros((8, primitive_config['rank'], 1, 3), dtype=torch.float32, device=self.device)
+            self.features_ac = torch.zeros((8, primitive_config['rank'], (sh_degree+1)**2-1, 3), dtype=torch.float32, device=self.device)
+            self.density_shift = 0
+            self.data += ['decoupoly_V', 'decoupoly_g', 'density', 'features_dc', 'features_ac']
+            self.param_names += ['decoupoly_V', 'decoupoly_g', 'density', 'features_dc', 'features_ac']
+
+        self.setup_functions()
+
+    def setup_functions(self):
+        self.density_activation = (lambda x: torch.exp(x - 2)) if self.primitive != 'trivec' else (lambda x: x)
+        self.opacity_activation = lambda x: torch.sigmoid(x - 6)
+        self.inverse_opacity_activation = lambda x: torch.log(x / (1 - x)) + 6
+        self.color_activation = lambda x: torch.sigmoid(x)
+
+    @property
+    def num_non_leaf_nodes(self):
+        return self.structure.shape[0]
+    
+    @property
+    def num_leaf_nodes(self):
+        return self.depth.shape[0]
+
+    @property
+    def cur_depth(self):
+        return self.depth.max().item()
+    
+    @property
+    def occupancy(self):
+        return self.num_leaf_nodes / 8 ** self.cur_depth
+    
+    @property
+    def get_xyz(self):
+        return self.position
+
+    @property
+    def get_depth(self):
+        return self.depth
+
+    @property
+    def get_density(self):
+        if self.primitive == 'voxel' and self.primitive_config.get('solid', False):
+            return torch.full((self.position.shape[0], 1), torch.finfo(torch.float32).max, dtype=torch.float32, device=self.device)
+        return self.density_activation(self.density)
+    
+    @property
+    def get_opacity(self):
+        return self.opacity_activation(self.density)
+
+    @property
+    def get_trivec(self):
+        return self.trivec
+
+    @property
+    def get_decoupoly(self):
+        return F.normalize(self.decoupoly_V, dim=-1), self.decoupoly_g
+
+    @property
+    def get_color(self):
+        return self.color_activation(self.colors)
+
+    @property
+    def get_features(self):
+        if self.sh_degree == 0:
+            return self.features_dc
+        return torch.cat([self.features_dc, self.features_ac], dim=-2)
+
+    def state_dict(self):
+        ret = {'structure': self.structure, 'position': self.position, 'depth': self.depth, 'sh_degree': self.sh_degree, 'active_sh_degree': self.active_sh_degree, 'primitive_config': self.primitive_config, 'primitive': self.primitive}
+        if hasattr(self, 'density_shift'):
+            ret['density_shift'] = self.density_shift
+        for data in set(self.data + self.param_names):
+            if not isinstance(getattr(self, data), nn.Module):
+                ret[data] = getattr(self, data)
+            else:
+                ret[data] = getattr(self, data).state_dict()
+        return ret
+
+    def load_state_dict(self, state_dict):
+        keys = list(set(self.data + self.param_names + list(state_dict.keys()) + ['structure', 'position', 'depth']))
+        for key in keys:
+            if key not in state_dict:
+                print(f"Warning: key {key} not found in the state_dict.")
+                continue
+            try:
+                if not isinstance(getattr(self, key), nn.Module):
+                    setattr(self, key, state_dict[key])
+                else:
+                    getattr(self, key).load_state_dict(state_dict[key])
+            except Exception as e:
+                print(e)
+                raise ValueError(f"Error loading key {key}.")
+
+    def gather_from_leaf_children(self, data):
+        """
+        Gather the data from the leaf children.
+
+        Args:
+            data (torch.Tensor): the data to gather. The first dimension should be the number of leaf nodes.
+        """
+        leaf_cnt = self.structure[:, 0]
+        leaf_cnt_masks = [leaf_cnt == i for i in range(1, 9)]
+        ret = torch.zeros((self.num_non_leaf_nodes,), dtype=data.dtype, device=self.device)
+        for i in range(8):
+            if leaf_cnt_masks[i].sum() == 0:
+                continue
+            start = self.structure[leaf_cnt_masks[i], 2]
+            for j in range(i+1):
+                ret[leaf_cnt_masks[i]] += data[start + j]
+        return ret
+
+    def gather_from_non_leaf_children(self, data):
+        """
+        Gather the data from the non-leaf children.
+
+        Args:
+            data (torch.Tensor): the data to gather. The first dimension should be the number of leaf nodes.
+        """
+        non_leaf_cnt = 8 - self.structure[:, 0]
+        non_leaf_cnt_masks = [non_leaf_cnt == i for i in range(1, 9)]
+        ret = torch.zeros_like(data, device=self.device)
+        for i in range(8):
+            if non_leaf_cnt_masks[i].sum() == 0:
+                continue
+            start = self.structure[non_leaf_cnt_masks[i], 1]
+            for j in range(i+1):
+                ret[non_leaf_cnt_masks[i]] += data[start + j]
+        return ret
+
+    def structure_control(self, mask):
+        """
+        Control the structure of the octree.
+
+        Args:
+            mask (torch.Tensor): the mask to control the structure. 1 for subdivide, -1 for merge, 0 for keep.
+        """
+        # Dont subdivide when the depth is the maximum.
+        mask[self.depth.squeeze() == self.max_depth] = torch.clamp_max(mask[self.depth.squeeze() == self.max_depth], 0)
+        # Dont merge when the depth is the minimum.
+        mask[self.depth.squeeze() == 1] = torch.clamp_min(mask[self.depth.squeeze() == 1], 0)
+
+        # Gather control mask
+        structre_ctrl = self.gather_from_leaf_children(mask)
+        structre_ctrl[structre_ctrl==-8] = -1
+
+        new_leaf_num = self.structure[:, 0].clone()
+        # Modify the leaf num.
+        structre_valid = structre_ctrl >= 0
+        new_leaf_num[structre_valid] -= structre_ctrl[structre_valid]                               # Add the new nodes.
+        structre_delete = structre_ctrl < 0
+        merged_nodes = self.gather_from_non_leaf_children(structre_delete.int())
+        new_leaf_num += merged_nodes                                                                # Delete the merged nodes.
+
+        # Update the structure array to allocate new nodes.
+        mem_offset = torch.zeros((self.num_non_leaf_nodes + 1,), dtype=torch.int32, device=self.device)
+        mem_offset.index_add_(0, self.structure[structre_valid, 1], structre_ctrl[structre_valid])  # Add the new nodes.
+        mem_offset[:-1] -= structre_delete.int()                                                    # Delete the merged nodes.
+        new_structre_idx = torch.arange(0, self.num_non_leaf_nodes + 1, dtype=torch.int32, device=self.device) + mem_offset.cumsum(0)
+        new_structure_length = new_structre_idx[-1].item()
+        new_structre_idx = new_structre_idx[:-1]
+        new_structure = torch.empty((new_structure_length, 3), dtype=torch.int32, device=self.device)
+        new_structure[new_structre_idx[structre_valid], 0] = new_leaf_num[structre_valid]
+
+        # Initialize the new nodes.
+        new_node_mask = torch.ones((new_structure_length,), dtype=torch.bool, device=self.device)
+        new_node_mask[new_structre_idx[structre_valid]] = False
+        new_structure[new_node_mask, 0] = 8                                                         # Initialize to all leaf nodes.
+        new_node_num = new_node_mask.sum().item()
+
+        # Rebuild child ptr.
+        non_leaf_cnt = 8 - new_structure[:, 0]
+        new_child_ptr = torch.cat([torch.zeros((1,), dtype=torch.int32, device=self.device), non_leaf_cnt.cumsum(0)[:-1]])
+        new_structure[:, 1] = new_child_ptr + 1
+
+        # Rebuild data ptr with old data.
+        leaf_cnt = torch.zeros((new_structure_length,), dtype=torch.int32, device=self.device)
+        leaf_cnt.index_add_(0, new_structre_idx, self.structure[:, 0])
+        old_data_ptr = torch.cat([torch.zeros((1,), dtype=torch.int32, device=self.device), leaf_cnt.cumsum(0)[:-1]])
+
+        # Update the data array
+        subdivide_mask = mask == 1
+        merge_mask = mask == -1
+        data_valid = ~(subdivide_mask | merge_mask)
+        mem_offset = torch.zeros((self.num_leaf_nodes + 1,), dtype=torch.int32, device=self.device)
+        mem_offset.index_add_(0, old_data_ptr[new_node_mask], torch.full((new_node_num,), 8, dtype=torch.int32, device=self.device))    # Add data array for new nodes
+        mem_offset[:-1] -= subdivide_mask.int()                                                                                         # Delete data elements for subdivide nodes
+        mem_offset[:-1] -= merge_mask.int()                                                                                             # Delete data elements for merge nodes
+        mem_offset.index_add_(0, self.structure[structre_valid, 2], merged_nodes[structre_valid])                                       # Add data elements for merge nodes
+        new_data_idx = torch.arange(0, self.num_leaf_nodes + 1, dtype=torch.int32, device=self.device) + mem_offset.cumsum(0)
+        new_data_length = new_data_idx[-1].item()
+        new_data_idx = new_data_idx[:-1]
+        new_data = {data: torch.empty((new_data_length,) + getattr(self, data).shape[1:], dtype=getattr(self, data).dtype, device=self.device) for data in self.data}
+        for data in self.data:
+            new_data[data][new_data_idx[data_valid]] = getattr(self, data)[data_valid]
+
+        # Rebuild data ptr
+        leaf_cnt = new_structure[:, 0]
+        new_data_ptr = torch.cat([torch.zeros((1,), dtype=torch.int32, device=self.device), leaf_cnt.cumsum(0)[:-1]])
+        new_structure[:, 2] = new_data_ptr
+
+        # Initialize the new data array
+        ## For subdivide nodes
+        if subdivide_mask.sum() > 0:
+            subdivide_data_ptr = new_structure[new_node_mask, 2]
+            for data in self.data:
+                for i in range(8):
+                    if data == 'position':
+                        offset = torch.tensor([i // 4, (i // 2) % 2, i % 2], dtype=torch.float32, device=self.device) - 0.5
+                        scale = 2 ** (-1.0 - self.depth[subdivide_mask])
+                        new_data['position'][subdivide_data_ptr + i] = self.position[subdivide_mask] + offset * scale
+                    elif data == 'depth':
+                        new_data['depth'][subdivide_data_ptr + i] = self.depth[subdivide_mask] + 1
+                    elif data == 'opacity':
+                        new_data['opacity'][subdivide_data_ptr + i] = self.inverse_opacity_activation(torch.sqrt(self.opacity_activation(self.opacity[subdivide_mask])))
+                    elif data == 'trivec':
+                        offset = torch.tensor([i // 4, (i // 2) % 2, i % 2], dtype=torch.float32, device=self.device) * 0.5
+                        coord = (torch.linspace(0, 0.5, self.trivec.shape[-1], dtype=torch.float32, device=self.device)[None] + offset[:, None]).reshape(1, 3, self.trivec.shape[-1], 1)
+                        axis = torch.linspace(0, 1, 3, dtype=torch.float32, device=self.device).reshape(1, 3, 1, 1).repeat(1, 1, self.trivec.shape[-1], 1)
+                        coord = torch.stack([coord, axis], dim=3).reshape(1, 3, self.trivec.shape[-1], 2).expand(self.trivec[subdivide_mask].shape[0], -1, -1, -1) * 2 - 1
+                        new_data['trivec'][subdivide_data_ptr + i] = F.grid_sample(self.trivec[subdivide_mask], coord, align_corners=True)
+                    else:
+                        new_data[data][subdivide_data_ptr + i] = getattr(self, data)[subdivide_mask]
+        ## For merge nodes
+        if merge_mask.sum() > 0:
+            merge_data_ptr = torch.empty((merged_nodes.sum().item(),), dtype=torch.int32, device=self.device)
+            merge_nodes_cumsum = torch.cat([torch.zeros((1,), dtype=torch.int32, device=self.device), merged_nodes.cumsum(0)[:-1]])
+            for i in range(8):
+                merge_data_ptr[merge_nodes_cumsum[merged_nodes > i] + i] = new_structure[new_structre_idx[merged_nodes > i], 2] + i
+            old_merge_data_ptr = self.structure[structre_delete, 2]
+            for data in self.data:
+                if data == 'position':
+                    scale = 2 ** (1.0 - self.depth[old_merge_data_ptr])
+                    new_data['position'][merge_data_ptr] = ((self.position[old_merge_data_ptr] + 0.5) / scale).floor() * scale + 0.5 * scale - 0.5
+                elif data == 'depth':
+                    new_data['depth'][merge_data_ptr] = self.depth[old_merge_data_ptr] - 1
+                elif data == 'opacity':
+                    new_data['opacity'][subdivide_data_ptr + i] = self.inverse_opacity_activation(self.opacity_activation(self.opacity[subdivide_mask])**2)
+                elif data == 'trivec':
+                    new_data['trivec'][merge_data_ptr] = self.trivec[old_merge_data_ptr]
+                else:
+                    new_data[data][merge_data_ptr] = getattr(self, data)[old_merge_data_ptr]
+
+        # Update the structure and data array
+        self.structure = new_structure
+        for data in self.data:
+            setattr(self, data, new_data[data])
+
+        # Save data array control temp variables
+        self.data_rearrange_buffer = {
+            'subdivide_mask': subdivide_mask,
+            'merge_mask': merge_mask,
+            'data_valid': data_valid,
+            'new_data_idx': new_data_idx,
+            'new_data_length': new_data_length,
+            'new_data': new_data
+        } 

trellis/representations/radiance_field/__init__.py

@@ -0,0 +1 @@
+from .strivec import Strivec

trellis/representations/radiance_field/strivec.py

@@ -0,0 +1,28 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from ..octree import DfsOctree as Octree
+
+
+class Strivec(Octree):
+    def __init__(
+        self,
+        resolution: int,
+        aabb: list,
+        sh_degree: int = 0,
+        rank: int = 8,
+        dim: int = 8,
+        device: str = "cuda",
+    ):
+        assert np.log2(resolution) % 1 == 0, "Resolution must be a power of 2"
+        self.resolution = resolution
+        depth = int(np.round(np.log2(resolution)))
+        super().__init__(
+            depth=depth,
+            aabb=aabb,
+            sh_degree=sh_degree,
+            primitive="trivec",
+            primitive_config={"rank": rank, "dim": dim},
+            device=device,
+        )