Spaces:
Running
on
Zero
Running
on
Zero
File size: 6,871 Bytes
168b252 ac25c40 168b252 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 |
# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto. Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.
import torch
from ...modules.sparse import SparseTensor
from easydict import EasyDict as edict
from .utils_cube import *
from .flexicube import FlexiCubes
import numpy as np
import trimesh
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.vertex_normal = self.comput_v_normals(vertices, faces)
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
def to_trimesh(self, transform_pose=False):
vertices = self.vertices.detach().cpu().numpy()
faces = self.faces.detach().cpu().numpy()
if transform_pose:
transform_matrix = np.array([
[1, 0, 0],
[0, 0, -1],
[0, 1, 0]
])
vertices = vertices @ transform_matrix
vertex_normals = self.vertex_normal.detach().cpu().numpy() @ transform_matrix
else:
vertex_normals = self.vertex_normal.detach().cpu().numpy()
# Create the trimesh mesh
mesh = trimesh.Trimesh(
vertices=vertices,
faces=faces,
face_normals=self.face_normal.detach().cpu().numpy(),
vertex_normals=vertex_normals
)
return mesh
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
|