Big Code Refactor
This commit is contained in:
Mr. For Example
@@ -0,0 +1,162 @@
|
||||
import random
|
||||
import tqdm
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from kiui.mesh_utils import clean_mesh, decimate_mesh
|
||||
from kiui.mesh_utils import laplacian_smooth_loss, normal_consistency
|
||||
from pytorch_msssim import SSIM, MS_SSIM
|
||||
|
||||
import comfy.utils
|
||||
|
||||
from .diff_mesh_renderer import DiffRastRenderer
|
||||
|
||||
from shared_utils.camera_utils import BaseCameraController
|
||||
from shared_utils.image_utils import prepare_torch_img
|
||||
|
||||
class DiffMeshCameraController(BaseCameraController):
|
||||
|
||||
def get_render_result(self, render_pose, bg_color, **kwargs):
|
||||
ref_cam = (render_pose, self.cam.perspective)
|
||||
return self.renderer.render(*ref_cam, self.cam.H, self.cam.W, ssaa=1, bg_color=bg_color, **kwargs) #ssaa = min(2.0, max(0.125, 2 * np.random.random()))
|
||||
|
||||
class DiffMesh:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
mesh,
|
||||
training_iterations,
|
||||
batch_size,
|
||||
texture_learning_rate,
|
||||
train_mesh_geometry,
|
||||
geometry_learning_rate,
|
||||
ms_ssim_loss_weight,
|
||||
remesh_after_n_iteration,
|
||||
invert_bg_prob,
|
||||
force_cuda_rasterize
|
||||
):
|
||||
self.device = torch.device("cuda")
|
||||
|
||||
self.train_mesh_geometry = train_mesh_geometry
|
||||
self.remesh_after_n_iteration = remesh_after_n_iteration
|
||||
|
||||
# prepare main components for optimization
|
||||
self.renderer = DiffRastRenderer(mesh, force_cuda_rasterize).to(self.device)
|
||||
|
||||
self.optimizer = torch.optim.Adam(self.renderer.get_params(texture_learning_rate, train_mesh_geometry, geometry_learning_rate))
|
||||
#self.ssim_loss = SSIM(data_range=1, size_average=True, channel=3)
|
||||
self.ms_ssim_loss = MS_SSIM(data_range=1, size_average=True, channel=3)
|
||||
self.lambda_ssim = ms_ssim_loss_weight
|
||||
|
||||
self.training_iterations = training_iterations
|
||||
|
||||
self.batch_size = batch_size
|
||||
|
||||
self.invert_bg_prob = invert_bg_prob
|
||||
|
||||
def prepare_training(self, reference_images, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy):
|
||||
self.ref_imgs_num = len(reference_images)
|
||||
|
||||
self.ref_size_H = reference_images[0].shape[0]
|
||||
self.ref_size_W = reference_images[0].shape[1]
|
||||
|
||||
# default camera settings
|
||||
self.cam_controller = DiffMeshCameraController(
|
||||
self.renderer, self.ref_size_W, self.ref_size_H, reference_orbit_camera_fovy, self.invert_bg_prob, None, self.device
|
||||
)
|
||||
|
||||
self.all_ref_cam_poses = reference_orbit_camera_poses
|
||||
|
||||
# prepare reference images and masks
|
||||
ref_imgs_torch_list = []
|
||||
ref_masks_torch_list = []
|
||||
for i in range(self.ref_imgs_num):
|
||||
ref_imgs_torch_list.append(prepare_torch_img(reference_images[i].unsqueeze(0), self.ref_size_H, self.ref_size_W, self.device))
|
||||
ref_masks_torch_list.append(prepare_torch_img(reference_masks[i].unsqueeze(2).unsqueeze(0), self.ref_size_H, self.ref_size_W, self.device))
|
||||
|
||||
self.ref_imgs_torch = torch.cat(ref_imgs_torch_list, dim=0)
|
||||
self.ref_masks_torch = torch.cat(ref_masks_torch_list, dim=0)
|
||||
|
||||
def training(self, decimate_target=5e4):
|
||||
starter = torch.cuda.Event(enable_timing=True)
|
||||
ender = torch.cuda.Event(enable_timing=True)
|
||||
starter.record()
|
||||
|
||||
ref_imgs_masked = []
|
||||
for i in range(self.ref_imgs_num):
|
||||
ref_imgs_masked.append((self.ref_imgs_torch[i] * self.ref_masks_torch[i]).unsqueeze(0))
|
||||
|
||||
ref_imgs_num_minus_1 = self.ref_imgs_num-1
|
||||
|
||||
comfy_pbar = comfy.utils.ProgressBar(self.training_iterations)
|
||||
|
||||
for step in tqdm.trange(self.training_iterations):
|
||||
|
||||
### calculate loss between reference and rendered image from known view
|
||||
loss = 0
|
||||
masked_rendered_img_batch = []
|
||||
masked_ref_img_batch = []
|
||||
for _ in range(self.batch_size):
|
||||
|
||||
i = random.randint(0, ref_imgs_num_minus_1)
|
||||
|
||||
out = self.cam_controller.render_at_pose(self.all_ref_cam_poses[i])
|
||||
|
||||
image = out["image"] # [H, W, 3] in [0, 1]
|
||||
image = image.permute(2, 0, 1).contiguous() # [3, H, W] in [0, 1]
|
||||
|
||||
image_masked = (image * self.ref_masks_torch[i]).unsqueeze(0)
|
||||
|
||||
masked_rendered_img_batch.append(image_masked)
|
||||
masked_ref_img_batch.append(ref_imgs_masked[i])
|
||||
|
||||
masked_rendered_img_batch_torch = torch.cat(masked_rendered_img_batch, dim=0)
|
||||
masked_ref_img_batch_torch = torch.cat(masked_ref_img_batch, dim=0)
|
||||
|
||||
# rgb loss
|
||||
loss += (1 - self.lambda_ssim) * F.mse_loss(masked_rendered_img_batch_torch, masked_ref_img_batch_torch)
|
||||
|
||||
# D-SSIM loss
|
||||
# [1, 3, H, W] in [0, 1]
|
||||
#loss += self.lambda_ssim * (1 - self.ssim_loss(X, Y))
|
||||
loss += self.lambda_ssim * (1 - self.ms_ssim_loss(masked_ref_img_batch_torch, masked_rendered_img_batch_torch))
|
||||
|
||||
# Regularization loss
|
||||
if self.train_mesh_geometry:
|
||||
current_v = self.renderer.mesh.v + self.renderer.v_offsets
|
||||
loss += 0.01 * laplacian_smooth_loss(current_v, self.renderer.mesh.f)
|
||||
loss += 0.001 * normal_consistency(current_v, self.renderer.mesh.f)
|
||||
loss += 0.1 * (self.renderer.v_offsets ** 2).sum(-1).mean()
|
||||
|
||||
# remesh periodically
|
||||
if step > 0 and step % self.remesh_after_n_iteration == 0:
|
||||
vertices = (self.renderer.mesh.v + self.renderer.v_offsets).detach().cpu().numpy()
|
||||
triangles = self.renderer.mesh.f.detach().cpu().numpy()
|
||||
vertices, triangles = clean_mesh(vertices, triangles, remesh=True, remesh_size=0.01)
|
||||
if triangles.shape[0] > decimate_target:
|
||||
vertices, triangles = decimate_mesh(vertices, triangles, decimate_target, optimalplacement=False)
|
||||
self.renderer.mesh.v = torch.from_numpy(vertices).contiguous().float().to(self.device)
|
||||
self.renderer.mesh.f = torch.from_numpy(triangles).contiguous().int().to(self.device)
|
||||
self.renderer.v_offsets = nn.Parameter(torch.zeros_like(self.renderer.mesh.v)).to(self.device)
|
||||
|
||||
# optimize step
|
||||
loss.backward()
|
||||
self.optimizer.step()
|
||||
self.optimizer.zero_grad()
|
||||
|
||||
comfy_pbar.update_absolute(step + 1)
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
self.need_update = True
|
||||
|
||||
print(f"Step: {step}")
|
||||
|
||||
self.renderer.update_mesh()
|
||||
|
||||
ender.record()
|
||||
#t = starter.elapsed_time(ender)
|
||||
|
||||
def get_mesh_and_texture(self):
|
||||
return (self.renderer.mesh, self.renderer.mesh.albedo, )
|
||||
@@ -0,0 +1,160 @@
|
||||
import os
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import nvdiffrast.torch as dr
|
||||
|
||||
from kiui.op import inverse_sigmoid
|
||||
|
||||
from mesh_processer.mesh import safe_normalize
|
||||
|
||||
def scale_img_nhwc(x, size, mag='bilinear', min='bilinear'):
|
||||
assert (x.shape[1] >= size[0] and x.shape[2] >= size[1]) or (x.shape[1] < size[0] and x.shape[2] < size[1]), "Trying to magnify image in one dimension and minify in the other"
|
||||
y = x.permute(0, 3, 1, 2) # NHWC -> NCHW
|
||||
if x.shape[1] > size[0] and x.shape[2] > size[1]: # Minification, previous size was bigger
|
||||
y = torch.nn.functional.interpolate(y, size, mode=min)
|
||||
else: # Magnification
|
||||
if mag == 'bilinear' or mag == 'bicubic':
|
||||
y = torch.nn.functional.interpolate(y, size, mode=mag, align_corners=True)
|
||||
else:
|
||||
y = torch.nn.functional.interpolate(y, size, mode=mag)
|
||||
return y.permute(0, 2, 3, 1).contiguous() # NCHW -> NHWC
|
||||
|
||||
def scale_img_hwc(x, size, mag='bilinear', min='bilinear'):
|
||||
return scale_img_nhwc(x[None, ...], size, mag, min)[0]
|
||||
|
||||
def scale_img_nhw(x, size, mag='bilinear', min='bilinear'):
|
||||
return scale_img_nhwc(x[..., None], size, mag, min)[..., 0]
|
||||
|
||||
def scale_img_hw(x, size, mag='bilinear', min='bilinear'):
|
||||
return scale_img_nhwc(x[None, ..., None], size, mag, min)[0, ..., 0]
|
||||
|
||||
def make_divisible(x, m=8):
|
||||
return int(math.ceil(x / m) * m)
|
||||
|
||||
class DiffRastRenderer(nn.Module):
|
||||
def __init__(self, mesh, force_cuda_rast):
|
||||
|
||||
super().__init__()
|
||||
|
||||
self.mesh = mesh
|
||||
|
||||
if force_cuda_rast or os.name != 'nt':
|
||||
self.glctx = dr.RasterizeCudaContext()
|
||||
else:
|
||||
self.glctx = dr.RasterizeGLContext()
|
||||
|
||||
# extract trainable parameters
|
||||
self.v_offsets = nn.Parameter(torch.zeros_like(self.mesh.v), requires_grad=True)
|
||||
self.raw_albedo = nn.Parameter(inverse_sigmoid(self.mesh.albedo), requires_grad=True)
|
||||
|
||||
self.train_geo = False
|
||||
|
||||
def get_params(self, texture_lr, train_geo, geom_lr):
|
||||
|
||||
params = [
|
||||
{'params': self.raw_albedo, 'lr': texture_lr},
|
||||
]
|
||||
|
||||
self.train_geo = train_geo
|
||||
if train_geo:
|
||||
params.append({'params': self.v_offsets, 'lr': geom_lr})
|
||||
|
||||
return params
|
||||
|
||||
def update_mesh(self):
|
||||
self.mesh.v = (self.mesh.v + self.v_offsets).detach()
|
||||
self.mesh.albedo = torch.sigmoid(self.raw_albedo.detach())
|
||||
|
||||
def render(self, pose, proj, h0, w0, ssaa=1, bg_color=1, texture_filter='linear',
|
||||
optional_render_types=['depth', 'normal']):
|
||||
|
||||
# do super-sampling
|
||||
if ssaa != 1:
|
||||
h = make_divisible(h0 * ssaa, 8)
|
||||
w = make_divisible(w0 * ssaa, 8)
|
||||
else:
|
||||
h, w = h0, w0
|
||||
|
||||
results = {}
|
||||
|
||||
# get v
|
||||
if self.train_geo:
|
||||
v = self.mesh.v + self.v_offsets # [N, 3]
|
||||
else:
|
||||
v = self.mesh.v
|
||||
|
||||
pose = torch.from_numpy(pose.astype(np.float32)).to(v.device)
|
||||
proj = torch.from_numpy(proj.astype(np.float32)).to(v.device)
|
||||
|
||||
# get v_clip and render rgb
|
||||
v_cam = torch.matmul(F.pad(v, pad=(0, 1), mode='constant', value=1.0), torch.inverse(pose).T).float().unsqueeze(0)
|
||||
v_clip = v_cam @ proj.T
|
||||
|
||||
rast, rast_db = dr.rasterize(self.glctx, v_clip, self.mesh.f, (h, w))
|
||||
|
||||
#alpha = (rast[0, ..., 3:] > 0).float() # [H, W, 1]
|
||||
alpha = torch.clamp(rast[..., -1:], 0, 1).contiguous() # [1, H, W, 1]
|
||||
alpha = dr.antialias(alpha, rast, v_clip, self.mesh.f).clamp(0, 1).squeeze(0) # [H, W, 1] important to enable gradients!
|
||||
|
||||
# render albedo
|
||||
texc, texc_db = dr.interpolate(self.mesh.vt.unsqueeze(0).contiguous(), rast, self.mesh.ft, rast_db=rast_db, diff_attrs='all')
|
||||
albedo = dr.texture(self.raw_albedo.unsqueeze(0), texc, uv_da=texc_db, filter_mode=texture_filter) # [1, H, W, 3]
|
||||
albedo = torch.sigmoid(albedo)
|
||||
|
||||
# render depth
|
||||
if 'depth' in optional_render_types:
|
||||
depth, _ = dr.interpolate(-v_cam[..., [2]], rast, self.mesh.f) # [1, H, W, 1]
|
||||
depth = depth.squeeze(0) # [H, W, 1]
|
||||
|
||||
# get vn and render normal
|
||||
if 'normal' in optional_render_types:
|
||||
if self.train_geo:
|
||||
i0, i1, i2 = self.mesh.f[:, 0].long(), self.mesh.f[:, 1].long(), self.mesh.f[:, 2].long()
|
||||
v0, v1, v2 = v[i0, :], v[i1, :], v[i2, :]
|
||||
|
||||
face_normals = torch.cross(v1 - v0, v2 - v0)
|
||||
face_normals = safe_normalize(face_normals)
|
||||
|
||||
vn = torch.zeros_like(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)
|
||||
|
||||
vn = torch.where(torch.sum(vn * vn, -1, keepdim=True) > 1e-20, vn, torch.tensor([0.0, 0.0, 1.0], dtype=torch.float32, device=vn.device))
|
||||
else:
|
||||
vn = self.mesh.vn
|
||||
|
||||
normal, _ = dr.interpolate(vn.unsqueeze(0).contiguous(), rast, self.mesh.fn)
|
||||
normal = safe_normalize(normal[0])
|
||||
|
||||
# rotated normal (where [0, 0, 1] always faces camera)
|
||||
rot_normal = normal @ pose[:3, :3]
|
||||
viewcos = rot_normal[..., [2]]
|
||||
|
||||
# antialias
|
||||
albedo = dr.antialias(albedo, rast, v_clip, self.mesh.f).squeeze(0).contiguous() # [H, W, 3]
|
||||
albedo = alpha * albedo + (1 - alpha) * bg_color
|
||||
|
||||
# ssaa
|
||||
if ssaa != 1:
|
||||
albedo = scale_img_hwc(albedo, (h0, w0))
|
||||
alpha = scale_img_hwc(alpha, (h0, w0))
|
||||
if 'depth' in optional_render_types:
|
||||
depth = scale_img_hwc(depth, (h0, w0))
|
||||
if 'normal' in optional_render_types:
|
||||
normal = scale_img_hwc(normal, (h0, w0))
|
||||
viewcos = scale_img_hwc(viewcos, (h0, w0))
|
||||
|
||||
results['image'] = albedo.clamp(0, 1)
|
||||
results['alpha'] = alpha
|
||||
if 'depth' in optional_render_types:
|
||||
results['depth'] = depth
|
||||
if 'normal' in optional_render_types:
|
||||
results['normal'] = (normal + 1) / 2
|
||||
results['viewcos'] = viewcos
|
||||
|
||||
return results
|
||||
@@ -0,0 +1,579 @@
|
||||
# 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 .tables import *
|
||||
|
||||
__all__ = [
|
||||
'FlexiCubes'
|
||||
]
|
||||
|
||||
|
||||
class FlexiCubes:
|
||||
"""
|
||||
This class implements the FlexiCubes method for extracting meshes from scalar fields.
|
||||
It maintains a series of lookup tables and indices to support the mesh extraction process.
|
||||
FlexiCubes, a differentiable variant of the Dual Marching Cubes (DMC) scheme, enhances
|
||||
the geometric fidelity and mesh quality of reconstructed meshes by dynamically adjusting
|
||||
the surface representation through gradient-based optimization.
|
||||
|
||||
During instantiation, the class loads DMC tables from a file and transforms them into
|
||||
PyTorch tensors on the specified device.
|
||||
|
||||
Attributes:
|
||||
device (str): Specifies the computational device (default is "cuda").
|
||||
dmc_table (torch.Tensor): Dual Marching Cubes (DMC) table that encodes the edges
|
||||
associated with each dual vertex in 256 Marching Cubes (MC) configurations.
|
||||
num_vd_table (torch.Tensor): Table holding the number of dual vertices in each of
|
||||
the 256 MC configurations.
|
||||
check_table (torch.Tensor): Table resolving ambiguity in cases C16 and C19
|
||||
of the DMC configurations.
|
||||
tet_table (torch.Tensor): Lookup table used in tetrahedralizing the isosurface.
|
||||
quad_split_1 (torch.Tensor): Indices for splitting a quad into two triangles
|
||||
along one diagonal.
|
||||
quad_split_2 (torch.Tensor): Alternative indices for splitting a quad into
|
||||
two triangles along the other diagonal.
|
||||
quad_split_train (torch.Tensor): Indices for splitting a quad into four triangles
|
||||
during training by connecting all edges to their midpoints.
|
||||
cube_corners (torch.Tensor): Defines the positions of a standard unit cube's
|
||||
eight corners in 3D space, ordered starting from the origin (0,0,0),
|
||||
moving along the x-axis, then y-axis, and finally z-axis.
|
||||
Used as a blueprint for generating a voxel grid.
|
||||
cube_corners_idx (torch.Tensor): Cube corners indexed as powers of 2, used
|
||||
to retrieve the case id.
|
||||
cube_edges (torch.Tensor): Edge connections in a cube, listed in pairs.
|
||||
Used to retrieve edge vertices in DMC.
|
||||
edge_dir_table (torch.Tensor): A mapping tensor that associates edge indices with
|
||||
their corresponding axis. For instance, edge_dir_table[0] = 0 indicates that the
|
||||
first edge is oriented along the x-axis.
|
||||
dir_faces_table (torch.Tensor): A tensor that maps the corresponding axis of shared edges
|
||||
across four adjacent cubes to the shared faces of these cubes. For instance,
|
||||
dir_faces_table[0] = [5, 4] implies that for four cubes sharing an edge along
|
||||
the x-axis, the first and second cubes share faces indexed as 5 and 4, respectively.
|
||||
This tensor is only utilized during isosurface tetrahedralization.
|
||||
adj_pairs (torch.Tensor):
|
||||
A tensor containing index pairs that correspond to neighboring cubes that share the same edge.
|
||||
qef_reg_scale (float):
|
||||
The scaling factor applied to the regularization loss to prevent issues with singularity
|
||||
when solving the QEF. This parameter is only used when a 'grad_func' is specified.
|
||||
weight_scale (float):
|
||||
The scale of weights in FlexiCubes. Should be between 0 and 1.
|
||||
"""
|
||||
|
||||
def __init__(self, device="cuda", qef_reg_scale=1e-3, weight_scale=0.99):
|
||||
|
||||
self.device = device
|
||||
self.dmc_table = torch.tensor(dmc_table, dtype=torch.long, device=device, requires_grad=False)
|
||||
self.num_vd_table = torch.tensor(num_vd_table,
|
||||
dtype=torch.long, device=device, requires_grad=False)
|
||||
self.check_table = torch.tensor(
|
||||
check_table,
|
||||
dtype=torch.long, device=device, requires_grad=False)
|
||||
|
||||
self.tet_table = torch.tensor(tet_table, dtype=torch.long, device=device, requires_grad=False)
|
||||
self.quad_split_1 = torch.tensor([0, 1, 2, 0, 2, 3], dtype=torch.long, device=device, requires_grad=False)
|
||||
self.quad_split_2 = torch.tensor([0, 1, 3, 3, 1, 2], dtype=torch.long, device=device, requires_grad=False)
|
||||
self.quad_split_train = torch.tensor(
|
||||
[0, 1, 1, 2, 2, 3, 3, 0], dtype=torch.long, device=device, requires_grad=False)
|
||||
|
||||
self.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.float, device=device)
|
||||
self.cube_corners_idx = torch.pow(2, torch.arange(8, requires_grad=False))
|
||||
self.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, device=device, requires_grad=False)
|
||||
|
||||
self.edge_dir_table = torch.tensor([0, 2, 0, 2, 0, 2, 0, 2, 1, 1, 1, 1],
|
||||
dtype=torch.long, device=device)
|
||||
self.dir_faces_table = torch.tensor([
|
||||
[[5, 4], [3, 2], [4, 5], [2, 3]],
|
||||
[[5, 4], [1, 0], [4, 5], [0, 1]],
|
||||
[[3, 2], [1, 0], [2, 3], [0, 1]]
|
||||
], dtype=torch.long, device=device)
|
||||
self.adj_pairs = torch.tensor([0, 1, 1, 3, 3, 2, 2, 0], dtype=torch.long, device=device)
|
||||
self.qef_reg_scale = qef_reg_scale
|
||||
self.weight_scale = weight_scale
|
||||
|
||||
def construct_voxel_grid(self, res):
|
||||
"""
|
||||
Generates a voxel grid based on the specified resolution.
|
||||
|
||||
Args:
|
||||
res (int or list[int]): The resolution of the voxel grid. If an integer
|
||||
is provided, it is used for all three dimensions. If a list or tuple
|
||||
of 3 integers is provided, they define the resolution for the x,
|
||||
y, and z dimensions respectively.
|
||||
|
||||
Returns:
|
||||
(torch.Tensor, torch.Tensor): Returns the vertices and the indices of the
|
||||
cube corners (index into vertices) of the constructed voxel grid.
|
||||
The vertices are centered at the origin, with the length of each
|
||||
dimension in the grid being one.
|
||||
"""
|
||||
base_cube_f = torch.arange(8).to(self.device)
|
||||
if isinstance(res, int):
|
||||
res = (res, res, res)
|
||||
voxel_grid_template = torch.ones(res, device=self.device)
|
||||
|
||||
res = torch.tensor([res], dtype=torch.float, device=self.device)
|
||||
coords = torch.nonzero(voxel_grid_template).float() / res # N, 3
|
||||
verts = (self.cube_corners.unsqueeze(0) / res + coords.unsqueeze(1)).reshape(-1, 3)
|
||||
cubes = (base_cube_f.unsqueeze(0) +
|
||||
torch.arange(coords.shape[0], device=self.device).unsqueeze(1) * 8).reshape(-1)
|
||||
|
||||
verts_rounded = torch.round(verts * 10**5) / (10**5)
|
||||
verts_unique, inverse_indices = torch.unique(verts_rounded, dim=0, return_inverse=True)
|
||||
cubes = inverse_indices[cubes.reshape(-1)].reshape(-1, 8)
|
||||
|
||||
return verts_unique - 0.5, cubes
|
||||
|
||||
def __call__(self, x_nx3, s_n, cube_fx8, res, beta_fx12=None, alpha_fx8=None,
|
||||
gamma_f=None, training=False, output_tetmesh=False, grad_func=None):
|
||||
r"""
|
||||
Main function for mesh extraction from scalar field using FlexiCubes. This function converts
|
||||
discrete signed distance fields, encoded on voxel grids and additional per-cube parameters,
|
||||
to triangle or tetrahedral meshes using a differentiable operation as described in
|
||||
`Flexible Isosurface Extraction for Gradient-Based Mesh Optimization`_. FlexiCubes enhances
|
||||
mesh quality and geometric fidelity by adjusting the surface representation based on gradient
|
||||
optimization. The output surface is differentiable with respect to the input vertex positions,
|
||||
scalar field values, and weight parameters.
|
||||
|
||||
If you intend to extract a surface mesh from a fixed Signed Distance Field without the
|
||||
optimization of parameters, it is suggested to provide the "grad_func" which should
|
||||
return the surface gradient at any given 3D position. When grad_func is provided, the process
|
||||
to determine the dual vertex position adapts to solve a Quadratic Error Function (QEF), as
|
||||
described in the `Manifold Dual Contouring`_ paper, and employs an smart splitting strategy.
|
||||
Please note, this approach is non-differentiable.
|
||||
|
||||
For more details and example usage in optimization, refer to the
|
||||
`Flexible Isosurface Extraction for Gradient-Based Mesh Optimization`_ SIGGRAPH 2023 paper.
|
||||
|
||||
Args:
|
||||
x_nx3 (torch.Tensor): Coordinates of the voxel grid vertices, can be deformed.
|
||||
s_n (torch.Tensor): Scalar field values at each vertex of the voxel grid. Negative values
|
||||
denote that the corresponding vertex resides inside the isosurface. This affects
|
||||
the directions of the extracted triangle faces and volume to be tetrahedralized.
|
||||
cube_fx8 (torch.Tensor): Indices of 8 vertices for each cube in the voxel grid.
|
||||
res (int or list[int]): The resolution of the voxel grid. If an integer is provided, it
|
||||
is used for all three dimensions. If a list or tuple of 3 integers is provided, they
|
||||
specify the resolution for the x, y, and z dimensions respectively.
|
||||
beta_fx12 (torch.Tensor, optional): Weight parameters for the cube edges to adjust dual
|
||||
vertices positioning. Defaults to uniform value for all edges.
|
||||
alpha_fx8 (torch.Tensor, optional): Weight parameters for the cube corners to adjust dual
|
||||
vertices positioning. Defaults to uniform value for all vertices.
|
||||
gamma_f (torch.Tensor, optional): Weight parameters to control the splitting of
|
||||
quadrilaterals into triangles. Defaults to uniform value for all cubes.
|
||||
training (bool, optional): If set to True, applies differentiable quad splitting for
|
||||
training. Defaults to False.
|
||||
output_tetmesh (bool, optional): If set to True, outputs a tetrahedral mesh, otherwise,
|
||||
outputs a triangular mesh. Defaults to False.
|
||||
grad_func (callable, optional): A function to compute the surface gradient at specified
|
||||
3D positions (input: Nx3 positions). The function should return gradients as an Nx3
|
||||
tensor. If None, the original FlexiCubes algorithm is utilized. Defaults to None.
|
||||
|
||||
Returns:
|
||||
(torch.Tensor, torch.LongTensor, torch.Tensor): Tuple containing:
|
||||
- Vertices for the extracted triangular/tetrahedral mesh.
|
||||
- Faces for the extracted triangular/tetrahedral mesh.
|
||||
- Regularizer L_dev, computed per dual vertex.
|
||||
|
||||
.. _Flexible Isosurface Extraction for Gradient-Based Mesh Optimization:
|
||||
https://research.nvidia.com/labs/toronto-ai/flexicubes/
|
||||
.. _Manifold Dual Contouring:
|
||||
https://people.engr.tamu.edu/schaefer/research/dualsimp_tvcg.pdf
|
||||
"""
|
||||
|
||||
surf_cubes, occ_fx8 = self._identify_surf_cubes(s_n, cube_fx8)
|
||||
if surf_cubes.sum() == 0:
|
||||
return torch.zeros(
|
||||
(0, 3),
|
||||
device=self.device), torch.zeros(
|
||||
(0, 4),
|
||||
dtype=torch.long, device=self.device) if output_tetmesh else torch.zeros(
|
||||
(0, 3),
|
||||
dtype=torch.long, device=self.device), torch.zeros(
|
||||
(0),
|
||||
device=self.device)
|
||||
beta_fx12, alpha_fx8, gamma_f = self._normalize_weights(beta_fx12, alpha_fx8, gamma_f, surf_cubes)
|
||||
|
||||
case_ids = self._get_case_id(occ_fx8, surf_cubes, res)
|
||||
|
||||
surf_edges, idx_map, edge_counts, surf_edges_mask = self._identify_surf_edges(s_n, cube_fx8, surf_cubes)
|
||||
|
||||
vd, L_dev, vd_gamma, vd_idx_map = self._compute_vd(
|
||||
x_nx3, cube_fx8[surf_cubes], surf_edges, s_n, case_ids, beta_fx12, alpha_fx8, gamma_f, idx_map, grad_func)
|
||||
vertices, faces, s_edges, edge_indices = self._triangulate(
|
||||
s_n, surf_edges, vd, vd_gamma, edge_counts, idx_map, vd_idx_map, surf_edges_mask, training, grad_func)
|
||||
if not output_tetmesh:
|
||||
return vertices, faces, L_dev
|
||||
else:
|
||||
vertices, tets = self._tetrahedralize(
|
||||
x_nx3, s_n, cube_fx8, vertices, faces, surf_edges, s_edges, vd_idx_map, case_ids, edge_indices,
|
||||
surf_cubes, training)
|
||||
return vertices, tets, L_dev
|
||||
|
||||
def _compute_reg_loss(self, vd, ue, edge_group_to_vd, vd_num_edges):
|
||||
"""
|
||||
Regularizer L_dev as in Equation 8
|
||||
"""
|
||||
dist = torch.norm(ue - torch.index_select(input=vd, index=edge_group_to_vd, dim=0), dim=-1)
|
||||
mean_l2 = torch.zeros_like(vd[:, 0])
|
||||
mean_l2 = (mean_l2).index_add_(0, edge_group_to_vd, dist) / vd_num_edges.squeeze(1).float()
|
||||
mad = (dist - torch.index_select(input=mean_l2, index=edge_group_to_vd, dim=0)).abs()
|
||||
return mad
|
||||
|
||||
def _normalize_weights(self, beta_fx12, alpha_fx8, gamma_f, surf_cubes):
|
||||
"""
|
||||
Normalizes the given weights to be non-negative. If input weights are None, it creates and returns a set of weights of ones.
|
||||
"""
|
||||
n_cubes = surf_cubes.shape[0]
|
||||
|
||||
if beta_fx12 is not None:
|
||||
beta_fx12 = (torch.tanh(beta_fx12) * self.weight_scale + 1)
|
||||
else:
|
||||
beta_fx12 = torch.ones((n_cubes, 12), dtype=torch.float, device=self.device)
|
||||
|
||||
if alpha_fx8 is not None:
|
||||
alpha_fx8 = (torch.tanh(alpha_fx8) * self.weight_scale + 1)
|
||||
else:
|
||||
alpha_fx8 = torch.ones((n_cubes, 8), dtype=torch.float, device=self.device)
|
||||
|
||||
if gamma_f is not None:
|
||||
gamma_f = torch.sigmoid(gamma_f) * self.weight_scale + (1 - self.weight_scale)/2
|
||||
else:
|
||||
gamma_f = torch.ones((n_cubes), dtype=torch.float, device=self.device)
|
||||
|
||||
return beta_fx12[surf_cubes], alpha_fx8[surf_cubes], gamma_f[surf_cubes]
|
||||
|
||||
@torch.no_grad()
|
||||
def _get_case_id(self, occ_fx8, surf_cubes, res):
|
||||
"""
|
||||
Obtains the ID of topology cases based on cell corner occupancy. This function resolves the
|
||||
ambiguity in the Dual Marching Cubes (DMC) configurations as described in Section 1.3 of the
|
||||
supplementary material. It should be noted that this function assumes a regular grid.
|
||||
"""
|
||||
case_ids = (occ_fx8[surf_cubes] * self.cube_corners_idx.to(self.device).unsqueeze(0)).sum(-1)
|
||||
|
||||
problem_config = self.check_table.to(self.device)[case_ids]
|
||||
to_check = problem_config[..., 0] == 1
|
||||
problem_config = problem_config[to_check]
|
||||
if not isinstance(res, (list, tuple)):
|
||||
res = [res, res, res]
|
||||
|
||||
# The 'problematic_configs' only contain configurations for surface cubes. Next, we construct a 3D array,
|
||||
# 'problem_config_full', to store configurations for all cubes (with default config for non-surface cubes).
|
||||
# This allows efficient checking on adjacent cubes.
|
||||
problem_config_full = torch.zeros(list(res) + [5], device=self.device, dtype=torch.long)
|
||||
vol_idx = torch.nonzero(problem_config_full[..., 0] == 0) # N, 3
|
||||
vol_idx_problem = vol_idx[surf_cubes][to_check]
|
||||
problem_config_full[vol_idx_problem[..., 0], vol_idx_problem[..., 1], vol_idx_problem[..., 2]] = problem_config
|
||||
vol_idx_problem_adj = vol_idx_problem + problem_config[..., 1:4]
|
||||
|
||||
within_range = (
|
||||
vol_idx_problem_adj[..., 0] >= 0) & (
|
||||
vol_idx_problem_adj[..., 0] < res[0]) & (
|
||||
vol_idx_problem_adj[..., 1] >= 0) & (
|
||||
vol_idx_problem_adj[..., 1] < res[1]) & (
|
||||
vol_idx_problem_adj[..., 2] >= 0) & (
|
||||
vol_idx_problem_adj[..., 2] < res[2])
|
||||
|
||||
vol_idx_problem = vol_idx_problem[within_range]
|
||||
vol_idx_problem_adj = vol_idx_problem_adj[within_range]
|
||||
problem_config = problem_config[within_range]
|
||||
problem_config_adj = problem_config_full[vol_idx_problem_adj[..., 0],
|
||||
vol_idx_problem_adj[..., 1], vol_idx_problem_adj[..., 2]]
|
||||
# If two cubes with cases C16 and C19 share an ambiguous face, both cases are inverted.
|
||||
to_invert = (problem_config_adj[..., 0] == 1)
|
||||
idx = torch.arange(case_ids.shape[0], device=self.device)[to_check][within_range][to_invert]
|
||||
case_ids.index_put_((idx,), problem_config[to_invert][..., -1])
|
||||
return case_ids
|
||||
|
||||
@torch.no_grad()
|
||||
def _identify_surf_edges(self, s_n, cube_fx8, surf_cubes):
|
||||
"""
|
||||
Identifies grid edges that intersect with the underlying surface by checking for opposite signs. As each edge
|
||||
can be shared by multiple cubes, this function also assigns a unique index to each surface-intersecting edge
|
||||
and marks the cube edges with this index.
|
||||
"""
|
||||
occ_n = s_n < 0
|
||||
all_edges = cube_fx8[surf_cubes][:, self.cube_edges].reshape(-1, 2)
|
||||
unique_edges, _idx_map, counts = torch.unique(all_edges, dim=0, return_inverse=True, return_counts=True)
|
||||
|
||||
unique_edges = unique_edges.long()
|
||||
mask_edges = occ_n[unique_edges.reshape(-1)].reshape(-1, 2).sum(-1) == 1
|
||||
|
||||
surf_edges_mask = mask_edges[_idx_map]
|
||||
counts = counts[_idx_map]
|
||||
|
||||
mapping = torch.ones((unique_edges.shape[0]), dtype=torch.long, device=cube_fx8.device) * -1
|
||||
mapping[mask_edges] = torch.arange(mask_edges.sum(), device=cube_fx8.device)
|
||||
# Shaped as [number of cubes x 12 edges per cube]. This is later used to map a cube edge to the unique index
|
||||
# for a surface-intersecting edge. Non-surface-intersecting edges are marked with -1.
|
||||
idx_map = mapping[_idx_map]
|
||||
surf_edges = unique_edges[mask_edges]
|
||||
return surf_edges, idx_map, counts, surf_edges_mask
|
||||
|
||||
@torch.no_grad()
|
||||
def _identify_surf_cubes(self, s_n, cube_fx8):
|
||||
"""
|
||||
Identifies grid cubes that intersect with the underlying surface by checking if the signs at
|
||||
all corners are not identical.
|
||||
"""
|
||||
occ_n = s_n < 0
|
||||
occ_fx8 = occ_n[cube_fx8.reshape(-1)].reshape(-1, 8)
|
||||
_occ_sum = torch.sum(occ_fx8, -1)
|
||||
surf_cubes = (_occ_sum > 0) & (_occ_sum < 8)
|
||||
return surf_cubes, occ_fx8
|
||||
|
||||
def _linear_interp(self, edges_weight, edges_x):
|
||||
"""
|
||||
Computes the location of zero-crossings on 'edges_x' using linear interpolation with 'edges_weight'.
|
||||
"""
|
||||
edge_dim = edges_weight.dim() - 2
|
||||
assert edges_weight.shape[edge_dim] == 2
|
||||
edges_weight = torch.cat([torch.index_select(input=edges_weight, index=torch.tensor(1, device=self.device), dim=edge_dim), -
|
||||
torch.index_select(input=edges_weight, index=torch.tensor(0, device=self.device), dim=edge_dim)], edge_dim)
|
||||
denominator = edges_weight.sum(edge_dim)
|
||||
ue = (edges_x * edges_weight).sum(edge_dim) / denominator
|
||||
return ue
|
||||
|
||||
def _solve_vd_QEF(self, p_bxnx3, norm_bxnx3, c_bx3=None):
|
||||
p_bxnx3 = p_bxnx3.reshape(-1, 7, 3)
|
||||
norm_bxnx3 = norm_bxnx3.reshape(-1, 7, 3)
|
||||
c_bx3 = c_bx3.reshape(-1, 3)
|
||||
A = norm_bxnx3
|
||||
B = ((p_bxnx3) * norm_bxnx3).sum(-1, keepdims=True)
|
||||
|
||||
A_reg = (torch.eye(3, device=p_bxnx3.device) * self.qef_reg_scale).unsqueeze(0).repeat(p_bxnx3.shape[0], 1, 1)
|
||||
B_reg = (self.qef_reg_scale * c_bx3).unsqueeze(-1)
|
||||
A = torch.cat([A, A_reg], 1)
|
||||
B = torch.cat([B, B_reg], 1)
|
||||
dual_verts = torch.linalg.lstsq(A, B).solution.squeeze(-1)
|
||||
return dual_verts
|
||||
|
||||
def _compute_vd(self, x_nx3, surf_cubes_fx8, surf_edges, s_n, case_ids, beta_fx12, alpha_fx8, gamma_f, idx_map, grad_func):
|
||||
"""
|
||||
Computes the location of dual vertices as described in Section 4.2
|
||||
"""
|
||||
alpha_nx12x2 = torch.index_select(input=alpha_fx8, index=self.cube_edges, dim=1).reshape(-1, 12, 2)
|
||||
surf_edges_x = torch.index_select(input=x_nx3, index=surf_edges.reshape(-1), dim=0).reshape(-1, 2, 3)
|
||||
surf_edges_s = torch.index_select(input=s_n, index=surf_edges.reshape(-1), dim=0).reshape(-1, 2, 1)
|
||||
zero_crossing = self._linear_interp(surf_edges_s, surf_edges_x)
|
||||
|
||||
idx_map = idx_map.reshape(-1, 12)
|
||||
num_vd = torch.index_select(input=self.num_vd_table, index=case_ids, dim=0)
|
||||
edge_group, edge_group_to_vd, edge_group_to_cube, vd_num_edges, vd_gamma = [], [], [], [], []
|
||||
|
||||
total_num_vd = 0
|
||||
vd_idx_map = torch.zeros((case_ids.shape[0], 12), dtype=torch.long, device=self.device, requires_grad=False)
|
||||
if grad_func is not None:
|
||||
normals = torch.nn.functional.normalize(grad_func(zero_crossing), dim=-1)
|
||||
vd = []
|
||||
for num in torch.unique(num_vd):
|
||||
cur_cubes = (num_vd == num) # consider cubes with the same numbers of vd emitted (for batching)
|
||||
curr_num_vd = cur_cubes.sum() * num
|
||||
curr_edge_group = self.dmc_table[case_ids[cur_cubes], :num].reshape(-1, num * 7)
|
||||
curr_edge_group_to_vd = torch.arange(
|
||||
curr_num_vd, device=self.device).unsqueeze(-1).repeat(1, 7) + total_num_vd
|
||||
total_num_vd += curr_num_vd
|
||||
curr_edge_group_to_cube = torch.arange(idx_map.shape[0], device=self.device)[
|
||||
cur_cubes].unsqueeze(-1).repeat(1, num * 7).reshape_as(curr_edge_group)
|
||||
|
||||
curr_mask = (curr_edge_group != -1)
|
||||
edge_group.append(torch.masked_select(curr_edge_group, curr_mask))
|
||||
edge_group_to_vd.append(torch.masked_select(curr_edge_group_to_vd.reshape_as(curr_edge_group), curr_mask))
|
||||
edge_group_to_cube.append(torch.masked_select(curr_edge_group_to_cube, curr_mask))
|
||||
vd_num_edges.append(curr_mask.reshape(-1, 7).sum(-1, keepdims=True))
|
||||
vd_gamma.append(torch.masked_select(gamma_f, cur_cubes).unsqueeze(-1).repeat(1, num).reshape(-1))
|
||||
|
||||
if grad_func is not None:
|
||||
with torch.no_grad():
|
||||
cube_e_verts_idx = idx_map[cur_cubes]
|
||||
curr_edge_group[~curr_mask] = 0
|
||||
|
||||
verts_group_idx = torch.gather(input=cube_e_verts_idx, dim=1, index=curr_edge_group)
|
||||
verts_group_idx[verts_group_idx == -1] = 0
|
||||
verts_group_pos = torch.index_select(
|
||||
input=zero_crossing, index=verts_group_idx.reshape(-1), dim=0).reshape(-1, num.item(), 7, 3)
|
||||
v0 = x_nx3[surf_cubes_fx8[cur_cubes][:, 0]].reshape(-1, 1, 1, 3).repeat(1, num.item(), 1, 1)
|
||||
curr_mask = curr_mask.reshape(-1, num.item(), 7, 1)
|
||||
verts_centroid = (verts_group_pos * curr_mask).sum(2) / (curr_mask.sum(2))
|
||||
|
||||
normals_bx7x3 = torch.index_select(input=normals, index=verts_group_idx.reshape(-1), dim=0).reshape(
|
||||
-1, num.item(), 7,
|
||||
3)
|
||||
curr_mask = curr_mask.squeeze(2)
|
||||
vd.append(self._solve_vd_QEF((verts_group_pos - v0) * curr_mask, normals_bx7x3 * curr_mask,
|
||||
verts_centroid - v0.squeeze(2)) + v0.reshape(-1, 3))
|
||||
edge_group = torch.cat(edge_group)
|
||||
edge_group_to_vd = torch.cat(edge_group_to_vd)
|
||||
edge_group_to_cube = torch.cat(edge_group_to_cube)
|
||||
vd_num_edges = torch.cat(vd_num_edges)
|
||||
vd_gamma = torch.cat(vd_gamma)
|
||||
|
||||
if grad_func is not None:
|
||||
vd = torch.cat(vd)
|
||||
L_dev = torch.zeros([1], device=self.device)
|
||||
else:
|
||||
vd = torch.zeros((total_num_vd, 3), device=self.device)
|
||||
beta_sum = torch.zeros((total_num_vd, 1), device=self.device)
|
||||
|
||||
idx_group = torch.gather(input=idx_map.reshape(-1), dim=0, index=edge_group_to_cube * 12 + edge_group)
|
||||
|
||||
x_group = torch.index_select(input=surf_edges_x, index=idx_group.reshape(-1), dim=0).reshape(-1, 2, 3)
|
||||
s_group = torch.index_select(input=surf_edges_s, index=idx_group.reshape(-1), dim=0).reshape(-1, 2, 1)
|
||||
|
||||
zero_crossing_group = torch.index_select(
|
||||
input=zero_crossing, index=idx_group.reshape(-1), dim=0).reshape(-1, 3)
|
||||
|
||||
alpha_group = torch.index_select(input=alpha_nx12x2.reshape(-1, 2), dim=0,
|
||||
index=edge_group_to_cube * 12 + edge_group).reshape(-1, 2, 1)
|
||||
ue_group = self._linear_interp(s_group * alpha_group, x_group)
|
||||
|
||||
beta_group = torch.gather(input=beta_fx12.reshape(-1), dim=0,
|
||||
index=edge_group_to_cube * 12 + edge_group).reshape(-1, 1)
|
||||
beta_sum = beta_sum.index_add_(0, index=edge_group_to_vd, source=beta_group)
|
||||
vd = vd.index_add_(0, index=edge_group_to_vd, source=ue_group * beta_group) / beta_sum
|
||||
L_dev = self._compute_reg_loss(vd, zero_crossing_group, edge_group_to_vd, vd_num_edges)
|
||||
|
||||
v_idx = torch.arange(vd.shape[0], device=self.device) # + total_num_vd
|
||||
|
||||
vd_idx_map = (vd_idx_map.reshape(-1)).scatter(dim=0, index=edge_group_to_cube *
|
||||
12 + edge_group, src=v_idx[edge_group_to_vd])
|
||||
|
||||
return vd, L_dev, vd_gamma, vd_idx_map
|
||||
|
||||
def _triangulate(self, s_n, surf_edges, vd, vd_gamma, edge_counts, idx_map, vd_idx_map, surf_edges_mask, training, grad_func):
|
||||
"""
|
||||
Connects four neighboring dual vertices to form a quadrilateral. The quadrilaterals are then split into
|
||||
triangles based on the gamma parameter, as described in Section 4.3.
|
||||
"""
|
||||
with torch.no_grad():
|
||||
group_mask = (edge_counts == 4) & surf_edges_mask # surface edges shared by 4 cubes.
|
||||
group = idx_map.reshape(-1)[group_mask]
|
||||
vd_idx = vd_idx_map[group_mask]
|
||||
edge_indices, indices = torch.sort(group, stable=True)
|
||||
quad_vd_idx = vd_idx[indices].reshape(-1, 4)
|
||||
|
||||
# Ensure all face directions point towards the positive SDF to maintain consistent winding.
|
||||
s_edges = s_n[surf_edges[edge_indices.reshape(-1, 4)[:, 0]].reshape(-1)].reshape(-1, 2)
|
||||
flip_mask = s_edges[:, 0] > 0
|
||||
quad_vd_idx = torch.cat((quad_vd_idx[flip_mask][:, [0, 1, 3, 2]],
|
||||
quad_vd_idx[~flip_mask][:, [2, 3, 1, 0]]))
|
||||
if grad_func is not None:
|
||||
# when grad_func is given, split quadrilaterals along the diagonals with more consistent gradients.
|
||||
with torch.no_grad():
|
||||
vd_gamma = torch.nn.functional.normalize(grad_func(vd), dim=-1)
|
||||
quad_gamma = torch.index_select(input=vd_gamma, index=quad_vd_idx.reshape(-1), dim=0).reshape(-1, 4, 3)
|
||||
gamma_02 = (quad_gamma[:, 0] * quad_gamma[:, 2]).sum(-1, keepdims=True)
|
||||
gamma_13 = (quad_gamma[:, 1] * quad_gamma[:, 3]).sum(-1, keepdims=True)
|
||||
else:
|
||||
quad_gamma = torch.index_select(input=vd_gamma, index=quad_vd_idx.reshape(-1), dim=0).reshape(-1, 4)
|
||||
gamma_02 = torch.index_select(input=quad_gamma, index=torch.tensor(
|
||||
0, device=self.device), dim=1) * torch.index_select(input=quad_gamma, index=torch.tensor(2, device=self.device), dim=1)
|
||||
gamma_13 = torch.index_select(input=quad_gamma, index=torch.tensor(
|
||||
1, device=self.device), dim=1) * torch.index_select(input=quad_gamma, index=torch.tensor(3, device=self.device), dim=1)
|
||||
if not training:
|
||||
mask = (gamma_02 > gamma_13).squeeze(1)
|
||||
faces = torch.zeros((quad_gamma.shape[0], 6), dtype=torch.long, device=quad_vd_idx.device)
|
||||
faces[mask] = quad_vd_idx[mask][:, self.quad_split_1]
|
||||
faces[~mask] = quad_vd_idx[~mask][:, self.quad_split_2]
|
||||
faces = faces.reshape(-1, 3)
|
||||
else:
|
||||
vd_quad = torch.index_select(input=vd, index=quad_vd_idx.reshape(-1), dim=0).reshape(-1, 4, 3)
|
||||
vd_02 = (torch.index_select(input=vd_quad, index=torch.tensor(0, device=self.device), dim=1) +
|
||||
torch.index_select(input=vd_quad, index=torch.tensor(2, device=self.device), dim=1)) / 2
|
||||
vd_13 = (torch.index_select(input=vd_quad, index=torch.tensor(1, device=self.device), dim=1) +
|
||||
torch.index_select(input=vd_quad, index=torch.tensor(3, device=self.device), dim=1)) / 2
|
||||
weight_sum = (gamma_02 + gamma_13) + 1e-8
|
||||
vd_center = ((vd_02 * gamma_02.unsqueeze(-1) + vd_13 * gamma_13.unsqueeze(-1)) /
|
||||
weight_sum.unsqueeze(-1)).squeeze(1)
|
||||
vd_center_idx = torch.arange(vd_center.shape[0], device=self.device) + vd.shape[0]
|
||||
vd = torch.cat([vd, vd_center])
|
||||
faces = quad_vd_idx[:, self.quad_split_train].reshape(-1, 4, 2)
|
||||
faces = torch.cat([faces, vd_center_idx.reshape(-1, 1, 1).repeat(1, 4, 1)], -1).reshape(-1, 3)
|
||||
return vd, faces, s_edges, edge_indices
|
||||
|
||||
def _tetrahedralize(
|
||||
self, x_nx3, s_n, cube_fx8, vertices, faces, surf_edges, s_edges, vd_idx_map, case_ids, edge_indices,
|
||||
surf_cubes, training):
|
||||
"""
|
||||
Tetrahedralizes the interior volume to produce a tetrahedral mesh, as described in Section 4.5.
|
||||
"""
|
||||
occ_n = s_n < 0
|
||||
occ_fx8 = occ_n[cube_fx8.reshape(-1)].reshape(-1, 8)
|
||||
occ_sum = torch.sum(occ_fx8, -1)
|
||||
|
||||
inside_verts = x_nx3[occ_n]
|
||||
mapping_inside_verts = torch.ones((occ_n.shape[0]), dtype=torch.long, device=self.device) * -1
|
||||
mapping_inside_verts[occ_n] = torch.arange(occ_n.sum(), device=self.device) + vertices.shape[0]
|
||||
"""
|
||||
For each grid edge connecting two grid vertices with different
|
||||
signs, we first form a four-sided pyramid by connecting one
|
||||
of the grid vertices with four mesh vertices that correspond
|
||||
to the grid edge and then subdivide the pyramid into two tetrahedra
|
||||
"""
|
||||
inside_verts_idx = mapping_inside_verts[surf_edges[edge_indices.reshape(-1, 4)[:, 0]].reshape(-1, 2)[
|
||||
s_edges < 0]]
|
||||
if not training:
|
||||
inside_verts_idx = inside_verts_idx.unsqueeze(1).expand(-1, 2).reshape(-1)
|
||||
else:
|
||||
inside_verts_idx = inside_verts_idx.unsqueeze(1).expand(-1, 4).reshape(-1)
|
||||
|
||||
tets_surface = torch.cat([faces, inside_verts_idx.unsqueeze(-1)], -1)
|
||||
"""
|
||||
For each grid edge connecting two grid vertices with the
|
||||
same sign, the tetrahedron is formed by the two grid vertices
|
||||
and two vertices in consecutive adjacent cells
|
||||
"""
|
||||
inside_cubes = (occ_sum == 8)
|
||||
inside_cubes_center = x_nx3[cube_fx8[inside_cubes].reshape(-1)].reshape(-1, 8, 3).mean(1)
|
||||
inside_cubes_center_idx = torch.arange(
|
||||
inside_cubes_center.shape[0], device=inside_cubes.device) + vertices.shape[0] + inside_verts.shape[0]
|
||||
|
||||
surface_n_inside_cubes = surf_cubes | inside_cubes
|
||||
edge_center_vertex_idx = torch.ones(((surface_n_inside_cubes).sum(), 13),
|
||||
dtype=torch.long, device=x_nx3.device) * -1
|
||||
surf_cubes = surf_cubes[surface_n_inside_cubes]
|
||||
inside_cubes = inside_cubes[surface_n_inside_cubes]
|
||||
edge_center_vertex_idx[surf_cubes, :12] = vd_idx_map.reshape(-1, 12)
|
||||
edge_center_vertex_idx[inside_cubes, 12] = inside_cubes_center_idx
|
||||
|
||||
all_edges = cube_fx8[surface_n_inside_cubes][:, self.cube_edges].reshape(-1, 2)
|
||||
unique_edges, _idx_map, counts = torch.unique(all_edges, dim=0, return_inverse=True, return_counts=True)
|
||||
unique_edges = unique_edges.long()
|
||||
mask_edges = occ_n[unique_edges.reshape(-1)].reshape(-1, 2).sum(-1) == 2
|
||||
mask = mask_edges[_idx_map]
|
||||
counts = counts[_idx_map]
|
||||
mapping = torch.ones((unique_edges.shape[0]), dtype=torch.long, device=self.device) * -1
|
||||
mapping[mask_edges] = torch.arange(mask_edges.sum(), device=self.device)
|
||||
idx_map = mapping[_idx_map]
|
||||
|
||||
group_mask = (counts == 4) & mask
|
||||
group = idx_map.reshape(-1)[group_mask]
|
||||
edge_indices, indices = torch.sort(group)
|
||||
cube_idx = torch.arange((_idx_map.shape[0] // 12), dtype=torch.long,
|
||||
device=self.device).unsqueeze(1).expand(-1, 12).reshape(-1)[group_mask]
|
||||
edge_idx = torch.arange((12), dtype=torch.long, device=self.device).unsqueeze(
|
||||
0).expand(_idx_map.shape[0] // 12, -1).reshape(-1)[group_mask]
|
||||
# Identify the face shared by the adjacent cells.
|
||||
cube_idx_4 = cube_idx[indices].reshape(-1, 4)
|
||||
edge_dir = self.edge_dir_table[edge_idx[indices]].reshape(-1, 4)[..., 0]
|
||||
shared_faces_4x2 = self.dir_faces_table[edge_dir].reshape(-1)
|
||||
cube_idx_4x2 = cube_idx_4[:, self.adj_pairs].reshape(-1)
|
||||
# Identify an edge of the face with different signs and
|
||||
# select the mesh vertex corresponding to the identified edge.
|
||||
case_ids_expand = torch.ones((surface_n_inside_cubes).sum(), dtype=torch.long, device=x_nx3.device) * 255
|
||||
case_ids_expand[surf_cubes] = case_ids
|
||||
cases = case_ids_expand[cube_idx_4x2]
|
||||
quad_edge = edge_center_vertex_idx[cube_idx_4x2, self.tet_table[cases, shared_faces_4x2]].reshape(-1, 2)
|
||||
mask = (quad_edge == -1).sum(-1) == 0
|
||||
inside_edge = mapping_inside_verts[unique_edges[mask_edges][edge_indices].reshape(-1)].reshape(-1, 2)
|
||||
tets_inside = torch.cat([quad_edge, inside_edge], -1)[mask]
|
||||
|
||||
tets = torch.cat([tets_surface, tets_inside])
|
||||
vertices = torch.cat([vertices, inside_verts, inside_cubes_center])
|
||||
return vertices, tets
|
||||
@@ -0,0 +1,74 @@
|
||||
import os
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
import nvdiffrast.torch as dr
|
||||
from kiui.cam import orbit_camera
|
||||
|
||||
from FlexiCubes import util
|
||||
|
||||
class FlexiCubesRenderer:
|
||||
def __init__(self, force_cuda_rast):
|
||||
|
||||
if force_cuda_rast or os.name != 'nt':
|
||||
self.glctx = dr.RasterizeCudaContext()
|
||||
else:
|
||||
self.glctx = dr.RasterizeGLContext()
|
||||
|
||||
def get_rotate_camera(self, itr, fovy = np.deg2rad(45), iter_res=[512,512], cam_near_far=[0.1, 1000.0], cam_radius=3.0, device="cuda"):
|
||||
proj_mtx = util.perspective(fovy, iter_res[1] / iter_res[0], cam_near_far[0], cam_near_far[1])
|
||||
|
||||
# Smooth rotation for display.
|
||||
ang = (itr / 10) * np.pi * 2
|
||||
mv = util.translate(0, 0, -cam_radius) @ (util.rotate_x(-0.4) @ util.rotate_y(ang))
|
||||
mvp = proj_mtx @ mv
|
||||
return mv.to(device), mvp.to(device)
|
||||
|
||||
def get_orbit_camera(self, azimuth=0, elevation=0, fovy=45, iter_res=[512,512], cam_near_far=[0.1, 1000.0], cam_radius=3.0, device="cuda"):
|
||||
aspect = iter_res[1] / iter_res[0]
|
||||
fovy = np.deg2rad(fovy)
|
||||
proj_mtx = util.perspective(fovy, aspect, cam_near_far[0], cam_near_far[1])
|
||||
#fovx = 2 * np.arctan(np.tan(fovy / 2) * aspect)
|
||||
#proj_mtx = get_projection_matrix(cam_near_far[0], cam_near_far[1], fovx, -fovy, z_sign=-1.0)
|
||||
|
||||
# Smooth rotation for display.
|
||||
#mv = util.translate(0, 0, -cam_radius) @ (util.rotate_x(-np.de2rad(elevation)) @ util.rotate_y(np.deg2rad(azimuth)))
|
||||
mv = torch.tensor(np.linalg.inv(orbit_camera(-elevation, azimuth, cam_radius)))
|
||||
mvp = proj_mtx @ mv
|
||||
return mv.to(device), mvp.to(device)
|
||||
|
||||
def render_mesh(self, mesh, mv, mvp, iter_res, depth_min=-5.5, depth_max=-0.5, return_types=["mask", "depth", "normal"], white_bg=False):
|
||||
'''
|
||||
The rendering function used to produce the results in the paper.
|
||||
'''
|
||||
v_pos_clip = util.xfm_points(mesh.vertices.unsqueeze(0), mvp) # Rotate it to camera coordinates
|
||||
rast, db = dr.rasterize(
|
||||
dr.RasterizeGLContext(), v_pos_clip, mesh.faces.int(), iter_res)
|
||||
|
||||
alpha_index = rast[..., -1:] > 0
|
||||
inv_alpha_index = rast[..., -1:] <= 0
|
||||
alpha = alpha_index.float()
|
||||
|
||||
out_dict = {}
|
||||
for type in return_types:
|
||||
if type == "mask" :
|
||||
img = dr.antialias(alpha, rast, v_pos_clip, mesh.faces.int())
|
||||
elif type == "depth":
|
||||
v_pos_cam = util.xfm_points(mesh.vertices.unsqueeze(0), mv)
|
||||
img, _ = util.interpolate(v_pos_cam[..., [2]].contiguous(), rast, mesh.faces.int())
|
||||
img_masked = img[alpha_index]
|
||||
img_masked = torch.clamp(img_masked, min=depth_min, max=depth_max)
|
||||
img_masked = (img_masked - depth_min) / (depth_max - depth_min)
|
||||
img[alpha_index] = img_masked
|
||||
img[inv_alpha_index] = 0
|
||||
elif type == "normal" :
|
||||
normal_indices = (torch.arange(0, mesh.nrm.shape[0], dtype=torch.int64, device='cuda')[:, None]).repeat(1, 3)
|
||||
img, _ = util.interpolate(mesh.nrm.unsqueeze(0).contiguous(), rast, normal_indices.int())
|
||||
elif type == "vertex_normal":
|
||||
img, _ = util.interpolate(mesh.v_nrm.unsqueeze(0).contiguous(), rast, mesh.faces.int())
|
||||
img = dr.antialias((img + 1) * 0.5, rast, v_pos_clip, mesh.faces.int())
|
||||
if white_bg:
|
||||
bg = torch.ones_like(img)
|
||||
img = torch.lerp(bg, img, alpha)
|
||||
out_dict[type] = img
|
||||
return out_dict
|
||||
@@ -0,0 +1,213 @@
|
||||
import os
|
||||
import imageio
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
import tqdm
|
||||
import comfy.utils
|
||||
|
||||
from .flexicubes_renderer import FlexiCubesRenderer
|
||||
from .flexicubes import FlexiCubes
|
||||
from .util import SimpleMesh
|
||||
from .loss import sdf_reg_loss
|
||||
|
||||
from shared_utils.camera_utils import OrbitCamera
|
||||
from mesh_processer.mesh import Mesh
|
||||
|
||||
|
||||
def lr_schedule(iter):
|
||||
return max(0.0, 10**(-(iter)*0.0002)) # Exponential falloff from [1.0, 0.1] over 5k epochs.
|
||||
|
||||
class FlexiCubesTrainer:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
training_iterations,
|
||||
batch_size,
|
||||
learning_rate,
|
||||
voxel_grids_resolution,
|
||||
depth_near=0.5,
|
||||
depth_far=5.5,
|
||||
mask_loss_weight=1.0,
|
||||
depth_loss_weight=100.0,
|
||||
normal_loss_weight=1.0,
|
||||
sdf_regularizer_weight=0.2,
|
||||
remove_floaters_weight=0.5,
|
||||
cube_stabilizer_weight=0.1,
|
||||
force_cuda_rast=False,
|
||||
device='cuda'
|
||||
):
|
||||
self.device = torch.device(device)
|
||||
|
||||
self.renderer = FlexiCubesRenderer(force_cuda_rast)
|
||||
|
||||
# Create and initialize FlexiCubes
|
||||
self.voxel_grid_res = voxel_grids_resolution
|
||||
self.fc = FlexiCubes(self.device)
|
||||
self.x_nx3, self.cube_fx8 = self.fc.construct_voxel_grid(self.voxel_grid_res)
|
||||
self.x_nx3 *= 2 # scale up the grid so that it's larger than the target object
|
||||
|
||||
self.sdf = torch.rand_like(self.x_nx3[:,0]) - 0.1 # randomly init SDF
|
||||
self.sdf = torch.nn.Parameter(self.sdf.clone().detach(), requires_grad=True)
|
||||
# set per-cube learnable weights to zeros
|
||||
self.weight = torch.zeros((self.cube_fx8.shape[0], 21), dtype=torch.float, device=self.device)
|
||||
self.weight = torch.nn.Parameter(self.weight.clone().detach(), requires_grad=True)
|
||||
self.deform = torch.nn.Parameter(torch.zeros_like(self.x_nx3), requires_grad=True)
|
||||
|
||||
# Retrieve all the edges of the voxel grid; these edges will be utilized to
|
||||
# compute the regularization loss in subsequent steps of the process.
|
||||
all_edges = self.cube_fx8[:, self.fc.cube_edges].reshape(-1, 2)
|
||||
self.grid_edges = torch.unique(all_edges, dim=0)
|
||||
|
||||
# Setup optimizer
|
||||
self.optimizer = torch.optim.Adam([self.sdf, self.weight, self.deform], lr=learning_rate)
|
||||
self.scheduler = torch.optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda=lambda x: lr_schedule(x))
|
||||
|
||||
self.training_iterations = training_iterations
|
||||
self.batch_size = batch_size
|
||||
|
||||
self.depth_near = -depth_near
|
||||
self.depth_far = -depth_far
|
||||
|
||||
self.mask_loss_weight = mask_loss_weight
|
||||
self.depth_loss_weight = depth_loss_weight
|
||||
self.normal_loss_weight = normal_loss_weight
|
||||
self.sdf_regularizer_weight = sdf_regularizer_weight
|
||||
self.remove_floaters_weight = remove_floaters_weight
|
||||
self.cube_stabilizer_weight = cube_stabilizer_weight
|
||||
|
||||
def prepare_training(self, reference_depth_images, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy, reference_normals=None):
|
||||
self.ref_imgs_num = len(reference_depth_images)
|
||||
|
||||
self.all_ref_cam_poses = reference_orbit_camera_poses
|
||||
self.ref_cam_fovy = reference_orbit_camera_fovy
|
||||
|
||||
self.ref_size_H = reference_depth_images[0].shape[0]
|
||||
self.ref_size_W = reference_depth_images[0].shape[1]
|
||||
|
||||
self.cam = OrbitCamera(self.ref_size_W, self.ref_size_H, fovy=reference_orbit_camera_fovy)
|
||||
|
||||
# prepare reference images and masks
|
||||
self.ref_depth_imgs_torch = reference_depth_images[:, :, :, 0].unsqueeze(3).to(self.device) # (N, H, W, C) -> (N, H, W, 1)
|
||||
self.ref_masks_torch = reference_masks.unsqueeze(3).to(self.device)
|
||||
|
||||
if reference_normals is not None:
|
||||
self.ref_normal_imgs_torch = (reference_normals * 2 - 1).to(self.device) # change value from [0, 1] -> [-1, 1]
|
||||
else:
|
||||
self.ref_normal_imgs_torch = None
|
||||
|
||||
# prepare reference camera projection matrix for all camera poses
|
||||
mv_all = []
|
||||
mvp_all = []
|
||||
for pose in self.all_ref_cam_poses:
|
||||
mv, mvp = self.renderer.get_orbit_camera(pose[2], pose[1], cam_radius=pose[0], device=self.device)
|
||||
mv_all.append(mv)
|
||||
mvp_all.append(mvp)
|
||||
self.mv_all = torch.stack(mv_all).to(self.device)
|
||||
self.mvp_all = torch.stack(mvp_all).to(self.device)
|
||||
self.camposes_len = len(self.all_ref_cam_poses)
|
||||
|
||||
def training(self):
|
||||
starter = torch.cuda.Event(enable_timing=True)
|
||||
ender = torch.cuda.Event(enable_timing=True)
|
||||
starter.record()
|
||||
|
||||
comfy_pbar = comfy.utils.ProgressBar(self.training_iterations)
|
||||
|
||||
if self.ref_normal_imgs_torch is not None:
|
||||
return_types = ["mask", "depth", "normal"]
|
||||
else:
|
||||
return_types = ["mask", "depth"]
|
||||
|
||||
for step in tqdm.trange(self.training_iterations):
|
||||
# sample random render & camera pose from multi-views
|
||||
batch_index = np.random.randint(0, self.camposes_len, size=self.batch_size)
|
||||
mv = self.mv_all[batch_index, :, :]
|
||||
mvp = self.mvp_all[batch_index, :, :]
|
||||
|
||||
target_mask = self.ref_masks_torch[batch_index, :, :, :]
|
||||
target_depth = self.ref_depth_imgs_torch[batch_index, :, :, :]
|
||||
|
||||
# extract and render FlexiCubes mesh
|
||||
grid_verts = self.x_nx3 + (2-1e-8) / (self.voxel_grid_res * 2) * torch.tanh(self.deform)
|
||||
vertices, faces, L_dev = self.fc(grid_verts, self.sdf, self.cube_fx8, self.voxel_grid_res, beta_fx12=self.weight[:,:12], alpha_fx8=self.weight[:,12:20],
|
||||
gamma_f=self.weight[:,20], training=True)
|
||||
flexicubes_mesh = SimpleMesh(vertices, faces)
|
||||
if self.ref_normal_imgs_torch is not None:
|
||||
flexicubes_mesh.auto_normals()
|
||||
buffers = self.renderer.render_mesh(flexicubes_mesh, mv, mvp, (self.ref_size_H, self.ref_size_W), self.depth_far, self.depth_near, return_types)
|
||||
|
||||
t_iter = step / self.training_iterations
|
||||
# evaluate reconstruction loss
|
||||
# mask & depth shape: (N, H, W, 1)
|
||||
mask_loss = (buffers['mask'] - target_mask).abs().mean() * self.mask_loss_weight
|
||||
depth_loss = (((((buffers['depth'] - target_depth)* target_mask)**2).sum(-1)+1e-8)).sqrt().mean() * self.depth_loss_weight
|
||||
total_loss = mask_loss + depth_loss
|
||||
if self.ref_normal_imgs_torch is not None:
|
||||
target_normal = self.ref_normal_imgs_torch[batch_index, :, :, :]
|
||||
normal_loss = (((((buffers['normal'] - (target_normal))* target_mask)**2).sum(-1)+1e-8)).sqrt().mean() * self.normal_loss_weight * t_iter
|
||||
total_loss += normal_loss
|
||||
|
||||
sdf_weight = self.sdf_regularizer_weight - (self.sdf_regularizer_weight - self.sdf_regularizer_weight/20)*min(1.0, 4.0 * t_iter)
|
||||
reg_loss = sdf_reg_loss(self.sdf, self.grid_edges).mean() * sdf_weight # Loss to eliminate internal floaters that are not visible
|
||||
reg_loss += L_dev.mean() * self.remove_floaters_weight
|
||||
reg_loss += (self.weight[:,:20]).abs().mean() * self.cube_stabilizer_weight
|
||||
total_loss += reg_loss
|
||||
|
||||
# optimize step
|
||||
total_loss.backward()
|
||||
self.optimizer.step()
|
||||
self.scheduler.step()
|
||||
self.optimizer.zero_grad()
|
||||
|
||||
comfy_pbar.update_absolute(step + 1)
|
||||
|
||||
#self.test_save(step, mv, mvp, grid_verts, total_loss, out_dir="C:\\Users\\reall\\Softwares\\ComfyUI_windows_portable\\ComfyUI\\output\\FlexiCubes_Output\\Test_Normals")
|
||||
|
||||
ender.record()
|
||||
|
||||
def get_mesh(self):
|
||||
grid_verts = self.x_nx3 + (2-1e-8) / (self.voxel_grid_res * 2) * torch.tanh(self.deform)
|
||||
vertices, faces, L_dev = self.fc(grid_verts, self.sdf, self.cube_fx8, self.voxel_grid_res, beta_fx12=self.weight[:,:12], alpha_fx8=self.weight[:,12:20],
|
||||
gamma_f=self.weight[:,20], training=False)
|
||||
|
||||
v = vertices.detach().contiguous().float().to(self.device)
|
||||
f = faces.detach().contiguous().float().to(self.device)
|
||||
mesh = Mesh(v=v, f=f, device=self.device)
|
||||
mesh.auto_normal()
|
||||
mesh.auto_uv()
|
||||
|
||||
# Trimesh seems output better mesh
|
||||
#import trimesh
|
||||
#import os
|
||||
#mesh_np = trimesh.Trimesh(vertices = vertices.detach().cpu().numpy(), faces=faces.detach().cpu().numpy(), process=False)
|
||||
#mesh_np.export(os.path.join("C:\\Users\\reall\\Softwares\\ComfyUI_windows_portable\\ComfyUI\\output\\FlexiCubes_Output", 'output_trimesh.obj'))
|
||||
|
||||
return mesh
|
||||
|
||||
def test_save(self, step, mv, mvp, grid_verts, total_loss, out_dir, save_interval=20, display_res=[1024, 1024]):
|
||||
if (step % save_interval == 0 or step == (self.training_iterations-1)): # save normal image for visualization
|
||||
with torch.no_grad():
|
||||
# extract mesh with training=False
|
||||
vertices, faces, L_dev = self.fc(grid_verts, self.sdf, self.cube_fx8, self.voxel_grid_res, beta_fx12=self.weight[:,:12], alpha_fx8=self.weight[:,12:20],
|
||||
gamma_f=self.weight[:,20], training=False)
|
||||
flexicubes_mesh = SimpleMesh(vertices, faces)
|
||||
flexicubes_mesh.auto_normals() # compute face normals for visualization
|
||||
|
||||
mv, mvp = self.renderer.get_rotate_camera(step//save_interval, iter_res=display_res, device=self.device)
|
||||
mv = mv.unsqueeze(0)
|
||||
mvp = mvp.unsqueeze(0)
|
||||
val_buffers = self.renderer.render_mesh(flexicubes_mesh, mv, mvp, display_res, return_types=["normal"], white_bg=True)
|
||||
val_image = ((val_buffers["normal"][0].detach().cpu().numpy()+1)/2*255).astype(np.uint8)
|
||||
#val_image = (val_buffers["depth"][0].detach().cpu().numpy()*255).astype(np.uint8)
|
||||
#val_repeat = np.repeat(val_image, 4, axis=2)
|
||||
#val_repeat[:, :, 3] = 255
|
||||
|
||||
#gt_image = (gt_image[0].detach().cpu().numpy()*255).astype(np.uint8)
|
||||
#gt_repeat = np.repeat(gt_image, 4, axis=2)
|
||||
#gt_repeat[:, :, 3] = 255
|
||||
|
||||
imageio.imwrite(os.path.join(out_dir, '{:04d}.png'.format(step)), val_image)
|
||||
#imageio.imwrite(os.path.join(out_dir, '{:04d}.png'.format(step)), np.concatenate([val_repeat, gt_repeat], 1))
|
||||
print(f"Optimization Step [{step}/{self.training_iterations}], Loss: {total_loss.item():.4f}")
|
||||
@@ -0,0 +1,95 @@
|
||||
# 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
|
||||
import torch_scatter
|
||||
|
||||
###############################################################################
|
||||
# Pytorch implementation of the developability regularizer introduced in paper
|
||||
# "Developability of Triangle Meshes" by Stein et al.
|
||||
###############################################################################
|
||||
def mesh_developable_reg(mesh):
|
||||
|
||||
verts = mesh.vertices
|
||||
tris = mesh.faces
|
||||
|
||||
device = verts.device
|
||||
V = verts.shape[0]
|
||||
F = tris.shape[0]
|
||||
|
||||
POS_EPS = 1e-6
|
||||
REL_EPS = 1e-6
|
||||
|
||||
def normalize(vecs):
|
||||
return vecs / (torch.linalg.norm(vecs, dim=-1, keepdim=True) + POS_EPS)
|
||||
|
||||
tri_pos = verts[tris]
|
||||
|
||||
vert_normal_covariance_sum = torch.zeros((V, 9), device=device)
|
||||
vert_area = torch.zeros(V, device=device)
|
||||
vert_degree = torch.zeros(V, dtype=torch.int32, device=device)
|
||||
|
||||
for iC in range(3): # loop over three corners of each triangle
|
||||
|
||||
# gather tri verts
|
||||
pRoot = tri_pos[:, iC, :]
|
||||
pA = tri_pos[:, (iC + 1) % 3, :]
|
||||
pB = tri_pos[:, (iC + 2) % 3, :]
|
||||
|
||||
# compute the corner angle & normal
|
||||
vA = pA - pRoot
|
||||
vAn = normalize(vA)
|
||||
vB = pB - pRoot
|
||||
vBn = normalize(vB)
|
||||
area_normal = torch.linalg.cross(vA, vB, dim=-1)
|
||||
face_area = 0.5 * torch.linalg.norm(area_normal, dim=-1)
|
||||
normal = normalize(area_normal)
|
||||
corner_angle = torch.acos(torch.clamp(torch.sum(vAn * vBn, dim=-1), min=-1., max=1.))
|
||||
|
||||
# add up the contribution to the covariance matrix
|
||||
outer = normal[:, :, None] @ normal[:, None, :]
|
||||
contrib = corner_angle[:, None] * outer.reshape(-1, 9)
|
||||
|
||||
# scatter the result to the appropriate matrices
|
||||
vert_normal_covariance_sum = torch_scatter.scatter_add(src=contrib,
|
||||
index=tris[:, iC],
|
||||
dim=-2,
|
||||
out=vert_normal_covariance_sum)
|
||||
|
||||
vert_area = torch_scatter.scatter_add(src=face_area / 3.,
|
||||
index=tris[:, iC],
|
||||
dim=-1,
|
||||
out=vert_area)
|
||||
|
||||
vert_degree = torch_scatter.scatter_add(src=torch.ones(F, dtype=torch.int32, device=device),
|
||||
index=tris[:, iC],
|
||||
dim=-1,
|
||||
out=vert_degree)
|
||||
|
||||
# The energy is the smallest eigenvalue of the outer-product matrix
|
||||
vert_normal_covariance_sum = vert_normal_covariance_sum.reshape(
|
||||
-1, 3, 3) # reshape to a batch of matrices
|
||||
vert_normal_covariance_sum = vert_normal_covariance_sum + torch.eye(
|
||||
3, device=device)[None, :, :] * REL_EPS
|
||||
|
||||
min_eigvals = torch.min(torch.linalg.eigvals(vert_normal_covariance_sum).abs(), dim=-1).values
|
||||
|
||||
# Mask out degree-3 vertices
|
||||
vert_area = torch.where(vert_degree == 3, torch.tensor(0, dtype=vert_area.dtype,device=vert_area.device), vert_area)
|
||||
|
||||
# Adjust the vertex area weighting so it is unit-less, and 1 on average
|
||||
vert_area = vert_area * (V / torch.sum(vert_area, dim=-1, keepdim=True))
|
||||
|
||||
return vert_area * min_eigvals
|
||||
|
||||
def sdf_reg_loss(sdf, all_edges):
|
||||
sdf_f1x6x2 = sdf[all_edges.reshape(-1)].reshape(-1,2)
|
||||
mask = torch.sign(sdf_f1x6x2[...,0]) != torch.sign(sdf_f1x6x2[...,1])
|
||||
sdf_f1x6x2 = sdf_f1x6x2[mask]
|
||||
sdf_diff = torch.nn.functional.binary_cross_entropy_with_logits(sdf_f1x6x2[...,0], (sdf_f1x6x2[...,1] > 0).float()) + \
|
||||
torch.nn.functional.binary_cross_entropy_with_logits(sdf_f1x6x2[...,1], (sdf_f1x6x2[...,0] > 0).float())
|
||||
return sdf_diff
|
||||
@@ -0,0 +1,791 @@
|
||||
# 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.
|
||||
dmc_table = [
|
||||
[[-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 8, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 9, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 3, 8, 9, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 7, 8, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 4, 7, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 9, -1, -1, -1, -1], [4, 7, 8, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 3, 4, 7, 9, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 5, 9, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 8, -1, -1, -1, -1], [4, 5, 9, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 4, 5, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 3, 4, 5, 8, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[5, 7, 8, 9, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 5, 7, 9, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 5, 7, 8, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 3, 5, 7, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 8, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 9, -1, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 2, 8, 9, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 7, 8, -1, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 4, 7, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 9, -1, -1, -1, -1], [4, 7, 8, -1, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 2, 4, 7, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 5, 9, -1, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 8, 11, -1, -1, -1], [4, 5, 9, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 4, 5, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 2, 4, 5, 8, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[5, 7, 8, 9, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 5, 7, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 5, 7, 8, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 2, 5, 7, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 8, -1, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 9, 10, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[2, 3, 8, 9, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 7, 8, -1, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 4, 7, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 9, 10, -1, -1, -1], [4, 7, 8, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[2, 3, 4, 7, 9, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 5, 9, -1, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 8, -1, -1, -1, -1], [4, 5, 9, -1, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 4, 5, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[2, 3, 4, 5, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[5, 7, 8, 9, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 5, 7, 9, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 5, 7, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[2, 3, 5, 7, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 3, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 8, 10, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 9, 10, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[8, 9, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 7, 8, -1, -1, -1, -1], [1, 3, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 4, 7, 10, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 9, 10, 11, -1, -1], [4, 7, 8, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 7, 9, 10, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 5, 9, -1, -1, -1, -1], [1, 3, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 8, 10, 11, -1, -1], [4, 5, 9, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 4, 5, 10, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 5, 8, 10, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[5, 7, 8, 9, -1, -1, -1], [1, 3, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 5, 7, 9, 10, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 5, 7, 8, 10, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[5, 7, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 8, -1, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 9, -1, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 3, 8, 9, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 6, 8, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 4, 6, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 9, -1, -1, -1, -1], [4, 6, 8, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 3, 4, 6, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 5, 9, -1, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 8, -1, -1, -1, -1], [4, 5, 9, -1, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 4, 5, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 3, 4, 5, 8, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[5, 6, 8, 9, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 5, 6, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 5, 6, 8, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 3, 5, 6, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[2, 3, 6, 7, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 6, 7, 8, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 9, -1, -1, -1, -1], [2, 3, 6, 7, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 2, 6, 7, 8, 9, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[2, 3, 4, 6, 8, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 4, 6, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 9, -1, -1, -1, -1], [2, 3, 4, 6, 8, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 2, 4, 6, 9, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 5, 9, -1, -1, -1, -1], [2, 3, 6, 7, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 6, 7, 8, -1, -1], [4, 5, 9, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 4, 5, -1, -1, -1], [2, 3, 6, 7, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 2, 4, 5, 6, 7, 8], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[2, 3, 5, 6, 8, 9, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 5, 6, 9, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 2, 3, 5, 6, 8], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 2, 5, 6, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 2, 10, -1, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 8, -1, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 9, 10, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[2, 3, 8, 9, 10, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 6, 8, 11, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 4, 6, 11, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 9, 10, -1, -1, -1], [4, 6, 8, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[2, 3, 4, 6, 9, 10, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 5, 9, -1, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 8, -1, -1, -1, -1], [4, 5, 9, -1, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1]],
|
||||
[[0, 2, 4, 5, 10, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[2, 3, 4, 5, 8, 10, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[5, 6, 8, 9, 11, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 5, 6, 9, 11, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 5, 6, 8, 10, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[2, 3, 5, 6, 10, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 3, 6, 7, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 6, 7, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 6, 7, 9, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[6, 7, 8, 9, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 3, 4, 6, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 4, 6, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 4, 6, 8, 9, 10], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 6, 9, 10, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 5, 9, -1, -1, -1, -1], [1, 3, 6, 7, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 6, 7, 8, 10, -1], [4, 5, 9, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 4, 5, 6, 7, 10], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 5, 6, 7, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 3, 5, 6, 8, 9, 10], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 5, 6, 9, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 8, -1, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 8, -1, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 9, -1, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 3, 8, 9, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 7, 8, -1, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 4, 7, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 9, -1, -1, -1, -1], [4, 7, 8, -1, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 3, 4, 7, 9, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 6, 9, 10, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 8, -1, -1, -1, -1], [4, 6, 9, 10, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 4, 6, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 3, 4, 6, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[6, 7, 8, 9, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 6, 7, 9, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 6, 7, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 3, 6, 7, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[2, 3, 11, -1, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 8, 11, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 9, -1, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 2, 8, 9, 11, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 7, 8, -1, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 4, 7, 11, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 9, -1, -1, -1, -1], [4, 7, 8, -1, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1]],
|
||||
[[1, 2, 4, 7, 9, 11, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 6, 9, 10, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 8, 11, -1, -1, -1], [4, 6, 9, 10, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 4, 6, 10, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 2, 4, 6, 8, 10, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[6, 7, 8, 9, 10, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 6, 7, 9, 10, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 6, 7, 8, 10, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 2, 6, 7, 10, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 2, 5, 6, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 8, -1, -1, -1, -1], [1, 2, 5, 6, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 5, 6, 9, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[2, 3, 5, 6, 8, 9, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 7, 8, -1, -1, -1, -1], [1, 2, 5, 6, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 4, 7, -1, -1, -1], [1, 2, 5, 6, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 5, 6, 9, -1, -1], [4, 7, 8, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[2, 3, 4, 5, 6, 7, 9], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 2, 4, 6, 9, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 8, -1, -1, -1, -1], [1, 2, 4, 6, 9, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 4, 6, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[2, 3, 4, 6, 8, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 2, 6, 7, 8, 9, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 2, 3, 6, 7, 9], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 6, 7, 8, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[2, 3, 6, 7, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 3, 5, 6, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 5, 6, 8, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 5, 6, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[5, 6, 8, 9, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 7, 8, -1, -1, -1, -1], [1, 3, 5, 6, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 4, 5, 6, 7, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 5, 6, 9, 11, -1], [4, 7, 8, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 5, 6, 7, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 3, 4, 6, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 4, 6, 8, 9, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 4, 6, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 6, 8, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 3, 6, 7, 8, 9, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 9, -1, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 6, 7, 8, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[5, 7, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 8, -1, -1, -1, -1], [5, 7, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 9, -1, -1, -1, -1], [5, 7, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 3, 8, 9, -1, -1, -1], [5, 7, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 5, 8, 10, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 4, 5, 10, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 9, -1, -1, -1, -1], [4, 5, 8, 10, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 3, 4, 5, 9, 10, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 7, 9, 10, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 8, -1, -1, -1, -1], [4, 7, 9, 10, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 4, 7, 10, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 3, 4, 7, 8, 10, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[8, 9, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 9, 10, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 8, 10, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 3, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[2, 3, 5, 7, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 5, 7, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 9, -1, -1, -1, -1], [2, 3, 5, 7, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 2, 5, 7, 8, 9, 10], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[2, 3, 4, 5, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 4, 5, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 9, -1, -1, -1, -1], [2, 3, 4, 5, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 2, 4, 5, 9, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[2, 3, 4, 7, 9, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 4, 7, 8, 9, 10], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 2, 3, 4, 7, 10], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 7, 8, -1, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[2, 3, 8, 9, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 9, 10, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 2, 3, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 2, 5, 7, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 8, -1, -1, -1, -1], [1, 2, 5, 7, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 5, 7, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[2, 3, 5, 7, 8, 9, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 2, 4, 5, 8, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 2, 3, 4, 5, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 4, 5, 8, 9, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 5, 9, -1, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 2, 4, 7, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 8, -1, -1, -1, -1], [1, 2, 4, 7, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 4, 7, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[2, 3, 4, 7, 8, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 2, 8, 9, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 2, 3, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 2, 8, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 3, 5, 7, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 5, 7, 8, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 5, 7, 9, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[5, 7, 8, 9, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 3, 4, 5, 8, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 4, 5, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 4, 5, 8, 9, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 5, 9, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 3, 4, 7, 9, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 4, 7, 8, 9, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 4, 7, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[4, 7, 8, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[1, 3, 8, 9, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 1, 9, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[0, 3, 8, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
|
||||
[[-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]]
|
||||
]
|
||||
num_vd_table = [0, 1, 1, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 2, 1, 3, 1, 2, 2,
|
||||
2, 1, 2, 1, 2, 1, 1, 2, 1, 1, 2, 2, 2, 1, 2, 3, 1, 1, 2, 2, 1, 1, 1, 1, 1, 1, 2,
|
||||
1, 2, 1, 2, 2, 1, 1, 2, 1, 1, 1, 1, 2, 2, 2, 1, 1, 2, 1, 2, 3, 2, 2, 1, 1, 1, 1,
|
||||
1, 1, 2, 1, 1, 1, 2, 1, 2, 2, 2, 1, 1, 1, 1, 1, 2, 3, 2, 2, 2, 2, 2, 1, 3, 4, 2,
|
||||
2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 1, 1, 1, 1, 2, 1, 1, 2, 2, 2, 2, 2,
|
||||
3, 2, 1, 2, 1, 1, 1, 1, 1, 1, 2, 2, 3, 2, 3, 2, 4, 2, 2, 2, 2, 1, 2, 1, 2, 1, 1,
|
||||
2, 1, 1, 2, 2, 2, 1, 1, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 2, 1, 1, 1, 1, 1,
|
||||
1, 2, 1, 1, 1, 2, 2, 2, 1, 1, 2, 1, 1, 2, 1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 1, 2,
|
||||
1, 1, 1, 1, 2, 1, 1, 1, 1, 1, 2, 1, 1, 1, 1, 1, 2, 1, 2, 1, 1, 1, 1, 1, 1, 1, 1,
|
||||
1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0]
|
||||
check_table = [
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[1, 1, 0, 0, 194],
|
||||
[1, -1, 0, 0, 193],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[1, 0, 1, 0, 164],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[1, 0, -1, 0, 161],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 1, 152],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 1, 145],
|
||||
[1, 0, 0, 1, 144],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[1, 0, 0, -1, 137],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[1, 0, 1, 0, 133],
|
||||
[1, 0, 1, 0, 132],
|
||||
[1, 1, 0, 0, 131],
|
||||
[1, 1, 0, 0, 130],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 1, 100],
|
||||
[0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 1, 98],
|
||||
[0, 0, 0, 0, 0],
|
||||
[1, 0, 0, 1, 96],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[1, 0, 1, 0, 88],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[1, 0, -1, 0, 82],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[1, 0, 1, 0, 74],
|
||||
[0, 0, 0, 0, 0],
|
||||
[1, 0, 1, 0, 72],
|
||||
[0, 0, 0, 0, 0],
|
||||
[1, 0, 0, -1, 70],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[1, -1, 0, 0, 67],
|
||||
[0, 0, 0, 0, 0],
|
||||
[1, -1, 0, 0, 65],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[1, 1, 0, 0, 56],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[1, -1, 0, 0, 52],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[1, 1, 0, 0, 44],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[1, 1, 0, 0, 40],
|
||||
[0, 0, 0, 0, 0],
|
||||
[1, 0, 0, -1, 38],
|
||||
[1, 0, -1, 0, 37],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[1, 0, -1, 0, 33],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[1, -1, 0, 0, 28],
|
||||
[0, 0, 0, 0, 0],
|
||||
[1, 0, -1, 0, 26],
|
||||
[1, 0, 0, -1, 25],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[1, -1, 0, 0, 20],
|
||||
[0, 0, 0, 0, 0],
|
||||
[1, 0, -1, 0, 18],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[1, 0, 0, -1, 9],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[1, 0, 0, -1, 6],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0]
|
||||
]
|
||||
tet_table = [
|
||||
[-1, -1, -1, -1, -1, -1],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[4, 4, 4, 4, 4, 4],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[4, 0, 0, 4, 4, -1],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[4, 4, 4, 4, 4, 4],
|
||||
[0, 4, 0, 4, 4, -1],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[5, 5, 5, 5, 5, 5],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[2, 2, 2, 2, 2, 2],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[2, 0, 2, -1, 0, 2],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[2, -1, 2, 4, 4, 2],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[2, 0, 2, 4, 4, 2],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[2, 4, 2, 4, 4, 2],
|
||||
[0, 4, 0, 4, 4, 0],
|
||||
[2, 0, 2, 0, 0, 2],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[2, 5, 2, 5, 5, 2],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[2, 0, 2, 0, 0, 2],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[0, 1, 1, -1, 0, 1],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[2, 2, 2, 2, 2, 2],
|
||||
[4, 1, 1, 4, 4, 1],
|
||||
[0, 1, 1, 0, 0, 1],
|
||||
[4, 0, 0, 4, 4, 0],
|
||||
[2, 2, 2, 2, 2, 2],
|
||||
[-1, 1, 1, 4, 4, 1],
|
||||
[0, 1, 1, 4, 4, 1],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[2, 2, 2, 2, 2, 2],
|
||||
[5, 1, 1, 5, 5, 1],
|
||||
[0, 1, 1, 0, 0, 1],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[2, 2, 2, 2, 2, 2],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[8, 8, 8, 8, 8, 8],
|
||||
[1, 1, 1, 4, 4, 1],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[4, 0, 0, 4, 4, 0],
|
||||
[4, 4, 4, 4, 4, 4],
|
||||
[1, 1, 1, 4, 4, 1],
|
||||
[0, 4, 0, 4, 4, 0],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[4, 4, 4, 4, 4, 4],
|
||||
[1, 1, 1, 5, 5, 1],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[5, 5, 5, 5, 5, 5],
|
||||
[6, 6, 6, 6, 6, 6],
|
||||
[6, -1, 0, 6, 0, 6],
|
||||
[6, 0, 0, 6, 0, 6],
|
||||
[6, 1, 1, 6, 1, 6],
|
||||
[4, 4, 4, 4, 4, 4],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[4, 0, 0, 4, 4, 4],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[6, 4, -1, 6, 4, 6],
|
||||
[6, 4, 0, 6, 4, 6],
|
||||
[6, 0, 0, 6, 0, 6],
|
||||
[6, 1, 1, 6, 1, 6],
|
||||
[5, 5, 5, 5, 5, 5],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[2, 2, 2, 2, 2, 2],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[2, 0, 2, 2, 0, 2],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[2, 2, 2, 2, 2, 2],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[2, 0, 2, 2, 2, 2],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[2, 4, 2, 2, 4, 2],
|
||||
[0, 4, 0, 4, 4, 0],
|
||||
[2, 0, 2, 2, 0, 2],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[2, 2, 2, 2, 2, 2],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[6, 1, 1, 6, -1, 6],
|
||||
[6, 1, 1, 6, 0, 6],
|
||||
[6, 0, 0, 6, 0, 6],
|
||||
[6, 2, 2, 6, 2, 6],
|
||||
[4, 1, 1, 4, 4, 1],
|
||||
[0, 1, 1, 0, 0, 1],
|
||||
[4, 0, 0, 4, 4, 4],
|
||||
[2, 2, 2, 2, 2, 2],
|
||||
[6, 1, 1, 6, 4, 6],
|
||||
[6, 1, 1, 6, 4, 6],
|
||||
[6, 0, 0, 6, 0, 6],
|
||||
[6, 2, 2, 6, 2, 6],
|
||||
[5, 1, 1, 5, 5, 1],
|
||||
[0, 1, 1, 0, 0, 1],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[2, 2, 2, 2, 2, 2],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[6, 6, 6, 6, 6, 6],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[4, 4, 4, 4, 4, 4],
|
||||
[1, 1, 1, 1, 4, 1],
|
||||
[0, 4, 0, 4, 4, 0],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[4, 4, 4, 4, 4, 4],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[0, 5, 0, 5, 0, 5],
|
||||
[5, 5, 5, 5, 5, 5],
|
||||
[5, 5, 5, 5, 5, 5],
|
||||
[0, 5, 0, 5, 0, 5],
|
||||
[-1, 5, 0, 5, 0, 5],
|
||||
[1, 5, 1, 5, 1, 5],
|
||||
[4, 5, -1, 5, 4, 5],
|
||||
[0, 5, 0, 5, 0, 5],
|
||||
[4, 5, 0, 5, 4, 5],
|
||||
[1, 5, 1, 5, 1, 5],
|
||||
[4, 4, 4, 4, 4, 4],
|
||||
[0, 4, 0, 4, 4, 4],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[6, 6, 6, 6, 6, 6],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[2, 5, 2, 5, -1, 5],
|
||||
[0, 5, 0, 5, 0, 5],
|
||||
[2, 5, 2, 5, 0, 5],
|
||||
[1, 5, 1, 5, 1, 5],
|
||||
[2, 5, 2, 5, 4, 5],
|
||||
[0, 5, 0, 5, 0, 5],
|
||||
[2, 5, 2, 5, 4, 5],
|
||||
[1, 5, 1, 5, 1, 5],
|
||||
[2, 4, 2, 4, 4, 2],
|
||||
[0, 4, 0, 4, 4, 4],
|
||||
[2, 0, 2, 0, 0, 2],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[2, 6, 2, 6, 6, 2],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[2, 0, 2, 0, 0, 2],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[0, 1, 1, 1, 0, 1],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[2, 2, 2, 2, 2, 2],
|
||||
[4, 1, 1, 1, 4, 1],
|
||||
[0, 1, 1, 1, 0, 1],
|
||||
[4, 0, 0, 4, 4, 0],
|
||||
[2, 2, 2, 2, 2, 2],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[0, 1, 1, 1, 1, 1],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[2, 2, 2, 2, 2, 2],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[2, 2, 2, 2, 2, 2],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[5, 5, 5, 5, 5, 5],
|
||||
[1, 1, 1, 1, 4, 1],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[4, 0, 0, 4, 4, 0],
|
||||
[4, 4, 4, 4, 4, 4],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[4, 4, 4, 4, 4, 4],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[6, 0, 0, 6, 0, 6],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[6, 6, 6, 6, 6, 6],
|
||||
[5, 5, 5, 5, 5, 5],
|
||||
[5, 5, 0, 5, 0, 5],
|
||||
[5, 5, 0, 5, 0, 5],
|
||||
[5, 5, 1, 5, 1, 5],
|
||||
[4, 4, 4, 4, 4, 4],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[4, 4, 0, 4, 4, 4],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[4, 4, 4, 4, 4, 4],
|
||||
[4, 4, 0, 4, 4, 4],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[8, 8, 8, 8, 8, 8],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[2, 2, 2, 2, 2, 2],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[2, 2, 2, 2, 0, 2],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[2, 2, 2, 2, 2, 2],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[2, 2, 2, 2, 2, 2],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[2, 2, 2, 2, 2, 2],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[4, 1, 1, 4, 4, 1],
|
||||
[2, 2, 2, 2, 2, 2],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[1, 1, 1, 1, 0, 1],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[2, 2, 2, 2, 2, 2],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[2, 4, 2, 4, 4, 2],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[2, 2, 2, 2, 2, 2],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[2, 2, 2, 2, 2, 2],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[5, 5, 5, 5, 5, 5],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[4, 4, 4, 4, 4, 4],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[4, 4, 4, 4, 4, 4],
|
||||
[1, 1, 1, 1, 1, 1],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[0, 0, 0, 0, 0, 0],
|
||||
[12, 12, 12, 12, 12, 12]
|
||||
]
|
||||
@@ -0,0 +1,93 @@
|
||||
# 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 numpy as np
|
||||
import torch
|
||||
import nvdiffrast.torch as dr
|
||||
|
||||
###############################################################################
|
||||
# Functions adapted from https://github.com/NVlabs/nvdiffrec
|
||||
###############################################################################
|
||||
|
||||
def dot(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
return torch.sum(x*y, -1, keepdim=True)
|
||||
|
||||
def length(x: torch.Tensor, eps: float =1e-8) -> torch.Tensor:
|
||||
return torch.sqrt(torch.clamp(dot(x,x), min=eps)) # Clamp to avoid nan gradients because grad(sqrt(0)) = NaN
|
||||
|
||||
def safe_normalize(x: torch.Tensor, eps: float =1e-8) -> torch.Tensor:
|
||||
return x / length(x, eps)
|
||||
|
||||
def perspective(fovy=0.7854, aspect=1.0, n=0.1, f=1000.0, device=None):
|
||||
y = np.tan(fovy / 2)
|
||||
return torch.tensor([[1/(y*aspect), 0, 0, 0],
|
||||
[ 0, 1/-y, 0, 0],
|
||||
[ 0, 0, -(f+n)/(f-n), -(2*f*n)/(f-n)],
|
||||
[ 0, 0, -1, 0]], dtype=torch.float32, device=device)
|
||||
|
||||
def translate(x, y, z, device=None):
|
||||
return torch.tensor([[1, 0, 0, x],
|
||||
[0, 1, 0, y],
|
||||
[0, 0, 1, z],
|
||||
[0, 0, 0, 1]], dtype=torch.float32, device=device)
|
||||
|
||||
@torch.no_grad()
|
||||
def random_rotation_translation(t, device=None):
|
||||
m = np.random.normal(size=[3, 3])
|
||||
m[1] = np.cross(m[0], m[2])
|
||||
m[2] = np.cross(m[0], m[1])
|
||||
m = m / np.linalg.norm(m, axis=1, keepdims=True)
|
||||
m = np.pad(m, [[0, 1], [0, 1]], mode='constant')
|
||||
m[3, 3] = 1.0
|
||||
m[:3, 3] = np.random.uniform(-t, t, size=[3])
|
||||
return torch.tensor(m, dtype=torch.float32, device=device)
|
||||
|
||||
def rotate_x(a, device=None):
|
||||
s, c = np.sin(a), np.cos(a)
|
||||
return torch.tensor([[1, 0, 0, 0],
|
||||
[0, c, s, 0],
|
||||
[0, -s, c, 0],
|
||||
[0, 0, 0, 1]], dtype=torch.float32, device=device)
|
||||
|
||||
def rotate_y(a, device=None):
|
||||
s, c = np.sin(a), np.cos(a)
|
||||
return torch.tensor([[ c, 0, s, 0],
|
||||
[ 0, 1, 0, 0],
|
||||
[-s, 0, c, 0],
|
||||
[ 0, 0, 0, 1]], dtype=torch.float32, device=device)
|
||||
|
||||
class SimpleMesh:
|
||||
def __init__(self, vertices, faces):
|
||||
self.vertices = vertices
|
||||
self.faces = faces
|
||||
|
||||
def auto_normals(self):
|
||||
v0 = self.vertices[self.faces[:, 0], :]
|
||||
v1 = self.vertices[self.faces[:, 1], :]
|
||||
v2 = self.vertices[self.faces[:, 2], :]
|
||||
nrm = safe_normalize(torch.cross(v1 - v0, v2 - v0, dim=-1))
|
||||
self.nrm = nrm
|
||||
|
||||
def xfm_points(points, matrix):
|
||||
'''Transform points.
|
||||
Args:
|
||||
points: Tensor containing 3D points with shape [minibatch_size, num_vertices, 3] or [1, num_vertices, 3]
|
||||
matrix: A 4x4 transform matrix with shape [minibatch_size, 4, 4]
|
||||
use_python: Use PyTorch's torch.matmul (for validation)
|
||||
Returns:
|
||||
Transformed points in homogeneous 4D with shape [minibatch_size, num_vertices, 4].
|
||||
'''
|
||||
out = torch.matmul(
|
||||
torch.nn.functional.pad(points, pad=(0, 1), mode='constant', value=1.0), torch.transpose(matrix, 1, 2))
|
||||
if torch.is_anomaly_enabled():
|
||||
assert torch.all(torch.isfinite(out)), "Output of xfm_points contains inf or NaN"
|
||||
return out
|
||||
|
||||
def interpolate(attr, rast, attr_idx, rast_db=None):
|
||||
return dr.interpolate(
|
||||
attr, rast, attr_idx, rast_db=rast_db,
|
||||
diff_attrs=None if rast_db is None else 'all')
|
||||
@@ -0,0 +1,232 @@
|
||||
import random
|
||||
import tqdm
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from pytorch_msssim import SSIM, MS_SSIM
|
||||
|
||||
import comfy.utils
|
||||
|
||||
from .main_3DGS_renderer import GaussianSplattingRenderer
|
||||
|
||||
from shared_utils.camera_utils import BaseCameraController, MiniCam, calculate_fovX, get_projection_matrix
|
||||
from shared_utils.image_utils import prepare_torch_img
|
||||
|
||||
class GSParams:
|
||||
def __init__(
|
||||
self,
|
||||
training_iterations=30_000,
|
||||
batch_size=1,
|
||||
lambda_ssim=0.2,
|
||||
lambda_alpha=3,
|
||||
lambda_offset=0,
|
||||
lambda_offset_opacity=0,
|
||||
invert_bg_prob = 0.5,
|
||||
feature_lr=0.0025,
|
||||
opacity_lr=0.05,
|
||||
scaling_lr=0.005,
|
||||
rotation_lr=0.001,
|
||||
position_lr_init=0.00016,
|
||||
position_lr_final=0.0000016,
|
||||
position_lr_delay_mult=0.01,
|
||||
position_lr_max_steps=30_000,
|
||||
num_pts=10_000,
|
||||
K=3,
|
||||
percent_dense=0.01,
|
||||
density_start_iter=500,
|
||||
density_end_iter=15_000,
|
||||
densification_interval=100,
|
||||
opacity_reset_interval=3000,
|
||||
densify_grad_threshold=0.0002,
|
||||
sh_degree=3
|
||||
):
|
||||
|
||||
# training params
|
||||
self.training_iterations = training_iterations
|
||||
self.batch_size = batch_size
|
||||
self.lambda_ssim = lambda_ssim
|
||||
self.lambda_alpha = lambda_alpha
|
||||
self.lambda_offset = lambda_offset
|
||||
self.lambda_offset_opacity = lambda_offset_opacity
|
||||
self.invert_bg_prob = invert_bg_prob
|
||||
|
||||
# learning params
|
||||
self.feature_lr = feature_lr
|
||||
self.opacity_lr = opacity_lr
|
||||
self.scaling_lr = scaling_lr
|
||||
self.rotation_lr = rotation_lr
|
||||
self.position_lr_init = position_lr_init
|
||||
self.position_lr_final = position_lr_final
|
||||
self.position_lr_delay_mult = position_lr_delay_mult
|
||||
self.position_lr_max_steps = position_lr_max_steps
|
||||
|
||||
# densify and prune params
|
||||
self.num_pts = num_pts
|
||||
self.K = K
|
||||
self.percent_dense = percent_dense
|
||||
self.density_start_iter=density_start_iter
|
||||
self.density_end_iter=density_end_iter
|
||||
self.densification_interval=densification_interval
|
||||
self.opacity_reset_interval=opacity_reset_interval
|
||||
self.densify_grad_threshold=densify_grad_threshold
|
||||
|
||||
# other gaussian params
|
||||
self.sh_degree = sh_degree
|
||||
|
||||
class GaussianSplattingCameraController(BaseCameraController):
|
||||
def get_render_result(self, render_pose, bg_color, **kwargs):
|
||||
render_cam = MiniCam(render_pose, self.cam.W, self.cam.H, self.cam.fovy, self.cam.fovx, self.cam.near, self.cam.far, self.projection_matrix)
|
||||
return self.renderer.render(render_cam, bg_color=bg_color, **kwargs)
|
||||
|
||||
def post_init(self):
|
||||
self.projection_matrix = get_projection_matrix(self.cam.near, self.cam.far, self.cam.fovx, self.cam.fovy).transpose(0, 1).cuda()
|
||||
|
||||
class GaussianSplatting3D:
|
||||
|
||||
def __init__(self, gs_params=None, init_input=None, device='cuda'):
|
||||
self.device = torch.device(device)
|
||||
|
||||
# prepare renderer for optimization
|
||||
if gs_params is None:
|
||||
gs_params = GSParams()
|
||||
|
||||
self.renderer = GaussianSplattingRenderer(sh_degree=gs_params.sh_degree)
|
||||
self.renderer.initialize(init_input, num_pts=gs_params.num_pts)
|
||||
|
||||
# setup training
|
||||
self.renderer.gaussians.training_setup(gs_params)
|
||||
# do not do progressive sh-level
|
||||
self.renderer.gaussians.active_sh_degree = self.renderer.gaussians.max_sh_degree
|
||||
self.optimizer = self.renderer.gaussians.optimizer
|
||||
#self.ssim_loss = SSIM(data_range=1, size_average=True, channel=3)
|
||||
self.ms_ssim_loss = MS_SSIM(data_range=1, size_average=True, channel=3)
|
||||
|
||||
self.gs_params = gs_params
|
||||
|
||||
def prepare_training(self, reference_images, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy):
|
||||
self.ref_imgs_num = len(reference_images)
|
||||
|
||||
self.ref_size_H = reference_images[0].shape[0]
|
||||
self.ref_size_W = reference_images[0].shape[1]
|
||||
|
||||
# default camera settings
|
||||
self.cam_controller = GaussianSplattingCameraController(
|
||||
self.renderer, self.ref_size_W, self.ref_size_H, reference_orbit_camera_fovy, self.gs_params.invert_bg_prob, None, self.device
|
||||
)
|
||||
|
||||
self.all_ref_cam_poses = reference_orbit_camera_poses
|
||||
|
||||
# prepare reference images and masks
|
||||
ref_imgs_torch_list = []
|
||||
ref_masks_torch_list = []
|
||||
for i in range(self.ref_imgs_num):
|
||||
ref_imgs_torch_list.append(prepare_torch_img(reference_images[i].unsqueeze(0), self.ref_size_H, self.ref_size_W, self.device))
|
||||
ref_masks_torch_list.append(prepare_torch_img(reference_masks[i].unsqueeze(2).unsqueeze(0), self.ref_size_H, self.ref_size_W, self.device))
|
||||
|
||||
self.ref_imgs_torch = torch.cat(ref_imgs_torch_list, dim=0)
|
||||
self.ref_masks_torch = torch.cat(ref_masks_torch_list, dim=0)
|
||||
|
||||
def training(self):
|
||||
starter = torch.cuda.Event(enable_timing=True)
|
||||
ender = torch.cuda.Event(enable_timing=True)
|
||||
starter.record()
|
||||
|
||||
ref_imgs_masked = []
|
||||
for i in range(self.ref_imgs_num):
|
||||
ref_imgs_masked.append((self.ref_imgs_torch[i] * self.ref_masks_torch[i]).unsqueeze(0))
|
||||
|
||||
ref_imgs_num_minus_1 = self.ref_imgs_num-1
|
||||
|
||||
comfy_pbar = comfy.utils.ProgressBar(self.gs_params.training_iterations)
|
||||
|
||||
for step in tqdm.trange(self.gs_params.training_iterations):
|
||||
|
||||
### calculate loss between reference and rendered image from known view
|
||||
loss = 0
|
||||
masked_rendered_img_batch = []
|
||||
masked_ref_img_batch = []
|
||||
masks_batch = []
|
||||
ref_masks_batch = []
|
||||
|
||||
step_ratio = min(1, step / self.gs_params.training_iterations)
|
||||
|
||||
# update lr
|
||||
self.renderer.gaussians.update_learning_rate(step)
|
||||
|
||||
loss = 0
|
||||
|
||||
for _ in range(self.gs_params.batch_size):
|
||||
### calculate loss between reference and rendered image from known view
|
||||
|
||||
i = random.randint(0, ref_imgs_num_minus_1)
|
||||
|
||||
out = self.cam_controller.render_at_pose(self.all_ref_cam_poses[i])
|
||||
|
||||
image = out["image"] # [3, H, W] in [0, 1]
|
||||
mask = out["alpha"] # [1, H, W] in [0, 1]
|
||||
|
||||
ref_mask = self.ref_masks_torch[i]
|
||||
image_masked = (image * ref_mask).unsqueeze(0)
|
||||
|
||||
masked_rendered_img_batch.append(image_masked)
|
||||
masked_ref_img_batch.append(ref_imgs_masked[i])
|
||||
masks_batch.append(mask.unsqueeze(0))
|
||||
ref_masks_batch.append(ref_mask.unsqueeze(0))
|
||||
|
||||
masked_rendered_img_batch_torch = torch.cat(masked_rendered_img_batch, dim=0)
|
||||
masked_ref_img_batch_torch = torch.cat(masked_ref_img_batch, dim=0)
|
||||
masks_batch_torch = torch.cat(masks_batch, dim=0)
|
||||
ref_masks_batch_torch = torch.cat(ref_masks_batch, dim=0)
|
||||
|
||||
#loss_scaler = self.gs_params.loss_scale * step_ratio
|
||||
|
||||
# rgb loss
|
||||
loss += (1 - self.gs_params.lambda_ssim) * F.l1_loss(masked_rendered_img_batch_torch, masked_ref_img_batch_torch)
|
||||
|
||||
# alpha loss
|
||||
loss += self.gs_params.lambda_alpha * F.mse_loss(masks_batch_torch, ref_masks_batch_torch)
|
||||
|
||||
# D-SSIM loss
|
||||
# [1, 3, H, W] in [0, 1]
|
||||
#loss += self.lambda_ssim * (1 - self.ssim_loss(X, Y))
|
||||
loss += self.gs_params.lambda_ssim * (1 - self.ms_ssim_loss(masked_ref_img_batch_torch, masked_rendered_img_batch_torch))
|
||||
|
||||
if self.gs_params.lambda_offset > 0:
|
||||
# Reference offset loss
|
||||
offset_norm = self.renderer.gaussians.get_xyz_offset.norm(dim=-1, keepdim=True)
|
||||
loss += self.gs_params.lambda_offset * torch.mean(offset_norm)
|
||||
|
||||
if self.gs_params.lambda_offset_opacity > 0:
|
||||
# Alpha penalty loss
|
||||
loss += self.gs_params.lambda_offset_opacity * torch.mean(offset_norm.detach() * self.renderer.gaussians.get_opacity)
|
||||
|
||||
|
||||
# optimize step
|
||||
loss.backward()
|
||||
self.optimizer.step()
|
||||
self.optimizer.zero_grad()
|
||||
|
||||
# densify and prune
|
||||
if step >= self.gs_params.density_start_iter and step <= self.gs_params.density_end_iter:
|
||||
viewspace_point_tensor, visibility_filter, radii = out["viewspace_points"], out["visibility_filter"], out["radii"]
|
||||
self.renderer.gaussians.max_radii2D[visibility_filter] = torch.max(self.renderer.gaussians.max_radii2D[visibility_filter], radii[visibility_filter])
|
||||
self.renderer.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter)
|
||||
|
||||
if step % self.gs_params.densification_interval == 0:
|
||||
self.renderer.gaussians.densify_and_prune(self.gs_params.densify_grad_threshold, min_opacity=0.005, extent=4, max_screen_size=1)
|
||||
#self.renderer.gaussians.densify_and_prune_by_compatness(self.gs_params.K, min_opacity=0.01, extent=4, max_screen_size=1)
|
||||
|
||||
#self.renderer.gaussians.densify_by_clone_and_split(self.gs_params.densify_grad_threshold, extent=4)
|
||||
#self.renderer.gaussians.densify_by_compatness(self.gs_params.K)
|
||||
#self.renderer.gaussians.prune(min_opacity=0.01, extent=4, max_screen_size=1, max_offset=0.1)
|
||||
|
||||
if step % self.gs_params.opacity_reset_interval == 0:
|
||||
self.renderer.gaussians.reset_opacity()
|
||||
|
||||
comfy_pbar.update_absolute(step + 1)
|
||||
|
||||
ender.record()
|
||||
torch.cuda.synchronize()
|
||||
#t = starter.elapsed_time(ender)
|
||||
|
||||
self.need_update = True
|
||||
@@ -0,0 +1,982 @@
|
||||
import os
|
||||
import math
|
||||
import random
|
||||
import numpy as np
|
||||
from typing import Optional
|
||||
from torchtyping import TensorType
|
||||
from plyfile import PlyData, PlyElement
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
pytorch3d_capable = True
|
||||
try:
|
||||
import pytorch3d
|
||||
from pytorch3d.ops import knn_points
|
||||
except ImportError:
|
||||
pytorch3d_capable = False
|
||||
|
||||
from kornia.geometry.conversions import (
|
||||
quaternion_to_rotation_matrix,
|
||||
)
|
||||
|
||||
from diff_gaussian_rasterization import (
|
||||
GaussianRasterizationSettings,
|
||||
GaussianRasterizer,
|
||||
)
|
||||
from simple_knn._C import distCUDA2
|
||||
|
||||
from kiui.op import inverse_sigmoid
|
||||
|
||||
from shared_utils.sh_utils import eval_sh, SH2RGB, RGB2SH
|
||||
from mesh_processer.mesh import Mesh, PointCloud
|
||||
from mesh_processer.mesh_utils import construct_list_of_gs_attributes, write_gs_ply, read_gs_ply
|
||||
|
||||
def get_expon_lr_func(
|
||||
lr_init, lr_final, lr_delay_steps=0, lr_delay_mult=1.0, max_steps=1000000
|
||||
):
|
||||
|
||||
def helper(step):
|
||||
if lr_init == lr_final:
|
||||
# constant lr, ignore other params
|
||||
return lr_init
|
||||
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 gaussian_3d_coeff(xyzs, covs):
|
||||
# xyzs: [N, 3]
|
||||
# covs: [N, 6]
|
||||
x, y, z = xyzs[:, 0], xyzs[:, 1], xyzs[:, 2]
|
||||
a, b, c, d, e, f = covs[:, 0], covs[:, 1], covs[:, 2], covs[:, 3], covs[:, 4], covs[:, 5]
|
||||
|
||||
# eps must be small enough !!!
|
||||
inv_det = 1 / (a * d * f + 2 * e * c * b - e**2 * a - c**2 * d - b**2 * f + 1e-24)
|
||||
inv_a = (d * f - e**2) * inv_det
|
||||
inv_b = (e * c - b * f) * inv_det
|
||||
inv_c = (e * b - c * d) * inv_det
|
||||
inv_d = (a * f - c**2) * inv_det
|
||||
inv_e = (b * c - e * a) * inv_det
|
||||
inv_f = (a * d - b**2) * inv_det
|
||||
|
||||
power = -0.5 * (x**2 * inv_a + y**2 * inv_d + z**2 * inv_f) - x * y * inv_b - x * z * inv_c - y * z * inv_e
|
||||
|
||||
power[power > 0] = -1e10 # abnormal values... make weights 0
|
||||
|
||||
return torch.exp(power)
|
||||
|
||||
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 random_point_in_triangle(v0, v1, v2):
|
||||
"""
|
||||
Given three vertices v0, v1, v2, sample a point uniformly in the triangle
|
||||
Algorithm Reference:
|
||||
https://math.stackexchange.com/questions/538458/how-to-sample-points-on-a-triangle-surface-in-3d
|
||||
https://stackoverflow.com/questions/47410054/generate-random-locations-within-a-triangular-domain
|
||||
"""
|
||||
|
||||
r1 = random.random()
|
||||
r2 = random.random()
|
||||
|
||||
s1 = math.sqrt(r1)
|
||||
return v0 * (1.0 - s1) + v1 * (1.0 - r2) * s1 + v2 * r2 * s1#
|
||||
|
||||
def find_points_within_radius(query_points, vertex_points, d):
|
||||
"""
|
||||
Finds vertex points within a given radius for each query point.
|
||||
|
||||
Args:
|
||||
query_points (torch.Tensor): Tensor of shape (P1, 3) representing query points.
|
||||
vertex_points (torch.Tensor): Tensor of shape (P2, 3) representing vertex points.
|
||||
d (float): Radius within which to search for vertex points.
|
||||
|
||||
Returns:
|
||||
List[List[int]]: A list of lists, where each inner list contains indices of vertex points
|
||||
within the radius for the corresponding query point.
|
||||
"""
|
||||
num_query_points = query_points.shape[0]
|
||||
|
||||
# Calculate pairwise distances between query points and vertex points
|
||||
distances = torch.norm(query_points[:, None] - vertex_points, dim=2)
|
||||
|
||||
# Create a mask indicating which vertex points are within the radius
|
||||
mask = distances <= d
|
||||
|
||||
# Collect indices of vertex points within the radius for each query point
|
||||
result = []
|
||||
for i in range(num_query_points):
|
||||
indices_within_radius = torch.nonzero(mask[i]).squeeze().tolist()
|
||||
result.append(indices_within_radius)
|
||||
|
||||
return result
|
||||
|
||||
@torch.no_grad()
|
||||
def K_nearest_neighbors(
|
||||
points: torch.Tensor,
|
||||
K: int,
|
||||
query: Optional[torch.Tensor] = None,
|
||||
return_dist=False,
|
||||
):
|
||||
if not pytorch3d_capable:
|
||||
raise ImportError("pytorch3d is not installed, which is required for KNN")
|
||||
|
||||
# query/points: Tensor of shape (N, P1/P2, D) giving a batch of N point clouds, each containing up to P1/P2 points of dimension D
|
||||
if query is None:
|
||||
query = points
|
||||
dist, idx, nn = knn_points(query[None, ...], points[None, ...], K=K, return_nn=True)
|
||||
|
||||
# idx: Tensor of shape (N, P1, K)
|
||||
# nn: Tensor of shape (N, P1, K, D)
|
||||
|
||||
# take the index 1 since index 0 is the point itself
|
||||
if not return_dist:
|
||||
return nn[0, :, 1:, :], idx[0, :, 1:]
|
||||
else:
|
||||
return nn[0, :, 1:, :], idx[0, :, 1:], dist[0, :, 1:]
|
||||
|
||||
|
||||
def distance_to_gaussian_surface(points, svec, rotmat, query):
|
||||
"""
|
||||
Calculate the radius of the gaussian along the direction, which determined by offset between points and query
|
||||
Calculation using Mahalanobis distance
|
||||
|
||||
Args:
|
||||
points (Tensor): gaussians position of shape (*, 3)
|
||||
svec (Tensor): gaussians scale vector of shape (*, 3)
|
||||
rotmat : the rotation matrix of shape (*, 3, 3)
|
||||
query (Tensor): query positions of shape (*, 3)
|
||||
|
||||
Returns:
|
||||
Distance of shape (*, 1):
|
||||
"""
|
||||
offset_dir = query - points
|
||||
|
||||
offset_dir = torch.einsum("bij,bj->bi", rotmat.transpose(-1, -2), offset_dir)
|
||||
offset_dir = F.normalize(offset_dir, dim=-1)
|
||||
z = offset_dir[..., 2]
|
||||
y = offset_dir[..., 1]
|
||||
x = offset_dir[..., 0]
|
||||
r_xy = torch.sqrt(x**2 + y**2 + 1e-10)
|
||||
cos_theta = z
|
||||
sin_theta = r_xy
|
||||
cos_phi = x / r_xy
|
||||
sin_phi = y / r_xy
|
||||
|
||||
d2 = svec[..., 0] ** 2 * cos_phi**2 + svec[..., 1] ** 2 * sin_phi**2
|
||||
#r2 = svec[..., 2] ** 2 * cos_theta**2 + d2**2 * sin_theta**2
|
||||
r2 = svec[..., 2] ** 2 * cos_theta**2 + d2 * sin_theta**2 # same as: squared_dist = np.dot(diff.T, np.dot(np.linalg.inv(covariance), diff))
|
||||
|
||||
"""
|
||||
Alternatively:
|
||||
|
||||
def gaussian_squared_distance(query_point, mean, covariance):
|
||||
# Calculate the difference vector
|
||||
diff = query_point - mean
|
||||
|
||||
# Compute the squared Mahalanobis distance
|
||||
inv_covariance = np.linalg.inv(covariance)
|
||||
squared_dist = np.dot(diff.T, np.dot(inv_covariance, diff))
|
||||
|
||||
return squared_dist
|
||||
|
||||
# Example usage
|
||||
mean = np.array([x0, y0, z0]) # Gaussian mean (origin position)
|
||||
covariance = np.diag([sigma_x**2, sigma_y**2, sigma_z**2]) # Covariance matrix
|
||||
|
||||
query_point = np.array([x_q, y_q, z_q]) # Query point
|
||||
squared_dist_to_gaussian = gaussian_squared_distance(query_point, mean, covariance)
|
||||
|
||||
print(f"Squared distance to Gaussian surface: {squared_dist_to_gaussian:.4f}")
|
||||
"""
|
||||
|
||||
return torch.sqrt(r2 + 1e-10)
|
||||
|
||||
def qvec2rotmat_batched(qvec: TensorType["N", 4]):
|
||||
return quaternion_to_rotation_matrix(qvec)
|
||||
|
||||
class GaussianModel:
|
||||
|
||||
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
|
||||
|
||||
self.scaling_activation = torch.exp
|
||||
self.scaling_inverse_activation = torch.log
|
||||
|
||||
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
|
||||
|
||||
|
||||
def __init__(self, sh_degree : int):
|
||||
self.active_sh_degree = 0
|
||||
self.max_sh_degree = sh_degree
|
||||
self.init_xyz = torch.empty(0)
|
||||
self._xyz = torch.empty(0)
|
||||
self._features_dc = torch.empty(0) # diffuse color
|
||||
self._features_rest = torch.empty(0) # spherical harmonic coefficients
|
||||
self._scaling = torch.empty(0)
|
||||
self._rotation = torch.empty(0) # (w, x, y, z)
|
||||
self._opacity = torch.empty(0)
|
||||
self.max_radii2D = torch.empty(0)
|
||||
self.xyz_gradient_accum = torch.empty(0)
|
||||
self.denom = torch.empty(0)
|
||||
self.optimizer = None
|
||||
self.percent_dense = 0
|
||||
self.spatial_lr_scale = 0
|
||||
self.setup_functions()
|
||||
|
||||
def capture(self):
|
||||
return (
|
||||
self.active_sh_degree,
|
||||
self._xyz,
|
||||
self._features_dc,
|
||||
self._features_rest,
|
||||
self._scaling,
|
||||
self._rotation,
|
||||
self._opacity,
|
||||
self.max_radii2D,
|
||||
self.xyz_gradient_accum,
|
||||
self.denom,
|
||||
self.optimizer.state_dict(),
|
||||
self.spatial_lr_scale,
|
||||
)
|
||||
|
||||
def restore(self, model_args, training_args):
|
||||
(self.active_sh_degree,
|
||||
self._xyz,
|
||||
self._features_dc,
|
||||
self._features_rest,
|
||||
self._scaling,
|
||||
self._rotation,
|
||||
self._opacity,
|
||||
self.max_radii2D,
|
||||
xyz_gradient_accum,
|
||||
denom,
|
||||
opt_dict,
|
||||
self.spatial_lr_scale) = model_args
|
||||
self.training_setup(training_args)
|
||||
self.xyz_gradient_accum = xyz_gradient_accum
|
||||
self.denom = denom
|
||||
self.optimizer.load_state_dict(opt_dict)
|
||||
|
||||
@property
|
||||
def get_gaussians_num(self):
|
||||
return self.init_xyz.shape[0]
|
||||
|
||||
@property
|
||||
def get_scaling(self):
|
||||
return self.scaling_activation(self._scaling)
|
||||
|
||||
@property
|
||||
def get_rotation(self):
|
||||
return self.rotation_activation(self._rotation)
|
||||
|
||||
@property
|
||||
def rotmat(self):
|
||||
return qvec2rotmat_batched(self._rotation)
|
||||
|
||||
@property
|
||||
def get_xyz_offset(self):
|
||||
return self.init_xyz - self._xyz
|
||||
|
||||
@property
|
||||
def get_xyz(self):
|
||||
return self._xyz
|
||||
|
||||
@property
|
||||
def get_features(self):
|
||||
features_dc = self._features_dc
|
||||
features_rest = self._features_rest
|
||||
return torch.cat((features_dc, features_rest), dim=1)
|
||||
|
||||
@property
|
||||
def get_opacity(self):
|
||||
return self.opacity_activation(self._opacity)
|
||||
|
||||
@torch.no_grad()
|
||||
def extract_fields(self, resolution=128, num_blocks=16, relax_ratio=1.5):
|
||||
# resolution: resolution of field
|
||||
|
||||
block_size = 2 / num_blocks
|
||||
|
||||
assert resolution % block_size == 0
|
||||
split_size = resolution // num_blocks
|
||||
|
||||
opacities = self.get_opacity
|
||||
|
||||
# pre-filter low opacity gaussians to save computation
|
||||
mask = (opacities > 0.005).squeeze(1)
|
||||
|
||||
opacities = opacities[mask]
|
||||
xyzs = self.get_xyz[mask]
|
||||
stds = self.get_scaling[mask]
|
||||
|
||||
# normalize to ~ [-1, 1]
|
||||
mn, mx = xyzs.amin(0), xyzs.amax(0)
|
||||
self.center = (mn + mx) / 2
|
||||
self.scale = 1.8 / (mx - mn).amax().item()
|
||||
|
||||
xyzs = (xyzs - self.center) * self.scale
|
||||
stds = stds * self.scale
|
||||
|
||||
covs = self.covariance_activation(stds, 1, self._rotation[mask])
|
||||
|
||||
# tile
|
||||
device = opacities.device
|
||||
occ = torch.zeros([resolution] * 3, dtype=torch.float32, device=device)
|
||||
|
||||
X = torch.linspace(-1, 1, resolution).split(split_size)
|
||||
Y = torch.linspace(-1, 1, resolution).split(split_size)
|
||||
Z = torch.linspace(-1, 1, resolution).split(split_size)
|
||||
|
||||
|
||||
# loop blocks (assume max size of gaussian is small than relax_ratio * block_size !!!)
|
||||
for xi, xs in enumerate(X):
|
||||
for yi, ys in enumerate(Y):
|
||||
for zi, zs in enumerate(Z):
|
||||
xx, yy, zz = torch.meshgrid(xs, ys, zs)
|
||||
# sample points [M, 3]
|
||||
pts = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1).to(device)
|
||||
# in-tile gaussians mask
|
||||
vmin, vmax = pts.amin(0), pts.amax(0)
|
||||
vmin -= block_size * relax_ratio
|
||||
vmax += block_size * relax_ratio
|
||||
mask = (xyzs < vmax).all(-1) & (xyzs > vmin).all(-1)
|
||||
# if hit no gaussian, continue to next block
|
||||
if not mask.any():
|
||||
continue
|
||||
mask_xyzs = xyzs[mask] # [L, 3]
|
||||
mask_covs = covs[mask] # [L, 6]
|
||||
mask_opas = opacities[mask].view(1, -1) # [L, 1] --> [1, L]
|
||||
|
||||
# query per point-gaussian pair.
|
||||
g_pts = pts.unsqueeze(1).repeat(1, mask_covs.shape[0], 1) - mask_xyzs.unsqueeze(0) # [M, L, 3]
|
||||
g_covs = mask_covs.unsqueeze(0).repeat(pts.shape[0], 1, 1) # [M, L, 6]
|
||||
|
||||
# batch on gaussian to avoid OOM
|
||||
batch_g = 1024
|
||||
val = 0
|
||||
for start in range(0, g_covs.shape[1], batch_g):
|
||||
end = min(start + batch_g, g_covs.shape[1])
|
||||
w = gaussian_3d_coeff(g_pts[:, start:end].reshape(-1, 3), g_covs[:, start:end].reshape(-1, 6)).reshape(pts.shape[0], -1) # [M, l]
|
||||
val += (mask_opas[:, start:end] * w).sum(-1)
|
||||
|
||||
occ[xi * split_size: xi * split_size + len(xs),
|
||||
yi * split_size: yi * split_size + len(ys),
|
||||
zi * split_size: zi * split_size + len(zs)] = val.reshape(len(xs), len(ys), len(zs))
|
||||
|
||||
return occ
|
||||
|
||||
def get_covariance(self, scaling_modifier = 1, gaussain_idx = None):
|
||||
if gaussain_idx is None:
|
||||
return self.covariance_activation(self.get_scaling, scaling_modifier, self._rotation)
|
||||
else:
|
||||
return self.covariance_activation(self.get_scaling[gaussain_idx], scaling_modifier, self._rotation[gaussain_idx])
|
||||
|
||||
def oneupSHdegree(self):
|
||||
if self.active_sh_degree < self.max_sh_degree:
|
||||
self.active_sh_degree += 1
|
||||
|
||||
def create_from_pcd(self, pcd : PointCloud, spatial_lr_scale : float = 1):
|
||||
self.spatial_lr_scale = spatial_lr_scale
|
||||
fused_point_cloud = torch.tensor(np.asarray(pcd.points)).float().cuda()
|
||||
fused_color = RGB2SH(torch.tensor(np.asarray(pcd.colors)).float().cuda())
|
||||
features = torch.zeros((fused_color.shape[0], 3, (self.max_sh_degree + 1) ** 2)).float().cuda()
|
||||
features[:, :3, 0 ] = fused_color
|
||||
features[:, 3:, 1:] = 0.0
|
||||
|
||||
print("Number of points at initialisation : ", fused_point_cloud.shape[0])
|
||||
|
||||
dist2 = torch.clamp_min(distCUDA2(torch.from_numpy(np.asarray(pcd.points)).float().cuda()), 0.0000001)
|
||||
scales = torch.log(torch.sqrt(dist2))[...,None].repeat(1, 3)
|
||||
rots = torch.zeros((fused_point_cloud.shape[0], 4), device="cuda")
|
||||
rots[:, 0] = 1
|
||||
|
||||
opacities = inverse_sigmoid(0.1 * torch.ones((fused_point_cloud.shape[0], 1), dtype=torch.float, device="cuda"))
|
||||
|
||||
self.init_xyz = fused_point_cloud.clone().detach()
|
||||
self._xyz = nn.Parameter(fused_point_cloud.requires_grad_(True))
|
||||
self._features_dc = nn.Parameter(features[:,:,0:1].transpose(1, 2).contiguous().requires_grad_(True))
|
||||
self._features_rest = nn.Parameter(features[:,:,1:].transpose(1, 2).contiguous().requires_grad_(True))
|
||||
self._scaling = nn.Parameter(scales.requires_grad_(True))
|
||||
self._rotation = nn.Parameter(rots.requires_grad_(True))
|
||||
self._opacity = nn.Parameter(opacities.requires_grad_(True))
|
||||
self.max_radii2D = torch.zeros((self.get_xyz.shape[0]), device="cuda")
|
||||
|
||||
def training_setup(self, training_args):
|
||||
self.percent_dense = training_args.percent_dense
|
||||
self.xyz_gradient_accum = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
|
||||
self.denom = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
|
||||
|
||||
l = [
|
||||
{'params': [self._xyz], 'lr': training_args.position_lr_init * self.spatial_lr_scale, "name": "xyz"},
|
||||
{'params': [self._features_dc], 'lr': training_args.feature_lr, "name": "f_dc"},
|
||||
{'params': [self._features_rest], 'lr': training_args.feature_lr / 20.0, "name": "f_rest"},
|
||||
{'params': [self._opacity], 'lr': training_args.opacity_lr, "name": "opacity"},
|
||||
{'params': [self._scaling], 'lr': training_args.scaling_lr, "name": "scaling"},
|
||||
{'params': [self._rotation], 'lr': training_args.rotation_lr, "name": "rotation"}
|
||||
]
|
||||
|
||||
self.optimizer = torch.optim.Adam(l, lr=0.0, eps=1e-15)
|
||||
self.xyz_scheduler_args = get_expon_lr_func(lr_init=training_args.position_lr_init*self.spatial_lr_scale,
|
||||
lr_final=training_args.position_lr_final*self.spatial_lr_scale,
|
||||
lr_delay_mult=training_args.position_lr_delay_mult,
|
||||
max_steps=training_args.position_lr_max_steps)
|
||||
|
||||
def update_learning_rate(self, iteration):
|
||||
''' Learning rate scheduling per step '''
|
||||
for param_group in self.optimizer.param_groups:
|
||||
if param_group["name"] == "xyz":
|
||||
lr = self.xyz_scheduler_args(iteration)
|
||||
param_group['lr'] = lr
|
||||
return lr
|
||||
|
||||
def reset_opacity(self):
|
||||
opacities_new = inverse_sigmoid(torch.min(self.get_opacity, torch.ones_like(self.get_opacity)*0.01))
|
||||
optimizable_tensors = self.replace_tensor_to_optimizer(opacities_new, "opacity")
|
||||
self._opacity = optimizable_tensors["opacity"]
|
||||
|
||||
def get_points_cloud(self):
|
||||
xyz = self._xyz.detach().cpu().numpy()
|
||||
normals = np.zeros_like(xyz)
|
||||
shs = self.get_features.detach().cpu().numpy()
|
||||
pcd = PointCloud(points=xyz, colors=SH2RGB(shs), normals=normals)
|
||||
return pcd
|
||||
|
||||
def to_ply(self):
|
||||
xyz = self._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()
|
||||
f_rest = self._features_rest.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
|
||||
opacities = self._opacity.detach().cpu().numpy()
|
||||
scale = self._scaling.detach().cpu().numpy()
|
||||
rotation = self._rotation.detach().cpu().numpy()
|
||||
|
||||
return write_gs_ply(xyz, normals, f_dc, f_rest, opacities, scale, rotation, construct_list_of_gs_attributes(self._features_dc, self._features_rest, self._scaling, self._rotation))
|
||||
|
||||
def create_from_ply(self, plydata):
|
||||
xyz, features_dc, features_extra, opacities, scales, rots = read_gs_ply(plydata)
|
||||
|
||||
self._xyz = nn.Parameter(torch.tensor(xyz, dtype=torch.float, device="cuda").requires_grad_(True))
|
||||
self.init_xyz = self._xyz.clone().detach()
|
||||
self._features_dc = nn.Parameter(torch.tensor(features_dc, dtype=torch.float, device="cuda").transpose(1, 2).contiguous().requires_grad_(True))
|
||||
self._features_rest = nn.Parameter(torch.tensor(features_extra, dtype=torch.float, device="cuda").transpose(1, 2).contiguous().requires_grad_(True))
|
||||
self._opacity = nn.Parameter(torch.tensor(opacities, dtype=torch.float, device="cuda").requires_grad_(True))
|
||||
self._scaling = nn.Parameter(torch.tensor(scales, dtype=torch.float, device="cuda").requires_grad_(True))
|
||||
self._rotation = nn.Parameter(torch.tensor(rots, dtype=torch.float, device="cuda").requires_grad_(True))
|
||||
self.max_radii2D = torch.zeros((self.get_xyz.shape[0]), device="cuda")
|
||||
|
||||
self.active_sh_degree = self.max_sh_degree
|
||||
|
||||
def create_from_mesh(self, mesh, num_pts):
|
||||
"""
|
||||
Sample gaussians uniformly for each face, this method can capture the structure of the 3D meshes well even for the fine structure e.g. fingers
|
||||
"""
|
||||
# In .obj file, first index is 1 instead of zero
|
||||
vertices = np.concatenate([[[0., 0., 0.]], mesh.v.detach().cpu().numpy()])
|
||||
#normals = np.concatenate([[[0., 0., 0.]], mesh.n.detach().cpu().numpy()])
|
||||
|
||||
all_points = []
|
||||
num_pts_per_face = math.ceil(num_pts / mesh.f.shape[0])
|
||||
for triangle_face in mesh.f.detach().cpu().numpy():
|
||||
v0 = vertices[triangle_face[0]]
|
||||
v1 = vertices[triangle_face[1]]
|
||||
v2 = vertices[triangle_face[2]]
|
||||
for i in range(num_pts_per_face):
|
||||
all_points.append(random_point_in_triangle(v0, v1, v2))
|
||||
|
||||
num_pts = len(all_points)
|
||||
|
||||
xyz = np.array(all_points)
|
||||
shs = np.random.random((num_pts, 3)) / 255.0
|
||||
pcd = PointCloud(
|
||||
points=xyz, colors=SH2RGB(shs), normals=np.zeros((num_pts, 3))
|
||||
)
|
||||
self.create_from_pcd(pcd, 10)
|
||||
|
||||
def create_from_uv_data(self, uv_grids_data):
|
||||
"""
|
||||
Idealy this method should capture the most information for the texture,
|
||||
but depends on the UV of 3D meshes, fine structure like fingers could have too little gaussians to form the proper structure,
|
||||
also the number of gaussians is too big even for 1K texture
|
||||
"""
|
||||
if uv_grids_data is not None:
|
||||
uv_coords, uv_coords_3d, uv_normals_3d = uv_grids_data
|
||||
|
||||
num_pts = len(uv_coords_3d)
|
||||
xyz = np.array(uv_coords_3d)
|
||||
shs = np.random.random((num_pts, 3)) / 255.0
|
||||
pcd = PointCloud(
|
||||
points=xyz, colors=SH2RGB(shs), normals=np.zeros((num_pts, 3))
|
||||
)
|
||||
self.create_from_pcd(pcd, 10)
|
||||
|
||||
def replace_tensor_to_optimizer(self, tensor, name):
|
||||
optimizable_tensors = {}
|
||||
for group in self.optimizer.param_groups:
|
||||
if group["name"] == name:
|
||||
stored_state = self.optimizer.state.get(group['params'][0], None)
|
||||
stored_state["exp_avg"] = torch.zeros_like(tensor)
|
||||
stored_state["exp_avg_sq"] = torch.zeros_like(tensor)
|
||||
|
||||
del self.optimizer.state[group['params'][0]]
|
||||
group["params"][0] = nn.Parameter(tensor.requires_grad_(True))
|
||||
self.optimizer.state[group['params'][0]] = stored_state
|
||||
|
||||
optimizable_tensors[group["name"]] = group["params"][0]
|
||||
return optimizable_tensors
|
||||
|
||||
def _prune_optimizer(self, mask):
|
||||
optimizable_tensors = {}
|
||||
for group in self.optimizer.param_groups:
|
||||
stored_state = self.optimizer.state.get(group['params'][0], None)
|
||||
if stored_state is not None:
|
||||
stored_state["exp_avg"] = stored_state["exp_avg"][mask]
|
||||
stored_state["exp_avg_sq"] = stored_state["exp_avg_sq"][mask]
|
||||
|
||||
del self.optimizer.state[group['params'][0]]
|
||||
group["params"][0] = nn.Parameter((group["params"][0][mask].requires_grad_(True)))
|
||||
self.optimizer.state[group['params'][0]] = stored_state
|
||||
|
||||
optimizable_tensors[group["name"]] = group["params"][0]
|
||||
else:
|
||||
group["params"][0] = nn.Parameter(group["params"][0][mask].requires_grad_(True))
|
||||
optimizable_tensors[group["name"]] = group["params"][0]
|
||||
return optimizable_tensors
|
||||
|
||||
def prune_points(self, mask):
|
||||
valid_points_mask = ~mask
|
||||
optimizable_tensors = self._prune_optimizer(valid_points_mask)
|
||||
|
||||
self.init_xyz = self.init_xyz[valid_points_mask]
|
||||
|
||||
self._xyz = optimizable_tensors["xyz"]
|
||||
self._features_dc = optimizable_tensors["f_dc"]
|
||||
self._features_rest = optimizable_tensors["f_rest"]
|
||||
self._opacity = optimizable_tensors["opacity"]
|
||||
self._scaling = optimizable_tensors["scaling"]
|
||||
self._rotation = optimizable_tensors["rotation"]
|
||||
|
||||
self.xyz_gradient_accum = self.xyz_gradient_accum[valid_points_mask]
|
||||
|
||||
self.denom = self.denom[valid_points_mask]
|
||||
self.max_radii2D = self.max_radii2D[valid_points_mask]
|
||||
|
||||
def cat_tensors_to_optimizer(self, tensors_dict):
|
||||
optimizable_tensors = {}
|
||||
for group in self.optimizer.param_groups:
|
||||
assert len(group["params"]) == 1
|
||||
extension_tensor = tensors_dict[group["name"]]
|
||||
stored_state = self.optimizer.state.get(group['params'][0], None)
|
||||
if stored_state is not None:
|
||||
|
||||
stored_state["exp_avg"] = torch.cat((stored_state["exp_avg"], torch.zeros_like(extension_tensor)), dim=0)
|
||||
stored_state["exp_avg_sq"] = torch.cat((stored_state["exp_avg_sq"], torch.zeros_like(extension_tensor)), dim=0)
|
||||
|
||||
del self.optimizer.state[group['params'][0]]
|
||||
group["params"][0] = nn.Parameter(torch.cat((group["params"][0], extension_tensor), dim=0).requires_grad_(True))
|
||||
self.optimizer.state[group['params'][0]] = stored_state
|
||||
|
||||
optimizable_tensors[group["name"]] = group["params"][0]
|
||||
else:
|
||||
group["params"][0] = nn.Parameter(torch.cat((group["params"][0], extension_tensor), dim=0).requires_grad_(True))
|
||||
optimizable_tensors[group["name"]] = group["params"][0]
|
||||
|
||||
return optimizable_tensors
|
||||
|
||||
def densification_postfix(self, new_xyz, new_features_dc, new_features_rest, new_opacities, new_scaling, new_rotation):
|
||||
new_params = {
|
||||
"xyz": new_xyz,
|
||||
"f_dc": new_features_dc,
|
||||
"f_rest": new_features_rest,
|
||||
"opacity": new_opacities,
|
||||
"scaling" : new_scaling,
|
||||
"rotation" : new_rotation
|
||||
}
|
||||
|
||||
self.densify_with_new_params(new_params)
|
||||
|
||||
def densify_with_new_params(self, new_params):
|
||||
optimizable_tensors = self.cat_tensors_to_optimizer(new_params)
|
||||
self._xyz = optimizable_tensors["xyz"]
|
||||
self._features_dc = optimizable_tensors["f_dc"]
|
||||
self._features_rest = optimizable_tensors["f_rest"]
|
||||
self._opacity = optimizable_tensors["opacity"]
|
||||
self._scaling = optimizable_tensors["scaling"]
|
||||
self._rotation = optimizable_tensors["rotation"]
|
||||
|
||||
self.xyz_gradient_accum = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
|
||||
self.denom = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
|
||||
self.max_radii2D = torch.zeros((self.get_xyz.shape[0]), device="cuda")
|
||||
|
||||
def densify_and_split(self, grads, grad_threshold, scene_extent, N=2):
|
||||
n_init_points = self.get_xyz.shape[0]
|
||||
# Extract points that satisfy the gradient condition
|
||||
padded_grad = torch.zeros((n_init_points), device="cuda")
|
||||
padded_grad[:grads.shape[0]] = grads.squeeze()
|
||||
selected_pts_mask = torch.where(padded_grad >= grad_threshold, True, False)
|
||||
selected_pts_mask = torch.logical_and(selected_pts_mask,
|
||||
torch.max(self.get_scaling, dim=1).values > self.percent_dense*scene_extent
|
||||
)
|
||||
|
||||
stds = self.get_scaling[selected_pts_mask].repeat(N,1)
|
||||
means =torch.zeros((stds.size(0), 3),device="cuda")
|
||||
samples = torch.normal(mean=means, std=stds)
|
||||
rots = build_rotation(self._rotation[selected_pts_mask]).repeat(N,1,1)
|
||||
new_xyz = torch.bmm(rots, samples.unsqueeze(-1)).squeeze(-1) + self.get_xyz[selected_pts_mask].repeat(N, 1)
|
||||
new_scaling = self.scaling_inverse_activation(self.get_scaling[selected_pts_mask].repeat(N,1) / (0.8*N))
|
||||
new_rotation = self._rotation[selected_pts_mask].repeat(N,1)
|
||||
new_features_dc = self._features_dc[selected_pts_mask].repeat(N,1,1)
|
||||
new_features_rest = self._features_rest[selected_pts_mask].repeat(N,1,1)
|
||||
new_opacity = self._opacity[selected_pts_mask].repeat(N,1)
|
||||
|
||||
# Update reference position for gaussians
|
||||
self.init_xyz = torch.cat((self.init_xyz, self.init_xyz[selected_pts_mask].repeat(N,1)), dim=0)
|
||||
|
||||
self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_opacity, new_scaling, new_rotation)
|
||||
|
||||
prune_filter = torch.cat((selected_pts_mask, torch.zeros(N * selected_pts_mask.sum(), device="cuda", dtype=bool)))
|
||||
self.prune_points(prune_filter)
|
||||
|
||||
def densify_and_clone(self, grads, grad_threshold, scene_extent):
|
||||
# Extract points that satisfy the gradient condition
|
||||
selected_pts_mask = torch.where(torch.norm(grads, dim=-1) >= grad_threshold, True, False)
|
||||
selected_pts_mask = torch.logical_and(selected_pts_mask,
|
||||
torch.max(self.get_scaling, dim=1).values <= self.percent_dense*scene_extent
|
||||
)
|
||||
|
||||
new_xyz = self._xyz[selected_pts_mask]
|
||||
new_features_dc = self._features_dc[selected_pts_mask]
|
||||
new_features_rest = self._features_rest[selected_pts_mask]
|
||||
new_opacities = self._opacity[selected_pts_mask]
|
||||
new_scaling = self._scaling[selected_pts_mask]
|
||||
new_rotation = self._rotation[selected_pts_mask]
|
||||
|
||||
# Update reference position for gaussians
|
||||
self.init_xyz = torch.cat((self.init_xyz, self.init_xyz[selected_pts_mask]), dim=0)
|
||||
|
||||
self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_opacities, new_scaling, new_rotation)
|
||||
|
||||
def densify_by_compatnes_with_idx(self, idx):
|
||||
nn_scaling = self._scaling[idx]
|
||||
nn_rotmat = self.rotmat[idx]
|
||||
nn_pos = self._xyz[idx]
|
||||
|
||||
nn_gaussian_surface_dist = distance_to_gaussian_surface(
|
||||
nn_pos, nn_scaling, nn_rotmat, self._xyz
|
||||
)
|
||||
gaussian_surface_dist = distance_to_gaussian_surface(
|
||||
self._xyz, self._scaling, self.rotmat, nn_pos
|
||||
)
|
||||
|
||||
dist_to_nn = torch.norm(nn_pos - self._xyz, dim=-1)
|
||||
selected_pts_mask = (gaussian_surface_dist + nn_gaussian_surface_dist) < dist_to_nn
|
||||
new_direction = (nn_pos - self._xyz) / dist_to_nn[..., None]
|
||||
new_xyz = (
|
||||
self._xyz
|
||||
+ new_direction * (dist_to_nn + gaussian_surface_dist - nn_gaussian_surface_dist)[..., None]
|
||||
/ 2.0
|
||||
)[selected_pts_mask]
|
||||
new_features_dc = self._features_dc[selected_pts_mask]
|
||||
new_features_rest = self._features_rest[selected_pts_mask]
|
||||
new_raw_alpha = self._opacity[selected_pts_mask]
|
||||
new_raw_svec = self.scaling_inverse_activation(
|
||||
torch.ones_like(self._scaling[selected_pts_mask])
|
||||
* (dist_to_nn - gaussian_surface_dist - nn_gaussian_surface_dist)[selected_pts_mask][..., None]
|
||||
/ 6.0
|
||||
)
|
||||
# print(torch.ones_like(self.svec.data[selected_pts_mask]).shape)
|
||||
# print(
|
||||
# (dist_to_nn - gaussian_surface_dist - nn_gaussian_surface_dist)[selected_pts_mask].shape
|
||||
# )
|
||||
new_qvec = self._rotation[selected_pts_mask]
|
||||
|
||||
# Update reference position for gaussians
|
||||
self.init_xyz = torch.cat((self.init_xyz, self.init_xyz[selected_pts_mask]), dim=0)
|
||||
|
||||
new_params = {
|
||||
"xyz": new_xyz,
|
||||
"f_dc": new_features_dc,
|
||||
"f_rest": new_features_rest,
|
||||
"opacity": new_raw_alpha,
|
||||
"scaling": new_raw_svec,
|
||||
"rotation": new_qvec
|
||||
}
|
||||
|
||||
return new_params
|
||||
|
||||
def densify_by_compatness(self, K=1):
|
||||
_, idx = K_nearest_neighbors(self._xyz, K=K+1)
|
||||
|
||||
new_params_list = []
|
||||
for i in range(K):
|
||||
new_params = self.densify_by_compatnes_with_idx(idx[:, i])
|
||||
new_params_list.append(new_params)
|
||||
new_params = {}
|
||||
for key in new_params_list[0].keys():
|
||||
new_params[key] = torch.cat([p[key] for p in new_params_list], dim=0)
|
||||
|
||||
self.densify_with_new_params(new_params)
|
||||
|
||||
#num_densified = new_params["xyz"].shape[0]
|
||||
|
||||
def densify_by_clone_and_split(self, max_grad, extent):
|
||||
grads = self.xyz_gradient_accum / self.denom
|
||||
grads[grads.isnan()] = 0.0
|
||||
|
||||
self.densify_and_clone(grads, max_grad, extent)
|
||||
self.densify_and_split(grads, max_grad, extent)
|
||||
|
||||
def densify_and_prune(self, max_grad, min_opacity, extent, max_screen_size):
|
||||
self.densify_by_clone_and_split(max_grad, extent)
|
||||
self.prune(min_opacity, extent, max_screen_size)
|
||||
|
||||
def densify_and_prune_by_compatness(self, K, min_opacity, extent, max_screen_size):
|
||||
self.densify_by_compatness(K)
|
||||
self.prune(min_opacity, extent, max_screen_size)
|
||||
|
||||
def add_densification_stats(self, viewspace_point_tensor, update_filter):
|
||||
self.xyz_gradient_accum[update_filter] += torch.norm(viewspace_point_tensor.grad[update_filter,:2], dim=-1, keepdim=True)
|
||||
self.denom[update_filter] += 1
|
||||
|
||||
def prune(self, min_opacity, extent, max_screen_size, max_offset=0):
|
||||
prune_mask = (self.get_opacity < min_opacity).squeeze()
|
||||
if max_screen_size:
|
||||
big_points_vs = self.max_radii2D > max_screen_size
|
||||
big_points_ws = self.get_scaling.max(dim=1).values > 0.1 * extent
|
||||
#big_points_os = self.get_xyz_offset.norm(dim=1) > max_offset
|
||||
#prune_mask = torch.logical_or(torch.logical_or(torch.logical_or(prune_mask, big_points_vs), big_points_ws), big_points_os)
|
||||
prune_mask = torch.logical_or(torch.logical_or(prune_mask, big_points_vs), big_points_ws)
|
||||
self.prune_points(prune_mask)
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
class GaussianSplattingRenderer:
|
||||
def __init__(self, sh_degree=3, white_background=True, radius=1):
|
||||
|
||||
self.sh_degree = sh_degree
|
||||
self.white_background = white_background
|
||||
self.radius = radius
|
||||
|
||||
self.gaussians = GaussianModel(sh_degree)
|
||||
|
||||
self.bg_color = torch.tensor(
|
||||
[1, 1, 1] if white_background else [0, 0, 0],
|
||||
dtype=torch.float32,
|
||||
device="cuda",
|
||||
)
|
||||
|
||||
def initialize(self, input=None, num_pts=5000, radius=0.5):
|
||||
# load checkpoint
|
||||
if isinstance(input, Mesh):
|
||||
# load from 3D mesh
|
||||
self.gaussians.create_from_mesh(input, num_pts)
|
||||
elif isinstance(input, PlyData):
|
||||
self.gaussians.create_from_ply(input)
|
||||
elif isinstance(input, PointCloud):
|
||||
# load from a provided pcd
|
||||
self.gaussians.create_from_pcd(input, 1)
|
||||
elif input is None:
|
||||
# init from random point cloud
|
||||
|
||||
phis = np.random.random((num_pts,)) * 2 * np.pi
|
||||
costheta = np.random.random((num_pts,)) * 2 - 1
|
||||
thetas = np.arccos(costheta)
|
||||
mu = np.random.random((num_pts,))
|
||||
radius = radius * np.cbrt(mu)
|
||||
x = radius * np.sin(thetas) * np.cos(phis)
|
||||
y = radius * np.sin(thetas) * np.sin(phis)
|
||||
z = radius * np.cos(thetas)
|
||||
xyz = np.stack((x, y, z), axis=1)
|
||||
# xyz = np.random.random((num_pts, 3)) * 2.6 - 1.3
|
||||
|
||||
shs = np.random.random((num_pts, 3)) / 255.0
|
||||
pcd = PointCloud(
|
||||
points=xyz, colors=SH2RGB(shs), normals=np.zeros((num_pts, 3))
|
||||
)
|
||||
self.gaussians.create_from_pcd(pcd, 10)
|
||||
else:
|
||||
self.gaussians.create_from_uv_data(input)
|
||||
|
||||
def render(
|
||||
self,
|
||||
viewpoint_camera,
|
||||
scaling_modifier=1.0,
|
||||
gaussain_idx=None,
|
||||
bg_color=None,
|
||||
override_color=None,
|
||||
compute_cov3D_python=False,
|
||||
convert_SHs_python=False,
|
||||
):
|
||||
# Set up rasterization configuration
|
||||
tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)
|
||||
tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)
|
||||
|
||||
raster_settings = GaussianRasterizationSettings(
|
||||
image_height=int(viewpoint_camera.image_height),
|
||||
image_width=int(viewpoint_camera.image_width),
|
||||
tanfovx=tanfovx,
|
||||
tanfovy=tanfovy,
|
||||
bg=self.bg_color if bg_color is None else bg_color,
|
||||
scale_modifier=scaling_modifier,
|
||||
viewmatrix=viewpoint_camera.world_view_transform,
|
||||
projmatrix=viewpoint_camera.full_proj_transform,
|
||||
sh_degree=self.gaussians.active_sh_degree,
|
||||
campos=viewpoint_camera.camera_center,
|
||||
prefiltered=False,
|
||||
debug=False,
|
||||
)
|
||||
|
||||
rasterizer = GaussianRasterizer(raster_settings=raster_settings)
|
||||
|
||||
gaussians_xyz = self.gaussians.get_xyz
|
||||
gaussians_features = self.gaussians.get_features
|
||||
gaussians_opacity = self.gaussians.get_opacity
|
||||
|
||||
# If precomputed 3d covariance is provided, use it. If not, then it will be computed from
|
||||
# scaling / rotation by the rasterizer.
|
||||
gaussians_scales = None
|
||||
gaussians_rotations = None
|
||||
cov3D_precomp = None
|
||||
if compute_cov3D_python:
|
||||
cov3D_precomp = self.gaussians.get_covariance(scaling_modifier, gaussain_idx)
|
||||
else:
|
||||
gaussians_scales = self.gaussians.get_scaling
|
||||
gaussians_rotations = self.gaussians.get_rotation
|
||||
|
||||
if gaussain_idx is not None:
|
||||
gaussians_xyz = gaussians_xyz[gaussain_idx]
|
||||
gaussians_features = gaussians_features[gaussain_idx]
|
||||
gaussians_opacity = gaussians_opacity[gaussain_idx]
|
||||
if cov3D_precomp is None:
|
||||
gaussians_scales = gaussians_scales[gaussain_idx]
|
||||
gaussians_rotations = gaussians_rotations[gaussain_idx]
|
||||
|
||||
# Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
|
||||
screenspace_points = (
|
||||
torch.zeros_like(
|
||||
gaussians_xyz,
|
||||
dtype=gaussians_xyz.dtype,
|
||||
requires_grad=True,
|
||||
device="cuda",
|
||||
)
|
||||
+ 0
|
||||
)
|
||||
try:
|
||||
screenspace_points.retain_grad()
|
||||
except:
|
||||
pass
|
||||
|
||||
# If precomputed colors are provided, use them. Otherwise, if it is desired to precompute colors
|
||||
# from SHs in Python, do it. If not, then SH -> RGB conversion will be done by rasterizer.
|
||||
shs = None
|
||||
colors_precomp = None
|
||||
if override_color is None:
|
||||
if convert_SHs_python:
|
||||
shs_view = self.gaussians_features.transpose(1, 2).view(
|
||||
-1, 3, (self.gaussians.max_sh_degree + 1) ** 2
|
||||
)
|
||||
dir_pp = gaussians_xyz - viewpoint_camera.camera_center.repeat(
|
||||
gaussians_features.shape[0], 1
|
||||
)
|
||||
dir_pp_normalized = dir_pp / dir_pp.norm(dim=1, keepdim=True)
|
||||
sh2rgb = eval_sh(
|
||||
self.gaussians.active_sh_degree, shs_view, dir_pp_normalized
|
||||
)
|
||||
colors_precomp = torch.clamp_min(sh2rgb + 0.5, 0.0)
|
||||
else:
|
||||
shs = gaussians_features
|
||||
else:
|
||||
colors_precomp = override_color
|
||||
|
||||
# Rasterize visible Gaussians to image, obtain their radii (on screen).
|
||||
rendered_image, radii, rendered_depth, rendered_alpha = rasterizer(
|
||||
means3D=gaussians_xyz,
|
||||
means2D=screenspace_points,
|
||||
shs=shs,
|
||||
colors_precomp=colors_precomp,
|
||||
opacities=gaussians_opacity,
|
||||
scales=gaussians_scales,
|
||||
rotations=gaussians_rotations,
|
||||
cov3D_precomp=cov3D_precomp,
|
||||
)
|
||||
|
||||
rendered_image = rendered_image.clamp(0, 1)
|
||||
|
||||
# Those Gaussians that were frustum culled or had a radius of 0 were not visible.
|
||||
# They will be excluded from value updates used in the splitting criteria.
|
||||
return {
|
||||
"image": rendered_image, # rgb
|
||||
"depth": rendered_depth,
|
||||
"alpha": rendered_alpha,
|
||||
"viewspace_points": screenspace_points,
|
||||
"visibility_filter": radii > 0,
|
||||
"radii": radii,
|
||||
}
|
||||
@@ -0,0 +1,205 @@
|
||||
import tqdm
|
||||
import random
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import nerfacc
|
||||
|
||||
import comfy.utils
|
||||
|
||||
from pytorch_msssim import SSIM, MS_SSIM
|
||||
|
||||
from kiui.op import safe_normalize
|
||||
from kiui.cam import orbit_camera
|
||||
from kiui.nn import MLP, trunc_exp
|
||||
|
||||
from shared_utils.image_utils import prepare_torch_img
|
||||
|
||||
class InstantNGP(nn.Module):
|
||||
def __init__(self, resolution=128, device="cuda"):
|
||||
super().__init__()
|
||||
from kiui.gridencoder import GridEncoder
|
||||
|
||||
self.device = torch.device(device)
|
||||
self.ref_size_H = resolution
|
||||
self.ref_size_W = resolution
|
||||
|
||||
self.render_step_size = 5e-3
|
||||
self.aabb = torch.tensor([-1.0, -1.0, -1.0, 1.0, 1.0, 1.0], device=self.device)
|
||||
self.estimator = nerfacc.OccGridEstimator(roi_aabb=self.aabb, resolution=64, levels=1)
|
||||
|
||||
self.encoder_density = GridEncoder(num_levels=12) # VMEncoder(output_dim=16, mode='sum')
|
||||
self.encoder = GridEncoder(num_levels=12)
|
||||
self.mlp_density = MLP(self.encoder_density.output_dim, 1, 32, 2, bias=False)
|
||||
self.mlp = MLP(self.encoder.output_dim, 3, 32, 2, bias=False)
|
||||
|
||||
def get_rays(self, pose, h, w, fovy, opengl=True):
|
||||
|
||||
x, y = torch.meshgrid(
|
||||
torch.arange(w, device=pose.device),
|
||||
torch.arange(h, device=pose.device),
|
||||
indexing="xy",
|
||||
)
|
||||
x = x.flatten()
|
||||
y = y.flatten()
|
||||
|
||||
cx = w * 0.5
|
||||
cy = h * 0.5
|
||||
|
||||
focal = h * 0.5 / np.tan(0.5 * np.deg2rad(fovy))
|
||||
|
||||
camera_dirs = F.pad(
|
||||
torch.stack(
|
||||
[
|
||||
(x - cx + 0.5) / focal,
|
||||
(y - cy + 0.5) / focal * (-1.0 if opengl else 1.0),
|
||||
],
|
||||
dim=-1,
|
||||
),
|
||||
(0, 1),
|
||||
value=(-1.0 if opengl else 1.0),
|
||||
) # [hw, 3]
|
||||
|
||||
rays_d = camera_dirs @ pose[:3, :3].transpose(0, 1) # [hw, 3]
|
||||
rays_o = pose[:3, 3].unsqueeze(0).expand_as(rays_d) # [hw, 3]
|
||||
|
||||
rays_o = rays_o.view(h, w, 3)
|
||||
rays_d = safe_normalize(rays_d).view(h, w, 3)
|
||||
|
||||
return rays_o, rays_d
|
||||
|
||||
def get_color(self, xs):
|
||||
return torch.sigmoid(self.mlp(self.encoder(xs.to(self.device))))
|
||||
|
||||
def get_density(self, xs):
|
||||
# xs: [..., 3]
|
||||
xs = xs.to(self.device)
|
||||
prefix = xs.shape[:-1]
|
||||
xs = xs.view(-1, 3)
|
||||
feats = self.encoder_density(xs)
|
||||
density = trunc_exp(self.mlp_density(feats))
|
||||
density = density.view(*prefix, 1)
|
||||
return density
|
||||
|
||||
def prepare_training(self, reference_images, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy):
|
||||
self.ref_imgs_num = len(reference_images)
|
||||
|
||||
self.all_ref_cam_poses = reference_orbit_camera_poses
|
||||
self.ref_cam_fovy = reference_orbit_camera_fovy
|
||||
|
||||
# prepare reference images and masks
|
||||
ref_imgs_torch_list = []
|
||||
ref_masks_torch_list = []
|
||||
for i in range(self.ref_imgs_num):
|
||||
ref_imgs_torch_list.append(prepare_torch_img(reference_images[i].unsqueeze(0), self.ref_size_H, self.ref_size_W, self.device))
|
||||
ref_masks_torch_list.append(prepare_torch_img(reference_masks[i].unsqueeze(2).unsqueeze(0), self.ref_size_H, self.ref_size_W, self.device))
|
||||
|
||||
self.ref_imgs_torch = torch.cat(ref_imgs_torch_list, dim=0) # [N, 3, H, W]
|
||||
self.ref_masks_torch = torch.cat(ref_masks_torch_list, dim=0).squeeze(1) # [N, H, W]
|
||||
|
||||
def render_nerf(self, pose, bg_color=1):
|
||||
|
||||
pose = torch.from_numpy(pose.astype(np.float32)).to(self.device)
|
||||
|
||||
# get rays
|
||||
rays_o, rays_d = self.get_rays(pose, self.ref_size_H, self.ref_size_W, self.ref_cam_fovy)
|
||||
hw = rays_o.shape[0] * rays_o.shape[1]
|
||||
rays_o = rays_o.view(hw, 3)
|
||||
rays_d = rays_d.view(hw, 3)
|
||||
|
||||
# update occ grid
|
||||
if self.training:
|
||||
def occ_eval_fn(xs):
|
||||
sigmas = self.get_density(xs)
|
||||
return self.render_step_size * sigmas
|
||||
|
||||
self.estimator.update_every_n_steps(self.render_step, occ_eval_fn=occ_eval_fn, occ_thre=0.01, n=8)
|
||||
self.render_step += 1
|
||||
|
||||
# render
|
||||
def sigma_fn(t_starts, t_ends, ray_indices):
|
||||
t_origins = rays_o[ray_indices]
|
||||
t_dirs = rays_d[ray_indices]
|
||||
xs = t_origins + t_dirs * (t_starts + t_ends)[:, None] / 2.0
|
||||
sigmas = self.get_density(xs)
|
||||
return sigmas.squeeze(-1)
|
||||
|
||||
with torch.no_grad():
|
||||
ray_indices, t_starts, t_ends = self.estimator.sampling(
|
||||
rays_o,
|
||||
rays_d,
|
||||
sigma_fn=sigma_fn,
|
||||
near_plane=0.01,
|
||||
far_plane=100,
|
||||
render_step_size=self.render_step_size,
|
||||
stratified=self.training,
|
||||
cone_angle=0,
|
||||
)
|
||||
|
||||
t_origins = rays_o[ray_indices]
|
||||
t_dirs = rays_d[ray_indices]
|
||||
xs = t_origins + t_dirs * (t_starts + t_ends)[:, None] / 2.0
|
||||
sigmas = self.get_density(xs).squeeze(-1)
|
||||
rgbs = torch.sigmoid(self.mlp(self.encoder(xs)))
|
||||
|
||||
n_rays=rays_o.shape[0]
|
||||
weights, t, alphas = nerfacc.render_weight_from_density(t_starts, t_ends, sigmas, ray_indices=ray_indices, n_rays=n_rays)
|
||||
color = nerfacc.accumulate_along_rays(weights, values=rgbs, ray_indices=ray_indices, n_rays=n_rays)
|
||||
alpha = nerfacc.accumulate_along_rays(weights, values=None, ray_indices=ray_indices, n_rays=n_rays)
|
||||
|
||||
color = color + (1.0 - alpha) * bg_color
|
||||
|
||||
color = color.view(self.ref_size_H, self.ref_size_W, 3).clamp(0, 1).permute(2, 0, 1).contiguous()
|
||||
alpha = alpha.view(self.ref_size_H, self.ref_size_W).clamp(0, 1).contiguous()
|
||||
|
||||
return color, alpha
|
||||
|
||||
def fit_nerf(self, iters=512, bg_color=1):
|
||||
|
||||
optimizer = torch.optim.Adam([
|
||||
{'params': self.encoder_density.parameters(), 'lr': 1e-2},
|
||||
{'params': self.encoder.parameters(), 'lr': 1e-2},
|
||||
{'params': self.mlp_density.parameters(), 'lr': 1e-3},
|
||||
{'params': self.mlp.parameters(), 'lr': 1e-3},
|
||||
])
|
||||
|
||||
print(f"[INFO] fitting nerf...")
|
||||
self.render_step = 0
|
||||
|
||||
ref_imgs_num_minus_1 = self.ref_imgs_num-1
|
||||
|
||||
comfy_pbar = comfy.utils.ProgressBar(iters)
|
||||
pbar = tqdm.trange(iters)
|
||||
for step in pbar:
|
||||
|
||||
i = random.randint(0, ref_imgs_num_minus_1)
|
||||
|
||||
radius, elevation, azimuth, center_X, center_Y, center_Z = self.all_ref_cam_poses[i]
|
||||
|
||||
orbit_target = np.array([center_X, center_Y, center_Z], dtype=np.float32)
|
||||
pose = orbit_camera(elevation, azimuth, radius, target=orbit_target)
|
||||
|
||||
image_gt = self.ref_imgs_torch[i] # [3, H, W]
|
||||
alpha_gt = self.ref_masks_torch[i] # [H, W]
|
||||
image_pred, alpha_pred = self.render_nerf(pose, bg_color)
|
||||
|
||||
# if i % 200 == 0:
|
||||
# kiui.vis.plot_image(image_gt, alpha_gt, image_pred, alpha_pred)
|
||||
|
||||
loss_mse = F.mse_loss(image_pred, image_gt) + 0.1 * F.mse_loss(alpha_pred, alpha_gt)
|
||||
loss = loss_mse #+ 0.1 * self.encoder_density.tv_loss() #+ 0.0001 * self.encoder_density.density_loss()
|
||||
#loss += self.lambda_ssim * (1 - self.ms_ssim_loss(image_gt, image_pred))
|
||||
|
||||
loss.backward()
|
||||
self.encoder_density.grad_total_variation(1e-8)
|
||||
|
||||
optimizer.step()
|
||||
optimizer.zero_grad()
|
||||
|
||||
pbar.set_description(f"NeRF Fitting Loss = {loss_mse.item():.6f}")
|
||||
comfy_pbar.update_absolute(step + 1)
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
print(f"[INFO] finished fitting nerf!")
|
||||
Reference in New Issue
Block a user