souging/TRELLIS_TextTo3D
0
1from typing import *2import numpy as np3import torch4import utils3d5import nvdiffrast.torch as dr6from tqdm import tqdm7import trimesh8import trimesh.visual9import xatlas10import pyvista as pv11from pymeshfix import _meshfix12import igraph13import cv214from PIL import Image15from .random_utils import sphere_hammersley_sequence16from .render_utils import render_multiview17from ..renderers import GaussianRenderer18from ..representations import Strivec, Gaussian, MeshExtractResult19 20 21@torch.no_grad()22def _fill_holes(23 verts,24 faces,25 max_hole_size=0.04,26 max_hole_nbe=32,27 resolution=128,28 num_views=500,29 debug=False,30 verbose=False31):32 """33 Rasterize a mesh from multiple views and remove invisible faces.34 Also includes postprocessing to:35 1. Remove connected components that are have low visibility.36 2. Mincut to remove faces at the inner side of the mesh connected to the outer side with a small hole.37 38 Args:39 verts (torch.Tensor): Vertices of the mesh. Shape (V, 3).40 faces (torch.Tensor): Faces of the mesh. Shape (F, 3).41 max_hole_size (float): Maximum area of a hole to fill.42 resolution (int): Resolution of the rasterization.43 num_views (int): Number of views to rasterize the mesh.44 verbose (bool): Whether to print progress.45 """46 # Construct cameras47 yaws = []48 pitchs = []49 for i in range(num_views):50 y, p = sphere_hammersley_sequence(i, num_views)51 yaws.append(y)52 pitchs.append(p)53 yaws = torch.tensor(yaws).cuda()54 pitchs = torch.tensor(pitchs).cuda()55 radius = 2.056 fov = torch.deg2rad(torch.tensor(40)).cuda()57 projection = utils3d.torch.perspective_from_fov_xy(fov, fov, 1, 3)58 views = []59 for (yaw, pitch) in zip(yaws, pitchs):60 orig = torch.tensor([61 torch.sin(yaw) * torch.cos(pitch),62 torch.cos(yaw) * torch.cos(pitch),63 torch.sin(pitch),64 ]).cuda().float() * radius65 view = utils3d.torch.view_look_at(orig, torch.tensor([0, 0, 0]).float().cuda(), torch.tensor([0, 0, 1]).float().cuda())66 views.append(view)67 views = torch.stack(views, dim=0)68 69 # Rasterize70 visblity = torch.zeros(faces.shape[0], dtype=torch.int32, device=verts.device)71 rastctx = utils3d.torch.RastContext(backend='cuda')72 for i in tqdm(range(views.shape[0]), total=views.shape[0], disable=not verbose, desc='Rasterizing'):73 view = views[i]74 buffers = utils3d.torch.rasterize_triangle_faces(75 rastctx, verts[None], faces, resolution, resolution, view=view, projection=projection76 )77 face_id = buffers['face_id'][0][buffers['mask'][0] > 0.95] - 178 face_id = torch.unique(face_id).long()79 visblity[face_id] += 180 visblity = visblity.float() / num_views81 82 # Mincut83 ## construct outer faces84 edges, face2edge, edge_degrees = utils3d.torch.compute_edges(faces)85 boundary_edge_indices = torch.nonzero(edge_degrees == 1).reshape(-1)86 connected_components = utils3d.torch.compute_connected_components(faces, edges, face2edge)87 outer_face_indices = torch.zeros(faces.shape[0], dtype=torch.bool, device=faces.device)88 for i in range(len(connected_components)):89 outer_face_indices[connected_components[i]] = visblity[connected_components[i]] > min(max(visblity[connected_components[i]].quantile(0.75).item(), 0.25), 0.5)90 outer_face_indices = outer_face_indices.nonzero().reshape(-1)91 92 ## construct inner faces93 inner_face_indices = torch.nonzero(visblity == 0).reshape(-1)94 if verbose:95 tqdm.write(f'Found {inner_face_indices.shape[0]} invisible faces')96 if inner_face_indices.shape[0] == 0:97 return verts, faces98 99 ## Construct dual graph (faces as nodes, edges as edges)100 dual_edges, dual_edge2edge = utils3d.torch.compute_dual_graph(face2edge)101 dual_edge2edge = edges[dual_edge2edge]102 dual_edges_weights = torch.norm(verts[dual_edge2edge[:, 0]] - verts[dual_edge2edge[:, 1]], dim=1)103 if verbose:104 tqdm.write(f'Dual graph: {dual_edges.shape[0]} edges')105 106 ## solve mincut problem107 ### construct main graph108 g = igraph.Graph()109 g.add_vertices(faces.shape[0])110 g.add_edges(dual_edges.cpu().numpy())111 g.es['weight'] = dual_edges_weights.cpu().numpy()112 113 ### source and target114 g.add_vertex('s')115 g.add_vertex('t')116 117 ### connect invisible faces to source118 g.add_edges([(f, 's') for f in inner_face_indices], attributes={'weight': torch.ones(inner_face_indices.shape[0], dtype=torch.float32).cpu().numpy()})119 120 ### connect outer faces to target121 g.add_edges([(f, 't') for f in outer_face_indices], attributes={'weight': torch.ones(outer_face_indices.shape[0], dtype=torch.float32).cpu().numpy()})122 123 ### solve mincut124 cut = g.mincut('s', 't', (np.array(g.es['weight']) * 1000).tolist())125 remove_face_indices = torch.tensor([v for v in cut.partition[0] if v < faces.shape[0]], dtype=torch.long, device=faces.device)126 if verbose:127 tqdm.write(f'Mincut solved, start checking the cut')128 129 ### check if the cut is valid with each connected component130 to_remove_cc = utils3d.torch.compute_connected_components(faces[remove_face_indices])131 if debug:132 tqdm.write(f'Number of connected components of the cut: {len(to_remove_cc)}')133 valid_remove_cc = []134 cutting_edges = []135 for cc in to_remove_cc:136 #### check if the connected component has low visibility137 visblity_median = visblity[remove_face_indices[cc]].median()138 if debug:139 tqdm.write(f'visblity_median: {visblity_median}')140 if visblity_median > 0.25:141 continue142 143 #### check if the cuting loop is small enough144 cc_edge_indices, cc_edges_degree = torch.unique(face2edge[remove_face_indices[cc]], return_counts=True)145 cc_boundary_edge_indices = cc_edge_indices[cc_edges_degree == 1]146 cc_new_boundary_edge_indices = cc_boundary_edge_indices[~torch.isin(cc_boundary_edge_indices, boundary_edge_indices)]147 if len(cc_new_boundary_edge_indices) > 0:148 cc_new_boundary_edge_cc = utils3d.torch.compute_edge_connected_components(edges[cc_new_boundary_edge_indices])149 cc_new_boundary_edges_cc_center = [verts[edges[cc_new_boundary_edge_indices[edge_cc]]].mean(dim=1).mean(dim=0) for edge_cc in cc_new_boundary_edge_cc]150 cc_new_boundary_edges_cc_area = []151 for i, edge_cc in enumerate(cc_new_boundary_edge_cc):152 _e1 = verts[edges[cc_new_boundary_edge_indices[edge_cc]][:, 0]] - cc_new_boundary_edges_cc_center[i]153 _e2 = verts[edges[cc_new_boundary_edge_indices[edge_cc]][:, 1]] - cc_new_boundary_edges_cc_center[i]154 cc_new_boundary_edges_cc_area.append(torch.norm(torch.cross(_e1, _e2, dim=-1), dim=1).sum() * 0.5)155 if debug:156 cutting_edges.append(cc_new_boundary_edge_indices)157 tqdm.write(f'Area of the cutting loop: {cc_new_boundary_edges_cc_area}')158 if any([l > max_hole_size for l in cc_new_boundary_edges_cc_area]):159 continue160 161 valid_remove_cc.append(cc)162 163 if debug:164 face_v = verts[faces].mean(dim=1).cpu().numpy()165 vis_dual_edges = dual_edges.cpu().numpy()166 vis_colors = np.zeros((faces.shape[0], 3), dtype=np.uint8)167 vis_colors[inner_face_indices.cpu().numpy()] = [0, 0, 255]168 vis_colors[outer_face_indices.cpu().numpy()] = [0, 255, 0]169 vis_colors[remove_face_indices.cpu().numpy()] = [255, 0, 255]170 if len(valid_remove_cc) > 0:171 vis_colors[remove_face_indices[torch.cat(valid_remove_cc)].cpu().numpy()] = [255, 0, 0]172 utils3d.io.write_ply('dbg_dual.ply', face_v, edges=vis_dual_edges, vertex_colors=vis_colors)173 174 vis_verts = verts.cpu().numpy()175 vis_edges = edges[torch.cat(cutting_edges)].cpu().numpy()176 utils3d.io.write_ply('dbg_cut.ply', vis_verts, edges=vis_edges)177 178 179 if len(valid_remove_cc) > 0:180 remove_face_indices = remove_face_indices[torch.cat(valid_remove_cc)]181 mask = torch.ones(faces.shape[0], dtype=torch.bool, device=faces.device)182 mask[remove_face_indices] = 0183 faces = faces[mask]184 faces, verts = utils3d.torch.remove_unreferenced_vertices(faces, verts)185 if verbose:186 tqdm.write(f'Removed {(~mask).sum()} faces by mincut')187 else:188 if verbose:189 tqdm.write(f'Removed 0 faces by mincut')190 191 mesh = _meshfix.PyTMesh()192 mesh.load_array(verts.cpu().numpy(), faces.cpu().numpy())193 mesh.fill_small_boundaries(nbe=max_hole_nbe, refine=True)194 verts, faces = mesh.return_arrays()195 verts, faces = torch.tensor(verts, device='cuda', dtype=torch.float32), torch.tensor(faces, device='cuda', dtype=torch.int32)196 197 return verts, faces198 199 200def postprocess_mesh(201 vertices: np.array,202 faces: np.array,203 simplify: bool = True,204 simplify_ratio: float = 0.9,205 fill_holes: bool = True,206 fill_holes_max_hole_size: float = 0.04,207 fill_holes_max_hole_nbe: int = 32,208 fill_holes_resolution: int = 1024,209 fill_holes_num_views: int = 1000,210 debug: bool = False,211 verbose: bool = False,212):213 """214 Postprocess a mesh by simplifying, removing invisible faces, and removing isolated pieces.215 216 Args:217 vertices (np.array): Vertices of the mesh. Shape (V, 3).218 faces (np.array): Faces of the mesh. Shape (F, 3).219 simplify (bool): Whether to simplify the mesh, using quadric edge collapse.220 simplify_ratio (float): Ratio of faces to keep after simplification.221 fill_holes (bool): Whether to fill holes in the mesh.222 fill_holes_max_hole_size (float): Maximum area of a hole to fill.223 fill_holes_max_hole_nbe (int): Maximum number of boundary edges of a hole to fill.224 fill_holes_resolution (int): Resolution of the rasterization.225 fill_holes_num_views (int): Number of views to rasterize the mesh.226 verbose (bool): Whether to print progress.227 """228 229 if verbose:230 tqdm.write(f'Before postprocess: {vertices.shape[0]} vertices, {faces.shape[0]} faces')231 232 # Simplify233 if simplify and simplify_ratio > 0:234 mesh = pv.PolyData(vertices, np.concatenate([np.full((faces.shape[0], 1), 3), faces], axis=1))235 mesh = mesh.decimate(simplify_ratio, progress_bar=verbose)236 vertices, faces = mesh.points, mesh.faces.reshape(-1, 4)[:, 1:]237 if verbose:238 tqdm.write(f'After decimate: {vertices.shape[0]} vertices, {faces.shape[0]} faces')239 240 # Remove invisible faces241 if fill_holes:242 vertices, faces = torch.tensor(vertices).cuda(), torch.tensor(faces.astype(np.int32)).cuda()243 vertices, faces = _fill_holes(244 vertices, faces,245 max_hole_size=fill_holes_max_hole_size,246 max_hole_nbe=fill_holes_max_hole_nbe,247 resolution=fill_holes_resolution,248 num_views=fill_holes_num_views,249 debug=debug,250 verbose=verbose,251 )252 vertices, faces = vertices.cpu().numpy(), faces.cpu().numpy()253 if verbose:254 tqdm.write(f'After remove invisible faces: {vertices.shape[0]} vertices, {faces.shape[0]} faces')255 256 return vertices, faces257 258 259def parametrize_mesh(vertices: np.array, faces: np.array):260 """261 Parametrize a mesh to a texture space, using xatlas.262 263 Args:264 vertices (np.array): Vertices of the mesh. Shape (V, 3).265 faces (np.array): Faces of the mesh. Shape (F, 3).266 """267 268 vmapping, indices, uvs = xatlas.parametrize(vertices, faces)269 270 vertices = vertices[vmapping]271 faces = indices272 273 return vertices, faces, uvs274 275 276def bake_texture(277 vertices: np.array,278 faces: np.array,279 uvs: np.array,280 observations: List[np.array],281 masks: List[np.array],282 extrinsics: List[np.array],283 intrinsics: List[np.array],284 texture_size: int = 2048,285 near: float = 0.1,286 far: float = 10.0,287 mode: Literal['fast', 'opt'] = 'opt',288 lambda_tv: float = 1e-2,289 verbose: bool = False,290):291 """292 Bake texture to a mesh from multiple observations.293 294 Args:295 vertices (np.array): Vertices of the mesh. Shape (V, 3).296 faces (np.array): Faces of the mesh. Shape (F, 3).297 uvs (np.array): UV coordinates of the mesh. Shape (V, 2).298 observations (List[np.array]): List of observations. Each observation is a 2D image. Shape (H, W, 3).299 masks (List[np.array]): List of masks. Each mask is a 2D image. Shape (H, W).300 extrinsics (List[np.array]): List of extrinsics. Shape (4, 4).301 intrinsics (List[np.array]): List of intrinsics. Shape (3, 3).302 texture_size (int): Size of the texture.303 near (float): Near plane of the camera.304 far (float): Far plane of the camera.305 mode (Literal['fast', 'opt']): Mode of texture baking.306 lambda_tv (float): Weight of total variation loss in optimization.307 verbose (bool): Whether to print progress.308 """309 vertices = torch.tensor(vertices).cuda()310 faces = torch.tensor(faces.astype(np.int32)).cuda()311 uvs = torch.tensor(uvs).cuda()312 observations = [torch.tensor(obs / 255.0).float().cuda() for obs in observations]313 masks = [torch.tensor(m>0).bool().cuda() for m in masks]314 views = [utils3d.torch.extrinsics_to_view(torch.tensor(extr).cuda()) for extr in extrinsics]315 projections = [utils3d.torch.intrinsics_to_perspective(torch.tensor(intr).cuda(), near, far) for intr in intrinsics]316 317 if mode == 'fast':318 texture = torch.zeros((texture_size * texture_size, 3), dtype=torch.float32).cuda()319 texture_weights = torch.zeros((texture_size * texture_size), dtype=torch.float32).cuda()320 rastctx = utils3d.torch.RastContext(backend='cuda')321 for observation, view, projection in tqdm(zip(observations, views, projections), total=len(observations), disable=not verbose, desc='Texture baking (fast)'):322 with torch.no_grad():323 rast = utils3d.torch.rasterize_triangle_faces(324 rastctx, vertices[None], faces, observation.shape[1], observation.shape[0], uv=uvs[None], view=view, projection=projection325 )326 uv_map = rast['uv'][0].detach().flip(0)327 mask = rast['mask'][0].detach().bool() & masks[0]328 329 # nearest neighbor interpolation330 uv_map = (uv_map * texture_size).floor().long()331 obs = observation[mask]332 uv_map = uv_map[mask]333 idx = uv_map[:, 0] + (texture_size - uv_map[:, 1] - 1) * texture_size334 texture = texture.scatter_add(0, idx.view(-1, 1).expand(-1, 3), obs)335 texture_weights = texture_weights.scatter_add(0, idx, torch.ones((obs.shape[0]), dtype=torch.float32, device=texture.device))336 337 mask = texture_weights > 0338 texture[mask] /= texture_weights[mask][:, None]339 texture = np.clip(texture.reshape(texture_size, texture_size, 3).cpu().numpy() * 255, 0, 255).astype(np.uint8)340 341 # inpaint342 mask = (texture_weights == 0).cpu().numpy().astype(np.uint8).reshape(texture_size, texture_size)343 texture = cv2.inpaint(texture, mask, 3, cv2.INPAINT_TELEA)344 345 elif mode == 'opt':346 rastctx = utils3d.torch.RastContext(backend='cuda')347 observations = [observations.flip(0) for observations in observations]348 masks = [m.flip(0) for m in masks]349 _uv = []350 _uv_dr = []351 for observation, view, projection in tqdm(zip(observations, views, projections), total=len(views), disable=not verbose, desc='Texture baking (opt): UV'):352 with torch.no_grad():353 rast = utils3d.torch.rasterize_triangle_faces(354 rastctx, vertices[None], faces, observation.shape[1], observation.shape[0], uv=uvs[None], view=view, projection=projection355 )356 _uv.append(rast['uv'].detach())357 _uv_dr.append(rast['uv_dr'].detach())358 359 texture = torch.nn.Parameter(torch.zeros((1, texture_size, texture_size, 3), dtype=torch.float32).cuda())360 optimizer = torch.optim.Adam([texture], betas=(0.5, 0.9), lr=1e-2)361 362 def exp_anealing(optimizer, step, total_steps, start_lr, end_lr):363 return start_lr * (end_lr / start_lr) ** (step / total_steps)364 365 def cosine_anealing(optimizer, step, total_steps, start_lr, end_lr):366 return end_lr + 0.5 * (start_lr - end_lr) * (1 + np.cos(np.pi * step / total_steps))367 368 def tv_loss(texture):369 return torch.nn.functional.l1_loss(texture[:, :-1, :, :], texture[:, 1:, :, :]) + \370 torch.nn.functional.l1_loss(texture[:, :, :-1, :], texture[:, :, 1:, :])371 372 total_steps = 2500373 with tqdm(total=total_steps, disable=not verbose, desc='Texture baking (opt): optimizing') as pbar:374 for step in range(total_steps):375 optimizer.zero_grad()376 selected = np.random.randint(0, len(views))377 uv, uv_dr, observation, mask = _uv[selected], _uv_dr[selected], observations[selected], masks[selected]378 render = dr.texture(texture, uv, uv_dr)[0]379 loss = torch.nn.functional.l1_loss(render[mask], observation[mask])380 if lambda_tv > 0:381 loss += lambda_tv * tv_loss(texture)382 loss.backward()383 optimizer.step()384 # annealing385 optimizer.param_groups[0]['lr'] = cosine_anealing(optimizer, step, total_steps, 1e-2, 1e-5)386 pbar.set_postfix({'loss': loss.item()})387 pbar.update()388 texture = np.clip(texture[0].flip(0).detach().cpu().numpy() * 255, 0, 255).astype(np.uint8)389 mask = 1 - utils3d.torch.rasterize_triangle_faces(390 rastctx, (uvs * 2 - 1)[None], faces, texture_size, texture_size391 )['mask'][0].detach().cpu().numpy().astype(np.uint8)392 texture = cv2.inpaint(texture, mask, 3, cv2.INPAINT_TELEA)393 else:394 raise ValueError(f'Unknown mode: {mode}')395 396 return texture397 398 399def to_glb(400 app_rep: Union[Strivec, Gaussian],401 mesh: MeshExtractResult,402 simplify: float = 0.95,403 fill_holes: bool = True,404 fill_holes_max_size: float = 0.04,405 texture_size: int = 1024,406 debug: bool = False,407 verbose: bool = True,408) -> trimesh.Trimesh:409 """410 Convert a generated asset to a glb file.411 412 Args:413 app_rep (Union[Strivec, Gaussian]): Appearance representation.414 mesh (MeshExtractResult): Extracted mesh.415 simplify (float): Ratio of faces to remove in simplification.416 fill_holes (bool): Whether to fill holes in the mesh.417 fill_holes_max_size (float): Maximum area of a hole to fill.418 texture_size (int): Size of the texture.419 debug (bool): Whether to print debug information.420 verbose (bool): Whether to print progress.421 """422 vertices = mesh.vertices.cpu().numpy()423 faces = mesh.faces.cpu().numpy()424 425 # mesh postprocess426 vertices, faces = postprocess_mesh(427 vertices, faces,428 simplify=simplify > 0,429 simplify_ratio=simplify,430 fill_holes=fill_holes,431 fill_holes_max_hole_size=fill_holes_max_size,432 fill_holes_max_hole_nbe=int(250 * np.sqrt(1-simplify)),433 fill_holes_resolution=1024,434 fill_holes_num_views=1000,435 debug=debug,436 verbose=verbose,437 )438 439 # parametrize mesh440 vertices, faces, uvs = parametrize_mesh(vertices, faces)441 442 # bake texture443 observations, extrinsics, intrinsics = render_multiview(app_rep, resolution=1024, nviews=100)444 masks = [np.any(observation > 0, axis=-1) for observation in observations]445 extrinsics = [extrinsics[i].cpu().numpy() for i in range(len(extrinsics))]446 intrinsics = [intrinsics[i].cpu().numpy() for i in range(len(intrinsics))]447 texture = bake_texture(448 vertices, faces, uvs,449 observations, masks, extrinsics, intrinsics,450 texture_size=texture_size, mode='opt',451 lambda_tv=0.01,452 verbose=verbose453 )454 texture = Image.fromarray(texture)455 456 # rotate mesh (from z-up to y-up)457 vertices = vertices @ np.array([[1, 0, 0], [0, 0, -1], [0, 1, 0]])458 material = trimesh.visual.material.PBRMaterial(459 roughnessFactor=1.0,460 baseColorTexture=texture,461 baseColorFactor=np.array([255, 255, 255, 255], dtype=np.uint8)462 )463 mesh = trimesh.Trimesh(vertices, faces, visual=trimesh.visual.TextureVisuals(uv=uvs, material=material))464 return mesh465 466 467def simplify_gs(468 gs: Gaussian,469 simplify: float = 0.95,470 verbose: bool = True,471):472 """473 Simplify 3D Gaussians474 NOTE: this function is not used in the current implementation for the unsatisfactory performance.475 476 Args:477 gs (Gaussian): 3D Gaussian.478 simplify (float): Ratio of Gaussians to remove in simplification.479 """480 if simplify <= 0:481 return gs482 483 # simplify484 observations, extrinsics, intrinsics = render_multiview(gs, resolution=1024, nviews=100)485 observations = [torch.tensor(obs / 255.0).float().cuda().permute(2, 0, 1) for obs in observations]486 487 # Following https://arxiv.org/pdf/2411.06019488 renderer = GaussianRenderer({489 "resolution": 1024,490 "near": 0.8,491 "far": 1.6,492 "ssaa": 1,493 "bg_color": (0,0,0),494 })495 new_gs = Gaussian(**gs.init_params)496 new_gs._features_dc = gs._features_dc.clone()497 new_gs._features_rest = gs._features_rest.clone() if gs._features_rest is not None else None498 new_gs._opacity = torch.nn.Parameter(gs._opacity.clone())499 new_gs._rotation = torch.nn.Parameter(gs._rotation.clone())500 new_gs._scaling = torch.nn.Parameter(gs._scaling.clone())501 new_gs._xyz = torch.nn.Parameter(gs._xyz.clone())502 503 start_lr = [1e-4, 1e-3, 5e-3, 0.025]504 end_lr = [1e-6, 1e-5, 5e-5, 0.00025]505 optimizer = torch.optim.Adam([506 {"params": new_gs._xyz, "lr": start_lr[0]},507 {"params": new_gs._rotation, "lr": start_lr[1]},508 {"params": new_gs._scaling, "lr": start_lr[2]},509 {"params": new_gs._opacity, "lr": start_lr[3]},510 ], lr=start_lr[0])511 512 def exp_anealing(optimizer, step, total_steps, start_lr, end_lr):513 return start_lr * (end_lr / start_lr) ** (step / total_steps)514 515 def cosine_anealing(optimizer, step, total_steps, start_lr, end_lr):516 return end_lr + 0.5 * (start_lr - end_lr) * (1 + np.cos(np.pi * step / total_steps))517 518 _zeta = new_gs.get_opacity.clone().detach().squeeze()519 _lambda = torch.zeros_like(_zeta)520 _delta = 1e-7521 _interval = 10522 num_target = int((1 - simplify) * _zeta.shape[0])523 524 with tqdm(total=2500, disable=not verbose, desc='Simplifying Gaussian') as pbar:525 for i in range(2500):526 # prune527 if i % 100 == 0:528 mask = new_gs.get_opacity.squeeze() > 0.05529 mask = torch.nonzero(mask).squeeze()530 new_gs._xyz = torch.nn.Parameter(new_gs._xyz[mask])531 new_gs._rotation = torch.nn.Parameter(new_gs._rotation[mask])532 new_gs._scaling = torch.nn.Parameter(new_gs._scaling[mask])533 new_gs._opacity = torch.nn.Parameter(new_gs._opacity[mask])534 new_gs._features_dc = new_gs._features_dc[mask]535 new_gs._features_rest = new_gs._features_rest[mask] if new_gs._features_rest is not None else None536 _zeta = _zeta[mask]537 _lambda = _lambda[mask]538 # update optimizer state539 for param_group, new_param in zip(optimizer.param_groups, [new_gs._xyz, new_gs._rotation, new_gs._scaling, new_gs._opacity]):540 stored_state = optimizer.state[param_group['params'][0]]541 if 'exp_avg' in stored_state:542 stored_state['exp_avg'] = stored_state['exp_avg'][mask]543 stored_state['exp_avg_sq'] = stored_state['exp_avg_sq'][mask]544 del optimizer.state[param_group['params'][0]]545 param_group['params'][0] = new_param546 optimizer.state[param_group['params'][0]] = stored_state547 548 opacity = new_gs.get_opacity.squeeze()549 550 # sparisfy551 if i % _interval == 0:552 _zeta = _lambda + opacity.detach()553 if opacity.shape[0] > num_target:554 index = _zeta.topk(num_target)[1]555 _m = torch.ones_like(_zeta, dtype=torch.bool)556 _m[index] = 0557 _zeta[_m] = 0558 _lambda = _lambda + opacity.detach() - _zeta559 560 # sample a random view561 view_idx = np.random.randint(len(observations))562 observation = observations[view_idx]563 extrinsic = extrinsics[view_idx]564 intrinsic = intrinsics[view_idx]565 566 color = renderer.render(new_gs, extrinsic, intrinsic)['color']567 rgb_loss = torch.nn.functional.l1_loss(color, observation)568 loss = rgb_loss + \569 _delta * torch.sum(torch.pow(_lambda + opacity - _zeta, 2))570 571 optimizer.zero_grad()572 loss.backward()573 optimizer.step()574 575 # update lr576 for j in range(len(optimizer.param_groups)):577 optimizer.param_groups[j]['lr'] = cosine_anealing(optimizer, i, 2500, start_lr[j], end_lr[j])578 579 pbar.set_postfix({'loss': rgb_loss.item(), 'num': opacity.shape[0], 'lambda': _lambda.mean().item()})580 pbar.update()581 582 new_gs._xyz = new_gs._xyz.data583 new_gs._rotation = new_gs._rotation.data584 new_gs._scaling = new_gs._scaling.data585 new_gs._opacity = new_gs._opacity.data586 587 return new_gs588