162 lines
6.8 KiB
Python
162 lines
6.8 KiB
Python
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, ) |