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_processor.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}")