# # Copyright (C) 2023, Inria # GRAPHDECO research group, https://team.inria.fr/graphdeco # All rights reserved. # # This software is free for non-commercial, research and evaluation use # under the terms of the LICENSE.md file. # # For inquiries contact george.drettakis@inria.fr # import os import random import json from utils.system_utils import searchForMaxIteration from scene.dataset_readers import sceneLoadTypeCallbacks from scene.gaussian_model import GaussianModel from arguments import ModelParams from utils.camera_utils import cameraList_from_camInfos, camera_to_JSON from torch.nn import functional as F import torch class Scene: gaussians : GaussianModel def __init__(self, args : ModelParams, gaussians : GaussianModel, load_iteration=None, shuffle=False, resolution_scales=[1.0], load_coarse=False, #for waymo bg_gaussians: GaussianModel=None, build_octree=False, replace_pcd_by_octree_center=False, build_grid=False, build_featgrid=False, ): """b :param path: Path to colmap scene main folder. """ self.model_path = args.model_path self.loaded_iter = None self.gaussians = gaussians # for waymo self.bg_gaussians = bg_gaussians if load_iteration: if load_iteration == -1: self.loaded_iter = searchForMaxIteration(os.path.join(self.model_path, "point_cloud")) else: self.loaded_iter = load_iteration print("Loading trained model at iteration {}".format(self.loaded_iter)) self.train_cameras = {} self.test_cameras = {} self.full_cameras = {} if os.path.exists(os.path.join(args.source_path, "sparse")): #scene_info = sceneLoadTypeCallbacks["Colmap"](args.source_path, args.images, args.eval) scene_info = sceneLoadTypeCallbacks["Colmap"](args.source_path, args.images, args.eval, args.object_path, n_views=args.n_views, random_init=args.random_init, train_split=args.train_split) elif os.path.exists(os.path.join(args.source_path, "transforms_train.json")): print("Found transforms_train.json file, assuming Blender data set!") scene_info = sceneLoadTypeCallbacks["Blender"](args.source_path, args.white_background, args.eval) elif os.path.exists(os.path.join(args.source_path,"frame_info.json")): print("Found frame_info.json file, assuming Waymo data set!") scene_info = sceneLoadTypeCallbacks["Waymo"](args.source_path, args.white_background, args.eval, use_bg_gs = bg_gaussians is not None, load_sky_mask = args.load_sky_mask, #False, load_panoptic_mask = args.load_panoptic_mask, #True, load_intrinsic = args.load_intrinsic, #False, load_c2w = args.load_c2w, #False, load_sam_mask = args.load_sam_mask, #False, load_dynamic_mask = args.load_dynamic_mask, #False, load_feat_map = args.load_feat_map, #False, start_time = args.start_time, #0, end_time = args.end_time, # 100, num_pts = args.num_pts, save_occ_grid = args.save_occ_grid, occ_voxel_size = args.occ_voxel_size, recompute_occ_grid = args.recompute_occ_grid, stride = args.stride, original_start_time = args.original_start_time, n_views = args.n_views, waymo_split_mode = args.waymo_split_mode, ) dataset_type="waymo" else: assert False, "Could not recognize scene type!" if not self.loaded_iter: #with open(scene_info.ply_path, 'rb') as src_file, open(os.path.join(self.model_path, "input.ply") , 'wb') as dest_file: # dest_file.write(src_file.read()) json_cams = [] camlist = [] if scene_info.test_cameras: camlist.extend(scene_info.test_cameras) if scene_info.train_cameras: camlist.extend(scene_info.train_cameras) for id, cam in enumerate(camlist): json_cams.append(camera_to_JSON(id, cam)) with open(os.path.join(self.model_path, "cameras.json"), 'w') as file: json.dump(json_cams, file) if shuffle: random.shuffle(scene_info.train_cameras) # Multi-res consistent random shuffling random.shuffle(scene_info.test_cameras) # Multi-res consistent random shuffling self.cameras_extent = scene_info.nerf_normalization["radius"] for resolution_scale in resolution_scales: print("Loading Training Cameras") self.train_cameras[resolution_scale] = cameraList_from_camInfos(scene_info.train_cameras, resolution_scale, args) print("Loading Test Cameras") self.test_cameras[resolution_scale] = cameraList_from_camInfos(scene_info.test_cameras, resolution_scale, args) print("Loading Full Cameras") self.full_cameras[resolution_scale] = cameraList_from_camInfos(scene_info.full_cameras, resolution_scale, args) if self.loaded_iter: self.gaussians.load_ply(os.path.join(self.model_path, "point_cloud", "iteration_" + str(self.loaded_iter), "point_cloud.ply")) if bg_gaussians is not None: self.bg_gaussians.load_ply(os.path.join(self.model_path, "point_cloud", "iteration_" + str(self.loaded_iter), "bg_point_cloud.ply")) else: self.gaussians.create_from_pcd(scene_info.point_cloud, self.cameras_extent) # for waymo if bg_gaussians is not None: self.bg_gaussians.create_from_pcd(scene_info.bg_point_cloud, self.cameras_extent) self.gaussians.aabb = scene_info.cam_frustum_aabb self.gaussians.aabb_tensor = torch.tensor(scene_info.cam_frustum_aabb, dtype=torch.float32).cuda() self.gaussians.nerf_normalization = scene_info.nerf_normalization self.gaussians.img_width = scene_info.train_cameras[0].width self.gaussians.img_height = scene_info.train_cameras[0].height if scene_info.occ_grid is not None: self.gaussians.occ_grid = torch.tensor(scene_info.occ_grid, dtype=torch.bool).cuda() else: self.gaussians.occ_grid = scene_info.occ_grid self.gaussians.occ_voxel_size = args.occ_voxel_size # check occ #import numpy as np #voxel_coords = np.floor((self.gaussians._xyz.cpu().detach().numpy() - scene_info.cam_frustum_aabb[0]) / args.occ_voxel_size).astype(int) #occ = scene_info.occ_grid[voxel_coords[:, 0], voxel_coords[:, 1], voxel_coords[:, 2]] #occ_mask = self.gaussians.get_gs_mask_in_occGrid() #assert all(occ == occ_mask), 'occ should be equal to occ_mask' if args.load_panoptic_mask: self.gaussians.num_panoptic_objects = scene_info.num_panoptic_objects self.gaussians.panoptic_object_ids = scene_info.panoptic_object_ids self.gaussians.panoptic_id_to_idx = scene_info.panoptic_id_to_idx # for deformation-field if hasattr(self.gaussians, '_deformation'): self.gaussians._deformation.deformation_net.set_aabb(scene_info.cam_frustum_aabb[1], scene_info.cam_frustum_aabb[0]) ## make one-hot gt label #gt_label = F.one_hot(torch.arange(self.gaussians.num_panoptic_objects)).float().cuda() ## set as nn.Embedding #self.gaussians.gt_label = torch.nn.Embedding.from_pretrained(gt_label, freeze=True) if build_octree: # forward : point cloud -> octree self.gaussians.build_octree(aabb= scene_info.cam_frustum_aabb, # use camera-extent aabb resolution=5, threshold=10) if replace_pcd_by_octree_center: self.gaussians.replace_pcd_by_octree_node() # to cuda self.gaussians.octree.to_cuda() # backward : octree -> point cloud : get point - octree.node correspondence self.node_list, self.rot_list, self.scale_list = self.gaussians.octree.get_point_node_list() # check if all nodes are leaf nodes assert all([node.is_leaf() for node in self.node_list]), 'all nodes should be leaf nodes' if build_grid: # 建立 dense-occ-grid 来表达高斯的分布, 优势在于 索引快速 self.gaussians.build_grid(aabb= scene_info.cam_frustum_aabb, # use camera-extent aabb res=[128, 128, 128]) if build_featgrid: self.gaussians.build_featgrid(aabb= scene_info.cam_frustum_aabb, # use camera-extent aabb res=[128, 128, 128]) def save(self, iteration, stage): if stage == "coarse": point_cloud_path = os.path.join(self.model_path, "point_cloud/coarse_iteration_{}".format(iteration)) else: point_cloud_path = os.path.join(self.model_path, "point_cloud/iteration_{}".format(iteration)) # if save_spilt: # pc_dynamic_path = os.path.join(point_cloud_path,"point_cloud_dynamic.ply") # pc_static_path = os.path.join(point_cloud_path,"point_cloud_static.ply") self.gaussians.save_ply(os.path.join(point_cloud_path, "point_cloud.ply")) self.gaussians.save_deformation(point_cloud_path) # def save(self, iteration): # point_cloud_path = os.path.join(self.model_path, "point_cloud/iteration_{}".format(iteration)) # self.gaussians.save_ply(os.path.join(point_cloud_path, "point_cloud.ply")) # # background # # if self.gaussians.bg_gaussians is not None: # # self.gaussians.bg_gs.save_ply(os.path.join(point_cloud_path, "bg_point_cloud.ply")) # if self.bg_gaussians is not None: # self.bg_gaussians.save_ply(os.path.join(point_cloud_path, "bg_point_cloud.ply")) def save_gridgs(self, iteration): point_cloud_path = os.path.join(self.model_path, "point_cloud/iteration_{}_grid".format(iteration)) self.gaussians.save_ply(os.path.join(point_cloud_path, "point_cloud.ply")) # background if self.bg_gaussians is not None: self.bg_gaussians.save_ply(os.path.join(point_cloud_path, "bg_point_cloud.ply")) def getTrainCameras(self, scale=1.0): return self.train_cameras[scale] def getTestCameras(self, scale=1.0): return self.test_cameras[scale] def getFullCameras(self, scale=1.0): return self.full_cameras[scale]