nodes/MrForExample-ComfyUI-3D-Pack
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add 3D Visualization & Switch Axis for 3DGS & 3D Mesh & Gaussian Splatting Orbit Renderer

Browse Source
This commit is contained in:
MrForExample
2024-02-04 19:31:31 +01:00
parent d9a79e4039
commit 16c685d98f
32 changed files with 3952 additions and 2665 deletions
Download Patch File Download Diff File Expand all files Collapse all files
algorithms/_Test.py
+5 -4
View File
@@ -1,8 +1,9 @@
import os
import numpy as np
import torch
grid_res = 128
p = os.path.join(os.path.dirname(__file__), f'../data/tets/{grid_res}_tets.npz')
print(p)
vector3_tensor = torch.tensor([[1, 2, 3], [3, 2, 1], [2, 3, 1]]) # shape (N, 3)
target_axis = (2, 0, 1) # or [2, 0, 1]
vector3_tensor[:, [0, 1, 2]] = vector3_tensor[:, target_axis]
print(vector3_tensor)
algorithms/diff_mesh_renderer.py
+1 -1
View File
@@ -37,7 +37,7 @@ def trunc_rev_sigmoid(x, eps=1e-6):
def make_divisible(x, m=8):
return int(math.ceil(x / m) * m)
class Renderer(nn.Module):
class DiffRastRenderer(nn.Module):
def __init__(self, mesh, force_cuda_rast):
super().__init__()
algorithms/diff_texturing.py
+5 -10
View File
@@ -8,8 +8,9 @@ from pytorch_msssim import SSIM, MS_SSIM
import comfy.utils
from .diff_mesh_renderer import Renderer
from .diff_mesh_renderer import DiffRastRenderer
from ..shared_utils.camera_utils import orbit_camera, OrbitCamera
from ..shared_utils.image_utils import prepare_torch_img
class DiffTextureBaker:
@@ -17,7 +18,7 @@ class DiffTextureBaker:
self.device = torch.device("cuda")
# prepare main components for optimization
self.renderer = Renderer(mesh, force_cuda_rasterize).to(self.device)
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)
@@ -28,12 +29,6 @@ class DiffTextureBaker:
self.batch_size = batch_size
def prepare_img(self, img):
# (H, W, C) -> (C, H, W)
img_new = img.permute(2, 0, 1).unsqueeze(0).to(self.device)
img_new = F.interpolate(img_new, (self.ref_size_H, self.ref_size_W), mode="bilinear", align_corners=False).contiguous()
return img_new
def prepare_training(self, reference_images, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy):
self.ref_imgs_num = len(reference_images)
@@ -49,8 +44,8 @@ class DiffTextureBaker:
ref_imgs_torch_list = []
ref_masks_torch_list = []
for i in range(self.ref_imgs_num):
ref_imgs_torch_list.append(self.prepare_img(reference_images[i]))
ref_masks_torch_list.append(self.prepare_img(reference_masks[i].unsqueeze(2)))
ref_imgs_torch_list.append(prepare_torch_img(reference_images[i], self.ref_size_H, self.ref_size_W, self.device))
ref_masks_torch_list.append(prepare_torch_img(reference_masks[i].unsqueeze(2), 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)
algorithms/dmtet.py
+1 -1
View File
@@ -229,7 +229,7 @@ class DMTetMesh:
def get_mesh(self):
verts, faces = self.pred_verts_faces()
mesh = Mesh(v=verts, f=faces)
mesh = Mesh(v=verts.detach(), f=faces.detach(), device=self.device)
mesh.auto_normal()
mesh.auto_uv()
return mesh
algorithms/main_3DGS.py
+52 -24
View File
@@ -8,8 +8,9 @@ from pytorch_msssim import SSIM, MS_SSIM
import comfy.utils
from .main_3DGS_renderer import Renderer
from .main_3DGS_renderer import GaussianSplattingRenderer
from ..shared_utils.camera_utils import orbit_camera, OrbitCamera, MiniCam, calculate_fovX, get_projection_matrix
from ..shared_utils.image_utils import prepare_torch_img
class GSParams:
def __init__(self,
@@ -72,6 +73,46 @@ class GSParams:
# other gaussian params
self.sh_degree = sh_degree
class GaussianSplattingCameraController:
def __init__(self, renderer, cam_size_W, cam_size_H, reference_orbit_camera_fovy, invert_bg_prob=1.0, static_bg=None, device='cuda'):
self.device = torch.device(device)
self.renderer = renderer
self.cam = OrbitCamera(cam_size_W, cam_size_H, fovy=reference_orbit_camera_fovy)
self.projection_matrix = get_projection_matrix(self.cam.near, self.cam.far, self.cam.fovx, self.cam.fovy).transpose(0, 1).cuda()
self.invert_bg_prob = invert_bg_prob
self.black_bg = torch.tensor([0, 0, 0], dtype=torch.float32, device=self.device)
self.white_bg = torch.tensor([1, 1, 1], dtype=torch.float32, device=self.device)
self.static_bg = None if static_bg is None else torch.tensor(static_bg, dtype=torch.float32, device=self.device)
def render_at_pose(self, cam_pose):
radius, elevation, azimuth, center_X, center_Y, center_Z = cam_pose
orbit_target = np.array([center_X, center_Y, center_Z], dtype=np.float32)
render_pose = orbit_camera(elevation, azimuth, radius, target=orbit_target)
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)
if self.static_bg is None:
bg_color = self.white_bg if np.random.rand() > self.invert_bg_prob else self.black_bg
else:
bg_color = self.static_bg
return self.renderer.render(render_cam, bg_color=bg_color)
def render_all_pose(self, all_cam_poses):
all_rendered_images, all_rendered_masks = [], []
for cam_pose in all_cam_poses:
out = self.render_at_pose(cam_pose)
image = out["image"] # [3, H, W] in [0, 1]
mask = out["alpha"] # [1, H, W] in [0, 1]
all_rendered_images.append(image)
all_rendered_masks.append(mask)
# [Number of Poses, 3, H, W], [Number of Poses, 1, H, W] both in [0, 1]
return torch.stack(all_rendered_images, dim=0), torch.stack(all_rendered_masks, dim=0)
class GaussianSplatting:
@@ -82,8 +123,7 @@ class GaussianSplatting:
if gs_params is None:
gs_params = GSParams()
self.renderer = Renderer(sh_degree=gs_params.sh_degree)
self.gaussain_scale_factor = 1
self.renderer = GaussianSplattingRenderer(sh_degree=gs_params.sh_degree)
self.renderer.initialize(init_input, num_pts=gs_params.num_pts)
# setup training
@@ -95,31 +135,26 @@ class GaussianSplatting:
self.ms_ssim_loss = MS_SSIM(data_range=1, size_average=True, channel=3)
self.gs_params = gs_params
def prepare_img(self, img):
img_new = img.permute(2, 0, 1).unsqueeze(0).to(self.device)
img_new = F.interpolate(img_new, (self.ref_size_H, self.ref_size_W), mode="bilinear", align_corners=False).contiguous()
return img_new
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
self.ref_size_H = reference_images[0].shape[0]
self.ref_size_W = reference_images[0].shape[1]
# default camera settings
self.cam = OrbitCamera(self.ref_size_W, self.ref_size_H, fovy=reference_orbit_camera_fovy)
self.projection_matrix = get_projection_matrix(self.cam.near, self.cam.far, self.cam.fovx, self.cam.fovy).transpose(0, 1).cuda()
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(self.prepare_img(reference_images[i]))
ref_masks_torch_list.append(self.prepare_img(reference_masks[i].unsqueeze(2)))
ref_imgs_torch_list.append(prepare_torch_img(reference_images[i], self.ref_size_H, self.ref_size_W, self.device))
ref_masks_torch_list.append(prepare_torch_img(reference_masks[i].unsqueeze(2), 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)
@@ -158,14 +193,7 @@ class GaussianSplatting:
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)
ref_pose = orbit_camera(elevation, azimuth, radius, target=orbit_target)
ref_cam = MiniCam(ref_pose, self.ref_size_W, self.ref_size_H, self.cam.fovy, self.cam.fovx, self.cam.near, self.cam.far, self.projection_matrix)
bg_color = torch.tensor([1, 1, 1] if np.random.rand() > self.gs_params.invert_bg_prob else [0, 0, 0], dtype=torch.float32, device=self.device)
out = self.renderer.render(ref_cam, bg_color=bg_color)
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]
algorithms/main_3DGS_renderer.py
+8 -31
View File
@@ -10,8 +10,6 @@ import torch
from torch import nn
import torch.nn.functional as F
from ..mesh_processer.mesh import PointCloud
pytorch3d_capable = True
try:
import pytorch3d
@@ -30,8 +28,8 @@ from diff_gaussian_rasterization import (
from simple_knn._C import distCUDA2
from ..shared_utils.sh_utils import eval_sh, SH2RGB, RGB2SH
from ..mesh_processer.mesh import Mesh
from ..mesh_processer.mesh_utils import read_gs_ply
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 inverse_sigmoid(x):
return torch.log(x/(1-x))
@@ -285,7 +283,7 @@ class GaussianModel:
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)
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)
@@ -505,20 +503,6 @@ class GaussianModel:
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 construct_list_of_attributes(self):
l = ['x', 'y', 'z', 'nx', 'ny', 'nz']
# All channels except the 3 DC
for i in range(self._features_dc.shape[1]*self._features_dc.shape[2]):
l.append('f_dc_{}'.format(i))
for i in range(self._features_rest.shape[1]*self._features_rest.shape[2]):
l.append('f_rest_{}'.format(i))
l.append('opacity')
for i in range(self._scaling.shape[1]):
l.append('scale_{}'.format(i))
for i in range(self._rotation.shape[1]):
l.append('rot_{}'.format(i))
return l
def get_points_cloud(self):
xyz = self._xyz.detach().cpu().numpy()
@@ -526,9 +510,6 @@ class GaussianModel:
shs = self.get_features.detach().cpu().numpy()
pcd = PointCloud(points=xyz, colors=SH2RGB(shs), normals=normals)
return pcd
def get_mesh(self):
xyz = self._xyz.detach().cpu().numpy()
def to_ply(self):
xyz = self._xyz.detach().cpu().numpy()
@@ -538,24 +519,20 @@ class GaussianModel:
opacities = self._opacity.detach().cpu().numpy()
scale = self._scaling.detach().cpu().numpy()
rotation = self._rotation.detach().cpu().numpy()
dtype_full = [(attribute, 'f4') for attribute in self.construct_list_of_attributes()]
elements = np.empty(xyz.shape[0], dtype=dtype_full)
attributes = np.concatenate((xyz, normals, f_dc, f_rest, opacities, scale, rotation), axis=1)
elements[:] = list(map(tuple, attributes))
el = PlyElement.describe(elements, 'vertex')
return PlyData([el])
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
@@ -842,7 +819,7 @@ class GaussianModel:
torch.cuda.empty_cache()
class Renderer:
class GaussianSplattingRenderer:
def __init__(self, sh_degree=3, white_background=True, radius=1):
self.sh_degree = sh_degree