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/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)