├── .gitignore ├── .gitmodules ├── README.md ├── arguments └── __init__.py ├── database.py ├── eval.py ├── gaussian_renderer └── __init__.py ├── lpipsPyTorch ├── __init__.py └── modules │ ├── lpips.py │ ├── networks.py │ └── utils.py ├── metrics.py ├── motion_model ├── dataset.py ├── gcn.py └── modules.py ├── options └── gaussian_option.py ├── requirements.txt ├── scene ├── __init__.py ├── cameras.py ├── colmap_loader.py ├── dataset_readers.py ├── deformable_field.py ├── gaussian_model.py ├── hyper_loader.py └── utils.py ├── scripts ├── eval │ ├── d-nerf │ │ ├── bouncingballs.sh │ │ ├── hellwarrior.sh │ │ ├── hook.sh │ │ ├── jumping.sh │ │ ├── mutant.sh │ │ ├── standup.sh │ │ └── trex.sh │ └── hyper │ │ ├── chickchicken.sh │ │ ├── lemon.sh │ │ ├── printer.sh │ │ └── torchocolate.sh ├── predict │ ├── d-nerf │ │ ├── bouncingballs.sh │ │ ├── hellwarrior.sh │ │ ├── hook.sh │ │ ├── jumping.sh │ │ ├── mutant.sh │ │ ├── standup.sh │ │ └── trex.sh │ └── hyper │ │ ├── chickchicken.sh │ │ ├── lemon.sh │ │ ├── printer.sh │ │ └── torchocolate.sh ├── train │ ├── d-nerf │ │ ├── bouncingballs.sh │ │ ├── hellwarrior.sh │ │ ├── hook.sh │ │ ├── jumping.sh │ │ ├── mutant.sh │ │ ├── standup.sh │ │ ├── train_eval.sh │ │ ├── train_predict.sh │ │ └── trex.sh │ └── hyper │ │ ├── chickchicken.sh │ │ ├── lemon.sh │ │ ├── printer.sh │ │ ├── torchocolate.sh │ │ ├── train_eval.sh │ │ └── train_predict.sh └── utils │ ├── colmap.sh │ ├── dnerf_show.sh │ ├── env.sh │ └── hyper_show.sh ├── show.py ├── submodules └── lib │ ├── __init__.py │ └── pointops │ ├── __init__.py │ ├── dist │ └── pointops-0.0.0-py3.7-linux-x86_64.egg │ ├── functions │ ├── __init__.py │ └── pointops.py │ ├── pointops.egg-info │ ├── PKG-INFO │ ├── SOURCES.txt │ ├── dependency_links.txt │ └── top_level.txt │ ├── setup.py │ └── src │ ├── __init__.py │ ├── aggregation │ ├── aggregation_cuda.cpp │ ├── aggregation_cuda_kernel.cu │ └── aggregation_cuda_kernel.h │ ├── cuda_utils.h │ ├── fast_sampling │ ├── fast_sampling_cuda.cpp │ ├── fast_sampling_cuda_kernel.cu │ └── fast_sampling_cuda_kernel.h │ ├── grouping │ ├── grouping_cuda.cpp │ ├── grouping_cuda_kernel.cu │ └── grouping_cuda_kernel.h │ ├── interpolation │ ├── interpolation_cuda.cpp │ ├── interpolation_cuda_kernel.cu │ └── interpolation_cuda_kernel.h │ ├── knnquery │ ├── knnquery_cuda.cpp │ ├── knnquery_cuda_kernel.cu │ └── knnquery_cuda_kernel.h │ ├── pointops_api.cpp │ ├── sampling │ ├── sampling_cuda.cpp │ ├── sampling_cuda_kernel.cu │ └── sampling_cuda_kernel.h │ └── subtraction │ ├── subtraction_cuda.cpp │ ├── subtraction_cuda_kernel.cu │ └── subtraction_cuda_kernel.h ├── train.py ├── train_GCN.py └── utils ├── camera_utils.py ├── fps.py ├── general_utils.py ├── graphics_utils.py ├── image_utils.py ├── loss_utils.py ├── prepare ├── downsample_points.py ├── hypernerf2colmap.py └── makeVideo.py ├── sh_utils.py ├── system_utils.py └── visualizer_utils.py /.gitignore: -------------------------------------------------------------------------------- 1 | *.pyc 2 | .vscode 3 | **/__pycache__ 4 | results/* 5 | datasets/* 6 | build/* 7 | logs/* -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "submodules/simple-knn"] 2 | path = submodules/simple-knn 3 | url = https://gitlab.inria.fr/bkerbl/simple-knn.git 4 | [submodule "submodules/diff-gaussian-rasterization-w-depth"] 5 | path = submodules/diff-gaussian-rasterization-w-depth 6 | url = https://github.com/BoMingZhao/diff-gaussian-rasterization-w-depth.git 7 | [submodule "submodules/FRNN"] 8 | path = submodules/FRNN 9 | url = https://github.com/lxxue/FRNN.git 10 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # GaussianPrediction [SIGGRAPH2024] 2 | ### [Project Page](https://zju3dv.github.io/gaussian-prediction/) | [arXiv](https://arxiv.org/abs/2405.19745) | [Supplementary](https://raw.githubusercontent.com/BoMingZhao/open_access_assets/main/GaussianPrediction/supp.pdf) 3 | This is the official pytorch implementation of **GaussianPrediction: Dynamic 3D Gaussian Prediction for Motion Extrapolation and Free View Synthesis** 4 | 5 | ## Installation 6 | ```bash 7 | git clone https://github.com/BoMingZhao/GaussianPrediction.git --recursive 8 | cd GaussianPrediction 9 | 10 | # conda env create -f environment.yml 11 | conda create -n gaussian_prediction python=3.7.13 12 | conda activate gaussian_prediction 13 | 14 | bash scripts/utils/env.sh 15 | ``` 16 | 17 | **Note:** We tested our code on PyTorch 1.13.1 with CUDA 11.7 and CUDA 11.8, and no issues were observed. However, we found that when using PyTorch 2.1.0 with CUDA 12.1, the overall reconstruction quality significantly deteriorates. Therefore, please ensure that your installation steps align with the ones we provided. Additionally, if you are aware of the reasons causing this issue, we welcome you to open an issue and share your findings with us. 18 | 19 | ## Data Preparation 20 | In our paper, We use synthetic dataset [D-NeRF](https://github.com/albertpumarola/D-NeRF?tab=readme-ov-file#download-dataset) and real-world dataset [HyperNeRF](https://github.com/google/hypernerf/releases/tag/v0.1). 21 | 22 | 1. Please download the dataset and organize them as follows: 23 | ``` 24 | GaussianPrediction 25 | ├── datasets 26 | │ ├── d-nerf 27 | │ │ │── data 28 | │ │ │ │── bouncingballs 29 | │ │ │ │ │── test 30 | │ │ │ │ │── train 31 | │ │ │ │ │── val 32 | │ │ │ │ │── transforms_test.json 33 | │ │ │ │ │── transforms_train.json 34 | │ │ │ │ │── transforms_val.json 35 | │ │ │ │── mutant 36 | │ │ │ │── trex 37 | │ │ │ │── ... 38 | │ ├── HyperNeRF 39 | │ │ │── cut-lemon 40 | │ │ │ │── rgb 41 | │ │ │ │── camera 42 | │ │ │ │── dataset.json 43 | │ │ │ │── metadata.json 44 | │ │ │ │── scene.json 45 | │ │ │── chickchicken 46 | │ │ │── vrig-3dprinter 47 | │ │ │── torchocolate 48 | ``` 49 | 2. For the HyperNeRF dataset, we follow the [4D-GS](https://github.com/hustvl/4DGaussians) and use the results of COLMAP MVS as the initialization point cloud for the Gaussian Splatting. Therefore, please ensure that COLMAP is installed on your computer and run the following command: 50 | ```bash 51 | bash scripts/utils/colmap.sh datasets/HyperNeRF/${Scene} 52 | ``` 53 | 54 | 55 | ## Usage 56 | You can find all training command in **scripts/train**. You can train a single scene by running: 57 | ```bash 58 | bash ./scripts/train/${Dataset}/${Scene}.sh 59 | ``` 60 | Next, you can train the GCN to obtain prediction results: 61 | ```bash 62 | bash ./scripts/predict/${Dataset}/${Scene}.sh 63 | ``` 64 | 65 | or evaluate the performance of NVS: 66 | ```bash 67 | bash ./scripts/eval/${Dataset}/${Scene}.sh 68 | ``` 69 | Alternatively, you can run the following command to train, predict, and evaluate multiple scenes at once: 70 | 71 | ```bash 72 | # For prediction 73 | bash ./scripts/train/${Dataset}/train_predict.sh 74 | # For evaluation 75 | bash ./scripts/train/${Dataset}/train_eval.sh 76 | ``` 77 | 78 | **Note**: If you want to test the prediction results, please set the `max_time` parameter to `0.8` (this is the value used in our paper, but you can also set it to any value less than 1). If you only want to compare the NVS results, please set max_time to `1.0`. 79 | 80 |
81 | Some important Command Line Arguments for training 82 | 87 | 88 | 89 | 94 | 95 | 96 | 101 | 106 | 107 | 112 | 113 |
114 | 115 | ## Acknowledgement 116 | We would like to thanks [Guanjun Wu](https://guanjunwu.github.io/) for providing the source code of [4D-GS](https://github.com/hustvl/4DGaussians). We adopted his code and scripts for processing the hypernerf dataset. 117 | We also want to thank [Ziyi Yang](https://github.com/ingra14m) for providing the code and training strategy. Some of our training strategies refer to his paper [Deformable-GS](https://github.com/ingra14m/Deformable-3D-Gaussians). 118 | 119 | ## Citing 120 | If you find this work helpful in your research. We would appreciate a citation via 121 | ``` 122 | @inproceedings{zhao2024gaussianprediction, 123 | title={Gaussianprediction: Dynamic 3d gaussian prediction for motion extrapolation and free view synthesis}, 124 | author={Zhao, Boming and Li, Yuan and Sun, Ziyu and Zeng, Lin and Shen, Yujun and Ma, Rui and Zhang, Yinda and Bao, Hujun and Cui, Zhaopeng}, 125 | booktitle={ACM SIGGRAPH 2024 Conference Papers}, 126 | pages={1--12}, 127 | year={2024} 128 | } 129 | ``` -------------------------------------------------------------------------------- /arguments/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | from argparse import ArgumentParser, Namespace 13 | import sys 14 | import os 15 | 16 | class GroupParams: 17 | pass 18 | 19 | class ParamGroup: 20 | def __init__(self, parser: ArgumentParser, name : str, fill_none = False): 21 | group = parser.add_argument_group(name) 22 | for key, value in vars(self).items(): 23 | shorthand = False 24 | if key.startswith("_"): 25 | shorthand = True 26 | key = key[1:] 27 | t = type(value) 28 | value = value if not fill_none else None 29 | if shorthand: 30 | if t == bool: 31 | group.add_argument("--" + key, ("-" + key[0:1]), default=value, action="store_true") 32 | else: 33 | group.add_argument("--" + key, ("-" + key[0:1]), default=value, type=t) 34 | else: 35 | if t == bool: 36 | group.add_argument("--" + key, default=value, action="store_true") 37 | else: 38 | group.add_argument("--" + key, default=value, type=t) 39 | 40 | def extract(self, args): 41 | group = GroupParams() 42 | for arg in vars(args).items(): 43 | if arg[0] in vars(self) or ("_" + arg[0]) in vars(self): 44 | setattr(group, arg[0], arg[1]) 45 | return group 46 | 47 | class ModelParams(ParamGroup): 48 | def __init__(self, parser, sentinel=False): 49 | self.sh_degree = 3 50 | self._source_path = "" 51 | self._model_path = "" 52 | self._images = "images" 53 | self._resolution = -1 54 | self._white_background = False 55 | self.max_time = 1.0 56 | self.data_device = "cuda" 57 | self.eval = False 58 | super().__init__(parser, "Loading Parameters", sentinel) 59 | 60 | def extract(self, args): 61 | g = super().extract(args) 62 | g.source_path = os.path.abspath(g.source_path) 63 | return g 64 | 65 | class PipelineParams(ParamGroup): 66 | def __init__(self, parser): 67 | self.convert_SHs_python = False 68 | self.compute_cov3D_python = False 69 | self.debug = False 70 | super().__init__(parser, "Pipeline Parameters") 71 | 72 | class OptimizationParams(ParamGroup): 73 | def __init__(self, parser): 74 | self.iterations = 30_000 75 | self.position_lr_init = 0.00016 76 | self.position_lr_final = 0.0000016 77 | self.position_lr_delay_mult = 0.01 78 | self.position_lr_max_steps = 30_000 79 | self.feature_lr = 0.0025 80 | self.opacity_lr = 0.05 81 | self.scaling_lr = 0.005 82 | self.rotation_lr = 0.001 83 | self.mfeature_lr = 0.0008 84 | self.mfeature_lr_final = 0.00008 85 | self.kpts_lr = 0.0008 86 | self.kpts_lr_final = 0.00008 87 | self.hash_lr = 0.005 88 | self.hash_lr_final = 0.00005 89 | self.mlp_lr = 0.0008 90 | self.percent_dense = 0.01 91 | self.lambda_dssim = 0.2 92 | self.densification_interval = 100 93 | self.opacity_reset_interval = 3000 94 | self.densify_from_iter = 500 95 | self.densify_until_iter = 15_000 96 | self.densify_grad_threshold = 0.0002 97 | self.xyz_freq = 10 98 | self.time_freq = 6 # 10 for real scene 99 | self.superKeypoints_reset_interval = 1000 100 | super().__init__(parser, "Optimization Parameters") 101 | 102 | def get_combined_args(parser : ArgumentParser): 103 | cmdlne_string = sys.argv[1:] 104 | cfgfile_string = "Namespace()" 105 | args_cmdline = parser.parse_args(cmdlne_string) 106 | 107 | try: 108 | cfgfilepath = os.path.join(args_cmdline.model_path, "cfg_args") 109 | print("Looking for config file in", cfgfilepath) 110 | with open(cfgfilepath) as cfg_file: 111 | print("Config file found: {}".format(cfgfilepath)) 112 | cfgfile_string = cfg_file.read() 113 | except TypeError: 114 | print("Config file not found at") 115 | pass 116 | start = cfgfile_string.find("\n") 117 | cfgfile_string = cfgfile_string[:start] 118 | args_cfgfile = eval(cfgfile_string) 119 | 120 | merged_dict = vars(args_cfgfile).copy() 121 | for k,v in vars(args_cmdline).items(): 122 | if v != None: 123 | merged_dict[k] = v 124 | return Namespace(**merged_dict) 125 | -------------------------------------------------------------------------------- /database.py: -------------------------------------------------------------------------------- 1 | # This script is based on an original implementation by True Price. 2 | # Created by liminghao 3 | import sys 4 | import numpy as np 5 | import sqlite3 6 | 7 | IS_PYTHON3 = sys.version_info[0] >= 3 8 | 9 | def array_to_blob(array): 10 | if IS_PYTHON3: 11 | return array.tostring() 12 | else: 13 | return np.getbuffer(array) 14 | 15 | def blob_to_array(blob, dtype, shape=(-1,)): 16 | if IS_PYTHON3: 17 | return np.fromstring(blob, dtype=dtype).reshape(*shape) 18 | else: 19 | return np.frombuffer(blob, dtype=dtype).reshape(*shape) 20 | 21 | class COLMAPDatabase(sqlite3.Connection): 22 | 23 | @staticmethod 24 | def connect(database_path): 25 | return sqlite3.connect(database_path, factory=COLMAPDatabase) 26 | 27 | def __init__(self, *args, **kwargs): 28 | super(COLMAPDatabase, self).__init__(*args, **kwargs) 29 | 30 | self.create_tables = lambda: self.executescript(CREATE_ALL) 31 | self.create_cameras_table = \ 32 | lambda: self.executescript(CREATE_CAMERAS_TABLE) 33 | self.create_descriptors_table = \ 34 | lambda: self.executescript(CREATE_DESCRIPTORS_TABLE) 35 | self.create_images_table = \ 36 | lambda: self.executescript(CREATE_IMAGES_TABLE) 37 | self.create_two_view_geometries_table = \ 38 | lambda: self.executescript(CREATE_TWO_VIEW_GEOMETRIES_TABLE) 39 | self.create_keypoints_table = \ 40 | lambda: self.executescript(CREATE_KEYPOINTS_TABLE) 41 | self.create_matches_table = \ 42 | lambda: self.executescript(CREATE_MATCHES_TABLE) 43 | self.create_name_index = lambda: self.executescript(CREATE_NAME_INDEX) 44 | 45 | def update_camera(self, model, width, height, params, camera_id): 46 | params = np.asarray(params, np.float64) 47 | cursor = self.execute( 48 | "UPDATE cameras SET model=?, width=?, height=?, params=?, prior_focal_length=True WHERE camera_id=?", 49 | (model, width, height, array_to_blob(params),camera_id)) 50 | return cursor.lastrowid 51 | 52 | def camTodatabase(): 53 | import os 54 | import argparse 55 | 56 | camModelDict = {'SIMPLE_PINHOLE': 0, 57 | 'PINHOLE': 1, 58 | 'SIMPLE_RADIAL': 2, 59 | 'RADIAL': 3, 60 | 'OPENCV': 4, 61 | 'FULL_OPENCV': 5, 62 | 'SIMPLE_RADIAL_FISHEYE': 6, 63 | 'RADIAL_FISHEYE': 7, 64 | 'OPENCV_FISHEYE': 8, 65 | 'FOV': 9, 66 | 'THIN_PRISM_FISHEYE': 10} 67 | parser = argparse.ArgumentParser() 68 | parser.add_argument("--database_path", type=str, default="database.db") 69 | parser.add_argument("--txt_path", type=str, default="colmap/sparse_cameras.txt") 70 | # breakpoint() 71 | args = parser.parse_args() 72 | if os.path.exists(args.database_path)==False: 73 | print("ERROR: database path dosen't exist -- please check database.db.") 74 | return 75 | # Open the database. 76 | db = COLMAPDatabase.connect(args.database_path) 77 | 78 | idList=list() 79 | modelList=list() 80 | widthList=list() 81 | heightList=list() 82 | paramsList=list() 83 | # Update real cameras from .txt 84 | with open(args.txt_path, "r") as cam: 85 | lines = cam.readlines() 86 | for i in range(0,len(lines),1): 87 | if lines[i][0]!='#': 88 | strLists = lines[i].split() 89 | cameraId=int(strLists[0]) 90 | cameraModel=camModelDict[strLists[1]] #SelectCameraModel 91 | width=int(strLists[2]) 92 | height=int(strLists[3]) 93 | paramstr=np.array(strLists[4:12]) 94 | params = paramstr.astype(np.float64) 95 | idList.append(cameraId) 96 | modelList.append(cameraModel) 97 | widthList.append(width) 98 | heightList.append(height) 99 | paramsList.append(params) 100 | camera_id = db.update_camera(cameraModel, width, height, params, cameraId) 101 | 102 | # Commit the data to the file. 103 | db.commit() 104 | # Read and check cameras. 105 | rows = db.execute("SELECT * FROM cameras") 106 | for i in range(0,len(idList),1): 107 | camera_id, model, width, height, params, prior = next(rows) 108 | params = blob_to_array(params, np.float64) 109 | assert camera_id == idList[i] 110 | assert model == modelList[i] and width == widthList[i] and height == heightList[i] 111 | assert np.allclose(params, paramsList[i]) 112 | 113 | # Close database.db. 114 | db.close() 115 | 116 | if __name__ == "__main__": 117 | import sys,os 118 | 119 | camTodatabase() -------------------------------------------------------------------------------- /gaussian_renderer/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | import torch 13 | import math 14 | from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer 15 | from scene.gaussian_model import GaussianModel 16 | from utils.sh_utils import eval_sh 17 | 18 | def render(viewpoint_camera, pc : GaussianModel, pipe, bg_color : torch.Tensor, scaling_modifier = 1.0, override_color = None, delta=None, 19 | time=None, it=1): 20 | """ 21 | Render the scene. 22 | 23 | Background tensor (bg_color) must be on GPU! 24 | """ 25 | 26 | # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means 27 | screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0 28 | try: 29 | screenspace_points.retain_grad() 30 | except: 31 | pass 32 | 33 | # Set up rasterization configuration 34 | tanfovx = math.tan(viewpoint_camera.FoVx * 0.5) 35 | tanfovy = math.tan(viewpoint_camera.FoVy * 0.5) 36 | 37 | raster_settings = GaussianRasterizationSettings( 38 | image_height=int(viewpoint_camera.image_height), 39 | image_width=int(viewpoint_camera.image_width), 40 | tanfovx=tanfovx, 41 | tanfovy=tanfovy, 42 | bg=bg_color, 43 | scale_modifier=scaling_modifier, 44 | viewmatrix=viewpoint_camera.world_view_transform, 45 | projmatrix=viewpoint_camera.full_proj_transform, 46 | sh_degree=pc.active_sh_degree, 47 | campos=viewpoint_camera.camera_center, 48 | prefiltered=False, 49 | # debug=pipe.debug 50 | ) 51 | 52 | rasterizer = GaussianRasterizer(raster_settings=raster_settings) 53 | 54 | opacity = pc.get_opacity 55 | if time is None: 56 | if delta is not None: 57 | means3D = pc.get_xyz + delta 58 | else: 59 | means3D = pc.get_xyz 60 | means2D = screenspace_points 61 | 62 | # If precomputed 3d covariance is provided, use it. If not, then it will be computed from 63 | # scaling / rotation by the rasterizer. 64 | scales = None 65 | rotations = None 66 | cov3D_precomp = None 67 | if pipe.compute_cov3D_python: 68 | cov3D_precomp = pc.get_covariance(scaling_modifier) 69 | else: 70 | scales = pc.get_scaling 71 | rotations = pc.get_rotation 72 | else: 73 | cov3D_precomp = None 74 | xyz_t, r_t, s_t, o_t = pc(time, it) 75 | means3D = xyz_t 76 | means2D = screenspace_points 77 | rotations = r_t 78 | scales = s_t 79 | opacity = o_t 80 | 81 | # If precomputed colors are provided, use them. Otherwise, if it is desired to precompute colors 82 | # from SHs in Python, do it. If not, then SH -> RGB conversion will be done by rasterizer. 83 | shs = None 84 | colors_precomp = None 85 | if override_color is None: 86 | if pipe.convert_SHs_python: 87 | shs_view = pc.get_features.transpose(1, 2).view(-1, 3, (pc.max_sh_degree+1)**2) 88 | dir_pp = ((pc.get_xyz + delta) - viewpoint_camera.camera_center.repeat(pc.get_features.shape[0], 1)) 89 | dir_pp_normalized = dir_pp/dir_pp.norm(dim=1, keepdim=True) 90 | sh2rgb = eval_sh(pc.active_sh_degree, shs_view, dir_pp_normalized) 91 | colors_precomp = torch.clamp_min(sh2rgb + 0.5, 0.0) 92 | else: 93 | shs = pc.get_features 94 | else: 95 | colors_precomp = override_color 96 | 97 | # Rasterize visible Gaussians to image, obtain their radii (on screen). 98 | rendered_image, radii, depth, tidx = rasterizer( 99 | means3D = means3D, 100 | means2D = means2D, 101 | shs = shs, 102 | colors_precomp = colors_precomp, 103 | opacities = opacity, 104 | scales = scales, 105 | rotations = rotations, 106 | cov3D_precomp = cov3D_precomp) 107 | 108 | # Those Gaussians that were frustum culled or had a radius of 0 were not visible. 109 | # They will be excluded from value updates used in the splitting criteria. 110 | return {"render": rendered_image, 111 | "viewspace_points": screenspace_points, 112 | "visibility_filter" : radii > 0, 113 | "radii": radii, 114 | "depth": depth, 115 | "tidx": tidx} 116 | 117 | def render_motion(viewpoint_camera, pc : GaussianModel, pipe, bg_color : torch.Tensor, scaling_modifier = 1.0, override_color = None, 118 | xyz_t=None, r_t=None, opacity=None): 119 | """ 120 | Render the scene. 121 | 122 | Background tensor (bg_color) must be on GPU! 123 | """ 124 | 125 | # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means 126 | screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0 127 | try: 128 | screenspace_points.retain_grad() 129 | except: 130 | pass 131 | 132 | # Set up rasterization configuration 133 | tanfovx = math.tan(viewpoint_camera.FoVx * 0.5) 134 | tanfovy = math.tan(viewpoint_camera.FoVy * 0.5) 135 | 136 | raster_settings = GaussianRasterizationSettings( 137 | image_height=int(viewpoint_camera.image_height), 138 | image_width=int(viewpoint_camera.image_width), 139 | tanfovx=tanfovx, 140 | tanfovy=tanfovy, 141 | bg=bg_color, 142 | scale_modifier=scaling_modifier, 143 | viewmatrix=viewpoint_camera.world_view_transform, 144 | projmatrix=viewpoint_camera.full_proj_transform, 145 | sh_degree=pc.active_sh_degree, 146 | campos=viewpoint_camera.camera_center, 147 | prefiltered=False, 148 | ) 149 | 150 | rasterizer = GaussianRasterizer(raster_settings=raster_settings) 151 | 152 | opacity = pc.get_opacity if opacity is None else opacity 153 | cov3D_precomp = None 154 | means3D = xyz_t 155 | means2D = screenspace_points 156 | rotations = r_t 157 | scales = pc.get_scaling 158 | 159 | # If precomputed colors are provided, use them. Otherwise, if it is desired to precompute colors 160 | # from SHs in Python, do it. If not, then SH -> RGB conversion will be done by rasterizer. 161 | shs = None 162 | colors_precomp = None 163 | if override_color is None: 164 | if pipe.convert_SHs_python: 165 | shs_view = pc.get_features.transpose(1, 2).view(-1, 3, (pc.max_sh_degree+1)**2) 166 | dir_pp = ((pc.get_xyz) - viewpoint_camera.camera_center.repeat(pc.get_features.shape[0], 1)) 167 | dir_pp_normalized = dir_pp/dir_pp.norm(dim=1, keepdim=True) 168 | sh2rgb = eval_sh(pc.active_sh_degree, shs_view, dir_pp_normalized) 169 | colors_precomp = torch.clamp_min(sh2rgb + 0.5, 0.0) 170 | else: 171 | shs = pc.get_features 172 | else: 173 | colors_precomp = override_color 174 | 175 | # Rasterize visible Gaussians to image, obtain their radii (on screen). 176 | rendered_image, radii, depth, tidx = rasterizer( 177 | means3D = means3D, 178 | means2D = means2D, 179 | shs = shs, 180 | colors_precomp = colors_precomp, 181 | opacities = opacity, 182 | scales = scales, 183 | rotations = rotations, 184 | cov3D_precomp = cov3D_precomp) 185 | 186 | # Those Gaussians that were frustum culled or had a radius of 0 were not visible. 187 | # They will be excluded from value updates used in the splitting criteria. 188 | return {"render": rendered_image, 189 | "viewspace_points": screenspace_points, 190 | "visibility_filter" : radii > 0, 191 | "radii": radii} -------------------------------------------------------------------------------- /lpipsPyTorch/__init__.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from .modules.lpips import LPIPS 4 | 5 | 6 | def lpips(x: torch.Tensor, 7 | y: torch.Tensor, 8 | net_type: str = 'alex', 9 | version: str = '0.1'): 10 | r"""Function that measures 11 | Learned Perceptual Image Patch Similarity (LPIPS). 12 | 13 | Arguments: 14 | x, y (torch.Tensor): the input tensors to compare. 15 | net_type (str): the network type to compare the features: 16 | 'alex' | 'squeeze' | 'vgg'. Default: 'alex'. 17 | version (str): the version of LPIPS. Default: 0.1. 18 | """ 19 | device = x.device 20 | criterion = LPIPS(net_type, version).to(device) 21 | return criterion(x, y) 22 | -------------------------------------------------------------------------------- /lpipsPyTorch/modules/lpips.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | from .networks import get_network, LinLayers 5 | from .utils import get_state_dict 6 | 7 | 8 | class LPIPS(nn.Module): 9 | r"""Creates a criterion that measures 10 | Learned Perceptual Image Patch Similarity (LPIPS). 11 | 12 | Arguments: 13 | net_type (str): the network type to compare the features: 14 | 'alex' | 'squeeze' | 'vgg'. Default: 'alex'. 15 | version (str): the version of LPIPS. Default: 0.1. 16 | """ 17 | def __init__(self, net_type: str = 'alex', version: str = '0.1'): 18 | 19 | assert version in ['0.1'], 'v0.1 is only supported now' 20 | 21 | super(LPIPS, self).__init__() 22 | 23 | # pretrained network 24 | self.net = get_network(net_type) 25 | 26 | # linear layers 27 | self.lin = LinLayers(self.net.n_channels_list) 28 | self.lin.load_state_dict(get_state_dict(net_type, version)) 29 | 30 | def forward(self, x: torch.Tensor, y: torch.Tensor): 31 | feat_x, feat_y = self.net(x), self.net(y) 32 | 33 | diff = [(fx - fy) ** 2 for fx, fy in zip(feat_x, feat_y)] 34 | res = [l(d).mean((2, 3), True) for d, l in zip(diff, self.lin)] 35 | 36 | return torch.sum(torch.cat(res, 0), 0, True) 37 | -------------------------------------------------------------------------------- /lpipsPyTorch/modules/networks.py: -------------------------------------------------------------------------------- 1 | from typing import Sequence 2 | 3 | from itertools import chain 4 | 5 | import torch 6 | import torch.nn as nn 7 | from torchvision import models 8 | 9 | from .utils import normalize_activation 10 | 11 | 12 | def get_network(net_type: str): 13 | if net_type == 'alex': 14 | return AlexNet() 15 | elif net_type == 'squeeze': 16 | return SqueezeNet() 17 | elif net_type == 'vgg': 18 | return VGG16() 19 | else: 20 | raise NotImplementedError('choose net_type from [alex, squeeze, vgg].') 21 | 22 | 23 | class LinLayers(nn.ModuleList): 24 | def __init__(self, n_channels_list: Sequence[int]): 25 | super(LinLayers, self).__init__([ 26 | nn.Sequential( 27 | nn.Identity(), 28 | nn.Conv2d(nc, 1, 1, 1, 0, bias=False) 29 | ) for nc in n_channels_list 30 | ]) 31 | 32 | for param in self.parameters(): 33 | param.requires_grad = False 34 | 35 | 36 | class BaseNet(nn.Module): 37 | def __init__(self): 38 | super(BaseNet, self).__init__() 39 | 40 | # register buffer 41 | self.register_buffer( 42 | 'mean', torch.Tensor([-.030, -.088, -.188])[None, :, None, None]) 43 | self.register_buffer( 44 | 'std', torch.Tensor([.458, .448, .450])[None, :, None, None]) 45 | 46 | def set_requires_grad(self, state: bool): 47 | for param in chain(self.parameters(), self.buffers()): 48 | param.requires_grad = state 49 | 50 | def z_score(self, x: torch.Tensor): 51 | return (x - self.mean) / self.std 52 | 53 | def forward(self, x: torch.Tensor): 54 | x = self.z_score(x) 55 | 56 | output = [] 57 | for i, (_, layer) in enumerate(self.layers._modules.items(), 1): 58 | x = layer(x) 59 | if i in self.target_layers: 60 | output.append(normalize_activation(x)) 61 | if len(output) == len(self.target_layers): 62 | break 63 | return output 64 | 65 | 66 | class SqueezeNet(BaseNet): 67 | def __init__(self): 68 | super(SqueezeNet, self).__init__() 69 | 70 | self.layers = models.squeezenet1_1(True).features 71 | self.target_layers = [2, 5, 8, 10, 11, 12, 13] 72 | self.n_channels_list = [64, 128, 256, 384, 384, 512, 512] 73 | 74 | self.set_requires_grad(False) 75 | 76 | 77 | class AlexNet(BaseNet): 78 | def __init__(self): 79 | super(AlexNet, self).__init__() 80 | 81 | self.layers = models.alexnet(True).features 82 | self.target_layers = [2, 5, 8, 10, 12] 83 | self.n_channels_list = [64, 192, 384, 256, 256] 84 | 85 | self.set_requires_grad(False) 86 | 87 | 88 | class VGG16(BaseNet): 89 | def __init__(self): 90 | super(VGG16, self).__init__() 91 | 92 | self.layers = models.vgg16(weights=models.VGG16_Weights.IMAGENET1K_V1).features 93 | self.target_layers = [4, 9, 16, 23, 30] 94 | self.n_channels_list = [64, 128, 256, 512, 512] 95 | 96 | self.set_requires_grad(False) 97 | -------------------------------------------------------------------------------- /lpipsPyTorch/modules/utils.py: -------------------------------------------------------------------------------- 1 | from collections import OrderedDict 2 | 3 | import torch 4 | 5 | 6 | def normalize_activation(x, eps=1e-10): 7 | norm_factor = torch.sqrt(torch.sum(x ** 2, dim=1, keepdim=True)) 8 | return x / (norm_factor + eps) 9 | 10 | 11 | def get_state_dict(net_type: str = 'alex', version: str = '0.1'): 12 | # build url 13 | url = 'https://raw.githubusercontent.com/richzhang/PerceptualSimilarity/' \ 14 | + f'master/lpips/weights/v{version}/{net_type}.pth' 15 | 16 | # download 17 | old_state_dict = torch.hub.load_state_dict_from_url( 18 | url, progress=True, 19 | map_location=None if torch.cuda.is_available() else torch.device('cpu') 20 | ) 21 | 22 | # rename keys 23 | new_state_dict = OrderedDict() 24 | for key, val in old_state_dict.items(): 25 | new_key = key 26 | new_key = new_key.replace('lin', '') 27 | new_key = new_key.replace('model.', '') 28 | new_state_dict[new_key] = val 29 | 30 | return new_state_dict 31 | -------------------------------------------------------------------------------- /metrics.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | from pathlib import Path 13 | import os 14 | from PIL import Image 15 | import numpy as np 16 | import torch 17 | import torchvision.transforms.functional as tf 18 | from utils.loss_utils import ssim 19 | from lpipsPyTorch import lpips 20 | import json 21 | from tqdm import tqdm 22 | from utils.image_utils import psnr 23 | from argparse import ArgumentParser 24 | from pytorch_msssim import ms_ssim 25 | 26 | def readImages(renders_dir, gt_dir, resize=False, resize_ratio=0.5): 27 | renders = [] 28 | gts = [] 29 | image_names = [] 30 | image_gt_names = [] 31 | 32 | render_fname_list = sorted(os.listdir(renders_dir)) 33 | gt_fname_list = sorted(os.listdir(gt_dir)) 34 | 35 | for fname in render_fname_list: 36 | if "depth" in fname: 37 | continue 38 | render = Image.open(os.path.join(renders_dir, fname)) 39 | size = render.size 40 | W = int(size[0] * resize_ratio) 41 | H = int(size[1] * resize_ratio) 42 | new_size = (W, H) 43 | if resize: 44 | render = render.resize(new_size) 45 | renders.append(tf.to_tensor(render).unsqueeze(0)[:, :3, :, :].cuda()) 46 | image_names.append(fname) 47 | for fname in gt_fname_list: 48 | gt = Image.open(os.path.join(gt_dir, fname)) 49 | size = gt.size 50 | W = int(size[0] * resize_ratio) 51 | H = int(size[1] * resize_ratio) 52 | new_size = (W, H) 53 | if resize: 54 | gt = gt.resize(new_size) 55 | gts.append(tf.to_tensor(gt).unsqueeze(0)[:, :3, :, :].cuda()) 56 | image_gt_names.append(fname) 57 | return renders, gts, image_names 58 | 59 | def evaluate_twodir(gt_dir, renders_dir): 60 | try: 61 | full_dict = {} 62 | per_view_dict = {} 63 | renders, gts, image_names = readImages(renders_dir, gt_dir) 64 | 65 | ssims = [] 66 | psnrs = [] 67 | lpipss = [] 68 | lpipsa = [] 69 | ms_ssims = [] 70 | Dssims = [] 71 | for idx in tqdm(range(len(renders)), desc="Metric evaluation progress"): 72 | ssims.append(ssim(renders[idx], gts[idx])) 73 | psnrs.append(psnr(renders[idx], gts[idx])) 74 | lpipss.append(lpips(renders[idx], gts[idx], net_type='vgg')) 75 | ms_ssims.append(ms_ssim(renders[idx], gts[idx],data_range=1, size_average=True )) 76 | lpipsa.append(lpips(renders[idx], gts[idx], net_type='alex')) 77 | Dssims.append((1-ms_ssims[-1])/2) 78 | 79 | print("Scene: ", renders_dir, "SSIM : {:>12.7f}".format(torch.tensor(ssims).mean(), ".5")) 80 | print("Scene: ", renders_dir, "PSNR : {:>12.7f}".format(torch.tensor(psnrs).mean(), ".5")) 81 | print("Scene: ", renders_dir, "LPIPS-vgg: {:>12.7f}".format(torch.tensor(lpipss).mean(), ".5")) 82 | print("Scene: ", renders_dir, "LPIPS-alex: {:>12.7f}".format(torch.tensor(lpipsa).mean(), ".5")) 83 | print("Scene: ", renders_dir, "MS-SSIM: {:>12.7f}".format(torch.tensor(ms_ssims).mean(), ".5")) 84 | print("Scene: ", renders_dir, "D-SSIM: {:>12.7f}".format(torch.tensor(Dssims).mean(), ".5")) 85 | 86 | full_dict.update({"SSIM": torch.tensor(ssims).mean().item(), 87 | "PSNR": torch.tensor(psnrs).mean().item(), 88 | "LPIPS-vgg": torch.tensor(lpipss).mean().item(), 89 | "LPIPS-alex": torch.tensor(lpipsa).mean().item(), 90 | "MS-SSIM": torch.tensor(ms_ssims).mean().item(), 91 | "D-SSIM": torch.tensor(Dssims).mean().item()}, 92 | 93 | ) 94 | per_view_dict.update({"SSIM": {name: ssim for ssim, name in zip(torch.tensor(ssims).tolist(), image_names)}, 95 | "PSNR": {name: psnr for psnr, name in zip(torch.tensor(psnrs).tolist(), image_names)}, 96 | "LPIPS-vgg": {name: lp for lp, name in zip(torch.tensor(lpipss).tolist(), image_names)}, 97 | "LPIPS-alex": {name: lp for lp, name in zip(torch.tensor(lpipsa).tolist(), image_names)}, 98 | "MS-SSIM": {name: lp for lp, name in zip(torch.tensor(ms_ssims).tolist(), image_names)}, 99 | "D-SSIM": {name: lp for lp, name in zip(torch.tensor(Dssims).tolist(), image_names)}, 100 | 101 | } 102 | ) 103 | 104 | with open(renders_dir + "/results.json", 'w') as fp: 105 | json.dump(full_dict, fp, indent=True) 106 | with open(renders_dir + "/per_view.json", 'w') as fp: 107 | json.dump(per_view_dict, fp, indent=True) 108 | 109 | except Exception as e: 110 | print("Unable to compute metrics for model", renders_dir) 111 | raise e 112 | 113 | def evaluate(model_paths, resize=False, resize_ratio=0.5): 114 | 115 | full_dict = {} 116 | per_view_dict = {} 117 | full_dict_polytopeonly = {} 118 | per_view_dict_polytopeonly = {} 119 | print("") 120 | 121 | for method in os.listdir(model_paths): 122 | if not os.path.isdir(os.path.join(model_paths, method)): 123 | continue 124 | try: 125 | print("Method:", method) 126 | 127 | full_dict[method] = {} 128 | per_view_dict[method] = {} 129 | full_dict_polytopeonly[method] = {} 130 | per_view_dict_polytopeonly[method] = {} 131 | 132 | method_dir = os.path.join(model_paths, method) 133 | gt_dir = os.path.join(method_dir, "gt") 134 | renders_dir = os.path.join(method_dir, "renders") 135 | os.makedirs(os.path.join(method_dir, "deltas"), exist_ok=True) 136 | renders, gts, image_names = readImages(renders_dir, gt_dir, resize=resize, resize_ratio=resize_ratio) 137 | 138 | ssims, psnrs, lpipss, lpipsa, ms_ssims, Dssims = [], [], [], [], [], [] 139 | for idx in tqdm(range(len(renders)), desc="Metric evaluation progress"): 140 | ssims.append(ssim(renders[idx], gts[idx])) 141 | psnrs.append(psnr(renders[idx], gts[idx])) 142 | lpipss.append(lpips(renders[idx], gts[idx], net_type='vgg')) 143 | ms_ssims.append(ms_ssim(renders[idx], gts[idx],data_range=1, size_average=True )) 144 | lpipsa.append(lpips(renders[idx], gts[idx], net_type='alex')) 145 | Dssims.append((1-ms_ssims[-1])/2) 146 | error = (renders[idx] - gts[idx])[0].permute(1, 2, 0).abs().cpu().numpy() 147 | error_img = Image.fromarray((error * 255).astype(np.uint8)) 148 | error_img.save(os.path.join(method_dir, "deltas", "{0:05d}.jpg".format(idx))) 149 | 150 | print("Scene: ", model_paths, "SSIM : {:>12.7f}".format(torch.tensor(ssims).mean(), ".5")) 151 | print("Scene: ", model_paths, "PSNR : {:>12.7f}".format(torch.tensor(psnrs).mean(), ".5")) 152 | print("Scene: ", model_paths, "LPIPS-vgg: {:>12.7f}".format(torch.tensor(lpipss).mean(), ".5")) 153 | print("Scene: ", model_paths, "LPIPS-alex: {:>12.7f}".format(torch.tensor(lpipsa).mean(), ".5")) 154 | print("Scene: ", model_paths, "MS-SSIM: {:>12.7f}".format(torch.tensor(ms_ssims).mean(), ".5")) 155 | print("Scene: ", model_paths, "D-SSIM: {:>12.7f}".format(torch.tensor(Dssims).mean(), ".5")) 156 | 157 | full_dict[method].update({"SSIM": torch.tensor(ssims).mean().item(), 158 | "PSNR": torch.tensor(psnrs).mean().item(), 159 | "LPIPS-vgg": torch.tensor(lpipss).mean().item(), 160 | "LPIPS-alex": torch.tensor(lpipsa).mean().item(), 161 | "MS-SSIM": torch.tensor(ms_ssims).mean().item(), 162 | "D-SSIM": torch.tensor(Dssims).mean().item()}) 163 | 164 | per_view_dict[method].update({"SSIM": {name: ssim for ssim, name in zip(torch.tensor(ssims).tolist(), image_names)}, 165 | "PSNR": {name: psnr for psnr, name in zip(torch.tensor(psnrs).tolist(), image_names)}, 166 | "LPIPS-vgg": {name: lp for lp, name in zip(torch.tensor(lpipss).tolist(), image_names)}, 167 | "LPIPS-alex": {name: lp for lp, name in zip(torch.tensor(lpipsa).tolist(), image_names)}, 168 | "MS-SSIM": {name: lp for lp, name in zip(torch.tensor(ms_ssims).tolist(), image_names)}, 169 | "D-SSIM": {name: lp for lp, name in zip(torch.tensor(Dssims).tolist(), image_names)}}) 170 | 171 | with open(model_paths + "/results.json", 'w') as fp: 172 | json.dump(full_dict[method], fp, indent=True) 173 | with open(model_paths + "/per_view.json", 'w') as fp: 174 | json.dump(per_view_dict[method], fp, indent=True) 175 | 176 | except Exception as e: 177 | print("Unable to compute metrics for model", model_paths) 178 | raise e 179 | 180 | if __name__ == "__main__": 181 | device = torch.device("cuda:0") 182 | torch.cuda.set_device(device) 183 | 184 | # Set up command line argument parser 185 | parser = ArgumentParser(description="Training script parameters") 186 | parser.add_argument('--model_paths', '-m', required=True, nargs="+", type=str, default=[]) 187 | parser.add_argument('--gt_paths', '-g', type=str, default="") 188 | parser.add_argument('--split', action='store_true', default=False) 189 | parser.add_argument('--resize', action='store_true', default=False) 190 | args = parser.parse_args() 191 | if args.split: 192 | evaluate_twodir(args.gt_paths, args.model_paths[0]) 193 | else: 194 | evaluate(args.model_paths) -------------------------------------------------------------------------------- /motion_model/dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch.utils.data as data 3 | from abc import ABC, abstractmethod 4 | import torch 5 | import torch.nn.functional as F 6 | from tqdm import tqdm, trange 7 | from scene.deformable_field import positional_encoding 8 | import json 9 | import numpy as np 10 | 11 | class GCNBaseDataset(data.Dataset, ABC): 12 | def __init__(self, gaussians, time_freq, iteration, model_path, source_path, max_time=0.8, input_size=20, output_size=5, split="train"): 13 | super().__init__() 14 | self.gaussians = gaussians 15 | self.time_freq = time_freq 16 | self.iteration = iteration 17 | self.split = split 18 | self.max_time = max_time 19 | assert self.max_time < 1.0 20 | self.input_size = input_size 21 | self.output_size = output_size 22 | self.model_path = model_path 23 | self.source_path = source_path 24 | 25 | self.train_data = [] 26 | self.test_data = [] 27 | self.val_data = [] 28 | 29 | self.process_data() 30 | 31 | def load_hyper_times(self): 32 | datadir = self.source_path 33 | with open(f'{datadir}/metadata.json', 'r') as f: 34 | meta_json = json.load(f) 35 | with open(f'{datadir}/dataset.json', 'r') as f: 36 | dataset_json = json.load(f) 37 | self.all_img = dataset_json['ids'] 38 | self.val_id = dataset_json['val_ids'] 39 | self.all_time = [meta_json[i]['warp_id'] for i in self.all_img] 40 | 41 | max_time = max(self.all_time) 42 | if self.max_time < 1.0: 43 | self.all_time, self.i_train, self.i_test = [], [], [] 44 | for idx, i in enumerate(self.all_img): 45 | time = meta_json[i]['warp_id']/max_time 46 | self.all_time += [time] 47 | if len(self.val_id) == 0: 48 | if idx % 4 ==0 and time < self.max_time: 49 | self.i_train += [idx] 50 | if (idx - 2) % 4 ==0 and time >= self.max_time: 51 | self.i_test += [idx] 52 | else: 53 | self.train_id = dataset_json['train_ids'] 54 | if i in self.val_id and time >= self.max_time: 55 | self.i_test.append(idx) 56 | if i in self.train_id and time < self.max_time: 57 | self.i_train.append(idx) 58 | self.i_test = np.array(self.i_test) 59 | self.i_train = np.array(self.i_train) 60 | np_all_time = np.array(self.all_time) 61 | test_time = np_all_time[self.i_test] 62 | train_time = np_all_time[self.i_train] 63 | assert test_time.max() >= self.max_time 64 | assert train_time.min() < self.max_time 65 | print("train:", self.i_train) 66 | print("test:", self.i_test) 67 | else: 68 | self.all_time = [meta_json[i]['warp_id']/max_time for i in self.all_img] 69 | 70 | self.test_times = [self.all_time[idx] for idx in self.i_test] 71 | self.train_times = [self.all_time[idx] for idx in self.i_train] 72 | if self.split == "test": 73 | print(f"Check!!!!!There are {len(self.test_times)} test data!!!!") 74 | 75 | 76 | def load_dnerf_times(self): 77 | with open(os.path.join(self.source_path, "transforms_train.json"), "r") as train_file: 78 | data = json.load(train_file)["frames"] 79 | self.train_times, self.test_times = [], [] 80 | for i in range(len(data)): 81 | if float(data[i]["time"]) < self.max_time: 82 | self.train_times.append(float(data[i]["time"])) 83 | else: 84 | self.test_times.append(float(data[i]["time"])) 85 | train_file.close() 86 | 87 | def load_times(self): 88 | if "d-nerf" in self.model_path or "white" in self.model_path: 89 | print("Find transforms_train.json file, assume it D-NeRF Dataset!") 90 | self.load_dnerf_times() 91 | else: 92 | print("Find metadata.json file, assume it HyperNeRF Dataset!") 93 | self.load_hyper_times() 94 | 95 | @abstractmethod 96 | def get_lens(self): 97 | pass 98 | 99 | @abstractmethod 100 | def prepare_item(self): 101 | pass 102 | 103 | def process_data(self): 104 | self.load_times() 105 | self.generate_data() 106 | self.get_lens() 107 | self.prepare_item() 108 | 109 | @property 110 | def nodes_num(self): 111 | return self.gaussians.super_gaussians.shape[0] 112 | 113 | def generate_data(self): 114 | kpts_xyz, kpts_r = [], [] 115 | self.kpts_xyz_motion_train, self.kpts_r_motion_train = [], [] 116 | self.kpts_xyz_motion_test, self.kpts_r_motion_test = [], [] 117 | for i in range(len(self.train_times)): 118 | time_ = torch.tensor([self.train_times[i]], dtype=torch.float32).cuda() 119 | _, _, _, _ = self.gaussians(time_, self.iteration) 120 | kpts_xyz += [self.gaussians.get_superGaussians + self.gaussians.kpts_xyz_motion] 121 | kpts_r += [self.gaussians.kpts_rotation_motion] 122 | self.kpts_xyz_motion_train += [self.gaussians.kpts_xyz_motion] 123 | self.kpts_r_motion_train += [self.gaussians.kpts_rotation_motion] 124 | self.kpts_xyz_train = torch.stack(kpts_xyz, dim=0) 125 | self.kpts_r_train = torch.stack(kpts_r, dim=0) 126 | kpts_xyz, kpts_r = [], [] 127 | for i in range(len(self.test_times)): 128 | time_ = torch.tensor([self.test_times[i]], dtype=torch.float32).cuda() 129 | _, _, _, _ = self.gaussians(time_, self.iteration) 130 | kpts_xyz += [self.gaussians.get_superGaussians + self.gaussians.kpts_xyz_motion] 131 | kpts_r += [self.gaussians.kpts_rotation_motion] 132 | self.kpts_xyz_motion_test += [self.gaussians.kpts_xyz_motion] 133 | self.kpts_r_motion_test += [self.gaussians.kpts_rotation_motion] 134 | self.kpts_xyz_test = torch.stack(kpts_xyz, dim=0) 135 | self.kpts_r_test = torch.stack(kpts_r, dim=0) 136 | 137 | def __len__(self): 138 | if self.split == "train": 139 | # return self.train_lens 140 | return len(self.train_data) 141 | elif self.split == "test": 142 | # return self.test_lens 143 | return len(self.test_data) 144 | else: 145 | # return self.val_lens 146 | return len(self.val_data) 147 | 148 | class GCN3DDataset(GCNBaseDataset): 149 | def get_lens(self): 150 | self.train_lens = len(self.kpts_xyz_train) - self.input_size - self.output_size 151 | self.test_lens = len(self.kpts_xyz_test) 152 | self.val_lens = 2 153 | 154 | def prepare_item(self): 155 | 156 | if self.split == "train": 157 | for i in trange(self.train_lens, desc="loading training data"): 158 | xyz_input = self.kpts_xyz_train[i:i+self.input_size] 159 | xyz_gt = self.kpts_xyz_train[i+self.input_size:i+self.input_size+self.output_size] 160 | time_interval = self.train_times[i+self.input_size] - self.train_times[i+self.input_size - 1] 161 | rotation_input = self.kpts_r_train[i:i+self.input_size] 162 | rotation_gt = self.kpts_r_train[i+self.input_size:i+self.input_size+self.output_size] 163 | self.train_data.append({"xyz_inputs": xyz_input, "xyz_gt": xyz_gt, 164 | "rotation_inputs": rotation_input, "rotation_gt": rotation_gt, "time": time_interval}) 165 | elif self.split == "test": 166 | concat_kpts_xyz_test = torch.cat([self.kpts_xyz_train[-self.input_size:], self.kpts_xyz_test], dim=0) 167 | concat_kpts_r_test = torch.cat([self.kpts_r_train[-self.input_size:], self.kpts_r_test], dim=0) 168 | concat_times = self.train_times[-self.input_size:] + self.test_times 169 | for i in trange(0, self.test_lens, self.output_size, desc="loading testing data"): 170 | xyz_input = concat_kpts_xyz_test[i:i+self.input_size] 171 | xyz_gt = concat_kpts_xyz_test[i+self.input_size:i+self.input_size+self.output_size] 172 | time_interval = concat_times[i+self.input_size] - concat_times[i+self.input_size - 1] 173 | rotation_input = concat_kpts_r_test[i:i+self.input_size] 174 | rotation_gt = concat_kpts_r_test[i+self.input_size:i+self.input_size+self.output_size] 175 | self.test_data.append({"xyz_inputs": xyz_input, "xyz_gt": xyz_gt, 176 | "rotation_inputs": rotation_input, "rotation_gt": rotation_gt, "time": time_interval}) 177 | else: 178 | idx = torch.randint(0, self.train_lens, [1]).item() 179 | self.val_data.append(self.train_data[idx]) 180 | last_test_xyz = self.kpts_xyz_test[-self.input_size:] 181 | last_test_r = self.kpts_r_test[-self.input_size:] 182 | self.val_data.append({"xyz_inputs": last_test_xyz, "xyz_gt": torch.zeros([self.output_size, *last_test_xyz.shape[1:]]).to(last_test_xyz), 183 | "rotation_inputs": last_test_r, "rotation_gt": torch.zeros([self.output_size, *last_test_r.shape[1:]]).to(last_test_r)}) 184 | 185 | def __getitem__(self, index): 186 | if self.split == "train": 187 | return self.train_data[index] 188 | elif self.split == "test": 189 | return self.test_data[index] 190 | else: 191 | return self.val_data[index] 192 | -------------------------------------------------------------------------------- /motion_model/gcn.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import torch 3 | import math 4 | from torch.nn.parameter import Parameter 5 | import torch.nn.functional as F 6 | import numpy as np 7 | 8 | class SemskeConv(nn.Module): 9 | 10 | def __init__(self, node_num=100, bias=True): 11 | super(SemskeConv, self).__init__() 12 | self.node_num = node_num 13 | 14 | A_sem = nn.Parameter(torch.zeros(node_num, node_num)) #Af 15 | self.A_sem = A_sem 16 | self.M = nn.Parameter(torch.zeros(node_num, node_num)) 17 | self.W = nn.Parameter(torch.zeros(node_num, node_num)) 18 | 19 | if bias: 20 | self.bias = nn.Parameter(torch.FloatTensor(node_num)) 21 | stdv = 1. / math.sqrt(self.M.size(1)) 22 | self.bias.data.uniform_(-stdv, stdv) 23 | 24 | else: 25 | self.register_parameter('bias', None) 26 | 27 | def forward(self, input): 28 | 29 | self.A_sem = nn.Parameter(torch.where(torch.isnan(self.A_sem), torch.full_like(self.A_sem, 0), self.A_sem)).cuda() 30 | self.W = nn.Parameter(torch.where(torch.isnan(self.W), torch.full_like(self.W, 0), self.W)) 31 | self.M = nn.Parameter(torch.where(torch.isnan(self.M), torch.full_like(self.M, 0), self.M)) 32 | Adj = self.A_sem 33 | Adj_W = torch.mul(Adj, self.W) 34 | support = torch.matmul(input, Adj_W) 35 | output = torch.matmul(support, self.M) 36 | 37 | if self.bias is not None: 38 | return output + self.bias 39 | else: 40 | return output 41 | 42 | class _GraphConv(nn.Module): 43 | def __init__(self, in_features, hidden_feature, node_num, p_dropout= 0): 44 | super(_GraphConv, self).__init__() 45 | 46 | self.gconv1 = SemskeConv(node_num) 47 | self.bn = nn.BatchNorm1d(node_num * hidden_feature) 48 | 49 | self.gconv2 = SemskeConv(node_num) 50 | 51 | self.tanh = nn.Tanh() 52 | 53 | if p_dropout is not None: 54 | self.dropout = nn.Dropout(p_dropout) 55 | else: 56 | self.dropout = None 57 | 58 | def forward(self, x): 59 | 60 | y = self.gconv1(x) 61 | 62 | y = self.tanh(y) 63 | if self.dropout is not None: 64 | y = self.dropout(y) 65 | y = self.gconv2(y) 66 | 67 | y = self.tanh(y) 68 | y = y + x 69 | 70 | return y 71 | 72 | 73 | class Generator(nn.Module): 74 | 75 | def __init__(self, input_size, hidden_size, output_size, node_num, batch_size, num_layers=2): 76 | super(Generator, self).__init__() 77 | 78 | 79 | self.hidden_prev = nn.Parameter(torch.zeros(num_layers, batch_size, hidden_size)) 80 | 81 | 82 | self.GRU = nn.GRU(input_size=input_size,hidden_size=hidden_size, 83 | num_layers=num_layers, dropout=0, batch_first = True) 84 | 85 | 86 | self.GCN = _GraphConv(1, 10, node_num) 87 | 88 | self.linear = nn.Linear(hidden_size, output_size) 89 | 90 | 91 | def forward(self, x, hidden_size): 92 | 93 | # GCN block 94 | x = x.permute(0, 2, 1) 95 | GCN_set = self.GCN(x) 96 | 97 | 98 | x = GCN_set.reshape(x.shape[0],x.shape[1],x.shape[2]) 99 | x = x.permute(0, 2, 1) 100 | 101 | 102 | out, h = self.GRU(x, self.hidden_prev) 103 | out = out.reshape(-1, hidden_size) 104 | out = self.linear(out) 105 | out = out.unsqueeze(dim=0) 106 | return out, h 107 | 108 | class GraphConvolution(nn.Module): 109 | """ 110 | adapted from : https://github.com/tkipf/gcn/blob/92600c39797c2bfb61a508e52b88fb554df30177/gcn/layers.py#L132 111 | """ 112 | 113 | def __init__(self, in_features, out_features, bias=True, node_n=48): 114 | super(GraphConvolution, self).__init__() 115 | self.in_features = in_features 116 | self.out_features = out_features 117 | self.weight = Parameter(torch.FloatTensor(in_features, out_features)) 118 | self.att = Parameter(torch.FloatTensor(node_n, node_n)) 119 | if bias: 120 | self.bias = Parameter(torch.FloatTensor(out_features)) 121 | else: 122 | self.register_parameter('bias', None) 123 | self.reset_parameters() 124 | 125 | def reset_parameters(self): 126 | stdv = 1. / math.sqrt(self.weight.size(1)) 127 | self.weight.data.uniform_(-stdv, stdv) 128 | self.att.data.uniform_(-stdv, stdv) 129 | if self.bias is not None: 130 | self.bias.data.uniform_(-stdv, stdv) 131 | 132 | def forward(self, input): 133 | support = torch.matmul(input, self.weight) 134 | output = torch.matmul(self.att, support) 135 | if self.bias is not None: 136 | return output + self.bias 137 | else: 138 | return output 139 | 140 | def __repr__(self): 141 | return self.__class__.__name__ + ' (' \ 142 | + str(self.in_features) + ' -> ' \ 143 | + str(self.out_features) + ')' 144 | 145 | 146 | class GC_Block(nn.Module): 147 | def __init__(self, in_features, p_dropout, bias=True, node_n=48): 148 | """ 149 | Define a residual block of GCN 150 | """ 151 | super(GC_Block, self).__init__() 152 | self.in_features = in_features 153 | self.out_features = in_features 154 | 155 | self.gc1 = GraphConvolution(in_features, in_features, node_n=node_n, bias=bias) 156 | self.bn1 = nn.BatchNorm1d(node_n * in_features) 157 | 158 | self.gc2 = GraphConvolution(in_features, in_features, node_n=node_n, bias=bias) 159 | self.bn2 = nn.BatchNorm1d(node_n * in_features) 160 | 161 | self.do = nn.Dropout(p_dropout) 162 | self.act_f = nn.Tanh() 163 | 164 | def forward(self, x): 165 | y = self.gc1(x) 166 | b, n, f = y.shape 167 | y = self.bn1(y.view(b, -1)).view(b, n, f) 168 | y = self.act_f(y) 169 | y = self.do(y) 170 | 171 | y = self.gc2(y) 172 | b, n, f = y.shape 173 | y = self.bn2(y.view(b, -1)).view(b, n, f) 174 | y = self.act_f(y) 175 | y = self.do(y) 176 | 177 | return y + x 178 | 179 | def __repr__(self): 180 | return self.__class__.__name__ + ' (' \ 181 | + str(self.in_features) + ' -> ' \ 182 | + str(self.out_features) + ')' 183 | 184 | 185 | class GCN(nn.Module): 186 | def __init__(self, input_feature, hidden_feature, output_feature, p_dropout, num_stage=1, node_n=48, no_mapping=False): 187 | """ 188 | 189 | :param input_feature: num of input feature 190 | :param hidden_feature: num of hidden feature 191 | :param p_dropout: drop out prob. 192 | :param num_stage: number of residual blocks 193 | :param node_n: number of nodes in graph 194 | """ 195 | super(GCN, self).__init__() 196 | self.num_stage = num_stage 197 | 198 | self.gc1 = GraphConvolution(input_feature, hidden_feature, node_n=node_n) 199 | self.bn1 = nn.BatchNorm1d(node_n * hidden_feature) 200 | 201 | self.gcbs = [] 202 | for i in range(num_stage): 203 | self.gcbs.append(GC_Block(hidden_feature, p_dropout=p_dropout, node_n=node_n)) 204 | 205 | self.gcbs = nn.ModuleList(self.gcbs) 206 | 207 | if no_mapping: 208 | self.gc_out = GraphConvolution(hidden_feature, output_feature, node_n=node_n) 209 | else: 210 | self.gc_out = nn.Sequential( 211 | nn.Linear(hidden_feature, hidden_feature), 212 | nn.ReLU(), 213 | nn.Linear(hidden_feature, output_feature) 214 | ) 215 | 216 | self.do = nn.Dropout(p_dropout) 217 | 218 | self.act_f = nn.Tanh() 219 | 220 | def forward(self, x): 221 | # print(x[0][0].cpu().detach().numpy()) 222 | y = self.gc1(x) 223 | b, n, f = y.shape 224 | y = self.bn1(y.view(b, -1)).view(b, n, f) 225 | y = self.act_f(y) 226 | y = self.do(y) 227 | 228 | for i in range(self.num_stage): 229 | y = self.gcbs[i](y) 230 | 231 | y = self.gc_out(y) 232 | # print(y[0][0].cpu().detach().numpy()) 233 | # y = y + x 234 | 235 | return y 236 | 237 | class Channel_GCN(nn.Module): 238 | def __init__(self, input_feature, hidden_feature, output_feature, p_dropout, num_stage=1, node_n=48, channel=3, no_mapping=False): 239 | super().__init__() 240 | self.channel = channel 241 | self.output_feature = output_feature 242 | # self.list = [] 243 | # for _ in range(channel): 244 | # self.list.append(GCN(input_feature, hidden_feature, output_feature, p_dropout, num_stage, node_n)) 245 | self.GCN = GCN(input_feature, hidden_feature, output_feature, p_dropout, num_stage, node_n*channel, no_mapping) 246 | 247 | # self.list = nn.ModuleList(self.list) 248 | 249 | def forward(self, x): 250 | ''' 251 | x: (B, C, nodes, input_feature) 252 | ''' 253 | B, C, N, _ = x.shape 254 | return self.GCN(x.reshape(x.shape[0], -1, x.shape[-1])).reshape((B, C, N, self.output_feature)) 255 | # output = [] 256 | # for i in range(self.channel): 257 | # output.append(self.list[i](x[:, i])) 258 | # return torch.stack(output, dim=1) 259 | 260 | class GCN_xyzr(nn.Module): 261 | def __init__(self, input_feature, hidden_feature, output_feature, p_dropout, num_stage=1, node_n=48, no_mapping=False): 262 | super().__init__() 263 | self.output_feature = output_feature 264 | self.GCN_xyz = Channel_GCN(input_feature, hidden_feature, output_feature, p_dropout, num_stage, node_n, 3, no_mapping) 265 | self.GCN_r = Channel_GCN(input_feature, hidden_feature, output_feature, p_dropout, num_stage, node_n, 4, no_mapping) 266 | 267 | def forward(self, x, r): 268 | ''' 269 | x: (B, 3, nodes, input_feature) 270 | r: (B, 4, nodes, input_feature) 271 | ''' 272 | x_out = self.GCN_xyz(x) 273 | r_out = F.normalize(self.GCN_r(r), dim=1) 274 | 275 | return x_out, r_out 276 | 277 | def get_dct_matrix(N): 278 | dct_m = np.eye(N) 279 | for k in np.arange(N): 280 | for i in np.arange(N): 281 | w = np.sqrt(2 / N) 282 | if k == 0: 283 | w = np.sqrt(1 / N) 284 | dct_m[k, i] = w * np.cos(np.pi * (i + 1 / 2) * k / N) 285 | idct_m = np.linalg.inv(dct_m) 286 | return dct_m, idct_m 287 | 288 | -------------------------------------------------------------------------------- /options/gaussian_option.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | class Gaussian_Options: 4 | def __init__(self, parser) -> None: 5 | self.parser = parser 6 | 7 | def get_parser(self): 8 | return self.parser 9 | 10 | def eval(self): 11 | self.parser.add_argument("--ckpt_iteration", default=30000, type=int) 12 | self.parser.add_argument("--render_train", action='store_true', default=False) 13 | self.parser.add_argument("--render_video", action='store_true', default=False) 14 | self.parser.add_argument("--train_view", nargs="+", type=int, default=[0]) 15 | self.parser.add_argument("--interpolation", type=int, default=5) 16 | 17 | def dynamic_training(self): 18 | # Adaptive strategy 19 | pass 20 | 21 | def gcn_training(self): 22 | self.parser.add_argument("--epoch", default=101, type=int) 23 | self.parser.add_argument("--batch_size", default=32, type=int) 24 | self.parser.add_argument("--no_mapping", action='store_true', default=False) 25 | self.parser.add_argument("--num_stage", default=4, type=int) 26 | self.parser.add_argument("--linear_size", default=128, type=int) 27 | self.parser.add_argument("--dropout", default=0., type=float) 28 | self.parser.add_argument("--Rscale_ratio", default=0.5, type=float) 29 | self.parser.add_argument("--input_size", default=10, type=int) 30 | self.parser.add_argument("--output_size", default=1, type=int) 31 | self.parser.add_argument("--predict_more", action='store_true', default=False) 32 | self.parser.add_argument("--metrics", action='store_true', default=False) 33 | self.parser.add_argument("--cam_id", default=0, type=int) 34 | self.parser.add_argument("--frames", default=150, type=int) 35 | self.parser.add_argument("--noise_init", default=0.1, type=float) 36 | self.parser.add_argument("--noise_step", default=100, type=int) 37 | self.parser.add_argument("--exp_name", default="", type=str) 38 | self.parser.add_argument("--evaluate", action='store_true', default=False) 39 | 40 | 41 | def initial(self): 42 | self.parser.add_argument("--test_iterations", nargs="+", type=int, default=[7_000, 30_000]) 43 | self.parser.add_argument("--save_iterations", nargs="+", type=int, default=[7_000, 30_000]) 44 | self.parser.add_argument("--checkpoint_iterations", nargs="+", type=int, default=[]) 45 | self.parser.add_argument("--start_checkpoint", type=str, default = None) 46 | self.parser.add_argument("--batch", type=int, default=1) 47 | self.parser.add_argument("--max_gaussian_size", type=int, default=2e5) 48 | 49 | # Train strategy 50 | self.parser.add_argument("--jointly_iteration", type=int, default=1000, help="warm up iter") 51 | self.parser.add_argument("--max_points", type=int, default=100, help="The numbers of Initialize key points") 52 | self.parser.add_argument("--adaptive_points_num", type=int, default=0, help="Upperbound of the increasing key points") 53 | self.parser.add_argument("--set_iter", type=int, default=30) # need to change name 54 | self.parser.add_argument("--d", type=int, default=4) 55 | self.parser.add_argument("--w", type=int, default=256) 56 | self.parser.add_argument("--feature_dim", type=int, default=32, help="motion feature dimensions") 57 | 58 | 59 | self.parser.add_argument("--use_time_decay", action='store_true', default=False, help="Smooth the temporary input") 60 | self.parser.add_argument("--time_noise_ratio", type=float, default=0.5, help="decay time noise ratio") 61 | self.parser.add_argument("--time_noise_iteration", type=int, default=10000, help="After this iteration, no noise will add to the temporary input") 62 | 63 | 64 | self.parser.add_argument("--xyz_spatial_noise", action='store_true', default=False) 65 | self.parser.add_argument("--xyz_noise_iteration", type=int, default=10000, help="After this iteration, no noise will add to the xyz position") 66 | 67 | 68 | self.parser.add_argument("--nearest_num", type=int, default=6) 69 | self.parser.add_argument("--step_opacity", action='store_true', default=False) 70 | self.parser.add_argument("--step_opacity_iteration", type=int, default=5000) 71 | self.parser.add_argument("--opacity_type", type=str, default="implicit") 72 | self.parser.add_argument("--beta", type=float, default=0.1) 73 | self.parser.add_argument("--feature_amplify", type=float, default=5) 74 | self.parser.add_argument("--resize", type=float, default=1) 75 | self.parser.add_argument("--knn_type", type=str, default="hybird") 76 | self.parser.add_argument("--norm_rotation", action='store_true', default=False) 77 | self.parser.add_argument("--ratio", type=float, default=0.5, help="HyperNeRF dataset resize") 78 | 79 | self.parser.add_argument("--densify_from_teaching", action='store_true', default=False) 80 | self.parser.add_argument('--densify_from_grad', type=str, choices=['True', 'False'], default="True") 81 | self.parser.add_argument("--teaching_threshold", type=float, default=0.2) 82 | 83 | 84 | self.parser.add_argument("--second_stage_iteration", type=int, default=30000) 85 | self.parser.add_argument("--third_stage_iteration", type=int, default=40000) 86 | 87 | self.parser.add_argument("--adaptive_from_iter", type=int, default=3000) 88 | self.parser.add_argument("--adaptive_end_iter", type=int, default=10000) 89 | self.parser.add_argument("--adaptive_interval", type=int, default=200) 90 | self.parser.add_argument("--seed", type=int, default=1) 91 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | trimesh==4.0.5 2 | plyfile==0.8.1 3 | tqdm 4 | tensorboard 5 | prettytable 6 | open3d 7 | einops 8 | kmeans_pytorch 9 | opencv_python 10 | pytorch_msssim 11 | imageio 12 | torch-scatter==2.1.1 -------------------------------------------------------------------------------- /scene/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | import os 13 | import random 14 | import json 15 | from utils.system_utils import searchForMaxIteration 16 | from scene.dataset_readers import sceneLoadTypeCallbacks 17 | from scene.gaussian_model import GaussianModel 18 | from arguments import ModelParams 19 | from utils.camera_utils import cameraList_from_camInfos, camera_to_JSON 20 | 21 | class Scene: 22 | 23 | gaussians : GaussianModel 24 | 25 | def __init__(self, args : ModelParams, gaussians : GaussianModel, load_iteration=None, shuffle=True, resolution_scales=[1.0], ratio=0.5): 26 | """b 27 | :param path: Path to colmap scene main folder. 28 | """ 29 | self.model_path = args.model_path 30 | self.loaded_iter = None 31 | self.gaussians = gaussians 32 | 33 | if load_iteration: 34 | if load_iteration == -1: 35 | self.loaded_iter = searchForMaxIteration(os.path.join(self.model_path, "point_cloud")) 36 | else: 37 | self.loaded_iter = load_iteration 38 | print("Loading trained model at iteration {}".format(self.loaded_iter)) 39 | 40 | self.train_cameras = {} 41 | self.test_cameras = {} 42 | self.render_cameras = {} 43 | 44 | if os.path.exists(os.path.join(args.source_path, "sparse")): 45 | scene_info = sceneLoadTypeCallbacks["Colmap"](args.source_path, args.images, args.eval) 46 | elif os.path.exists(os.path.join(args.source_path, "transforms_train.json")): 47 | print("Found transforms_train.json file, assuming Blender data set!") 48 | scene_info = sceneLoadTypeCallbacks["Blender"](args.source_path, args.white_background, args.eval, max_time=args.max_time) 49 | elif os.path.exists(os.path.join(args.source_path, "points.npy")): 50 | print("Found .bin file, assuming HyperNeRF data set!") 51 | scene_info = sceneLoadTypeCallbacks["Hyper"](args.source_path, args.max_time, ratio) 52 | else: 53 | assert False, "Could not recognize scene type!" 54 | 55 | self.total_frame = scene_info.total_frame 56 | 57 | if not self.loaded_iter: 58 | with open(scene_info.ply_path, 'rb') as src_file, open(os.path.join(self.model_path, "input.ply") , 'wb') as dest_file: 59 | dest_file.write(src_file.read()) 60 | json_cams = [] 61 | camlist = [] 62 | if scene_info.test_cameras: 63 | camlist.extend(scene_info.test_cameras) 64 | if scene_info.train_cameras: 65 | camlist.extend(scene_info.train_cameras) 66 | for id, cam in enumerate(camlist): 67 | if isinstance(cam, dict): 68 | cam = cam["cam"] 69 | json_cams.append(camera_to_JSON(id, cam)) 70 | with open(os.path.join(self.model_path, "cameras.json"), 'w') as file: 71 | json.dump(json_cams, file) 72 | 73 | if shuffle: 74 | random.shuffle(scene_info.train_cameras) # Multi-res consistent random shuffling 75 | random.shuffle(scene_info.test_cameras) # Multi-res consistent random shuffling 76 | 77 | self.cameras_extent = scene_info.nerf_normalization["radius"] 78 | 79 | for resolution_scale in resolution_scales: 80 | print("Loading Training Cameras") 81 | self.train_cameras[resolution_scale] = cameraList_from_camInfos(scene_info.train_cameras, resolution_scale, args) 82 | print("Loading Test Cameras") 83 | self.test_cameras[resolution_scale] = cameraList_from_camInfos(scene_info.test_cameras, resolution_scale, args) 84 | if scene_info.render_cameras: 85 | print("Loading Render Cameras") 86 | self.render_cameras[resolution_scale] = cameraList_from_camInfos(scene_info.render_cameras, resolution_scale, args) 87 | 88 | if self.loaded_iter: 89 | self.gaussians.load_ply(os.path.join(self.model_path, 90 | "point_cloud", 91 | "iteration_" + str(self.loaded_iter), 92 | "point_cloud.ply")) 93 | else: 94 | self.gaussians.create_from_pcd(scene_info.point_cloud, self.cameras_extent) 95 | 96 | def save(self, iteration): 97 | point_cloud_path = os.path.join(self.model_path, "point_cloud/iteration_{}".format(iteration)) 98 | self.gaussians.save_ply(os.path.join(point_cloud_path, "point_cloud.ply")) 99 | 100 | def getTrainCameras(self, scale=1.0): 101 | return self.train_cameras[scale] 102 | 103 | def getTestCameras(self, scale=1.0): 104 | return self.test_cameras[scale] 105 | 106 | def getRenderCameras(self, scale=1.0): 107 | return self.render_cameras[scale] -------------------------------------------------------------------------------- /scene/cameras.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | import torch 13 | from torch import nn 14 | import numpy as np 15 | from utils.graphics_utils import getWorld2View2, getProjectionMatrix 16 | 17 | class Camera(nn.Module): 18 | def __init__(self, colmap_id, R, T, FoVx, FoVy, image, gt_alpha_mask, 19 | image_name, uid, 20 | trans=np.array([0.0, 0.0, 0.0]), scale=1.0, data_device = "cuda", time=0., 21 | fwd_flow=None, bwd_flow=None, fwd_flow_mask=None, bwd_flow_mask=None,): 22 | super(Camera, self).__init__() 23 | 24 | self.uid = uid 25 | self.colmap_id = colmap_id 26 | self.R = R 27 | self.T = T 28 | self.FoVx = FoVx 29 | self.FoVy = FoVy 30 | self.image_name = image_name 31 | self.time = time 32 | self.fwd_flow = fwd_flow 33 | self.fwd_flow_mask = fwd_flow_mask 34 | self.bwd_flow = bwd_flow 35 | self.bwd_flow_mask = bwd_flow_mask 36 | 37 | try: 38 | self.data_device = torch.device(data_device) 39 | except Exception as e: 40 | print(e) 41 | print(f"[Warning] Custom device {data_device} failed, fallback to default cuda device" ) 42 | self.data_device = torch.device("cuda") 43 | 44 | self.original_image = image.clamp(0.0, 1.0).to(self.data_device) 45 | self.image_width = self.original_image.shape[2] 46 | self.image_height = self.original_image.shape[1] 47 | 48 | if gt_alpha_mask is not None: 49 | self.original_image *= gt_alpha_mask.to(self.data_device) 50 | else: 51 | self.original_image *= torch.ones((1, self.image_height, self.image_width), device=self.data_device) 52 | 53 | self.zfar = 100.0 54 | self.znear = 0.01 55 | 56 | self.trans = trans 57 | self.scale = scale 58 | 59 | self.world_view_transform = torch.tensor(getWorld2View2(R, T, trans, scale)).transpose(0, 1).cuda() 60 | self.projection_matrix = getProjectionMatrix(znear=self.znear, zfar=self.zfar, fovX=self.FoVx, fovY=self.FoVy).transpose(0,1).cuda() 61 | self.full_proj_transform = (self.world_view_transform.unsqueeze(0).bmm(self.projection_matrix.unsqueeze(0))).squeeze(0) 62 | self.camera_center = self.world_view_transform.inverse()[3, :3] 63 | 64 | @classmethod 65 | def from_Camera(cls, Camera): 66 | colmap_id, R, T, FoVx, FoVy, image, gt_alpha_mask, image_name, uid, \ 67 | trans, scale, data_device, time, fwd_flow, bwd_flow, fwd_flow_mask, bwd_flow_mask = \ 68 | Camera.colmap_id, Camera.R, Camera.T, Camera.FoVx, Camera.FoVy, Camera.original_image, None, Camera.image_name, Camera.uid, \ 69 | Camera.trans, Camera.scale, Camera.data_device, Camera.time, Camera.fwd_flow, Camera.bwd_flow, Camera.fwd_flow_mask, Camera.bwd_flow_mask 70 | return cls(colmap_id, R, T, FoVx, FoVy, image, gt_alpha_mask, 71 | image_name, uid, trans, scale, data_device, time, 72 | fwd_flow, bwd_flow, fwd_flow_mask, bwd_flow_mask) 73 | 74 | class MiniCam: 75 | def __init__(self, width, height, fovy, fovx, znear, zfar, world_view_transform, full_proj_transform): 76 | self.image_width = width 77 | self.image_height = height 78 | self.FoVy = fovy 79 | self.FoVx = fovx 80 | self.znear = znear 81 | self.zfar = zfar 82 | self.world_view_transform = world_view_transform 83 | self.full_proj_transform = full_proj_transform 84 | view_inv = torch.inverse(self.world_view_transform) 85 | self.camera_center = view_inv[3][:3] 86 | 87 | -------------------------------------------------------------------------------- /scene/deformable_field.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | import torch.nn.functional as F 4 | import numpy as np 5 | from torch.autograd import Variable 6 | from inspect import isfunction 7 | from einops import rearrange, repeat 8 | from torch import einsum 9 | 10 | def exists(val): 11 | return val is not None 12 | 13 | def default(val, d): 14 | if exists(val): 15 | return val 16 | return d() if isfunction(d) else d 17 | 18 | class CrossAttention(nn.Module): 19 | def __init__(self, query_dim, context_dim=None, output_dim=None, heads=8, dim_head=64, dropout=0.): 20 | super().__init__() 21 | inner_dim = dim_head * heads 22 | context_dim = default(context_dim, query_dim) 23 | output_dim = default(output_dim, query_dim) 24 | 25 | self.scale = dim_head ** -0.5 26 | self.heads = heads 27 | 28 | self.to_q = nn.Linear(query_dim, inner_dim, bias=False) 29 | self.to_k = nn.Linear(context_dim, inner_dim, bias=False) 30 | self.to_v = nn.Linear(context_dim, inner_dim, bias=False) 31 | 32 | self.to_out = nn.Sequential( 33 | nn.Linear(inner_dim, query_dim), 34 | nn.Dropout(dropout) 35 | ) 36 | 37 | def forward(self, x, context=None, mask=None): 38 | h = self.heads 39 | 40 | q = self.to_q(x) 41 | context = default(context, x) 42 | k = self.to_k(context) 43 | v = self.to_v(context) 44 | 45 | q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v)) 46 | 47 | sim = einsum('b i d, b j d -> b i j', q, k) * self.scale 48 | 49 | if exists(mask): 50 | mask = mask>0 51 | mask = rearrange(mask, 'b ... -> b (...)') 52 | max_neg_value = -torch.finfo(sim.dtype).max 53 | mask = repeat(mask, 'b j -> (b h) () j', h=h) 54 | sim.masked_fill_(~mask, max_neg_value) 55 | 56 | # attention, what we cannot get enough of 57 | attn = sim.softmax(dim=-1) 58 | 59 | out = einsum('b i j, b j d -> b i d', attn, v) 60 | out = rearrange(out, '(b h) n d -> b n (h d)', h=h) 61 | return self.to_out(out) 62 | 63 | def positional_encoding(positions, freqs, ori=False): 64 | freq_bands = (2**torch.arange(freqs).float()).to(positions.device) 65 | ori_c = positions.shape[-1] 66 | pts = (positions[..., None] * freq_bands).reshape(positions.shape[:-1] + 67 | (freqs * positions.shape[-1], )) 68 | if ori: 69 | pts = torch.cat([positions, torch.sin(pts), torch.cos(pts)], dim=-1).reshape(pts.shape[:-1] + (pts.shape[-1] * 2 + ori_c, )) 70 | else: 71 | pts = torch.stack([torch.sin(pts), torch.cos(pts)], dim=-1).reshape(pts.shape[:-1] + (pts.shape[-1] * 2,)) 72 | return pts 73 | 74 | class Deformable_Field(nn.Module): 75 | def __init__(self, input_dim, output_dim=10, d=8, w=256, use_softmax=False, split_xyz=False): 76 | super(Deformable_Field, self).__init__() 77 | self.input_dim = input_dim 78 | self.d = d 79 | self.w = w 80 | self.use_softmax = use_softmax 81 | self.output_dim = output_dim 82 | 83 | self.split_xyz = split_xyz 84 | if self.split_xyz: 85 | self.output_times = self.output_dim 86 | self.output_dim = 1 87 | self.mlp = {} 88 | self.feature_to_deformation = {} 89 | for times in range(self.output_times): 90 | mlp = [] 91 | for i in range(d): 92 | if i == 0: 93 | mlp.append(nn.Linear(self.input_dim, self.w)) 94 | else: 95 | mlp.append(nn.Linear(self.w, self.w)) 96 | mlp.append(nn.ReLU()) 97 | self.mlp[f"mlp{times:d}"] = nn.Sequential(*mlp) 98 | self.feature_to_deformation[f"feature_to_deformation{times:d}"] = nn.Sequential(nn.Linear(w, self.output_dim)) 99 | self.mlp = nn.ModuleDict(self.mlp) 100 | self.feature_to_deformation = nn.ModuleDict(self.feature_to_deformation) 101 | else: 102 | mlp = [] 103 | for i in range(d): 104 | if i == 0: 105 | mlp.append(nn.Linear(self.input_dim, self.w)) 106 | else: 107 | mlp.append(nn.Linear(self.w, self.w)) 108 | mlp.append(nn.ReLU()) 109 | self.mlp = nn.Sequential(*mlp) 110 | self.feature_to_deformation = nn.Sequential(nn.Linear(w, self.output_dim)) 111 | 112 | def forward(self, x): 113 | if self.split_xyz: 114 | raw_list = [] 115 | for i in range(self.output_times): 116 | feature = self.mlp[f"mlp{i:d}"](x) 117 | raw_output = self.feature_to_deformation[f"feature_to_deformation{i:d}"](feature) 118 | if self.use_softmax: 119 | raw_output = nn.Softmax(dim=-1)(raw_output) 120 | raw_list += [raw_output] 121 | raw_output = torch.cat(raw_list, dim=-1) 122 | else: 123 | feature = self.mlp(x) 124 | raw_output = self.feature_to_deformation(feature) 125 | if self.use_softmax: 126 | raw_output = nn.Softmax(dim=-1)(raw_output) 127 | return raw_output 128 | 129 | class get_model(nn.Module): 130 | def __init__(self, num_class): 131 | super(get_model, self).__init__() 132 | self.k = num_class 133 | self.feat = PointNetEncoder(global_feat=False, feature_transform=False, channel=3) 134 | self.conv1 = torch.nn.Conv1d(1088, 512, 1) 135 | self.conv2 = torch.nn.Conv1d(512, 256, 1) 136 | self.conv3 = torch.nn.Conv1d(256, 128, 1) 137 | self.conv4 = torch.nn.Conv1d(128, self.k, 1) 138 | self.bn1 = nn.InstanceNorm1d(512) 139 | self.bn2 = nn.InstanceNorm1d(256) 140 | self.bn3 = nn.InstanceNorm1d(128) 141 | 142 | def forward(self, x): 143 | batchsize = x.size()[0] 144 | n_pts = x.size()[2] 145 | x, trans, trans_feat = self.feat(x) 146 | x = F.relu(self.bn1(self.conv1(x))) 147 | x = F.relu(self.bn2(self.conv2(x))) 148 | x = F.relu(self.bn3(self.conv3(x))) 149 | x = self.conv4(x) 150 | x = x.transpose(2,1).contiguous() 151 | x = F.softmax(x.view(-1,self.k), dim=-1) 152 | x = x.view(batchsize, n_pts, self.k) 153 | return x 154 | # return x, trans_feat 155 | 156 | 157 | class STN3d(nn.Module): 158 | def __init__(self, channel): 159 | super(STN3d, self).__init__() 160 | self.conv1 = torch.nn.Conv1d(channel, 64, 1) 161 | self.conv2 = torch.nn.Conv1d(64, 128, 1) 162 | self.conv3 = torch.nn.Conv1d(128, 1024, 1) 163 | self.fc1 = nn.Linear(1024, 512) 164 | self.fc2 = nn.Linear(512, 256) 165 | self.fc3 = nn.Linear(256, 9) 166 | self.relu = nn.ReLU() 167 | 168 | self.bn1 = nn.InstanceNorm1d(64) 169 | self.bn2 = nn.InstanceNorm1d(128) 170 | self.bn3 = nn.InstanceNorm1d(1024) 171 | self.bn4 = nn.InstanceNorm1d(512) 172 | self.bn5 = nn.InstanceNorm1d(256) 173 | 174 | def forward(self, x): 175 | batchsize = x.size()[0] 176 | x = F.relu(self.bn1(self.conv1(x))) 177 | x = F.relu(self.bn2(self.conv2(x))) 178 | x = F.relu(self.bn3(self.conv3(x))) 179 | x = torch.max(x, 2, keepdim=True)[0] 180 | x = x.view(-1, 1024) 181 | 182 | x = F.relu(self.bn4(self.fc1(x))) 183 | x = F.relu(self.bn5(self.fc2(x))) 184 | x = self.fc3(x) 185 | 186 | iden = Variable(torch.from_numpy(np.array([1, 0, 0, 0, 1, 0, 0, 0, 1]).astype(np.float32))).view(1, 9).repeat( 187 | batchsize, 1) 188 | if x.is_cuda: 189 | iden = iden.cuda() 190 | x = x + iden 191 | x = x.view(-1, 3, 3) 192 | return x 193 | 194 | 195 | class STNkd(nn.Module): 196 | def __init__(self, k=64): 197 | super(STNkd, self).__init__() 198 | self.conv1 = torch.nn.Conv1d(k, 64, 1) 199 | self.conv2 = torch.nn.Conv1d(64, 128, 1) 200 | self.conv3 = torch.nn.Conv1d(128, 1024, 1) 201 | self.fc1 = nn.Linear(1024, 512) 202 | self.fc2 = nn.Linear(512, 256) 203 | self.fc3 = nn.Linear(256, k * k) 204 | self.relu = nn.ReLU() 205 | 206 | self.bn1 = nn.InstanceNorm1d(64) 207 | self.bn2 = nn.InstanceNorm1d(128) 208 | self.bn3 = nn.InstanceNorm1d(1024) 209 | self.bn4 = nn.InstanceNorm1d(512) 210 | self.bn5 = nn.InstanceNorm1d(256) 211 | 212 | self.k = k 213 | 214 | def forward(self, x): 215 | batchsize = x.size()[0] 216 | x = F.relu(self.bn1(self.conv1(x))) 217 | x = F.relu(self.bn2(self.conv2(x))) 218 | x = F.relu(self.bn3(self.conv3(x))) 219 | x = torch.max(x, 2, keepdim=True)[0] 220 | x = x.view(-1, 1024) 221 | 222 | x = F.relu(self.bn4(self.fc1(x))) 223 | x = F.relu(self.bn5(self.fc2(x))) 224 | x = self.fc3(x) 225 | 226 | iden = Variable(torch.from_numpy(np.eye(self.k).flatten().astype(np.float32))).view(1, self.k * self.k).repeat( 227 | batchsize, 1) 228 | if x.is_cuda: 229 | iden = iden.cuda() 230 | x = x + iden 231 | x = x.view(-1, self.k, self.k) 232 | return x 233 | 234 | 235 | class PointNetEncoder(nn.Module): 236 | def __init__(self, global_feat=True, feature_transform=False, channel=3): 237 | super(PointNetEncoder, self).__init__() 238 | self.stn = STN3d(channel) 239 | self.conv1 = torch.nn.Conv1d(channel, 64, 1) 240 | self.conv2 = torch.nn.Conv1d(64, 128, 1) 241 | self.conv3 = torch.nn.Conv1d(128, 1024, 1) 242 | self.bn1 = nn.InstanceNorm1d(64) 243 | self.bn2 = nn.InstanceNorm1d(128) 244 | self.bn3 = nn.InstanceNorm1d(1024) 245 | self.global_feat = global_feat 246 | self.feature_transform = feature_transform 247 | if self.feature_transform: 248 | self.fstn = STNkd(k=64) 249 | 250 | def forward(self, x): 251 | B, D, N = x.size() 252 | trans = self.stn(x) 253 | x = x.transpose(2, 1) 254 | if D > 3: 255 | feature = x[:, :, 3:] 256 | x = x[:, :, :3] 257 | x = torch.bmm(x, trans) 258 | if D > 3: 259 | x = torch.cat([x, feature], dim=2) 260 | x = x.transpose(2, 1) 261 | x = F.relu(self.bn1(self.conv1(x))) 262 | 263 | if self.feature_transform: 264 | trans_feat = self.fstn(x) 265 | x = x.transpose(2, 1) 266 | x = torch.bmm(x, trans_feat) 267 | x = x.transpose(2, 1) 268 | else: 269 | trans_feat = None 270 | 271 | pointfeat = x 272 | x = F.relu(self.bn2(self.conv2(x))) 273 | x = self.bn3(self.conv3(x)) 274 | x = torch.max(x, 2, keepdim=True)[0] 275 | x = x.view(-1, 1024) 276 | if self.global_feat: 277 | return x, trans, trans_feat 278 | else: 279 | x = x.view(-1, 1024, 1).repeat(1, 1, N) 280 | return torch.cat([x, pointfeat], 1), trans, trans_feat 281 | 282 | 283 | def feature_transform_reguliarzer(trans): 284 | d = trans.size()[1] 285 | I = torch.eye(d)[None, :, :] 286 | if trans.is_cuda: 287 | I = I.cuda() 288 | loss = torch.mean(torch.norm(torch.bmm(trans, trans.transpose(2, 1)) - I, dim=(1, 2))) 289 | return loss -------------------------------------------------------------------------------- /scene/hyper_loader.py: -------------------------------------------------------------------------------- 1 | import warnings 2 | 3 | warnings.filterwarnings("ignore") 4 | 5 | import json 6 | import os 7 | import random 8 | 9 | import numpy as np 10 | import torch 11 | from PIL import Image 12 | import math 13 | from tqdm import tqdm 14 | from scene.utils import Camera 15 | from typing import NamedTuple 16 | from torch.utils.data import Dataset 17 | from utils.general_utils import PILtoTorch 18 | # from scene.dataset_readers import 19 | from utils.graphics_utils import getWorld2View2, focal2fov, fov2focal 20 | import copy 21 | class CameraInfo(NamedTuple): 22 | uid: int 23 | R: np.array 24 | T: np.array 25 | FovY: np.array 26 | FovX: np.array 27 | image: np.array 28 | image_path: str 29 | image_name: str 30 | width: int 31 | height: int 32 | time: np.array 33 | 34 | 35 | class Load_hyper_data(Dataset): 36 | def __init__(self, 37 | datadir, 38 | ratio=1.0, 39 | split="train", 40 | thres_time=1.0 41 | ): 42 | 43 | from .utils import Camera 44 | datadir = os.path.expanduser(datadir) 45 | self.datadir = datadir 46 | with open(f'{datadir}/scene.json', 'r') as f: 47 | scene_json = json.load(f) 48 | with open(f'{datadir}/metadata.json', 'r') as f: 49 | meta_json = json.load(f) 50 | with open(f'{datadir}/dataset.json', 'r') as f: 51 | dataset_json = json.load(f) 52 | 53 | self.near = scene_json['near'] 54 | self.far = scene_json['far'] 55 | self.coord_scale = scene_json['scale'] 56 | self.scene_center = scene_json['center'] 57 | 58 | self.all_img = dataset_json['ids'] 59 | self.val_id = dataset_json['val_ids'] 60 | self.split = split 61 | if len(self.val_id) == 0: 62 | self.i_train = np.array([i for i in np.arange(len(self.all_img)) if 63 | (i%4 == 0)]) 64 | self.i_test = self.i_train + 2 65 | self.i_test = self.i_test[:-1,] 66 | else: 67 | self.train_id = dataset_json['train_ids'] 68 | self.i_test = [] 69 | self.i_train = [] 70 | for i in range(len(self.all_img)): 71 | id = self.all_img[i] 72 | if id in self.val_id: 73 | self.i_test.append(i) 74 | if id in self.train_id: 75 | self.i_train.append(i) 76 | 77 | 78 | self.all_cam = [meta_json[i]['camera_id'] for i in self.all_img] 79 | self.all_time = [meta_json[i]['warp_id'] for i in self.all_img] 80 | max_time = max(self.all_time) 81 | if thres_time < 1.0: 82 | self.all_time, self.i_train, self.i_test = [], [], [] 83 | for idx, i in enumerate(self.all_img): 84 | time = meta_json[i]['warp_id']/max_time 85 | self.all_time += [time] 86 | if len(self.val_id) == 0: 87 | if idx % 4 ==0 and time < thres_time: 88 | self.i_train += [idx] 89 | if (idx - 2) % 4 ==0 and time >= thres_time: 90 | self.i_test += [idx] 91 | else: 92 | if i in self.val_id and time >= thres_time: 93 | self.i_test.append(idx) 94 | if i in self.train_id and time < thres_time: 95 | self.i_train.append(idx) 96 | self.i_test = np.array(self.i_test) 97 | self.i_train = np.array(self.i_train) 98 | np_all_time = np.array(self.all_time) 99 | test_time = np_all_time[self.i_test] 100 | train_time = np_all_time[self.i_train] 101 | assert test_time.max() >= thres_time 102 | assert train_time.min() < thres_time 103 | print("train:", self.i_train) 104 | print("test:", self.i_test) 105 | else: 106 | self.all_time = [meta_json[i]['warp_id']/max_time for i in self.all_img] 107 | 108 | self.selected_time = set(self.all_time) 109 | self.ratio = ratio 110 | self.max_time = max(self.all_time) 111 | self.min_time = min(self.all_time) 112 | self.i_video = [i for i in range(len(self.all_img))] 113 | self.i_video.sort() 114 | # all poses 115 | self.all_cam_params = [] 116 | for im in self.all_img: 117 | camera = Camera.from_json(f'{datadir}/camera/{im}.json') 118 | self.all_cam_params.append(camera) 119 | 120 | self.all_img = [f'{datadir}/rgb/{int(1/ratio)}x/{i}.png' for i in self.all_img] 121 | self.h, self.w = self.all_cam_params[0].image_shape 122 | self.map = {} 123 | self.image_one = Image.open(self.all_img[0]) 124 | self.image_one_torch = PILtoTorch(self.image_one, None).to(torch.float32) 125 | 126 | def __getitem__(self, index): 127 | if self.split == "train": 128 | return self.load_raw(self.i_train[index]) 129 | elif self.split == "test": 130 | return self.load_raw(self.i_test[index]) 131 | elif self.split == "video": 132 | return self.load_video(self.i_video[index]) 133 | 134 | def __len__(self): 135 | if self.split == "train": 136 | return len(self.i_train) 137 | elif self.split == "test": 138 | return len(self.i_test) 139 | elif self.split == "video": 140 | # return len(self.i_video) 141 | return len(self.video_v2) 142 | 143 | def load_video(self, idx): 144 | if idx in self.map.keys(): 145 | return self.map[idx] 146 | camera = self.all_cam_params[idx] 147 | w = self.image_one.size[0] 148 | h = self.image_one.size[1] 149 | # image = PILtoTorch(image,None) 150 | # image = image.to(torch.float32) 151 | time = self.all_time[idx] 152 | R = camera.orientation.T 153 | T = - camera.position @ R 154 | FovY = focal2fov(camera.focal_length, self.h) 155 | FovX = focal2fov(camera.focal_length, self.w) 156 | image_path = "/".join(self.all_img[idx].split("/")[:-1]) 157 | image_name = self.all_img[idx].split("/")[-1] 158 | caminfo = CameraInfo(uid=idx, R=R, T=T, FovY=FovY, FovX=FovX, image=self.image_one_torch, 159 | image_path=image_path, image_name=image_name, width=w, height=h, time=time, 160 | ) 161 | self.map[idx] = caminfo 162 | return caminfo 163 | 164 | def load_raw(self, idx): 165 | if idx in self.map.keys(): 166 | return self.map[idx] 167 | camera = self.all_cam_params[idx] 168 | image = Image.open(self.all_img[idx]) 169 | image_name = os.path.basename(self.all_img[idx]) 170 | 171 | w = image.size[0] 172 | h = image.size[1] 173 | time = np.array([self.all_time[idx]]).astype(np.float32) 174 | R = camera.orientation.T 175 | T = - camera.position @ R 176 | FovY = focal2fov(camera.focal_length, self.h) 177 | FovX = focal2fov(camera.focal_length, self.w) 178 | image_path = "/".join(self.all_img[idx].split("/")[:-1]) 179 | image_name = self.all_img[idx].split("/")[-1] 180 | caminfo = CameraInfo(uid=idx, R=R, T=T, FovY=FovY, FovX=FovX, image=image, 181 | image_path=image_path, image_name=image_name, width=w, height=h, time=time 182 | ) 183 | self.map[idx] = caminfo 184 | return caminfo 185 | 186 | def format_hyper_data(data_class, split): 187 | if split == "train": 188 | data_idx = data_class.i_train 189 | elif split == "test": 190 | data_idx = data_class.i_test 191 | 192 | cam_infos = [] 193 | for uid, index in tqdm(enumerate(data_idx)): 194 | camera = data_class.all_cam_params[index] 195 | time = data_class.all_time[index] 196 | R = camera.orientation.T 197 | T = - camera.position @ R 198 | FovY = focal2fov(camera.focal_length, data_class.h) 199 | FovX = focal2fov(camera.focal_length, data_class.w) 200 | image_path = "/".join(data_class.all_img[index].split("/")[:-1]) 201 | image_name = data_class.all_img[index].split("/")[-1] 202 | cam_info = CameraInfo(uid=uid, R=R, T=T, FovY=FovY, FovX=FovX, image=None, 203 | image_path=image_path, image_name=image_name, width=int(data_class.w), height=int(data_class.h), time=time 204 | ) 205 | cam_infos.append(cam_info) 206 | return cam_infos 207 | -------------------------------------------------------------------------------- /scripts/eval/d-nerf/bouncingballs.sh: -------------------------------------------------------------------------------- 1 | scene_name="bouncingballs" 2 | max_keypoints=100 3 | adaptive_points_num=100 4 | time_freq=6 5 | nearest_num=6 6 | max_time=0.8 7 | feature_amplify=0.5 8 | 9 | model_path="./results/d-nerf_${max_time}/${scene_name}/finalVersion/" 10 | 11 | CUDA_VISIBLE_DEVICES=0 python eval.py -m $model_path --max_points $max_keypoints --max_time $max_time \ 12 | --adaptive_points_num $adaptive_points_num --ckpt_iteration 60000 --time_freq $time_freq \ 13 | --nearest_num $nearest_num --feature_amplify $feature_amplify --norm_rotation -------------------------------------------------------------------------------- /scripts/eval/d-nerf/hellwarrior.sh: -------------------------------------------------------------------------------- 1 | scene_name="hellwarrior" 2 | max_keypoints=100 3 | adaptive_points_num=100 4 | time_freq=6 5 | nearest_num=6 6 | max_time=0.8 7 | feature_amplify=0.5 8 | 9 | model_path="./results/d-nerf_${max_time}/${scene_name}/finalVersion/" 10 | 11 | CUDA_VISIBLE_DEVICES=0 python eval.py -m $model_path --max_points $max_keypoints --max_time $max_time \ 12 | --adaptive_points_num $adaptive_points_num --ckpt_iteration 60000 --time_freq $time_freq \ 13 | --nearest_num $nearest_num --feature_amplify $feature_amplify --norm_rotation -------------------------------------------------------------------------------- /scripts/eval/d-nerf/hook.sh: -------------------------------------------------------------------------------- 1 | scene_name="hook" 2 | max_keypoints=100 3 | adaptive_points_num=100 4 | time_freq=6 5 | nearest_num=6 6 | max_time=0.8 7 | feature_amplify=0.5 8 | 9 | model_path="./results/d-nerf_${max_time}/${scene_name}/finalVersion/" 10 | 11 | CUDA_VISIBLE_DEVICES=0 python eval.py -m $model_path --max_points $max_keypoints --max_time $max_time \ 12 | --adaptive_points_num $adaptive_points_num --ckpt_iteration 60000 --time_freq $time_freq \ 13 | --nearest_num $nearest_num --feature_amplify $feature_amplify --norm_rotation -------------------------------------------------------------------------------- /scripts/eval/d-nerf/jumping.sh: -------------------------------------------------------------------------------- 1 | scene_name="jumpingjacks" 2 | max_keypoints=100 3 | adaptive_points_num=100 4 | time_freq=6 5 | nearest_num=6 6 | max_time=0.8 7 | feature_amplify=0.5 8 | 9 | model_path="./results/d-nerf_${max_time}/${scene_name}/finalVersion/" 10 | 11 | CUDA_VISIBLE_DEVICES=0 python eval.py -m $model_path --max_points $max_keypoints --max_time $max_time \ 12 | --adaptive_points_num $adaptive_points_num --ckpt_iteration 60000 --time_freq $time_freq \ 13 | --nearest_num $nearest_num --feature_amplify $feature_amplify --norm_rotation 14 | -------------------------------------------------------------------------------- /scripts/eval/d-nerf/mutant.sh: -------------------------------------------------------------------------------- 1 | scene_name="mutant" 2 | max_keypoints=100 3 | adaptive_points_num=100 4 | time_freq=6 5 | nearest_num=6 6 | max_time=0.8 7 | feature_amplify=0.5 8 | 9 | model_path="./results/d-nerf_${max_time}/${scene_name}/finalVersion/" 10 | 11 | CUDA_VISIBLE_DEVICES=0 python eval.py -m $model_path --max_points $max_keypoints --max_time $max_time \ 12 | --adaptive_points_num $adaptive_points_num --ckpt_iteration 60000 --time_freq $time_freq \ 13 | --nearest_num $nearest_num --feature_amplify $feature_amplify --norm_rotation -------------------------------------------------------------------------------- /scripts/eval/d-nerf/standup.sh: -------------------------------------------------------------------------------- 1 | scene_name="standup" 2 | max_keypoints=100 3 | adaptive_points_num=100 4 | time_freq=6 5 | nearest_num=6 6 | max_time=0.8 7 | feature_amplify=0.5 8 | 9 | model_path="./results/d-nerf_${max_time}/${scene_name}/finalVersion/" 10 | 11 | CUDA_VISIBLE_DEVICES=0 python eval.py -m $model_path --max_points $max_keypoints --max_time $max_time \ 12 | --adaptive_points_num $adaptive_points_num --ckpt_iteration 60000 --time_freq $time_freq \ 13 | --nearest_num $nearest_num --feature_amplify $feature_amplify --norm_rotation -------------------------------------------------------------------------------- /scripts/eval/d-nerf/trex.sh: -------------------------------------------------------------------------------- 1 | scene_name="trex" 2 | max_keypoints=100 3 | adaptive_points_num=100 4 | time_freq=6 5 | nearest_num=6 6 | max_time=0.8 7 | feature_amplify=0.5 8 | 9 | model_path="./results/d-nerf_${max_time}/${scene_name}/finalVersion/" 10 | 11 | CUDA_VISIBLE_DEVICES=0 python eval.py -m $model_path --max_points $max_keypoints --max_time $max_time \ 12 | --adaptive_points_num $adaptive_points_num --ckpt_iteration 60000 --time_freq $time_freq \ 13 | --nearest_num $nearest_num --feature_amplify $feature_amplify --norm_rotation -------------------------------------------------------------------------------- /scripts/eval/hyper/chickchicken.sh: -------------------------------------------------------------------------------- 1 | scene_name="chickchicken" 2 | time_freq=8 3 | max_keypoints=100 4 | adaptive_points_num=100 5 | nearest_num=6 6 | max_time=1.0 7 | feature_amplify=5 8 | model_path="./results/HyperNeRF_${max_time}/${scene_name}/finalVersion" 9 | 10 | CUDA_VISIBLE_DEVICES=0 python eval.py -m $model_path --max_points $max_keypoints --time_freq $time_freq \ 11 | --adaptive_points_num $adaptive_points_num --ckpt_iteration 70000 --feature_amplify $feature_amplify \ 12 | --nearest_num $nearest_num --max_time $max_time --norm_rotation -------------------------------------------------------------------------------- /scripts/eval/hyper/lemon.sh: -------------------------------------------------------------------------------- 1 | scene_name="cut-lemon-new" 2 | time_freq=10 3 | max_keypoints=100 4 | adaptive_points_num=200 5 | nearest_num=6 6 | max_time=1.0 7 | feature_amplify=5 8 | model_path="./results/HyperNeRF_${max_time}/${scene_name}/finalVersion/" 9 | 10 | CUDA_VISIBLE_DEVICES=0 python eval.py -m $model_path --max_points $max_keypoints --time_freq $time_freq \ 11 | --adaptive_points_num $adaptive_points_num --ckpt_iteration 70000 --feature_amplify $feature_amplify \ 12 | --nearest_num $nearest_num --max_time $max_time --norm_rotation --step_opacity -------------------------------------------------------------------------------- /scripts/eval/hyper/printer.sh: -------------------------------------------------------------------------------- 1 | scene_name="vrig-3dprinter-new" 2 | time_freq=8 3 | max_keypoints=150 4 | adaptive_points_num=100 5 | nearest_num=6 6 | max_time=1.0 7 | feature_amplify=5 8 | model_path="./results/HyperNeRF_${max_time}/${scene_name}/finalVersion/" 9 | 10 | CUDA_VISIBLE_DEVICES=0 python eval.py -m $model_path --max_points $max_keypoints --time_freq $time_freq \ 11 | --adaptive_points_num $adaptive_points_num --ckpt_iteration 70000 --feature_amplify $feature_amplify \ 12 | --nearest_num $nearest_num --max_time $max_time --norm_rotation -------------------------------------------------------------------------------- /scripts/eval/hyper/torchocolate.sh: -------------------------------------------------------------------------------- 1 | scene_name="torchocolate" 2 | time_freq=8 3 | max_keypoints=50 4 | adaptive_points_num=100 5 | nearest_num=6 6 | max_time=1.0 7 | feature_amplify=5 8 | model_path="./results/HyperNeRF_${max_time}/${scene_name}/finalVersion/" 9 | 10 | CUDA_VISIBLE_DEVICES=0 python eval.py -m $model_path --max_points $max_keypoints --time_freq $time_freq \ 11 | --adaptive_points_num $adaptive_points_num --ckpt_iteration 70000 --feature_amplify $feature_amplify \ 12 | --nearest_num $nearest_num --max_time $max_time --norm_rotation -------------------------------------------------------------------------------- /scripts/predict/d-nerf/bouncingballs.sh: -------------------------------------------------------------------------------- 1 | scene_name="bouncingballs" 2 | max_keypoints=100 3 | adaptive_points_num=100 4 | max_time=0.8 5 | time_freq=6 6 | nearest_num=6 7 | feature_amplify=0.5 8 | 9 | # gcn param 10 | num_stage=6 11 | noise_init=0 12 | noise_step=100 13 | input_size=10 14 | epoch=2001 15 | exp_name="input${input_size}_stage${num_stage}_noise${noise_init}_epoch${epoch}" 16 | 17 | model_path="./results/d-nerf_${max_time}/${scene_name}/finalVersion/" 18 | 19 | CUDA_VISIBLE_DEVICES=0 python train_GCN.py -m $model_path --max_points $max_keypoints \ 20 | --adaptive_points_num $adaptive_points_num --ckpt_iteration 60000 --time_freq $time_freq \ 21 | --max_time $max_time --nearest_num $nearest_num --feature_amplify $feature_amplify \ 22 | --num_stage $num_stage --noise_step $noise_step --noise_init $noise_init --input_size $input_size \ 23 | --epoch $epoch --exp_name $exp_name --predict_more --metrics --norm_rotation -------------------------------------------------------------------------------- /scripts/predict/d-nerf/hellwarrior.sh: -------------------------------------------------------------------------------- 1 | scene_name="hellwarrior" 2 | max_keypoints=100 3 | max_time=0.8 4 | time_freq=6 5 | adaptive_points_num=100 6 | nearest_num=6 7 | feature_amplify=0.5 8 | 9 | # gcn param 10 | num_stage=6 11 | noise_init=0 12 | noise_step=100 13 | input_size=10 14 | epoch=2001 15 | exp_name="input${input_size}_stage${num_stage}_noise${noise_init}_epoch${epoch}" 16 | 17 | model_path="./results/d-nerf_${max_time}/${scene_name}/finalVersion/" 18 | 19 | CUDA_VISIBLE_DEVICES=0 python train_GCN.py -m $model_path --max_points $max_keypoints \ 20 | --adaptive_points_num $adaptive_points_num --ckpt_iteration 60000 --time_freq $time_freq \ 21 | --max_time $max_time --nearest_num $nearest_num --feature_amplify $feature_amplify \ 22 | --num_stage $num_stage --noise_step $noise_step --noise_init $noise_init --input_size $input_size \ 23 | --epoch $epoch --exp_name $exp_name --predict_more --metrics --norm_rotation -------------------------------------------------------------------------------- /scripts/predict/d-nerf/hook.sh: -------------------------------------------------------------------------------- 1 | scene_name="hook" 2 | max_keypoints=100 3 | max_time=0.8 4 | time_freq=6 5 | adaptive_points_num=100 6 | nearest_num=6 7 | feature_amplify=0.5 8 | 9 | # gcn param 10 | num_stage=6 11 | noise_init=0 12 | noise_step=100 13 | input_size=10 14 | epoch=2001 15 | exp_name="input${input_size}_stage${num_stage}_noise${noise_init}_epoch${epoch}" 16 | 17 | model_path="./results/d-nerf_${max_time}/${scene_name}/finalVersion/" 18 | 19 | CUDA_VISIBLE_DEVICES=0 python train_GCN.py -m $model_path --max_points $max_keypoints \ 20 | --adaptive_points_num $adaptive_points_num --ckpt_iteration 60000 --time_freq $time_freq \ 21 | --max_time $max_time --nearest_num $nearest_num --feature_amplify $feature_amplify \ 22 | --num_stage $num_stage --noise_step $noise_step --noise_init $noise_init --input_size $input_size \ 23 | --epoch $epoch --exp_name $exp_name --predict_more --metrics --norm_rotation -------------------------------------------------------------------------------- /scripts/predict/d-nerf/jumping.sh: -------------------------------------------------------------------------------- 1 | scene_name="jumpingjacks" 2 | max_keypoints=100 3 | max_time=0.8 4 | time_freq=6 5 | adaptive_points_num=100 6 | nearest_num=6 7 | feature_amplify=0.5 8 | 9 | # gcn param 10 | num_stage=6 11 | noise_init=0 12 | noise_step=100 13 | input_size=10 14 | epoch=2001 15 | exp_name="input${input_size}_stage${num_stage}_noise${noise_init}_epoch${epoch}" 16 | 17 | model_path="./results/d-nerf_${max_time}/${scene_name}/finalVersion/" 18 | 19 | CUDA_VISIBLE_DEVICES=0 python train_GCN.py -m $model_path --max_points $max_keypoints \ 20 | --adaptive_points_num $adaptive_points_num --ckpt_iteration 60000 --time_freq $time_freq \ 21 | --max_time $max_time --nearest_num $nearest_num --feature_amplify $feature_amplify \ 22 | --num_stage $num_stage --noise_step $noise_step --noise_init $noise_init --input_size $input_size \ 23 | --epoch $epoch --exp_name $exp_name --predict_more --metrics --norm_rotation -------------------------------------------------------------------------------- /scripts/predict/d-nerf/mutant.sh: -------------------------------------------------------------------------------- 1 | scene_name="mutant" 2 | max_keypoints=100 3 | max_time=0.8 4 | time_freq=6 5 | adaptive_points_num=100 6 | nearest_num=6 7 | feature_amplify=0.5 8 | 9 | # gcn param 10 | num_stage=6 11 | noise_init=0 12 | noise_step=100 13 | input_size=10 14 | epoch=2001 15 | exp_name="input${input_size}_stage${num_stage}_noise${noise_init}_epoch${epoch}" 16 | 17 | model_path="./results/d-nerf_${max_time}/${scene_name}/finalVersion/" 18 | 19 | CUDA_VISIBLE_DEVICES=0 python train_GCN.py -m $model_path --max_points $max_keypoints \ 20 | --adaptive_points_num $adaptive_points_num --ckpt_iteration 60000 --time_freq $time_freq \ 21 | --max_time $max_time --nearest_num $nearest_num --feature_amplify $feature_amplify \ 22 | --num_stage $num_stage --noise_step $noise_step --noise_init $noise_init --input_size $input_size \ 23 | --epoch $epoch --exp_name $exp_name --predict_more --metrics --norm_rotation -------------------------------------------------------------------------------- /scripts/predict/d-nerf/standup.sh: -------------------------------------------------------------------------------- 1 | scene_name="standup" 2 | max_keypoints=100 3 | max_time=0.8 4 | time_freq=6 5 | adaptive_points_num=100 6 | nearest_num=6 7 | feature_amplify=0.5 8 | 9 | # gcn param 10 | num_stage=6 11 | noise_init=0 12 | noise_step=100 13 | input_size=10 14 | epoch=2001 15 | exp_name="input${input_size}_stage${num_stage}_noise${noise_init}_epoch${epoch}" 16 | 17 | model_path="./results/d-nerf_${max_time}/${scene_name}/finalVersion/" 18 | 19 | CUDA_VISIBLE_DEVICES=0 python train_GCN.py -m $model_path --max_points $max_keypoints \ 20 | --adaptive_points_num $adaptive_points_num --ckpt_iteration 60000 --time_freq $time_freq \ 21 | --max_time $max_time --nearest_num $nearest_num --feature_amplify $feature_amplify \ 22 | --num_stage $num_stage --noise_step $noise_step --noise_init $noise_init --input_size $input_size \ 23 | --epoch $epoch --exp_name $exp_name --predict_more --metrics --norm_rotation -------------------------------------------------------------------------------- /scripts/predict/d-nerf/trex.sh: -------------------------------------------------------------------------------- 1 | scene_name="trex" 2 | max_keypoints=100 3 | max_time=0.8 4 | time_freq=6 5 | adaptive_points_num=100 6 | nearest_num=6 7 | feature_amplify=0.5 8 | 9 | num_stage=6 10 | noise_init=0 11 | noise_step=100 12 | input_size=10 13 | epoch=2001 14 | exp_name="input${input_size}_stage${num_stage}_noise${noise_init}_epoch${epoch}" 15 | 16 | model_path="./results/d-nerf_${max_time}/${scene_name}/finalVersion/" 17 | 18 | CUDA_VISIBLE_DEVICES=0 python train_GCN.py -m $model_path --max_points $max_keypoints \ 19 | --adaptive_points_num $adaptive_points_num --ckpt_iteration 60000 --time_freq $time_freq \ 20 | --max_time $max_time --nearest_num $nearest_num --feature_amplify $feature_amplify \ 21 | --num_stage $num_stage --noise_step $noise_step --noise_init $noise_init --input_size $input_size \ 22 | --epoch $epoch --exp_name $exp_name --predict_more --metrics --norm_rotation -------------------------------------------------------------------------------- /scripts/predict/hyper/chickchicken.sh: -------------------------------------------------------------------------------- 1 | scene_name="chickchicken" 2 | time_freq=8 3 | max_keypoints=100 4 | adaptive_points_num=100 5 | max_time=0.8 6 | noise_init=5.0 7 | R_ratio=0.5 8 | 9 | num_stage=16 10 | input_size=20 11 | noise_step=300 12 | epoch=2001 13 | linear_size=256 14 | exp_name="input${input_size}_stage${num_stage}_hidden${linear_size}_noise${noise_init}_RNoise${R_ratio}_noiseStep${noise_step}_epoch${epoch}" 15 | 16 | model_path="./results/HyperNeRF_${max_time}/${scene_name}/finalVersion/" 17 | 18 | CUDA_VISIBLE_DEVICES=0 python train_GCN.py -m $model_path --max_points $max_keypoints --time_freq $time_freq \ 19 | --adaptive_points_num $adaptive_points_num --ckpt_iteration 70000 --max_time $max_time \ 20 | --num_stage $num_stage --noise_step $noise_step --noise_init $noise_init --linear_size $linear_size --input_size $input_size \ 21 | --Rscale_ratio $R_ratio --batch_size 32 --epoch $epoch --exp_name $exp_name --predict_more --metrics --norm_rotation 22 | -------------------------------------------------------------------------------- /scripts/predict/hyper/lemon.sh: -------------------------------------------------------------------------------- 1 | scene_name="cut-lemon" 2 | time_freq=10 3 | max_keypoints=100 4 | adaptive_points_num=200 5 | max_time=0.8 6 | noise_init=0.5 7 | R_ratio=1.0 8 | 9 | num_stage=16 10 | input_size=10 11 | noise_step=500 12 | epoch=1501 13 | linear_size=128 14 | exp_name="input${input_size}_stage${num_stage}_hidden${linear_size}_noise${noise_init}_RNoise${R_ratio}_noiseStep${noise_step}_epoch${epoch}" 15 | 16 | 17 | model_path="./results/HyperNeRF_${max_time}/${scene_name}/finalVersion/" 18 | 19 | CUDA_VISIBLE_DEVICES=0 python train_GCN.py -m $model_path --max_points $max_keypoints --time_freq $time_freq \ 20 | --adaptive_points_num $adaptive_points_num --ckpt_iteration 70000 --max_time $max_time \ 21 | --num_stage $num_stage --noise_step $noise_step --noise_init $noise_init --linear_size $linear_size --input_size $input_size \ 22 | --Rscale_ratio $R_ratio --batch_size 32 --epoch $epoch --exp_name $exp_name --predict_more --metrics --norm_rotation --step_opacity --cam_id=-16 -------------------------------------------------------------------------------- /scripts/predict/hyper/printer.sh: -------------------------------------------------------------------------------- 1 | scene_name="vrig-3dprinter" 2 | time_freq=8 3 | max_keypoints=150 4 | adaptive_points_num=100 5 | max_time=0.8 6 | noise_init=5.0 7 | R_ratio=0.5 8 | 9 | num_stage=16 10 | input_size=10 11 | noise_step=300 12 | epoch=2001 13 | linear_size=128 14 | exp_name="input${input_size}_stage${num_stage}_hidden${linear_size}_noise${noise_init}_RNoise${R_ratio}_noiseStep${noise_step}_epoch${epoch}" 15 | 16 | model_path="./results/HyperNeRF${max_time}/${scene_name}/finalVersion/" 17 | 18 | CUDA_VISIBLE_DEVICES=0 python train_GCN.py -m $model_path --max_points $max_keypoints --time_freq $time_freq \ 19 | --adaptive_points_num $adaptive_points_num --ckpt_iteration 70000 --max_time $max_time \ 20 | --num_stage $num_stage --noise_step $noise_step --noise_init $noise_init --linear_size $linear_size --input_size $input_size \ 21 | --Rscale_ratio $R_ratio --batch_size 32 --epoch $epoch --exp_name $exp_name --predict_more --metrics --norm_rotation -------------------------------------------------------------------------------- /scripts/predict/hyper/torchocolate.sh: -------------------------------------------------------------------------------- 1 | scene_name="torchocolate" 2 | time_freq=8 3 | max_keypoints=50 4 | adaptive_points_num=100 5 | max_time=0.8 6 | noise_init=5.0 7 | R_ratio=0.5 8 | 9 | num_stage=16 10 | input_size=10 11 | noise_step=100 12 | epoch=2001 13 | linear_size=256 14 | exp_name="input${input_size}_stage${num_stage}_hidden${linear_size}_noise${noise_init}_RNoise${R_ratio}_noiseStep${noise_step}_epoch${epoch}" 15 | 16 | model_path="./results/HyperNeRF_${max_time}/${scene_name}/finalVersion/" 17 | 18 | CUDA_VISIBLE_DEVICES=0 python train_GCN.py -m $model_path --max_points $max_keypoints --time_freq $time_freq \ 19 | --adaptive_points_num $adaptive_points_num --ckpt_iteration 70000 --max_time $max_time \ 20 | --num_stage $num_stage --noise_step $noise_step --noise_init $noise_init --linear_size $linear_size --input_size $input_size \ 21 | --Rscale_ratio $R_ratio --batch_size 32 --epoch $epoch --exp_name $exp_name --predict_more --metrics --norm_rotation -------------------------------------------------------------------------------- /scripts/train/d-nerf/bouncingballs.sh: -------------------------------------------------------------------------------- 1 | 2 | if [ $# -eq 0 ]; then 3 | echo "Seed: 1*2024" 4 | seed=1 5 | else 6 | seed=$1 7 | fi 8 | scene_name="bouncingballs" 9 | max_keypoints=100 10 | adaptive_points_num=100 11 | time_freq=6 12 | nearest_num=6 13 | source_path="./datasets/d-nerf/data/${scene_name}/" 14 | max_time=1.0 15 | model_path="./results/d-nerf_${max_time}/${scene_name}/finalVersion/" 16 | position_lr_max_steps=40000 17 | adaptive_from_iter=3000 18 | feature_amplify=0.5 19 | 20 | # Train 21 | CUDA_VISIBLE_DEVICES=0 python train.py -s $source_path \ 22 | -m $model_path --max_points $max_keypoints --adaptive_points_num $adaptive_points_num \ 23 | --iterations 60000 --test_iterations 60000 --jointly_iteration 1000 --time_freq $time_freq \ 24 | --densify_from_iter 3000 --densify_until_iter 20000 --norm_rotation \ 25 | --save_iterations 29999 60000 --checkpoint_iterations 29999 60000 --max_time $max_time \ 26 | --position_lr_max_steps $position_lr_max_steps --adaptive_from_iter $adaptive_from_iter --seed $seed \ 27 | --eval --nearest_num $nearest_num --adaptive_interval 500 --feature_amplify $feature_amplify -------------------------------------------------------------------------------- /scripts/train/d-nerf/hellwarrior.sh: -------------------------------------------------------------------------------- 1 | 2 | if [ $# -eq 0 ]; then 3 | echo "Seed: 1*2024" 4 | seed=1 5 | else 6 | seed=$1 7 | fi 8 | scene_name="hellwarrior" 9 | max_keypoints=100 10 | adaptive_points_num=100 11 | time_freq=6 12 | nearest_num=6 13 | source_path="./datasets/d-nerf/data/${scene_name}/" 14 | max_time=1.0 15 | model_path="./results/d-nerf_${max_time}/${scene_name}/finalVersion/" 16 | position_lr_max_steps=40000 17 | adaptive_from_iter=3000 18 | feature_amplify=0.5 19 | 20 | # Train 21 | CUDA_VISIBLE_DEVICES=0 python train.py -s $source_path \ 22 | -m $model_path --max_points $max_keypoints --adaptive_points_num $adaptive_points_num \ 23 | --iterations 60000 --test_iterations 60000 --jointly_iteration 1000 --time_freq $time_freq \ 24 | --densify_from_iter 3000 --densify_until_iter 20000 --norm_rotation \ 25 | --save_iterations 29999 60000 --checkpoint_iterations 29999 60000 --max_time $max_time \ 26 | --position_lr_max_steps $position_lr_max_steps --adaptive_from_iter $adaptive_from_iter --seed $seed \ 27 | --eval --nearest_num $nearest_num --adaptive_interval 500 --feature_amplify $feature_amplify -------------------------------------------------------------------------------- /scripts/train/d-nerf/hook.sh: -------------------------------------------------------------------------------- 1 | 2 | if [ $# -eq 0 ]; then 3 | echo "Seed: 1*2024" 4 | seed=1 5 | else 6 | seed=$1 7 | fi 8 | scene_name="hook" 9 | max_keypoints=100 10 | adaptive_points_num=100 11 | time_freq=6 12 | nearest_num=6 13 | source_path="./datasets/d-nerf/data/${scene_name}/" 14 | max_time=1.0 15 | model_path="./results/d-nerf_${max_time}/${scene_name}/finalVersion/" 16 | position_lr_max_steps=40000 17 | adaptive_from_iter=3000 18 | feature_amplify=0.5 19 | 20 | # Train 21 | CUDA_VISIBLE_DEVICES=0 python train.py -s $source_path \ 22 | -m $model_path --max_points $max_keypoints --adaptive_points_num $adaptive_points_num \ 23 | --iterations 60000 --test_iterations 60000 --jointly_iteration 1000 --time_freq $time_freq \ 24 | --densify_from_iter 3000 --densify_until_iter 20000 --norm_rotation \ 25 | --save_iterations 29999 60000 --checkpoint_iterations 29999 60000 --max_time $max_time \ 26 | --position_lr_max_steps $position_lr_max_steps --adaptive_from_iter $adaptive_from_iter --seed $seed \ 27 | --eval --nearest_num $nearest_num --adaptive_interval 500 --feature_amplify $feature_amplify -------------------------------------------------------------------------------- /scripts/train/d-nerf/jumping.sh: -------------------------------------------------------------------------------- 1 | if [ $# -eq 0 ]; then 2 | echo "Seed: 1*2024" 3 | seed=1 4 | else 5 | seed=$1 6 | fi 7 | 8 | scene_name="jumpingjacks" 9 | max_keypoints=100 10 | adaptive_points_num=100 11 | time_freq=6 12 | nearest_num=6 13 | source_path="./datasets/d-nerf/data/${scene_name}/" 14 | max_time=1.0 15 | model_path="./results/d-nerf_${max_time}/${scene_name}/finalVersion/" 16 | position_lr_max_steps=40000 17 | adaptive_from_iter=3000 18 | feature_amplify=0.5 19 | 20 | # Train 21 | CUDA_VISIBLE_DEVICES=0 python train.py -s $source_path \ 22 | -m $model_path --max_points $max_keypoints --adaptive_points_num $adaptive_points_num \ 23 | --iterations 60000 --test_iterations 60000 --jointly_iteration 1000 --time_freq $time_freq \ 24 | --densify_from_iter 3000 --densify_until_iter 20000 --norm_rotation \ 25 | --save_iterations 29999 60000 --checkpoint_iterations 29999 60000 --max_time $max_time \ 26 | --position_lr_max_steps $position_lr_max_steps --adaptive_from_iter $adaptive_from_iter --seed $seed \ 27 | --eval --nearest_num $nearest_num --adaptive_interval 500 --feature_amplify $feature_amplify -------------------------------------------------------------------------------- /scripts/train/d-nerf/mutant.sh: -------------------------------------------------------------------------------- 1 | if [ $# -eq 0 ]; then 2 | echo "Seed: 1*2024" 3 | seed=1 4 | else 5 | seed=$1 6 | fi 7 | scene_name="mutant" 8 | max_keypoints=100 9 | adaptive_points_num=100 10 | time_freq=6 11 | nearest_num=6 12 | source_path="./datasets/d-nerf/data/${scene_name}/" 13 | max_time=1.0 14 | model_path="./results/d-nerf_${max_time}/${scene_name}/finalVersion/" 15 | position_lr_max_steps=40000 16 | adaptive_from_iter=3000 17 | feature_amplify=0.5 18 | 19 | # Train 20 | CUDA_VISIBLE_DEVICES=0 python train.py -s $source_path \ 21 | -m $model_path --max_points $max_keypoints --adaptive_points_num $adaptive_points_num \ 22 | --iterations 60000 --test_iterations 60000 --jointly_iteration 1000 --time_freq $time_freq \ 23 | --densify_from_iter 3000 --densify_until_iter 20000 --norm_rotation \ 24 | --save_iterations 29999 60000 --checkpoint_iterations 29999 60000 --max_time $max_time \ 25 | --position_lr_max_steps $position_lr_max_steps --adaptive_from_iter $adaptive_from_iter --seed $seed \ 26 | --eval --nearest_num $nearest_num --adaptive_interval 500 --feature_amplify $feature_amplify -------------------------------------------------------------------------------- /scripts/train/d-nerf/standup.sh: -------------------------------------------------------------------------------- 1 | if [ $# -eq 0 ]; then 2 | echo "Seed: 1*2024" 3 | seed=1 4 | else 5 | seed=$1 6 | fi 7 | scene_name="standup" 8 | max_keypoints=100 9 | adaptive_points_num=100 10 | time_freq=6 11 | nearest_num=6 12 | source_path="./datasets/d-nerf/data/${scene_name}/" 13 | max_time=1.0 14 | model_path="./results/d-nerf_${max_time}/${scene_name}/finalVersion/" 15 | position_lr_max_steps=40000 16 | adaptive_from_iter=3000 17 | feature_amplify=0.5 18 | 19 | # Train 20 | CUDA_VISIBLE_DEVICES=0 python train.py -s $source_path \ 21 | -m $model_path --max_points $max_keypoints --adaptive_points_num $adaptive_points_num \ 22 | --iterations 60000 --test_iterations 60000 --jointly_iteration 1000 --time_freq $time_freq \ 23 | --densify_from_iter 3000 --densify_until_iter 20000 --norm_rotation \ 24 | --save_iterations 29999 60000 --checkpoint_iterations 29999 60000 --max_time $max_time \ 25 | --position_lr_max_steps $position_lr_max_steps --adaptive_from_iter $adaptive_from_iter --seed $seed \ 26 | --eval --nearest_num $nearest_num --adaptive_interval 500 --feature_amplify $feature_amplify -------------------------------------------------------------------------------- /scripts/train/d-nerf/train_eval.sh: -------------------------------------------------------------------------------- 1 | echo "Start training and evaluation process" 2 | echo "Please make sure your scripts [max_time] is 1.0 " 3 | 4 | read -p "Do you want to continue? (yes/no): " user_input 5 | 6 | # 判断输入 7 | if [ "$user_input" = "yes" ]; then 8 | echo "Continuing execution..." 9 | # 在这里继续执行你的脚本逻辑 10 | else 11 | echo "Exiting..." 12 | exit 1 13 | fi 14 | 15 | # Train 16 | ./scripts/train/d-nerf/bouncingballs.sh 17 | wait 18 | ./scripts/eval/d-nerf/bouncingballs.sh 19 | wait 20 | 21 | ./scripts/train/d-nerf/hellwarrior.sh 22 | wait 23 | ./scripts/eval/d-nerf/hellwarrior.sh 24 | wait 25 | 26 | ./scripts/train/d-nerf/hook.sh 27 | wait 28 | ./scripts/eval/d-nerf/hook.sh 29 | wait 30 | 31 | ./scripts/train/d-nerf/mutant.sh 32 | wait 33 | ./scripts/eval/d-nerf/mutant.sh 34 | wait 35 | 36 | ./scripts/train/d-nerf/standup.sh 37 | wait 38 | ./scripts/eval/d-nerf/standup.sh 39 | wait 40 | 41 | ./scripts/train/d-nerf/jumping.sh 42 | wait 43 | ./scripts/eval/d-nerf/jumping.sh 44 | wait 45 | 46 | ./scripts/train/d-nerf/trex.sh 47 | wait 48 | ./scripts/eval/d-nerf/trex.sh 49 | wait 50 | 51 | ./scripts/utils/dnerf_show.sh -------------------------------------------------------------------------------- /scripts/train/d-nerf/train_predict.sh: -------------------------------------------------------------------------------- 1 | echo "Start training and prediction process" 2 | echo "Please make sure your scripts [max_time] is 0.8" 3 | 4 | read -p "Do you want to continue? (yes/no): " user_input 5 | 6 | # 判断输入 7 | if [ "$user_input" = "yes" ]; then 8 | echo "Continuing execution..." 9 | # 在这里继续执行你的脚本逻辑 10 | else 11 | echo "Exiting..." 12 | exit 1 13 | fi 14 | 15 | # Train 16 | ./scripts/train/d-nerf/bouncingballs.sh 17 | wait 18 | ./scripts/predict/d-nerf/bouncingballs.sh 19 | wait 20 | 21 | ./scripts/train/d-nerf/hellwarrior.sh 22 | wait 23 | ./scripts/predict/d-nerf/hellwarrior.sh 24 | wait 25 | 26 | ./scripts/train/d-nerf/hook.sh 27 | wait 28 | ./scripts/predict/d-nerf/hook.sh 29 | wait 30 | 31 | ./scripts/train/d-nerf/mutant.sh 32 | wait 33 | ./scripts/predict/d-nerf/mutant.sh 34 | wait 35 | 36 | ./scripts/train/d-nerf/standup.sh 37 | wait 38 | ./scripts/predict/d-nerf/standup.sh 39 | wait 40 | 41 | ./scripts/train/d-nerf/jumping.sh 42 | wait 43 | ./scripts/predict/d-nerf/jumping.sh 44 | wait 45 | 46 | ./scripts/train/d-nerf/trex.sh 47 | wait 48 | ./scripts/predict/d-nerf/trex.sh 49 | wait 50 | 51 | -------------------------------------------------------------------------------- /scripts/train/d-nerf/trex.sh: -------------------------------------------------------------------------------- 1 | if [ $# -eq 0 ]; then 2 | echo "Seed: 1*2024" 3 | seed=1 4 | else 5 | seed=$1 6 | fi 7 | scene_name="trex" 8 | max_keypoints=100 9 | adaptive_points_num=100 10 | time_freq=6 11 | nearest_num=6 12 | source_path="./datasets/d-nerf/data/${scene_name}/" 13 | max_time=1.0 14 | model_path="./results/d-nerf_${max_time}/${scene_name}/finalVersion/" 15 | position_lr_max_steps=40000 16 | adaptive_from_iter=3000 17 | feature_amplify=0.5 18 | 19 | # Train 20 | CUDA_VISIBLE_DEVICES=0 python train.py -s $source_path \ 21 | -m $model_path --max_points $max_keypoints --adaptive_points_num $adaptive_points_num \ 22 | --iterations 60000 --test_iterations 60000 --jointly_iteration 1000 --time_freq $time_freq \ 23 | --densify_from_iter 3000 --densify_until_iter 20000 --norm_rotation \ 24 | --save_iterations 29999 60000 --checkpoint_iterations 29999 60000 --max_time $max_time \ 25 | --position_lr_max_steps $position_lr_max_steps --adaptive_from_iter $adaptive_from_iter --seed $seed \ 26 | --eval --nearest_num $nearest_num --adaptive_interval 500 --feature_amplify $feature_amplify -------------------------------------------------------------------------------- /scripts/train/hyper/chickchicken.sh: -------------------------------------------------------------------------------- 1 | scene_name="chickchicken" 2 | time_freq=8 3 | feature_amplify=5 4 | max_keypoints=100 5 | adaptive_points_num=100 6 | nearest_num=6 7 | max_time=1.0 8 | source_path="./datasets/HyperNeRF/${scene_name}/" 9 | model_path="./results/HyperNeRF_${max_time}/${scene_name}/finalVersion" 10 | position_lr_max_steps=40000 11 | data_device="cpu" 12 | 13 | CUDA_VISIBLE_DEVICES=1 python train.py -s $source_path \ 14 | -m $model_path --max_points $max_keypoints --adaptive_points_num $adaptive_points_num \ 15 | --iterations 70000 --test_iterations 70000 --jointly_iteration 1000 --time_freq $time_freq\ 16 | --densify_from_iter 5000 --save_iterations 29998 70000 --checkpoint_iterations 29998 70000 \ 17 | --densify_until_iter 15000 --opacity_reset_interval 3000000 --max_time $max_time --norm_rotation \ 18 | --position_lr_max_steps $position_lr_max_steps --data_device $data_device \ 19 | --eval --use_time_decay --nearest_num $nearest_num --adaptive_interval 1000 --feature_amplify $feature_amplify -------------------------------------------------------------------------------- /scripts/train/hyper/lemon.sh: -------------------------------------------------------------------------------- 1 | scene_name="cut-lemon" 2 | time_freq=10 3 | feature_amplify=5 4 | max_keypoints=100 5 | adaptive_points_num=200 6 | nearest_num=6 7 | max_time=1.0 8 | source_path="./datasets/HyperNeRF/${scene_name}/" 9 | model_path="./results/HyperNeRF_${max_time}/${scene_name}/finalVersion" 10 | position_lr_max_steps=40000 11 | data_device="cpu" 12 | step_opacity_iteration=5000 13 | 14 | # Train 15 | CUDA_VISIBLE_DEVICES=1 python train.py -s $source_path \ 16 | -m $model_path --max_points $max_keypoints --adaptive_points_num $adaptive_points_num \ 17 | --iterations 70000 --test_iterations 70000 --jointly_iteration 1000 --time_freq $time_freq\ 18 | --densify_from_iter 5000 --save_iterations 29998 70000 --checkpoint_iterations 29998 70000 \ 19 | --densify_until_iter 15000 --opacity_reset_interval 3000000 --max_time $max_time --norm_rotation --step_opacity \ 20 | --position_lr_max_steps $position_lr_max_steps --data_device $data_device --step_opacity_iteration $step_opacity_iteration \ 21 | --eval --use_time_decay --nearest_num $nearest_num --adaptive_interval 1000 --feature_amplify $feature_amplify -------------------------------------------------------------------------------- /scripts/train/hyper/printer.sh: -------------------------------------------------------------------------------- 1 | scene_name="vrig-3dprinter" 2 | time_freq=8 3 | feature_amplify=5 4 | max_keypoints=150 5 | adaptive_points_num=100 6 | nearest_num=6 7 | max_time=1.0 8 | source_path="./datasets/HyperNeRF/${scene_name}/" 9 | model_path="./results/HyperNeRF_${max_time}/${scene_name}/finalVersion/" 10 | position_lr_max_steps=40000 11 | data_device="cpu" 12 | 13 | # Train 14 | CUDA_VISIBLE_DEVICES=1 python train.py -s $source_path \ 15 | -m $model_path --max_points $max_keypoints --adaptive_points_num $adaptive_points_num \ 16 | --iterations 70000 --test_iterations 70000 --jointly_iteration 1000 --time_freq $time_freq\ 17 | --densify_from_iter 5000 --save_iterations 29998 70000 --checkpoint_iterations 29998 70000 \ 18 | --densify_until_iter 15000 --opacity_reset_interval 3000000 --max_time $max_time --norm_rotation \ 19 | --position_lr_max_steps $position_lr_max_steps --data_device $data_device \ 20 | --eval --use_time_decay --nearest_num $nearest_num --adaptive_interval 1000 --feature_amplify $feature_amplify -------------------------------------------------------------------------------- /scripts/train/hyper/torchocolate.sh: -------------------------------------------------------------------------------- 1 | scene_name="torchocolate" 2 | time_freq=8 3 | feature_amplify=5 4 | max_keypoints=50 5 | adaptive_points_num=100 6 | nearest_num=6 7 | max_time=1.0 8 | source_path="./datasets/HyperNeRF/${scene_name}/" 9 | model_path="./results/HyperNeRF_${max_time}/${scene_name}/finalVersion/" 10 | position_lr_max_steps=40000 11 | data_device="cpu" 12 | 13 | CUDA_VISIBLE_DEVICES=1 python train.py -s $source_path \ 14 | -m $model_path --max_points $max_keypoints --adaptive_points_num $adaptive_points_num \ 15 | --iterations 70000 --test_iterations 70000 --jointly_iteration 1000 --time_freq $time_freq \ 16 | --densify_from_iter 5000 --save_iterations 29998 70000 --checkpoint_iterations 29998 70000 --norm_rotation \ 17 | --densify_until_iter 15000 --opacity_reset_interval 3000000 --max_time $max_time \ 18 | --position_lr_max_steps $position_lr_max_steps --data_device $data_device \ 19 | --eval --use_time_decay --nearest_num $nearest_num --adaptive_interval 1000 --feature_amplify $feature_amplify -------------------------------------------------------------------------------- /scripts/train/hyper/train_eval.sh: -------------------------------------------------------------------------------- 1 | echo "Start training and evaluation process" 2 | echo "Please make sure your scripts [max_time] is 1.0 " 3 | 4 | read -p "Do you want to continue? (yes/no): " user_input 5 | 6 | # 判断输入 7 | if [ "$user_input" = "yes" ]; then 8 | echo "Continuing execution..." 9 | # 在这里继续执行你的脚本逻辑 10 | else 11 | echo "Exiting..." 12 | exit 1 13 | fi 14 | 15 | ./scripts/train/hyper/lemon.sh 16 | wait 17 | ./scripts/eval/hyper/lemon.sh 18 | wait 19 | 20 | ./scripts/train/hyper/chickchicken.sh 21 | wait 22 | ./scripts/eval/hyper/chickchicken.sh 23 | wait 24 | 25 | ./scripts/train/hyper/printer.sh 26 | wait 27 | ./scripts/eval/hyper/printer.sh 28 | wait 29 | 30 | ./scripts/train/hyper/torchocolate.sh 31 | wait 32 | ./scripts/eval/hyper/torchocolate.sh 33 | wait 34 | 35 | ./scripts/utils/hyper_show.sh -------------------------------------------------------------------------------- /scripts/train/hyper/train_predict.sh: -------------------------------------------------------------------------------- 1 | echo "Start training and prediction process" 2 | echo "Please make sure your scripts [max_time] is 0.8" 3 | 4 | read -p "Do you want to continue? (yes/no): " user_input 5 | 6 | # 判断输入 7 | if [ "$user_input" = "yes" ]; then 8 | echo "Continuing execution..." 9 | # 在这里继续执行你的脚本逻辑 10 | else 11 | echo "Exiting..." 12 | exit 1 13 | fi 14 | 15 | ./scripts/train/hyper/lemon.sh 16 | wait 17 | ./scripts/predict/hyper/lemon.sh 18 | wait 19 | 20 | ./scripts/train/hyper/chickchicken.sh 21 | wait 22 | ./scripts/predict/hyper/chickchicken.sh 23 | wait 24 | 25 | ./scripts/train/hyper/printer.sh 26 | wait 27 | ./scripts/predict/hyper/printer.sh 28 | wait 29 | 30 | ./scripts/train/hyper/torchocolate.sh 31 | wait 32 | ./scripts/predict/hyper/torchocolate.sh 33 | wait 34 | 35 | echo "Done" -------------------------------------------------------------------------------- /scripts/utils/colmap.sh: -------------------------------------------------------------------------------- 1 | workdir=$1 2 | datatype="hypernerf" 3 | export CUDA_VISIBLE_DEVICES=0 4 | rm -rf $workdir/sparse_ 5 | rm -rf $workdir/image_colmap 6 | python utils/prepare/"$datatype"2colmap.py $workdir 7 | rm -rf $workdir/colmap 8 | rm -rf $workdir/colmap/sparse/0 9 | 10 | mkdir $workdir/colmap 11 | cp -r $workdir/image_colmap $workdir/colmap/images 12 | cp -r $workdir/sparse_ $workdir/colmap/sparse_custom 13 | colmap feature_extractor --database_path $workdir/colmap/database.db --image_path $workdir/colmap/images --SiftExtraction.max_image_size 4096 --SiftExtraction.max_num_features 16384 --SiftExtraction.estimate_affine_shape 1 --SiftExtraction.domain_size_pooling 1 14 | python database.py --database_path $workdir/colmap/database.db --txt_path $workdir/colmap/sparse_custom/cameras.txt 15 | colmap exhaustive_matcher --database_path $workdir/colmap/database.db 16 | mkdir -p $workdir/colmap/sparse/0 17 | 18 | colmap point_triangulator --database_path $workdir/colmap/database.db --image_path $workdir/colmap/images --input_path $workdir/colmap/sparse_custom --output_path $workdir/colmap/sparse/0 --clear_points 1 19 | 20 | mkdir -p $workdir/colmap/dense/workspace 21 | colmap image_undistorter --image_path $workdir/colmap/images --input_path $workdir/colmap/sparse/0 --output_path $workdir/colmap/dense/workspace 22 | colmap patch_match_stereo --workspace_path $workdir/colmap/dense/workspace 23 | colmap stereo_fusion --workspace_path $workdir/colmap/dense/workspace --output_path $workdir/colmap/dense/workspace/fused.ply 24 | python utils/prepare/downsample_points.py $workdir/colmap/dense/workspace/fused.ply $workdir/points3D_downsample.ply -------------------------------------------------------------------------------- /scripts/utils/dnerf_show.sh: -------------------------------------------------------------------------------- 1 | root="./results" 2 | model_path="finalVersion" 3 | 4 | ### Max time == 1 (reconsturction results) 5 | max_time=1.0 6 | dataset="d-nerf_${max_time}" 7 | python show.py -r $root -d $dataset -m $model_path 8 | 9 | ### Max time < 1 (prediction results) 10 | # max_time=0.8 11 | # dataset="d-nerf_${max_time}" 12 | # eval_path="input10_stage6_noise0_epoch2001/[metrics]Predicted_by_GCN_on_test_views" 13 | # python show.py -r $root -d $dataset -m $model_path --eval_path $eval_path -------------------------------------------------------------------------------- /scripts/utils/env.sh: -------------------------------------------------------------------------------- 1 | conda install -c nvidia cudatoolkit=11.7 2 | conda install -c conda-forge cudatoolkit-dev=11.7 3 | 4 | pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 --extra-index-url https://download.pytorch.org/whl/cu117 5 | pip install -r requirements.txt 6 | 7 | pip install git+https://github.com/BoMingZhao/tiny-cuda-nn-float32/#subdirectory=bindings/torch 8 | pip install submodules/diff-gaussian-rasterization-w-depth 9 | pip install submodules/simple-knn 10 | pip install submodules/FRNN/external/prefix_sum 11 | pip install submodules/FRNN 12 | cd submodules/lib/pointops 13 | pip install . 14 | pip install "git+https://github.com/facebookresearch/pytorch3d.git" -------------------------------------------------------------------------------- /scripts/utils/hyper_show.sh: -------------------------------------------------------------------------------- 1 | root="./results" 2 | model_path="finalVersion" 3 | dataset="HyperNeRF_1.0" 4 | 5 | python show.py -r $root -d $dataset -m $model_path -------------------------------------------------------------------------------- /show.py: -------------------------------------------------------------------------------- 1 | # To vis the final evaluation results in a beautiful table 2 | 3 | import os 4 | import json 5 | from prettytable import PrettyTable 6 | from argparse import ArgumentParser 7 | 8 | if __name__ == '__main__': 9 | parser = ArgumentParser(description="Visualization parameters") 10 | parser.add_argument('--root', '-r', type=str, default="./results/") 11 | parser.add_argument('--dataset', '-d', type=str, required=True) 12 | parser.add_argument('--scene', '-s', nargs='*', help="If this list is empty, we will search for all scenes in the root dir") 13 | parser.add_argument('--model_path', '-m', type=str, required=True) 14 | parser.add_argument('--eval_path', '-e', type=str, default="eval/test") 15 | 16 | args = parser.parse_args() 17 | root_dir = os.path.join(args.root, args.dataset) 18 | 19 | if args.scene is not None: 20 | scenes = args.scene 21 | else: 22 | scenes = os.listdir(root_dir) 23 | 24 | info = [] 25 | cnt, psnr_avg, ssim_avg, msssim_avg, lpivgg_avg, lpialex_avg = 0, 0., 0., 0., 0., 0. 26 | for scene in scenes: 27 | json_path = os.path.join(root_dir, scene, args.model_path, args.eval_path, "results.json") 28 | if os.path.exists(json_path): 29 | with open(json_path, "r") as f: 30 | cnt += 1 31 | data_ = json.load(f) 32 | # for key in data_.keys(): 33 | temp_data = {} 34 | temp_data["name"] = os.path.join(scene) 35 | temp_data["psnr"] = round(data_["PSNR"], 2) 36 | psnr_avg += round(data_["PSNR"], 2) 37 | temp_data["ssim"] = round(data_["SSIM"], 4) 38 | ssim_avg += round(data_["SSIM"], 4) 39 | temp_data["msssim"] = round(data_["MS-SSIM"], 4) 40 | msssim_avg += round(data_["MS-SSIM"], 4) 41 | temp_data["lpivgg"] = round(data_["LPIPS-vgg"], 4) 42 | lpivgg_avg += round(data_["LPIPS-vgg"], 4) 43 | temp_data["lpialex"] = round(data_["LPIPS-alex"], 4) 44 | lpialex_avg += round(data_["LPIPS-alex"], 4) 45 | 46 | info.append(temp_data) 47 | 48 | info = sorted(info, key=lambda x: x["name"]) 49 | info.append({"name": "Average", 50 | "psnr": round(psnr_avg / cnt, 2), 51 | "ssim": round(ssim_avg / cnt, 4), 52 | "msssim": round(msssim_avg / cnt, 4), 53 | "lpivgg": round(lpivgg_avg / cnt, 4), 54 | "lpialex": round(lpialex_avg / cnt, 4)}) 55 | 56 | table = PrettyTable() 57 | table.field_names = ["Scene", "PSNR", "SSIM", "MS-SSIMS", "LPIPS-vgg", "LPIPS-alex"] 58 | 59 | for data in info: 60 | table.add_row([data["name"], data["psnr"], data["ssim"], data["msssim"], data["lpivgg"], data["lpialex"]]) 61 | 62 | print("+++++++++++++++++++++++++++++++++++++") 63 | print(f"Visualizing {args.dataset} dataset: {args.model_path}") 64 | print("+++++++++++++++++++++++++++++++++++++") 65 | print(table) -------------------------------------------------------------------------------- /submodules/lib/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BoMingZhao/GaussianPrediction/301f3d552b8380a8540f7eb2a2bca22bafb00b50/submodules/lib/__init__.py -------------------------------------------------------------------------------- /submodules/lib/pointops/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BoMingZhao/GaussianPrediction/301f3d552b8380a8540f7eb2a2bca22bafb00b50/submodules/lib/pointops/__init__.py -------------------------------------------------------------------------------- /submodules/lib/pointops/dist/pointops-0.0.0-py3.7-linux-x86_64.egg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BoMingZhao/GaussianPrediction/301f3d552b8380a8540f7eb2a2bca22bafb00b50/submodules/lib/pointops/dist/pointops-0.0.0-py3.7-linux-x86_64.egg -------------------------------------------------------------------------------- /submodules/lib/pointops/functions/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BoMingZhao/GaussianPrediction/301f3d552b8380a8540f7eb2a2bca22bafb00b50/submodules/lib/pointops/functions/__init__.py -------------------------------------------------------------------------------- /submodules/lib/pointops/pointops.egg-info/PKG-INFO: -------------------------------------------------------------------------------- 1 | Metadata-Version: 2.1 2 | Name: pointops 3 | Version: 0.0.0 4 | Summary: UNKNOWN 5 | Home-page: UNKNOWN 6 | Author: Hengshuang Zhao 7 | License: UNKNOWN 8 | Platform: UNKNOWN 9 | 10 | UNKNOWN 11 | 12 | -------------------------------------------------------------------------------- /submodules/lib/pointops/pointops.egg-info/SOURCES.txt: -------------------------------------------------------------------------------- 1 | setup.py 2 | pointops.egg-info/PKG-INFO 3 | pointops.egg-info/SOURCES.txt 4 | pointops.egg-info/dependency_links.txt 5 | pointops.egg-info/top_level.txt 6 | src/pointops_api.cpp 7 | src/aggregation/aggregation_cuda.cpp 8 | src/aggregation/aggregation_cuda_kernel.cu 9 | src/fast_sampling/fast_sampling_cuda.cpp 10 | src/fast_sampling/fast_sampling_cuda_kernel.cu 11 | src/grouping/grouping_cuda.cpp 12 | src/grouping/grouping_cuda_kernel.cu 13 | src/interpolation/interpolation_cuda.cpp 14 | src/interpolation/interpolation_cuda_kernel.cu 15 | src/knnquery/knnquery_cuda.cpp 16 | src/knnquery/knnquery_cuda_kernel.cu 17 | src/sampling/sampling_cuda.cpp 18 | src/sampling/sampling_cuda_kernel.cu 19 | src/subtraction/subtraction_cuda.cpp 20 | src/subtraction/subtraction_cuda_kernel.cu -------------------------------------------------------------------------------- /submodules/lib/pointops/pointops.egg-info/dependency_links.txt: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /submodules/lib/pointops/pointops.egg-info/top_level.txt: -------------------------------------------------------------------------------- 1 | pointops_cuda 2 | -------------------------------------------------------------------------------- /submodules/lib/pointops/setup.py: -------------------------------------------------------------------------------- 1 | #python3 setup.py install 2 | from setuptools import setup 3 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 4 | import os 5 | from distutils.sysconfig import get_config_vars 6 | 7 | (opt,) = get_config_vars('OPT') 8 | os.environ['OPT'] = " ".join( 9 | flag for flag in opt.split() if flag != '-Wstrict-prototypes' 10 | ) 11 | 12 | setup( 13 | name='pointops', 14 | author='Hengshuang Zhao', 15 | ext_modules=[ 16 | CUDAExtension('pointops_cuda', [ 17 | 'src/pointops_api.cpp', 18 | 'src/knnquery/knnquery_cuda.cpp', 19 | 'src/knnquery/knnquery_cuda_kernel.cu', 20 | 'src/sampling/sampling_cuda.cpp', 21 | 'src/sampling/sampling_cuda_kernel.cu', 22 | 'src/fast_sampling/fast_sampling_cuda.cpp', 23 | 'src/fast_sampling/fast_sampling_cuda_kernel.cu', 24 | 'src/grouping/grouping_cuda.cpp', 25 | 'src/grouping/grouping_cuda_kernel.cu', 26 | 'src/interpolation/interpolation_cuda.cpp', 27 | 'src/interpolation/interpolation_cuda_kernel.cu', 28 | 'src/subtraction/subtraction_cuda.cpp', 29 | 'src/subtraction/subtraction_cuda_kernel.cu', 30 | 'src/aggregation/aggregation_cuda.cpp', 31 | 'src/aggregation/aggregation_cuda_kernel.cu', 32 | ], 33 | extra_compile_args={'cxx': ['-g'], 'nvcc': ['-O2']} 34 | ) 35 | ], 36 | cmdclass={'build_ext': BuildExtension} 37 | ) 38 | -------------------------------------------------------------------------------- /submodules/lib/pointops/src/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BoMingZhao/GaussianPrediction/301f3d552b8380a8540f7eb2a2bca22bafb00b50/submodules/lib/pointops/src/__init__.py -------------------------------------------------------------------------------- /submodules/lib/pointops/src/aggregation/aggregation_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | // #include 3 | #include 4 | #include 5 | #include "aggregation_cuda_kernel.h" 6 | 7 | 8 | void aggregation_forward_cuda(int n, int nsample, int c, int w_c, at::Tensor input_tensor, at::Tensor position_tensor, at::Tensor weight_tensor, at::Tensor idx_tensor, at::Tensor output_tensor) 9 | { 10 | const float *input = input_tensor.data_ptr(); 11 | const float *position = position_tensor.data_ptr(); 12 | const float *weight = weight_tensor.data_ptr(); 13 | const int *idx = idx_tensor.data_ptr(); 14 | float *output = output_tensor.data_ptr(); 15 | aggregation_forward_cuda_launcher(n, nsample, c, w_c, input, position, weight, idx, output); 16 | } 17 | 18 | void aggregation_backward_cuda(int n, int nsample, int c, int w_c, at::Tensor input_tensor, at::Tensor position_tensor, at::Tensor weight_tensor, at::Tensor idx_tensor, at::Tensor grad_output_tensor, at::Tensor grad_input_tensor, at::Tensor grad_position_tensor, at::Tensor grad_weight_tensor) 19 | { 20 | const float *input = input_tensor.data_ptr(); 21 | const float *position = position_tensor.data_ptr(); 22 | const float *weight = weight_tensor.data_ptr(); 23 | const int *idx = idx_tensor.data_ptr(); 24 | const float *grad_output = grad_output_tensor.data_ptr(); 25 | float *grad_input = grad_input_tensor.data_ptr(); 26 | float *grad_position = grad_position_tensor.data_ptr(); 27 | float *grad_weight = grad_weight_tensor.data_ptr(); 28 | aggregation_backward_cuda_launcher(n, nsample, c, w_c, input, position, weight, idx, grad_output, grad_input, grad_position, grad_weight); 29 | } 30 | -------------------------------------------------------------------------------- /submodules/lib/pointops/src/aggregation/aggregation_cuda_kernel.cu: -------------------------------------------------------------------------------- 1 | #include "../cuda_utils.h" 2 | #include "aggregation_cuda_kernel.h" 3 | 4 | 5 | __global__ void aggregation_forward_cuda_kernel(int n, int nsample, int c, int w_c, const float *input, const float *position, const float *weight, const int *idx, float *output) { 6 | // input: input: (n, c), position: (n, nsample, c), weight: (n, nsample, w_c), idx: (n, nsample), output: (n, c) 7 | int index = blockIdx.x * blockDim.x + threadIdx.x; 8 | if (index >= n * c) return; 9 | const int c_idx = index % c; 10 | const int n_idx = index / c; 11 | const int w_c_idx = c_idx % w_c; 12 | for (int nsample_idx = 0; nsample_idx < nsample; nsample_idx++) 13 | { 14 | int idx_idx = n_idx * nsample + nsample_idx; 15 | int input_idx = idx[idx_idx] * c + c_idx; 16 | int position_idx = n_idx * nsample * c + nsample_idx * c + c_idx; 17 | int weight_idx = n_idx * nsample * w_c + nsample_idx * w_c + w_c_idx; 18 | output[index] += (input[input_idx] + position[position_idx]) * weight[weight_idx]; 19 | } 20 | } 21 | 22 | __global__ void aggregation_backward_cuda_kernel(int n, int nsample, int c, int w_c, const float *input, const float *position, const float *weight, const int *idx, const float *grad_output, float *grad_input, float *grad_position, float *grad_weight) { 23 | // input: grad_output: (n, c), output: grad_input: (n, c), grad_position: (n, nsample, c), grad_weight: (n, nsample, w_c) 24 | int index = blockIdx.x * blockDim.x + threadIdx.x; 25 | if (index >= n * c) return; 26 | const int c_idx = index % c; 27 | const int n_idx = index / c; 28 | const int w_c_idx = c_idx % w_c; 29 | for (int nsample_idx = 0; nsample_idx < nsample; nsample_idx++) 30 | { 31 | int idx_idx = n_idx * nsample + nsample_idx; 32 | int input_idx = idx[idx_idx] * c + c_idx; 33 | int position_idx = n_idx * nsample * c + nsample_idx * c + c_idx; 34 | int weight_idx = n_idx * nsample * w_c + nsample_idx * w_c + w_c_idx; 35 | atomicAdd(grad_input + input_idx, grad_output[index] * weight[weight_idx]); 36 | grad_position[position_idx] = grad_output[index] * weight[weight_idx]; 37 | atomicAdd(grad_weight + weight_idx, grad_output[index] * (input[input_idx] + position[position_idx])); 38 | } 39 | } 40 | 41 | void aggregation_forward_cuda_launcher(int n, int nsample, int c, int w_c, const float *input, const float *position, const float *weight, const int *idx, float *output) { 42 | // input: input: (n, c), position: (n, nsample, c), weight: (n, nsample, w_c), idx: (n, nsample), output: (n, c) 43 | dim3 blocks(DIVUP(n * c, THREADS_PER_BLOCK)); 44 | dim3 threads(THREADS_PER_BLOCK); 45 | aggregation_forward_cuda_kernel<<>>(n, nsample, c, w_c, input, position, weight, idx, output); 46 | } 47 | 48 | void aggregation_backward_cuda_launcher(int n, int nsample, int c, int w_c, const float *input, const float *position, const float *weight, const int *idx, const float *grad_output, float *grad_input, float *grad_position, float *grad_weight) { 49 | // input: grad_output: (n, c), output: grad_input: (n, c), grad_position: (n, nsample, c), grad_weight: (n, nsample, w_c) 50 | dim3 blocks(DIVUP(n * c, THREADS_PER_BLOCK)); 51 | dim3 threads(THREADS_PER_BLOCK); 52 | aggregation_backward_cuda_kernel<<>>(n, nsample, c, w_c, input, position, weight, idx, grad_output, grad_input, grad_position, grad_weight); 53 | } 54 | -------------------------------------------------------------------------------- /submodules/lib/pointops/src/aggregation/aggregation_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _AGGREGATION_CUDA_KERNEL 2 | #define _AGGREGATION_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void aggregation_forward_cuda(int n, int nsample, int c, int w_c, at::Tensor input_tensor, at::Tensor position_tensor, at::Tensor weight_tensor, at::Tensor idx_tensor, at::Tensor output_tensor); 8 | void aggregation_backward_cuda(int n, int nsample, int c, int w_c, at::Tensor input_tensor, at::Tensor position_tensor, at::Tensor weight_tensor, at::Tensor idx_tensor, at::Tensor grad_output_tensor, at::Tensor grad_input_tensor, at::Tensor grad_position_tensor, at::Tensor grad_weight_tensor); 9 | 10 | #ifdef __cplusplus 11 | extern "C" { 12 | #endif 13 | 14 | void aggregation_forward_cuda_launcher(int n, int nsample, int c, int w_c, const float *input, const float *position, const float *weight, const int *idx, float *output); 15 | void aggregation_backward_cuda_launcher(int n, int nsample, int c, int w_c, const float *input, const float *position, const float *weight, const int *idx, const float *grad_output, float *grad_input, float *grad_position, float *grad_weight); 16 | 17 | #ifdef __cplusplus 18 | } 19 | #endif 20 | #endif 21 | -------------------------------------------------------------------------------- /submodules/lib/pointops/src/cuda_utils.h: -------------------------------------------------------------------------------- 1 | #ifndef _CUDA_UTILS_H 2 | #define _CUDA_UTILS_H 3 | #include 4 | #include 5 | #include 6 | 7 | #define TOTAL_THREADS 1024 8 | #define THREADS_PER_BLOCK 256 9 | #define DIVUP(m, n) ((m) / (n) + ((m) % (n) > 0)) 10 | 11 | inline void HandleError( cudaError_t err, 12 | const char *file, 13 | int line ) { 14 | if (err != cudaSuccess) { 15 | printf( "%s in %s at line %d\n", cudaGetErrorString( err ), 16 | file, line ); 17 | exit( EXIT_FAILURE ); 18 | } 19 | } 20 | #define HANDLE_ERROR( err ) (HandleError( err, __FILE__, __LINE__ )) 21 | 22 | 23 | inline int opt_n_threads(int work_size) { 24 | const int pow_2 = std::ceil(std::log(static_cast(work_size)) / std::log(2.0)); 25 | // const int pow_2 = std::log(static_cast(work_size)) / std::log(2.0); 26 | return std::max(std::min(1 << pow_2, TOTAL_THREADS), 1); 27 | } 28 | 29 | inline dim3 opt_block_config(int x, int y) { 30 | const int x_threads = opt_n_threads(x); 31 | const int y_threads = std::max(std::min(opt_n_threads(y), TOTAL_THREADS / x_threads), 1); 32 | dim3 block_config(x_threads, y_threads, 1); 33 | return block_config; 34 | } 35 | 36 | #endif 37 | -------------------------------------------------------------------------------- /submodules/lib/pointops/src/fast_sampling/fast_sampling_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | // #include 3 | #include 4 | #include 5 | #include "fast_sampling_cuda_kernel.h" 6 | 7 | 8 | void farthestsampling_cuda(int b, int n, int m, at::Tensor xyz_tensor, at::Tensor offset_tensor, at::Tensor new_offset_tensor, at::Tensor tmp_tensor, at::Tensor idx_tensor) 9 | { 10 | const float *xyz = xyz_tensor.data_ptr(); 11 | const int *offset = offset_tensor.data_ptr(); 12 | const int *new_offset = new_offset_tensor.data_ptr(); 13 | float *tmp = tmp_tensor.data_ptr(); 14 | int *idx = idx_tensor.data_ptr(); 15 | farthestsampling_cuda_launcher(b, n, m, xyz, offset, new_offset, tmp, idx); 16 | } 17 | -------------------------------------------------------------------------------- /submodules/lib/pointops/src/fast_sampling/fast_sampling_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _FAST_SAMPLING_CUDA_KERNEL 2 | #define _FAST_SAMPLING_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | #define MAX_BLOCKS 82 8 | 9 | void farthestsampling_cuda(int b, int n, int m, at::Tensor xyz_tensor, at::Tensor offset_tensor, at::Tensor new_offset_tensor, at::Tensor tmp_tensor, at::Tensor idx_tensor); 10 | 11 | #ifdef __cplusplus 12 | extern "C" { 13 | #endif 14 | 15 | void farthestsampling_cuda_launcher(int b, int n, int m, const float *xyz, const int *offset, const int *new_offset, float *tmp, int *idx); 16 | 17 | #ifdef __cplusplus 18 | } 19 | #endif 20 | #endif 21 | -------------------------------------------------------------------------------- /submodules/lib/pointops/src/grouping/grouping_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | // #include 3 | #include 4 | #include 5 | #include "grouping_cuda_kernel.h" 6 | 7 | 8 | void grouping_forward_cuda(int m, int nsample, int c, at::Tensor input_tensor, at::Tensor idx_tensor, at::Tensor output_tensor) 9 | { 10 | const float *input = input_tensor.data_ptr(); 11 | const int *idx = idx_tensor.data_ptr(); 12 | float *output = output_tensor.data_ptr(); 13 | grouping_forward_cuda_launcher(m, nsample, c, input, idx, output); 14 | } 15 | 16 | void grouping_backward_cuda(int m, int nsample, int c, at::Tensor grad_output_tensor, at::Tensor idx_tensor, at::Tensor grad_input_tensor) 17 | { 18 | const float *grad_output = grad_output_tensor.data_ptr(); 19 | const int *idx = idx_tensor.data_ptr(); 20 | float *grad_input = grad_input_tensor.data_ptr(); 21 | grouping_backward_cuda_launcher(m, nsample, c, grad_output, idx, grad_input); 22 | } 23 | -------------------------------------------------------------------------------- /submodules/lib/pointops/src/grouping/grouping_cuda_kernel.cu: -------------------------------------------------------------------------------- 1 | #include "../cuda_utils.h" 2 | #include "grouping_cuda_kernel.h" 3 | 4 | 5 | __global__ void grouping_forward_cuda_kernel(int m, int nsample, int c, const float *__restrict__ input, const int *__restrict__ idx, float *__restrict__ output) { 6 | // input: input: (n, c), idx: (m, nsample), output: (m, nsample, c) 7 | int index = blockIdx.x * blockDim.x + threadIdx.x; 8 | if (index >= m * nsample * c) return; 9 | const int c_idx = index % c; 10 | const int nsample_idx = (index / c) % nsample; 11 | const int m_idx = index / nsample / c; 12 | const int input_idx = idx[m_idx * nsample + nsample_idx] * c + c_idx; 13 | output[index] = input[input_idx]; 14 | } 15 | 16 | __global__ void grouping_backward_cuda_kernel(int m, int nsample, int c, const float *__restrict__ grad_output, const int *__restrict__ idx, float *__restrict__ grad_input) { 17 | // input: grad_output: (m, nsample, c), idx: (m, nsample), output: grad_input: (n, c) 18 | int index = blockIdx.x * blockDim.x + threadIdx.x; 19 | if (index >= m * nsample * c) return; 20 | const int c_idx = index % c; 21 | const int nsample_idx = (index / c) % nsample; 22 | const int m_idx = index / nsample / c; 23 | const int input_idx = idx[m_idx * nsample + nsample_idx] * c + c_idx; 24 | atomicAdd(grad_input + input_idx, grad_output[index]); 25 | } 26 | 27 | void grouping_forward_cuda_launcher(int m, int nsample, int c, const float *input, const int *idx, float *output) { 28 | // input: input: (n, c), idx: (m, nsample), output: (m, nsample, c) 29 | dim3 blocks(DIVUP(m * nsample * c, THREADS_PER_BLOCK)); 30 | dim3 threads(THREADS_PER_BLOCK); 31 | grouping_forward_cuda_kernel<<>>(m, nsample, c, input, idx, output); 32 | } 33 | 34 | void grouping_backward_cuda_launcher(int m, int nsample, int c, const float *grad_output, const int *idx, float *grad_input) 35 | { 36 | // input: grad_output: (m, nsample, c), idx: (m, nsample), output: grad_input: (n, c) 37 | dim3 blocks(DIVUP(m * nsample * c, THREADS_PER_BLOCK)); 38 | dim3 threads(THREADS_PER_BLOCK); 39 | grouping_backward_cuda_kernel<<>>(m, nsample, c, grad_output, idx, grad_input); 40 | } 41 | -------------------------------------------------------------------------------- /submodules/lib/pointops/src/grouping/grouping_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _GROUPING_CUDA_KERNEL 2 | #define _GROUPING_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void grouping_forward_cuda(int m, int nsample, int c, at::Tensor input_tensor, at::Tensor idx_tensor, at::Tensor output_tensor); 8 | void grouping_backward_cuda(int m, int nsample, int c, at::Tensor grad_output_tensor, at::Tensor idx_tensor, at::Tensor grad_input_tensor); 9 | 10 | #ifdef __cplusplus 11 | extern "C" { 12 | #endif 13 | 14 | void grouping_forward_cuda_launcher(int m, int nsample, int c, const float *input, const int *idx, float *output); 15 | void grouping_backward_cuda_launcher(int m, int nsample, int c, const float *grad_output, const int *idx, float *grad_input); 16 | 17 | #ifdef __cplusplus 18 | } 19 | #endif 20 | #endif 21 | -------------------------------------------------------------------------------- /submodules/lib/pointops/src/interpolation/interpolation_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | // #include 3 | #include 4 | #include 5 | #include "interpolation_cuda_kernel.h" 6 | 7 | 8 | void interpolation_forward_cuda(int n, int c, int k, at::Tensor input_tensor, at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor output_tensor) 9 | { 10 | const float *input = input_tensor.data_ptr(); 11 | const int *idx = idx_tensor.data_ptr(); 12 | const float *weight = weight_tensor.data_ptr(); 13 | float *output = output_tensor.data_ptr(); 14 | interpolation_forward_cuda_launcher(n, c, k, input, idx, weight, output); 15 | } 16 | 17 | void interpolation_backward_cuda(int n, int c, int k, at::Tensor grad_output_tensor, at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor grad_input_tensor) 18 | { 19 | const float *grad_output = grad_output_tensor.data_ptr(); 20 | const int *idx = idx_tensor.data_ptr(); 21 | const float *weight = weight_tensor.data_ptr(); 22 | float *grad_input = grad_input_tensor.data_ptr(); 23 | interpolation_backward_cuda_launcher(n, c, k, grad_output, idx, weight, grad_input); 24 | } 25 | -------------------------------------------------------------------------------- /submodules/lib/pointops/src/interpolation/interpolation_cuda_kernel.cu: -------------------------------------------------------------------------------- 1 | #include "../cuda_utils.h" 2 | #include "interpolation_cuda_kernel.h" 3 | 4 | 5 | __global__ void interpolation_forward_cuda_kernel(int n, int c, int k, const float *input, const int *idx, const float *weight, float *output) 6 | { 7 | // input: input: (m, c), idx: (n, k), weight: (n, k), output: output (n, c) 8 | int index = blockIdx.x * blockDim.x + threadIdx.x; 9 | if (index >= n * c) return; 10 | int c_idx = index % c; 11 | int n_idx = index / c; 12 | for (int i = 0; i < k; i++) 13 | { 14 | int idx_idx = n_idx * k + i; 15 | int input_idx = idx[idx_idx] * c + c_idx; 16 | output[index] += input[input_idx] * weight[idx_idx]; 17 | } 18 | } 19 | 20 | __global__ void interpolation_backward_cuda_kernel(int n, int c, int k, const float *grad_output, const int *idx, const float *weight, float *grad_input) 21 | { 22 | // input: grad_output: (n, c), idx: (n, k), weight: (n, k), output: grad_input (m, c) 23 | int index = blockIdx.x * blockDim.x + threadIdx.x; 24 | if (index >= n * c) return; 25 | int c_idx = index % c; 26 | int n_idx = index / c; 27 | for (int i = 0; i < k; i++) 28 | { 29 | int idx_idx = n_idx * k + i; 30 | int input_idx = idx[idx_idx] * c + c_idx; 31 | atomicAdd(grad_input + input_idx, grad_output[index] * weight[idx_idx]); 32 | } 33 | } 34 | 35 | void interpolation_forward_cuda_launcher(int n, int c, int k, const float *input, const int *idx, const float *weight, float *output) { 36 | // input: input: (m, c), idx: (n, k), weight: (n, k), output: output (n, c) 37 | dim3 blocks(DIVUP(n * c, THREADS_PER_BLOCK)); 38 | dim3 threads(THREADS_PER_BLOCK); 39 | interpolation_forward_cuda_kernel<<>>(n, c, k, input, idx, weight, output); 40 | } 41 | 42 | void interpolation_backward_cuda_launcher(int n, int c, int k, const float *grad_output, const int *idx, const float *weight, float *grad_input) { 43 | // input: grad_output: (n, c), idx: (n, k), weight: (n, k), output: grad_input (m, c) 44 | dim3 blocks(DIVUP(n * c, THREADS_PER_BLOCK)); 45 | dim3 threads(THREADS_PER_BLOCK); 46 | interpolation_backward_cuda_kernel<<>>(n, c, k, grad_output, idx, weight, grad_input); 47 | } 48 | -------------------------------------------------------------------------------- /submodules/lib/pointops/src/interpolation/interpolation_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _INTERPOLATION_CUDA_KERNEL 2 | #define _INTERPOLATION_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void interpolation_forward_cuda(int n, int c, int k, at::Tensor input_tensor, at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor output_tensor); 8 | void interpolation_backward_cuda(int n, int c, int k, at::Tensor grad_output_tensor, at::Tensor idx_tensor, at::Tensor weight_tensor, at::Tensor grad_input_tensor); 9 | 10 | #ifdef __cplusplus 11 | extern "C" { 12 | #endif 13 | 14 | void interpolation_forward_cuda_launcher(int n, int c, int k, const float *input, const int *idx, const float *weight, float *output); 15 | void interpolation_backward_cuda_launcher(int n, int c, int k, const float *grad_output, const int *idx, const float *weight, float *grad_input); 16 | 17 | #ifdef __cplusplus 18 | } 19 | #endif 20 | #endif 21 | -------------------------------------------------------------------------------- /submodules/lib/pointops/src/knnquery/knnquery_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | // #include 3 | #include 4 | #include 5 | #include "knnquery_cuda_kernel.h" 6 | 7 | 8 | void knnquery_cuda(int m, int nsample, at::Tensor xyz_tensor, at::Tensor new_xyz_tensor, at::Tensor offset_tensor, at::Tensor new_offset_tensor, at::Tensor idx_tensor, at::Tensor dist2_tensor) 9 | { 10 | const float *xyz = xyz_tensor.data_ptr(); 11 | const float *new_xyz = new_xyz_tensor.data_ptr(); 12 | const int *offset = offset_tensor.data_ptr(); 13 | const int *new_offset = new_offset_tensor.data_ptr(); 14 | int *idx = idx_tensor.data_ptr(); 15 | float *dist2 = dist2_tensor.data_ptr(); 16 | knnquery_cuda_launcher(m, nsample, xyz, new_xyz, offset, new_offset, idx, dist2); 17 | } 18 | -------------------------------------------------------------------------------- /submodules/lib/pointops/src/knnquery/knnquery_cuda_kernel.cu: -------------------------------------------------------------------------------- 1 | #include "../cuda_utils.h" 2 | #include "knnquery_cuda_kernel.h" 3 | 4 | 5 | __device__ void swap_float(float *x, float *y) 6 | { 7 | float tmp = *x; 8 | *x = *y; 9 | *y = tmp; 10 | } 11 | 12 | 13 | __device__ void swap_int(int *x, int *y) 14 | { 15 | int tmp = *x; 16 | *x = *y; 17 | *y = tmp; 18 | } 19 | 20 | 21 | __device__ void reheap(float *dist, int *idx, int k) 22 | { 23 | int root = 0; 24 | int child = root * 2 + 1; 25 | while (child < k) 26 | { 27 | if(child + 1 < k && dist[child+1] > dist[child]) 28 | child++; 29 | if(dist[root] > dist[child]) 30 | return; 31 | swap_float(&dist[root], &dist[child]); 32 | swap_int(&idx[root], &idx[child]); 33 | root = child; 34 | child = root * 2 + 1; 35 | } 36 | } 37 | 38 | 39 | __device__ void heap_sort(float *dist, int *idx, int k) 40 | { 41 | int i; 42 | for (i = k - 1; i > 0; i--) 43 | { 44 | swap_float(&dist[0], &dist[i]); 45 | swap_int(&idx[0], &idx[i]); 46 | reheap(dist, idx, i); 47 | } 48 | } 49 | 50 | 51 | __device__ int get_bt_idx(int idx, const int *offset) 52 | { 53 | int i = 0; 54 | while (1) 55 | { 56 | if (idx < offset[i]) 57 | break; 58 | else 59 | i++; 60 | } 61 | return i; 62 | } 63 | 64 | 65 | __global__ void knnquery_cuda_kernel(int m, int nsample, const float *__restrict__ xyz, const float *__restrict__ new_xyz, const int *__restrict__ offset, const int *__restrict__ new_offset, int *__restrict__ idx, float *__restrict__ dist2) { 66 | // input: xyz (n, 3) new_xyz (m, 3) 67 | // output: idx (m, nsample) dist2 (m, nsample) 68 | int pt_idx = blockIdx.x * blockDim.x + threadIdx.x; 69 | if (pt_idx >= m) return; 70 | 71 | new_xyz += pt_idx * 3; 72 | idx += pt_idx * nsample; 73 | dist2 += pt_idx * nsample; 74 | int bt_idx = get_bt_idx(pt_idx, new_offset); 75 | int start; 76 | if (bt_idx == 0) 77 | start = 0; 78 | else 79 | start = offset[bt_idx - 1]; 80 | int end = offset[bt_idx]; 81 | 82 | float new_x = new_xyz[0]; 83 | float new_y = new_xyz[1]; 84 | float new_z = new_xyz[2]; 85 | 86 | float best_dist[100]; 87 | int best_idx[100]; 88 | for(int i = 0; i < nsample; i++){ 89 | best_dist[i] = 1e10; 90 | best_idx[i] = start; 91 | } 92 | for(int i = start; i < end; i++){ 93 | float x = xyz[i * 3 + 0]; 94 | float y = xyz[i * 3 + 1]; 95 | float z = xyz[i * 3 + 2]; 96 | float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) + (new_z - z) * (new_z - z); 97 | if (d2 < best_dist[0]){ 98 | best_dist[0] = d2; 99 | best_idx[0] = i; 100 | reheap(best_dist, best_idx, nsample); 101 | } 102 | } 103 | heap_sort(best_dist, best_idx, nsample); 104 | for(int i = 0; i < nsample; i++){ 105 | idx[i] = best_idx[i]; 106 | dist2[i] = best_dist[i]; 107 | } 108 | } 109 | 110 | 111 | void knnquery_cuda_launcher(int m, int nsample, const float *xyz, const float *new_xyz, const int *offset, const int *new_offset, int *idx, float *dist2) { 112 | // input: new_xyz: (m, 3), xyz: (n, 3), idx: (m, nsample) 113 | dim3 blocks(DIVUP(m, THREADS_PER_BLOCK)); 114 | dim3 threads(THREADS_PER_BLOCK); 115 | knnquery_cuda_kernel<<>>(m, nsample, xyz, new_xyz, offset, new_offset, idx, dist2); 116 | } 117 | -------------------------------------------------------------------------------- /submodules/lib/pointops/src/knnquery/knnquery_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _KNNQUERY_CUDA_KERNEL 2 | #define _KNNQUERY_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void knnquery_cuda(int m, int nsample, at::Tensor xyz_tensor, at::Tensor new_xyz_tensor, at::Tensor offset_tensor, at::Tensor new_offset_tensor, at::Tensor idx_tensor, at::Tensor dist2_tensor); 8 | 9 | #ifdef __cplusplus 10 | extern "C" { 11 | #endif 12 | 13 | void knnquery_cuda_launcher(int m, int nsample, const float *xyz, const float *new_xyz, const int *offset, const int *new_offset, int *idx, float *dist2); 14 | 15 | #ifdef __cplusplus 16 | } 17 | #endif 18 | #endif 19 | -------------------------------------------------------------------------------- /submodules/lib/pointops/src/pointops_api.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include "knnquery/knnquery_cuda_kernel.h" 5 | #include "sampling/sampling_cuda_kernel.h" 6 | #include "fast_sampling/fast_sampling_cuda_kernel.h" 7 | #include "grouping/grouping_cuda_kernel.h" 8 | #include "interpolation/interpolation_cuda_kernel.h" 9 | #include "aggregation/aggregation_cuda_kernel.h" 10 | #include "subtraction/subtraction_cuda_kernel.h" 11 | 12 | 13 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 14 | m.def("knnquery_cuda", &knnquery_cuda, "knnquery_cuda"); 15 | m.def("furthestsampling_cuda", &furthestsampling_cuda, "furthestsampling_cuda"); 16 | m.def("farthestsampling_cuda", &farthestsampling_cuda, "farthestsampling_cuda"); 17 | m.def("grouping_forward_cuda", &grouping_forward_cuda, "grouping_forward_cuda"); 18 | m.def("grouping_backward_cuda", &grouping_backward_cuda, "grouping_backward_cuda"); 19 | m.def("interpolation_forward_cuda", &interpolation_forward_cuda, "interpolation_forward_cuda"); 20 | m.def("interpolation_backward_cuda", &interpolation_backward_cuda, "interpolation_backward_cuda"); 21 | m.def("subtraction_forward_cuda", &subtraction_forward_cuda, "subtraction_forward_cuda"); 22 | m.def("subtraction_backward_cuda", &subtraction_backward_cuda, "subtraction_backward_cuda"); 23 | m.def("aggregation_forward_cuda", &aggregation_forward_cuda, "aggregation_forward_cuda"); 24 | m.def("aggregation_backward_cuda", &aggregation_backward_cuda, "aggregation_backward_cuda"); 25 | } 26 | -------------------------------------------------------------------------------- /submodules/lib/pointops/src/sampling/sampling_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | // #include 3 | #include 4 | #include 5 | #include "sampling_cuda_kernel.h" 6 | 7 | 8 | void furthestsampling_cuda(int b, int n, at::Tensor xyz_tensor, at::Tensor offset_tensor, at::Tensor new_offset_tensor, at::Tensor tmp_tensor, at::Tensor idx_tensor) 9 | { 10 | const float *xyz = xyz_tensor.data_ptr(); 11 | const int *offset = offset_tensor.data_ptr(); 12 | const int *new_offset = new_offset_tensor.data_ptr(); 13 | float *tmp = tmp_tensor.data_ptr(); 14 | int *idx = idx_tensor.data_ptr(); 15 | furthestsampling_cuda_launcher(b, n, xyz, offset, new_offset, tmp, idx); 16 | } 17 | -------------------------------------------------------------------------------- /submodules/lib/pointops/src/sampling/sampling_cuda_kernel.cu: -------------------------------------------------------------------------------- 1 | #include "../cuda_utils.h" 2 | #include "sampling_cuda_kernel.h" 3 | 4 | 5 | __device__ void __update(float *dists, int *dists_i, int idx1, int idx2) { 6 | const float v1 = dists[idx1], v2 = dists[idx2]; 7 | const int i1 = dists_i[idx1], i2 = dists_i[idx2]; 8 | dists[idx1] = max(v1, v2); 9 | dists_i[idx1] = v2 > v1 ? i2 : i1; 10 | } 11 | 12 | // input xyz: (n, 3), tmp: (b, n_max) 13 | // ouput idx (m) 14 | template 15 | __global__ void furthestsampling_cuda_kernel(const float *xyz, const int *offset, const int *new_offset, float *tmp, int *idx) 16 | { 17 | __shared__ float dists[block_size]; 18 | __shared__ int dists_i[block_size]; 19 | 20 | int bid = blockIdx.x; 21 | int start_n, end_n, start_m, end_m, old; 22 | if (bid == 0) { 23 | start_n = 0; 24 | end_n = offset[0]; 25 | start_m = 0; 26 | end_m = new_offset[0]; 27 | old = 0; 28 | } 29 | else { 30 | start_n = offset[bid - 1]; 31 | end_n = offset[bid]; 32 | start_m = new_offset[bid - 1]; 33 | end_m = new_offset[bid]; 34 | old = offset[bid - 1]; 35 | } 36 | 37 | const int stride = block_size; 38 | int tid = threadIdx.x; 39 | if (tid == 0) idx[start_m] = start_n; 40 | 41 | __syncthreads(); 42 | for (int j = start_m + 1; j < end_m; j++) 43 | { 44 | int besti = start_n; 45 | float best = -1; 46 | float x1 = xyz[old * 3 + 0]; 47 | float y1 = xyz[old * 3 + 1]; 48 | float z1 = xyz[old * 3 + 2]; 49 | for (int k = start_n + tid; k < end_n; k += stride) 50 | { 51 | float x2 = xyz[k * 3 + 0]; 52 | float y2 = xyz[k * 3 + 1]; 53 | float z2 = xyz[k * 3 + 2]; 54 | float d = (x2 - x1) * (x2 - x1) + (y2 - y1) * (y2 - y1) + (z2 - z1) * (z2 - z1); 55 | float d2 = min(d, tmp[k]); 56 | tmp[k] = d2; 57 | besti = d2 > best ? k : besti; 58 | best = d2 > best ? d2 : best; 59 | } 60 | dists[tid] = best; 61 | dists_i[tid] = besti; 62 | __syncthreads(); 63 | 64 | if (block_size >= 1024) { 65 | if (tid < 512) { 66 | __update(dists, dists_i, tid, tid + 512); 67 | } 68 | __syncthreads(); 69 | } 70 | if (block_size >= 512) { 71 | if (tid < 256) { 72 | __update(dists, dists_i, tid, tid + 256); 73 | } 74 | __syncthreads(); 75 | } 76 | if (block_size >= 256) { 77 | if (tid < 128) { 78 | __update(dists, dists_i, tid, tid + 128); 79 | } 80 | __syncthreads(); 81 | } 82 | if (block_size >= 128) { 83 | if (tid < 64) { 84 | __update(dists, dists_i, tid, tid + 64); 85 | } 86 | __syncthreads(); 87 | } 88 | if (block_size >= 64) { 89 | if (tid < 32) { 90 | __update(dists, dists_i, tid, tid + 32); 91 | } 92 | __syncthreads(); 93 | } 94 | if (block_size >= 32) { 95 | if (tid < 16) { 96 | __update(dists, dists_i, tid, tid + 16); 97 | } 98 | __syncthreads(); 99 | } 100 | if (block_size >= 16) { 101 | if (tid < 8) { 102 | __update(dists, dists_i, tid, tid + 8); 103 | } 104 | __syncthreads(); 105 | } 106 | if (block_size >= 8) { 107 | if (tid < 4) { 108 | __update(dists, dists_i, tid, tid + 4); 109 | } 110 | __syncthreads(); 111 | } 112 | if (block_size >= 4) { 113 | if (tid < 2) { 114 | __update(dists, dists_i, tid, tid + 2); 115 | } 116 | __syncthreads(); 117 | } 118 | if (block_size >= 2) { 119 | if (tid < 1) { 120 | __update(dists, dists_i, tid, tid + 1); 121 | } 122 | __syncthreads(); 123 | } 124 | 125 | old = dists_i[0]; 126 | if (tid == 0) 127 | idx[j] = old; 128 | } 129 | } 130 | 131 | void furthestsampling_cuda_launcher(int b, int n, const float *xyz, const int *offset, const int *new_offset, float *tmp, int *idx) 132 | { 133 | unsigned int n_threads = opt_n_threads(n); 134 | switch (n_threads) { 135 | case 1024: 136 | furthestsampling_cuda_kernel<1024><<>>(xyz, offset, new_offset, tmp, idx); 137 | break; 138 | case 512: 139 | furthestsampling_cuda_kernel<512><<>>(xyz, offset, new_offset, tmp, idx); 140 | break; 141 | case 256: 142 | furthestsampling_cuda_kernel<256><<>>(xyz, offset, new_offset, tmp, idx); 143 | break; 144 | case 128: 145 | furthestsampling_cuda_kernel<128><<>>(xyz, offset, new_offset, tmp, idx); 146 | break; 147 | case 64: 148 | furthestsampling_cuda_kernel<64><<>>(xyz, offset, new_offset, tmp, idx); 149 | break; 150 | case 32: 151 | furthestsampling_cuda_kernel<32><<>>(xyz, offset, new_offset, tmp, idx); 152 | break; 153 | case 16: 154 | furthestsampling_cuda_kernel<16><<>>(xyz, offset, new_offset, tmp, idx); 155 | break; 156 | case 8: 157 | furthestsampling_cuda_kernel<8><<>>(xyz, offset, new_offset, tmp, idx); 158 | break; 159 | case 4: 160 | furthestsampling_cuda_kernel<4><<>>(xyz, offset, new_offset, tmp, idx); 161 | break; 162 | case 2: 163 | furthestsampling_cuda_kernel<2><<>>(xyz, offset, new_offset, tmp, idx); 164 | break; 165 | case 1: 166 | furthestsampling_cuda_kernel<1><<>>(xyz, offset, new_offset, tmp, idx); 167 | break; 168 | default: 169 | furthestsampling_cuda_kernel<512><<>>(xyz, offset, new_offset, tmp, idx); 170 | } 171 | } 172 | -------------------------------------------------------------------------------- /submodules/lib/pointops/src/sampling/sampling_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _SAMPLING_CUDA_KERNEL 2 | #define _SAMPLING_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void furthestsampling_cuda(int b, int n, at::Tensor xyz_tensor, at::Tensor offset_tensor, at::Tensor new_offset_tensor, at::Tensor tmp_tensor, at::Tensor idx_tensor); 8 | 9 | #ifdef __cplusplus 10 | extern "C" { 11 | #endif 12 | 13 | void furthestsampling_cuda_launcher(int b, int n, const float *xyz, const int *offset, const int *new_offset, float *tmp, int *idx); 14 | 15 | #ifdef __cplusplus 16 | } 17 | #endif 18 | #endif 19 | -------------------------------------------------------------------------------- /submodules/lib/pointops/src/subtraction/subtraction_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | // #include 3 | #include 4 | #include 5 | #include "subtraction_cuda_kernel.h" 6 | 7 | 8 | void subtraction_forward_cuda(int n, int nsample, int c, at::Tensor input1_tensor, at::Tensor input2_tensor, at::Tensor idx_tensor, at::Tensor output_tensor) 9 | { 10 | const float *input1 = input1_tensor.data_ptr(); 11 | const float *input2 = input2_tensor.data_ptr(); 12 | const int *idx = idx_tensor.data_ptr(); 13 | float *output = output_tensor.data_ptr(); 14 | subtraction_forward_cuda_launcher(n, nsample, c, input1, input2, idx, output); 15 | } 16 | 17 | void subtraction_backward_cuda(int n, int nsample, int c, at::Tensor idx_tensor, at::Tensor grad_output_tensor, at::Tensor grad_input1_tensor, at::Tensor grad_input2_tensor) 18 | { 19 | const int *idx = idx_tensor.data_ptr(); 20 | const float *grad_output = grad_output_tensor.data_ptr(); 21 | float *grad_input1 = grad_input1_tensor.data_ptr(); 22 | float *grad_input2 = grad_input2_tensor.data_ptr(); 23 | subtraction_backward_cuda_launcher(n, nsample, c, idx, grad_output, grad_input1, grad_input2); 24 | } 25 | -------------------------------------------------------------------------------- /submodules/lib/pointops/src/subtraction/subtraction_cuda_kernel.cu: -------------------------------------------------------------------------------- 1 | #include "../cuda_utils.h" 2 | #include "subtraction_cuda_kernel.h" 3 | 4 | 5 | __global__ void subtraction_forward_cuda_kernel(int n, int nsample, int c, const float *input1, const float *input2, const int *idx, float *output) { 6 | // input: input1: (n, c), input2: (n, c), idx: (n, nsample), output: (n, nsample, c) 7 | int index = blockIdx.x * blockDim.x + threadIdx.x; 8 | if (index >= n * nsample * c) return; 9 | const int c_idx = index % c; 10 | const int nsample_idx = (index / c) % nsample; 11 | const int n_idx = index / nsample / c; 12 | const int idx_idx = n_idx * nsample + nsample_idx; 13 | const int input1_idx = n_idx * c + c_idx; 14 | const int input2_idx = idx[idx_idx] * c + c_idx; 15 | output[index] = input1[input1_idx] - input2[input2_idx]; 16 | } 17 | 18 | __global__ void subtraction_backward_cuda_kernel(int n, int nsample, int c, const int *idx, const float *grad_output, float *grad_input1, float *grad_input2) { 19 | // input: grad_output: (n, nsample, c), output: grad_input1: (n, c), grad_input2: (n, c) 20 | int index = blockIdx.x * blockDim.x + threadIdx.x; 21 | if (index >= n * nsample * c) return; 22 | const int c_idx = index % c; 23 | const int nsample_idx = (index / c) % nsample; 24 | const int n_idx = index / nsample / c; 25 | const int idx_idx = n_idx * nsample + nsample_idx; 26 | const int input1_idx = n_idx * c + c_idx; 27 | const int input2_idx = idx[idx_idx] * c + c_idx; 28 | atomicAdd(grad_input1 + input1_idx, grad_output[index]); 29 | atomicAdd(grad_input2 + input2_idx, -grad_output[index]); 30 | } 31 | 32 | void subtraction_forward_cuda_launcher(int n, int nsample, int c, const float *input1, const float *input2, const int *idx, float *output) { 33 | // input: input1: (n, c), input2: (n, c), idx: (n, nsample), output: (n, nsample, c) 34 | dim3 blocks(DIVUP(n * nsample * c, THREADS_PER_BLOCK)); 35 | dim3 threads(THREADS_PER_BLOCK); 36 | subtraction_forward_cuda_kernel<<>>(n, nsample, c, input1, input2, idx, output); 37 | } 38 | 39 | void subtraction_backward_cuda_launcher(int n, int nsample, int c, const int *idx, const float *grad_output, float *grad_input1, float *grad_input2) { 40 | // input: grad_output: (n, nsample, c), output: grad_input1: (n, c), grad_input2: (n, c) 41 | dim3 blocks(DIVUP(n * nsample * c, THREADS_PER_BLOCK)); 42 | dim3 threads(THREADS_PER_BLOCK); 43 | subtraction_backward_cuda_kernel<<>>(n, nsample, c, idx, grad_output, grad_input1, grad_input2); 44 | } 45 | -------------------------------------------------------------------------------- /submodules/lib/pointops/src/subtraction/subtraction_cuda_kernel.h: -------------------------------------------------------------------------------- 1 | #ifndef _SUBTRACTION_CUDA_KERNEL 2 | #define _SUBTRACTION_CUDA_KERNEL 3 | #include 4 | #include 5 | #include 6 | 7 | void subtraction_forward_cuda(int n, int nsample, int c, at::Tensor input1_tensor, at::Tensor input2_tensor, at::Tensor idx_tensor, at::Tensor output_tensor); 8 | void subtraction_backward_cuda(int n, int nsample, int c, at::Tensor idx_tensor, at::Tensor grad_output_tensor, at::Tensor grad_input1_tensor, at::Tensor grad_input2_tensor); 9 | 10 | #ifdef __cplusplus 11 | extern "C" { 12 | #endif 13 | 14 | void subtraction_forward_cuda_launcher(int n, int nsample, int c, const float *input1, const float *input2, const int *idx, float *output); 15 | void subtraction_backward_cuda_launcher(int n, int nsample, int c, const int *idx, const float *grad_output, float *grad_input1, float *grad_input2); 16 | 17 | #ifdef __cplusplus 18 | } 19 | #endif 20 | #endif 21 | -------------------------------------------------------------------------------- /utils/camera_utils.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | from scene.cameras import Camera 13 | import numpy as np 14 | import torch 15 | from utils.general_utils import PILtoTorch 16 | from utils.graphics_utils import fov2focal 17 | from pytorch3d.transforms import matrix_to_quaternion, quaternion_to_matrix 18 | 19 | WARNED = False 20 | def lerp(t, v0, v1, is_tensor): 21 | x = (1 - t) * v0 + t * v1 22 | if is_tensor: 23 | x = torch.from_numpy(x).cuda() 24 | return x 25 | 26 | def slerp(t, v0, v1, DOT_THRESHOLD=0.9995): 27 | ''' 28 | Spherical linear interpolation 29 | Args: 30 | t (float/np.ndarray): Float value between 0.0 and 1.0 31 | v0 (np.ndarray): Starting vector 32 | v1 (np.ndarray): Final vector 33 | DOT_THRESHOLD (float): Threshold for considering the two vectors as 34 | colineal. Not recommended to alter this. 35 | Returns: 36 | v2 (np.ndarray): Interpolation vector between v0 and v1 37 | ''' 38 | c = False 39 | if not isinstance(v0,np.ndarray): 40 | c = True 41 | v0 = v0.detach().cpu().numpy() 42 | if not isinstance(v1,np.ndarray): 43 | c = True 44 | v1 = v1.detach().cpu().numpy() 45 | # Copy the vectors to reuse them later 46 | v0_copy = np.copy(v0) 47 | v1_copy = np.copy(v1) 48 | # Normalize the vectors to get the directions and angles 49 | v0 = v0 / np.linalg.norm(v0) 50 | v1 = v1 / np.linalg.norm(v1) 51 | # Dot product with the normalized vectors (can't use np.dot in W) 52 | dot = np.sum(v0 * v1) 53 | # If absolute value of dot product is almost 1, vectors are ~colineal, so use lerp 54 | if np.abs(dot) > DOT_THRESHOLD: 55 | return lerp(t, v0_copy, v1_copy, c) 56 | # Calculate initial angle between v0 and v1 57 | theta_0 = np.arccos(dot) 58 | sin_theta_0 = np.sin(theta_0) 59 | # Angle at timestep t 60 | theta_t = theta_0 * t 61 | sin_theta_t = np.sin(theta_t) 62 | # Finish the slerp algorithm 63 | s0 = np.sin(theta_0 - theta_t) / sin_theta_0 64 | s1 = sin_theta_t / sin_theta_0 65 | v2 = s0 * v0_copy + s1 * v1_copy 66 | if c: 67 | res = torch.from_numpy(v2).to("cuda") 68 | else: 69 | res = v2 70 | return res 71 | 72 | def loadCam(args, id, cam_info, resolution_scale): 73 | orig_w, orig_h = cam_info.image.size 74 | 75 | if args.resolution in [1, 2, 4, 8]: 76 | resolution = round(orig_w/(resolution_scale * args.resolution)), round(orig_h/(resolution_scale * args.resolution)) 77 | else: # should be a type that converts to float 78 | if args.resolution == -1: 79 | if orig_w > 1600: 80 | global WARNED 81 | if not WARNED: 82 | print("[ INFO ] Encountered quite large input images (>1.6K pixels width), rescaling to 1.6K.\n " 83 | "If this is not desired, please explicitly specify '--resolution/-r' as 1") 84 | WARNED = True 85 | global_down = orig_w / 1600 86 | else: 87 | global_down = 1 88 | else: 89 | global_down = orig_w / args.resolution 90 | 91 | scale = float(global_down) * float(resolution_scale) 92 | resolution = (int(orig_w / scale), int(orig_h / scale)) 93 | 94 | resized_image_rgb = PILtoTorch(cam_info.image, resolution) 95 | 96 | gt_image = resized_image_rgb[:3, ...] 97 | loaded_mask = None 98 | 99 | if resized_image_rgb.shape[1] == 4: 100 | loaded_mask = resized_image_rgb[3:4, ...] 101 | 102 | fwd_flow = None 103 | fwd_flow_mask = None 104 | bwd_flow = None 105 | bwd_flow_mask = None 106 | if "fwd_flow" in cam_info._fields: 107 | fwd_flow = cam_info.fwd_flow 108 | fwd_flow_mask = cam_info.fwd_flow_mask 109 | bwd_flow = cam_info.bwd_flow 110 | bwd_flow_mask = cam_info.bwd_flow_mask 111 | return Camera(colmap_id=cam_info.uid, R=cam_info.R, T=cam_info.T, 112 | FoVx=cam_info.FovX, FoVy=cam_info.FovY, 113 | image=gt_image, gt_alpha_mask=loaded_mask, 114 | image_name=cam_info.image_name, uid=id, data_device=args.data_device, 115 | time=cam_info.time, fwd_flow=fwd_flow, bwd_flow=bwd_flow, fwd_flow_mask=fwd_flow_mask, bwd_flow_mask=bwd_flow_mask) 116 | 117 | def cameraList_from_camInfos(cam_infos, resolution_scale, args): 118 | camera_list = [] 119 | 120 | for id, c in enumerate(cam_infos): 121 | if isinstance(c, dict): 122 | cam = loadCam(args, id, c["cam"], resolution_scale) 123 | fwd_cam = None 124 | bwd_cam = None 125 | if c["fwd_cam"] is not None: 126 | fwd_cam = loadCam(args, id, c["fwd_cam"], resolution_scale) 127 | if c["bwd_cam"] is not None: 128 | bwd_cam = loadCam(args, id, c["bwd_cam"], resolution_scale) 129 | data = {"cam": cam, "fwd_cam": fwd_cam, "bwd_cam": bwd_cam} 130 | camera_list.append(data) 131 | else: 132 | camera_list.append(loadCam(args, id, c, resolution_scale)) 133 | 134 | return camera_list 135 | 136 | def camera_to_JSON(id, camera : Camera): 137 | Rt = np.zeros((4, 4)) 138 | Rt[:3, :3] = camera.R.transpose() 139 | Rt[:3, 3] = camera.T 140 | Rt[3, 3] = 1.0 141 | 142 | W2C = np.linalg.inv(Rt) 143 | pos = W2C[:3, 3] 144 | rot = W2C[:3, :3] 145 | serializable_array_2d = [x.tolist() for x in rot] 146 | camera_entry = { 147 | 'id' : id, 148 | 'img_name' : camera.image_name, 149 | 'width' : camera.width, 150 | 'height' : camera.height, 151 | 'position': pos.tolist(), 152 | 'rotation': serializable_array_2d, 153 | 'fy' : fov2focal(camera.FovY, camera.height), 154 | 'fx' : fov2focal(camera.FovX, camera.width) 155 | } 156 | return camera_entry 157 | 158 | def quat_mul(q1, q2): 159 | ''' 160 | all quaternion in (w, x, y, z) 161 | q1: (N, 4) 162 | q2: (N, 4) 163 | return q1q2 164 | ''' 165 | q1q2 = torch.zeros_like(q1) 166 | q1q2[..., 0] = q2[..., 0]*q1[..., 0] - q2[..., 1]*q1[..., 1] - q2[..., 2]*q1[..., 2] - q2[..., 3]*q1[..., 3] 167 | q1q2[..., 1] = q2[..., 1]*q1[..., 0] + q2[..., 0]*q1[..., 1] + q2[..., 3]*q1[..., 2] - q2[..., 2]*q1[..., 3] 168 | q1q2[..., 2] = q2[..., 2]*q1[..., 0] - q2[..., 3]*q1[..., 1] + q2[..., 0]*q1[..., 2] + q2[..., 1]*q1[..., 3] 169 | q1q2[..., 3] = q2[..., 3]*q1[..., 0] + q2[..., 2]*q1[..., 1] - q2[..., 1]*q1[..., 2] + q2[..., 0]*q1[..., 3] 170 | return q1q2 171 | 172 | def quat_to_axis_theta(q): 173 | theta = 2 * torch.arccos(q[0]) 174 | v = 1 / torch.sqrt(1 - q[0]**2) 175 | v = v * q[1:] 176 | return v, theta 177 | 178 | def axis_theta_to_quat(axis, theta): 179 | w = torch.cos(theta / 2).view([1]) 180 | xyz = axis * torch.sqrt(1 - w**2) 181 | 182 | return torch.cat([w, xyz.view([-1])], dim=0) 183 | 184 | def slerp_tensor(q0, q1, t): 185 | dot = (q0 * q1).sum(dim=-1) 186 | theta = torch.acos(dot) 187 | 188 | sin_theta = torch.sin(theta) 189 | s0 = torch.sin((1 - t) * theta) / (sin_theta + 1e-8) 190 | s1 = torch.sin(t * theta) / (sin_theta + 1e-8) 191 | 192 | q = s0[..., None] * q0 + s1[..., None] * q1 193 | return q 194 | 195 | def transpose(R, t, X): 196 | """ 197 | Pytorch batch version of computing transform of the 3D points 198 | :param R: rotation matrix in dimension of (N, 3, 3) or (3, 3) 199 | :param t: translation vector could be (N, 3, 1) or (3, 1) 200 | :param X: points with 3D position, a 2D array with dimension of (N, num_points, 3) or (num_points, 3) 201 | :return: transformed 3D points 202 | """ 203 | keep_dim_n = False 204 | keep_dim_hw = False 205 | if R.dim() == 2: 206 | keep_dim_n = True 207 | R = R.unsqueeze(0) 208 | t = t.unsqueeze(0) 209 | if X.dim() == 2: 210 | X = X.unsqueeze(0) 211 | 212 | if X.dim() == 4: 213 | assert X.size(3) == 3 214 | keep_dim_hw = True 215 | N, H, W = X.shape[:3] 216 | X = X.view(N, H*W, 3) 217 | 218 | N = R.shape[0] 219 | M = X.shape[1] 220 | X_after_R = torch.bmm(R, torch.transpose(X, 1, 2)) 221 | X_after_R = torch.transpose(X_after_R, 1, 2) 222 | trans_X = X_after_R + t.view(N, 1, 3).expand(N, M, 3) 223 | 224 | if keep_dim_hw: 225 | trans_X = trans_X.view(N, H, W, 3) 226 | if keep_dim_n: 227 | trans_X = trans_X.squeeze(0) 228 | 229 | return trans_X 230 | 231 | def pi(K, X): 232 | """ 233 | Projecting the X in camera coordinates to the image plane 234 | :param K: camera intrinsic matrix tensor (N, 3, 3) or (3, 3) 235 | :param X: point position in 3D camera coordinates system, is a 3D array with dimension of (N, num_points, 3), or (num_points, 3) 236 | :return: N projected 2D pixel position u (N, num_points, 2) and the depth X (N, num_points, 1) 237 | """ 238 | keep_dim_n = False 239 | keep_dim_hw = False 240 | if K.dim() == 2: 241 | keep_dim_n = True 242 | K = K.unsqueeze(0) # make dim (1, 3, 3) 243 | if X.dim() == 2: 244 | X = X.unsqueeze(0) # make dim (1, num_points, 3) 245 | if X.dim() == 4: 246 | assert X.size(3) == 3 247 | keep_dim_hw = True 248 | N, H, W = X.shape[:3] 249 | X = X.view(N, H*W, 3) 250 | 251 | assert K.size(0) == X.size(0) 252 | N = K.shape[0] 253 | 254 | fx, fy, cx, cy = K[:, 0:1, 0:1], K[:, 1:2, 1:2], K[:, 0:1, 2:3], K[:, 1:2, 2:3] 255 | u_x = fx * X[:, :, 0:1] / (X[:, :, 2:3] + 0.0001) + cx 256 | u_y = fy * X[:, :, 1:2] / (X[:, :, 2:3] + 0.0001) + cy 257 | u = torch.cat([u_x, u_y], dim=-1) 258 | d = X[:, :, 2:3] 259 | 260 | if keep_dim_hw: 261 | u = u.view(N, H, W, 2) 262 | d = d.view(N, H, W) 263 | if keep_dim_n: 264 | u = u.squeeze(0) 265 | d = d.squeeze(0) 266 | 267 | return u, d 268 | 269 | def interpolation_pose(view, previous_view, interpolation_ratio): 270 | t, pre_t, R, pre_R = view.T, previous_view.T, view.R, previous_view.R 271 | new_t = pre_t + (t - pre_t) * interpolation_ratio 272 | quat = matrix_to_quaternion(torch.from_numpy(R)) 273 | pre_quat = matrix_to_quaternion(torch.from_numpy(pre_R)) 274 | new_quat = slerp(interpolation_ratio, pre_quat, quat) 275 | new_R = quaternion_to_matrix(new_quat[None]).squeeze().cpu().numpy() 276 | return new_t, new_R -------------------------------------------------------------------------------- /utils/general_utils.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | import torch 13 | import sys 14 | from datetime import datetime 15 | import numpy as np 16 | import random 17 | 18 | def inverse_sigmoid(x): 19 | return torch.log(x/(1-x)) 20 | 21 | def PILtoTorch(pil_image, resolution): 22 | resized_image_PIL = pil_image 23 | resized_image = torch.from_numpy(np.array(resized_image_PIL)) / 255.0 24 | if len(resized_image.shape) == 3: 25 | return resized_image.permute(2, 0, 1) 26 | else: 27 | return resized_image.unsqueeze(dim=-1).permute(2, 0, 1) 28 | 29 | def get_expon_lr_func( 30 | lr_init, lr_final, lr_delay_steps=0, lr_delay_mult=1.0, max_steps=1000000 31 | ): 32 | """ 33 | Copied from Plenoxels 34 | 35 | Continuous learning rate decay function. Adapted from JaxNeRF 36 | The returned rate is lr_init when step=0 and lr_final when step=max_steps, and 37 | is log-linearly interpolated elsewhere (equivalent to exponential decay). 38 | If lr_delay_steps>0 then the learning rate will be scaled by some smooth 39 | function of lr_delay_mult, such that the initial learning rate is 40 | lr_init*lr_delay_mult at the beginning of optimization but will be eased back 41 | to the normal learning rate when steps>lr_delay_steps. 42 | :param conf: config subtree 'lr' or similar 43 | :param max_steps: int, the number of steps during optimization. 44 | :return HoF which takes step as input 45 | """ 46 | 47 | def helper(step): 48 | if step < 0 or (lr_init == 0.0 and lr_final == 0.0): 49 | # Disable this parameter 50 | return 0.0 51 | if lr_delay_steps > 0: 52 | # A kind of reverse cosine decay. 53 | delay_rate = lr_delay_mult + (1 - lr_delay_mult) * np.sin( 54 | 0.5 * np.pi * np.clip(step / lr_delay_steps, 0, 1) 55 | ) 56 | else: 57 | delay_rate = 1.0 58 | t = np.clip(step / max_steps, 0, 1) 59 | log_lerp = np.exp(np.log(lr_init) * (1 - t) + np.log(lr_final) * t) 60 | return delay_rate * log_lerp 61 | 62 | return helper 63 | 64 | def strip_lowerdiag(L): 65 | uncertainty = torch.zeros((L.shape[0], 6), dtype=torch.float, device="cuda") 66 | 67 | uncertainty[:, 0] = L[:, 0, 0] 68 | uncertainty[:, 1] = L[:, 0, 1] 69 | uncertainty[:, 2] = L[:, 0, 2] 70 | uncertainty[:, 3] = L[:, 1, 1] 71 | uncertainty[:, 4] = L[:, 1, 2] 72 | uncertainty[:, 5] = L[:, 2, 2] 73 | return uncertainty 74 | 75 | def strip_symmetric(sym): 76 | return strip_lowerdiag(sym) 77 | 78 | def build_rotation(r): 79 | norm = torch.sqrt(r[:,0]*r[:,0] + r[:,1]*r[:,1] + r[:,2]*r[:,2] + r[:,3]*r[:,3]) 80 | 81 | q = r / norm[:, None] 82 | 83 | R = torch.zeros((q.size(0), 3, 3), device='cuda') 84 | 85 | r = q[:, 0] 86 | x = q[:, 1] 87 | y = q[:, 2] 88 | z = q[:, 3] 89 | 90 | R[:, 0, 0] = 1 - 2 * (y*y + z*z) 91 | R[:, 0, 1] = 2 * (x*y - r*z) 92 | R[:, 0, 2] = 2 * (x*z + r*y) 93 | R[:, 1, 0] = 2 * (x*y + r*z) 94 | R[:, 1, 1] = 1 - 2 * (x*x + z*z) 95 | R[:, 1, 2] = 2 * (y*z - r*x) 96 | R[:, 2, 0] = 2 * (x*z - r*y) 97 | R[:, 2, 1] = 2 * (y*z + r*x) 98 | R[:, 2, 2] = 1 - 2 * (x*x + y*y) 99 | return R 100 | 101 | def build_scaling_rotation(s, r): 102 | L = torch.zeros((s.shape[0], 3, 3), dtype=torch.float, device="cuda") 103 | R = build_rotation(r) 104 | 105 | L[:,0,0] = s[:,0] 106 | L[:,1,1] = s[:,1] 107 | L[:,2,2] = s[:,2] 108 | 109 | L = R @ L 110 | return L 111 | 112 | def safe_state(silent): 113 | old_f = sys.stdout 114 | class F: 115 | def __init__(self, silent): 116 | self.silent = silent 117 | 118 | def write(self, x): 119 | if not self.silent: 120 | if x.endswith("\n"): 121 | old_f.write(x.replace("\n", " [{}]\n".format(str(datetime.now().strftime("%d/%m %H:%M:%S"))))) 122 | else: 123 | old_f.write(x) 124 | 125 | def flush(self): 126 | old_f.flush() 127 | 128 | sys.stdout = F(silent) 129 | 130 | random.seed(0) 131 | np.random.seed(0) 132 | torch.manual_seed(0) 133 | torch.cuda.set_device(torch.device("cuda:0")) 134 | -------------------------------------------------------------------------------- /utils/graphics_utils.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | import torch 13 | import math 14 | import numpy as np 15 | from typing import NamedTuple 16 | 17 | class BasicPointCloud(NamedTuple): 18 | points : np.array 19 | colors : np.array 20 | normals : np.array 21 | 22 | def geom_transform_points(points, transf_matrix): 23 | P, _ = points.shape 24 | ones = torch.ones(P, 1, dtype=points.dtype, device=points.device) 25 | points_hom = torch.cat([points, ones], dim=1) 26 | points_out = torch.matmul(points_hom, transf_matrix.unsqueeze(0)) 27 | 28 | denom = points_out[..., 3:] + 0.0000001 29 | return (points_out[..., :3] / denom).squeeze(dim=0) 30 | 31 | def getWorld2View(R, t): 32 | Rt = np.zeros((4, 4)) 33 | Rt[:3, :3] = R.transpose() 34 | Rt[:3, 3] = t 35 | Rt[3, 3] = 1.0 36 | return np.float32(Rt) 37 | 38 | def getWorld2View2(R, t, translate=np.array([.0, .0, .0]), scale=1.0): 39 | Rt = np.zeros((4, 4)) 40 | Rt[:3, :3] = R.transpose() 41 | Rt[:3, 3] = t 42 | Rt[3, 3] = 1.0 43 | 44 | C2W = np.linalg.inv(Rt) 45 | cam_center = C2W[:3, 3] 46 | cam_center = (cam_center + translate) * scale 47 | C2W[:3, 3] = cam_center 48 | Rt = np.linalg.inv(C2W) 49 | return np.float32(Rt) 50 | 51 | def getProjectionMatrix(znear, zfar, fovX, fovY): 52 | tanHalfFovY = math.tan((fovY / 2)) 53 | tanHalfFovX = math.tan((fovX / 2)) 54 | 55 | top = tanHalfFovY * znear 56 | bottom = -top 57 | right = tanHalfFovX * znear 58 | left = -right 59 | 60 | P = torch.zeros(4, 4) 61 | 62 | z_sign = 1.0 63 | 64 | P[0, 0] = 2.0 * znear / (right - left) 65 | P[1, 1] = 2.0 * znear / (top - bottom) 66 | P[0, 2] = (right + left) / (right - left) 67 | P[1, 2] = (top + bottom) / (top - bottom) 68 | P[3, 2] = z_sign 69 | P[2, 2] = z_sign * zfar / (zfar - znear) 70 | P[2, 3] = -(zfar * znear) / (zfar - znear) 71 | return P 72 | 73 | def fov2focal(fov, pixels): 74 | return pixels / (2 * math.tan(fov / 2)) 75 | 76 | def focal2fov(focal, pixels): 77 | return 2*math.atan(pixels/(2*focal)) -------------------------------------------------------------------------------- /utils/image_utils.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | import torch 13 | import numpy as np 14 | 15 | def mse(img1, img2): 16 | return (((img1 - img2)) ** 2).view(img1.shape[0], -1).mean(1, keepdim=True) 17 | 18 | def psnr(img1, img2): 19 | mse = (((img1 - img2)) ** 2).view(img1.shape[0], -1).mean(1, keepdim=True) 20 | return 20 * torch.log10(1.0 / torch.sqrt(mse)) 21 | 22 | def psnr_diff(img1, img2): 23 | # image size (B, 3, H, W) 24 | mse = (((img1 - img2)) ** 2).mean(1, keepdim=True).view([1, img1.shape[-2], img1.shape[-1]]) 25 | return mse.cpu().numpy() 26 | -------------------------------------------------------------------------------- /utils/loss_utils.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | import torch 13 | import torch.nn.functional as F 14 | from torch.autograd import Variable 15 | from math import exp 16 | 17 | def motion_loss(gt_xyz, pred_xyz): 18 | loss = 0. 19 | for i in range(gt_xyz.shape[1]): 20 | gt = gt_xyz[:, i, ...] # (16, 100, 3) 21 | pred = pred_xyz[:, i, ...] 22 | # loss += l2_loss(pred, gt) 23 | loss += l1_loss(pred, gt) / (i + 5) 24 | # loss = l1_loss(pred, gt) 25 | return loss / gt_xyz.shape[1] 26 | 27 | def iso_loss(ut, canonical_ut, k=20, lamda=2000): 28 | if len(ut.shape) == 2: 29 | N = ut.shape[0] 30 | ut = ut.unsqueeze(0) 31 | else: 32 | N = ut.shape[1] 33 | B = ut.shape[0] 34 | # knn = KNN(k = k + 1, transpose_mode=True) 35 | 36 | dists, uj_idx, _ = knn_points(ut.unsqueeze(0), ut.unsqueeze(0), K=k+1) 37 | # _, uj_idx, _, _ = frnn.frnn_grid_points(ut, ut, K=k+1, r=0.08) 38 | uj = [] 39 | for b in range(B): 40 | uj += [ut[b, uj_idx[b, ..., 1:].contiguous().view([-1])]] # [N*k, 3] 41 | uj = torch.stack(uj) 42 | first_ut = canonical_ut 43 | _, first_uj_idx, _ = knn_points(first_ut.unsqueeze(0), first_ut.unsqueeze(0), K=k+1) 44 | # _, first_uj_idx, _, _ = frnn.frnn_grid_points(first_ut.unsqueeze(0), first_ut.unsqueeze(0), K=k+1, r=0.08) 45 | first_uj = first_ut[first_uj_idx[..., 1:].contiguous().view([-1])] 46 | 47 | # weights = torch.exp(-lamda * torch.norm(first_uj - first_ut.repeat(k, 1))**2) 48 | '''l_iso = (torch.abs(torch.norm(first_uj - first_ut.repeat(k, 1)) 49 | - torch.norm(uj - ut.repeat(k, 1))) * weights) / (k * N)''' 50 | l_iso = (torch.abs(torch.norm(first_uj - first_ut.repeat(k, 1)) - torch.norm(uj - ut.repeat(1, k, 1)))) / (k * N) 51 | 52 | return l_iso 53 | 54 | def l1_loss(network_output, gt): 55 | return torch.abs((network_output - gt)).mean() 56 | 57 | def l2_loss(network_output, gt): 58 | return ((network_output - gt) ** 2).mean() 59 | 60 | def gaussian(window_size, sigma): 61 | gauss = torch.Tensor([exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) 62 | return gauss / gauss.sum() 63 | 64 | def create_window(window_size, channel): 65 | _1D_window = gaussian(window_size, 1.5).unsqueeze(1) 66 | _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) 67 | window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) 68 | return window 69 | 70 | def ssim(img1, img2, window_size=11, size_average=True): 71 | channel = img1.size(-3) 72 | window = create_window(window_size, channel) 73 | 74 | if img1.is_cuda: 75 | window = window.cuda(img1.get_device()) 76 | window = window.type_as(img1) 77 | 78 | return _ssim(img1, img2, window, window_size, channel, size_average) 79 | 80 | def _ssim(img1, img2, window, window_size, channel, size_average=True): 81 | mu1 = F.conv2d(img1, window, padding=window_size // 2, groups=channel) 82 | mu2 = F.conv2d(img2, window, padding=window_size // 2, groups=channel) 83 | 84 | mu1_sq = mu1.pow(2) 85 | mu2_sq = mu2.pow(2) 86 | mu1_mu2 = mu1 * mu2 87 | 88 | sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size // 2, groups=channel) - mu1_sq 89 | sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size // 2, groups=channel) - mu2_sq 90 | sigma12 = F.conv2d(img1 * img2, window, padding=window_size // 2, groups=channel) - mu1_mu2 91 | 92 | C1 = 0.01 ** 2 93 | C2 = 0.03 ** 2 94 | 95 | ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) 96 | 97 | if size_average: 98 | return ssim_map.mean() 99 | else: 100 | return ssim_map.mean(1).mean(1).mean(1) -------------------------------------------------------------------------------- /utils/prepare/downsample_points.py: -------------------------------------------------------------------------------- 1 | import open3d as o3d 2 | import sys 3 | def process_ply_file(input_file, output_file): 4 | # 读取输入的ply文件 5 | pcd = o3d.io.read_point_cloud(input_file) 6 | print(f"Total points: {len(pcd.points)}") 7 | 8 | # 通过点云下采样将输入的点云减少 9 | voxel_size=0.02 10 | while len(pcd.points) > 40000: 11 | pcd = pcd.voxel_down_sample(voxel_size=voxel_size) 12 | print(f"Downsampled points: {len(pcd.points)}") 13 | voxel_size+=0.01 14 | 15 | # 将结果保存到输入的路径中 16 | o3d.io.write_point_cloud(output_file, pcd) 17 | 18 | # 使用函数 19 | process_ply_file(sys.argv[1], sys.argv[2]) -------------------------------------------------------------------------------- /utils/prepare/hypernerf2colmap.py: -------------------------------------------------------------------------------- 1 | 2 | import os 3 | import numpy as np 4 | from ipdb import set_trace 5 | import sys 6 | import json 7 | from PIL import Image 8 | from tqdm import tqdm 9 | import shutil 10 | def rotmat2qvec(R): 11 | Rxx, Ryx, Rzx, Rxy, Ryy, Rzy, Rxz, Ryz, Rzz = R.flat 12 | K = np.array([ 13 | [Rxx - Ryy - Rzz, 0, 0, 0], 14 | [Ryx + Rxy, Ryy - Rxx - Rzz, 0, 0], 15 | [Rzx + Rxz, Rzy + Ryz, Rzz - Rxx - Ryy, 0], 16 | [Ryz - Rzy, Rzx - Rxz, Rxy - Ryx, Rxx + Ryy + Rzz]]) / 3.0 17 | eigvals, eigvecs = np.linalg.eigh(K) 18 | qvec = eigvecs[[3, 0, 1, 2], np.argmax(eigvals)] 19 | if qvec[0] < 0: 20 | qvec *= -1 21 | return qvec 22 | 23 | root_dir = sys.argv[1] 24 | colmap_dir = os.path.join(root_dir,"sparse_") 25 | if not os.path.exists(colmap_dir): 26 | os.makedirs(colmap_dir) 27 | imagecolmap_dir = os.path.join(root_dir,"image_colmap") 28 | if not os.path.exists(imagecolmap_dir): 29 | os.makedirs(imagecolmap_dir) 30 | 31 | with open(f'{root_dir}/dataset.json', 'r') as f: 32 | dataset_json = json.load(f) 33 | 34 | all_img = dataset_json['ids'] 35 | val_id = dataset_json['val_ids'] 36 | if len(val_id) == 0: 37 | i_train = np.array([i for i in np.arange(len(all_img)) if 38 | (i%4 == 0)]) 39 | else: 40 | train_id = dataset_json['train_ids'] 41 | i_train = [] 42 | for i in range(len(all_img)): 43 | id = all_img[i] 44 | if id in train_id: 45 | i_train.append(i) 46 | 47 | # TODO read train img 48 | image_dir = os.path.join(root_dir,"rgb","2x") 49 | # images = os.listdir(image_dir) 50 | # images.sort() 51 | images = [f'{i}.png' for i in all_img] 52 | camera_dir = os.path.join(root_dir, "camera") 53 | # cameras = os.listdir(camera_dir) 54 | # cameras.sort() 55 | cameras = [f'{i}.json' for i in all_img] 56 | cameras = cameras[:399] 57 | images = images[:399] 58 | cams = [] 59 | for jsonfile in tqdm(cameras): 60 | with open (os.path.join(camera_dir,jsonfile)) as f: 61 | cams.append(json.load(f)) 62 | image_size = cams[0]['image_size'] 63 | image = Image.open(os.path.join(image_dir,images[0])) 64 | size = image.size 65 | # breakpoint() 66 | object_images_file = open(os.path.join(colmap_dir,"images.txt"),"w") 67 | object_cameras_file = open(os.path.join(colmap_dir,"cameras.txt"),"w") 68 | 69 | idx=0 70 | cnt=0 71 | sizes=1 72 | # while len(cams)//sizes > 200: 73 | # sizes += 1 74 | cameras = cameras[:100] 75 | images = images[:100] 76 | for cam, image in zip(cams, images): 77 | cnt+=1 78 | if cnt % sizes != 0: 79 | continue 80 | R = np.array(cam['orientation']).T 81 | T = -np.array(cam['position']) @ R 82 | 83 | T = [str(i) for i in T] 84 | qevc = [str(i) for i in rotmat2qvec(R.T)] 85 | print(idx+1," ".join(qevc)," ".join(T),1,image,"\n",file=object_images_file) 86 | 87 | print(idx,"SIMPLE_PINHOLE",image_size[0]/2,image_size[1]/2,cam['focal_length']/2,cam['principal_point'][0]/2,cam['principal_point'][1]/2,file=object_cameras_file) 88 | idx+=1 89 | shutil.copy(os.path.join(image_dir,image),os.path.join(imagecolmap_dir,image)) 90 | print(idx) 91 | # print(1,"SIMPLE_PINHOLE",image_size[0],image_size[1],focal[0],image_sizep0/2,image_size[1]/2,file=object_cameras_file) 92 | object_point_file = open(os.path.join(colmap_dir,"points3D.txt"),"w") 93 | 94 | object_cameras_file.close() 95 | object_images_file.close() 96 | object_point_file.close() 97 | -------------------------------------------------------------------------------- /utils/prepare/makeVideo.py: -------------------------------------------------------------------------------- 1 | #this file will split the pac-nerf datasets into per-frame format 2 | import os 3 | import imageio.v2 as imageio 4 | import numpy as np 5 | from tqdm import tqdm 6 | import cv2 7 | from scipy.ndimage import binary_dilation 8 | 9 | def make_video(path, image_num, output_name="video.mp4", fps=60, concat_delta=False, resize=1, bg_color="white"): 10 | render_rgbs = [] 11 | for i in tqdm(range(image_num)): 12 | if i in [497, 498, 499]: 13 | continue 14 | render_rgb = imageio.imread(os.path.join(path, f"{i:05d}.png")) 15 | if resize > 1: 16 | new_width = render_rgb.shape[1] // resize 17 | new_height = render_rgb.shape[0] // resize 18 | render_rgb = cv2.resize(render_rgb, (new_width, new_height)) 19 | if concat_delta: 20 | delta_img = imageio.imread(os.path.join(path.replace("/renders", "/deltas"), f"{i:05d}.jpg")) 21 | render_rgb = np.concatenate([render_rgb, delta_img], axis=1) 22 | if render_rgb.sum(axis=-1).max() == 255 * 3 and bg_color == "black": 23 | mask = (render_rgb.sum(axis=-1) >= 700) 24 | mask = binary_dilation(mask, structure=np.ones((7, 7))).astype(mask.dtype) 25 | render_rgb[mask, 0] = 0. 26 | render_rgb[mask, 1] = 0. 27 | render_rgb[mask, 2] = 0. 28 | render_rgbs.append(render_rgb) 29 | render_rgbs = np.stack(render_rgbs, axis=0) 30 | imageio.mimwrite(os.path.join(path, output_name), render_rgbs, fps=fps) 31 | 32 | -------------------------------------------------------------------------------- /utils/sh_utils.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The PlenOctree Authors. 2 | # Redistribution and use in source and binary forms, with or without 3 | # modification, are permitted provided that the following conditions are met: 4 | # 5 | # 1. Redistributions of source code must retain the above copyright notice, 6 | # this list of conditions and the following disclaimer. 7 | # 8 | # 2. Redistributions in binary form must reproduce the above copyright notice, 9 | # this list of conditions and the following disclaimer in the documentation 10 | # and/or other materials provided with the distribution. 11 | # 12 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 13 | # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 14 | # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE 15 | # ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE 16 | # LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 17 | # CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF 18 | # SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 19 | # INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN 20 | # CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) 21 | # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 22 | # POSSIBILITY OF SUCH DAMAGE. 23 | 24 | import torch 25 | 26 | C0 = 0.28209479177387814 27 | C1 = 0.4886025119029199 28 | C2 = [ 29 | 1.0925484305920792, 30 | -1.0925484305920792, 31 | 0.31539156525252005, 32 | -1.0925484305920792, 33 | 0.5462742152960396 34 | ] 35 | C3 = [ 36 | -0.5900435899266435, 37 | 2.890611442640554, 38 | -0.4570457994644658, 39 | 0.3731763325901154, 40 | -0.4570457994644658, 41 | 1.445305721320277, 42 | -0.5900435899266435 43 | ] 44 | C4 = [ 45 | 2.5033429417967046, 46 | -1.7701307697799304, 47 | 0.9461746957575601, 48 | -0.6690465435572892, 49 | 0.10578554691520431, 50 | -0.6690465435572892, 51 | 0.47308734787878004, 52 | -1.7701307697799304, 53 | 0.6258357354491761, 54 | ] 55 | 56 | 57 | def eval_sh(deg, sh, dirs): 58 | """ 59 | Evaluate spherical harmonics at unit directions 60 | using hardcoded SH polynomials. 61 | Works with torch/np/jnp. 62 | ... Can be 0 or more batch dimensions. 63 | Args: 64 | deg: int SH deg. Currently, 0-3 supported 65 | sh: jnp.ndarray SH coeffs [..., C, (deg + 1) ** 2] 66 | dirs: jnp.ndarray unit directions [..., 3] 67 | Returns: 68 | [..., C] 69 | """ 70 | assert deg <= 4 and deg >= 0 71 | coeff = (deg + 1) ** 2 72 | assert sh.shape[-1] >= coeff 73 | 74 | result = C0 * sh[..., 0] 75 | if deg > 0: 76 | x, y, z = dirs[..., 0:1], dirs[..., 1:2], dirs[..., 2:3] 77 | result = (result - 78 | C1 * y * sh[..., 1] + 79 | C1 * z * sh[..., 2] - 80 | C1 * x * sh[..., 3]) 81 | 82 | if deg > 1: 83 | xx, yy, zz = x * x, y * y, z * z 84 | xy, yz, xz = x * y, y * z, x * z 85 | result = (result + 86 | C2[0] * xy * sh[..., 4] + 87 | C2[1] * yz * sh[..., 5] + 88 | C2[2] * (2.0 * zz - xx - yy) * sh[..., 6] + 89 | C2[3] * xz * sh[..., 7] + 90 | C2[4] * (xx - yy) * sh[..., 8]) 91 | 92 | if deg > 2: 93 | result = (result + 94 | C3[0] * y * (3 * xx - yy) * sh[..., 9] + 95 | C3[1] * xy * z * sh[..., 10] + 96 | C3[2] * y * (4 * zz - xx - yy)* sh[..., 11] + 97 | C3[3] * z * (2 * zz - 3 * xx - 3 * yy) * sh[..., 12] + 98 | C3[4] * x * (4 * zz - xx - yy) * sh[..., 13] + 99 | C3[5] * z * (xx - yy) * sh[..., 14] + 100 | C3[6] * x * (xx - 3 * yy) * sh[..., 15]) 101 | 102 | if deg > 3: 103 | result = (result + C4[0] * xy * (xx - yy) * sh[..., 16] + 104 | C4[1] * yz * (3 * xx - yy) * sh[..., 17] + 105 | C4[2] * xy * (7 * zz - 1) * sh[..., 18] + 106 | C4[3] * yz * (7 * zz - 3) * sh[..., 19] + 107 | C4[4] * (zz * (35 * zz - 30) + 3) * sh[..., 20] + 108 | C4[5] * xz * (7 * zz - 3) * sh[..., 21] + 109 | C4[6] * (xx - yy) * (7 * zz - 1) * sh[..., 22] + 110 | C4[7] * xz * (xx - 3 * yy) * sh[..., 23] + 111 | C4[8] * (xx * (xx - 3 * yy) - yy * (3 * xx - yy)) * sh[..., 24]) 112 | return result 113 | 114 | def RGB2SH(rgb): 115 | return (rgb - 0.5) / C0 116 | 117 | def SH2RGB(sh): 118 | return sh * C0 + 0.5 -------------------------------------------------------------------------------- /utils/system_utils.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2023, Inria 3 | # GRAPHDECO research group, https://team.inria.fr/graphdeco 4 | # All rights reserved. 5 | # 6 | # This software is free for non-commercial, research and evaluation use 7 | # under the terms of the LICENSE.md file. 8 | # 9 | # For inquiries contact george.drettakis@inria.fr 10 | # 11 | 12 | from errno import EEXIST 13 | from os import makedirs, path 14 | import os 15 | 16 | def mkdir_p(folder_path): 17 | # Creates a directory. equivalent to using mkdir -p on the command line 18 | try: 19 | makedirs(folder_path) 20 | except OSError as exc: # Python >2.5 21 | if exc.errno == EEXIST and path.isdir(folder_path): 22 | pass 23 | else: 24 | raise 25 | 26 | def searchForMaxIteration(folder): 27 | saved_iters = [int(fname.split("_")[-1]) for fname in os.listdir(folder)] 28 | return max(saved_iters) 29 | -------------------------------------------------------------------------------- /utils/visualizer_utils.py: -------------------------------------------------------------------------------- 1 | # import open3d as o3d 2 | import numpy as np 3 | import torch 4 | # import open3d 5 | import matplotlib.pyplot as plt 6 | from sklearn.decomposition import PCA 7 | from sklearn.cluster import KMeans, MeanShift 8 | import trimesh 9 | from kmeans_pytorch import kmeans 10 | from torch_scatter import scatter 11 | import open3d as o3d 12 | 13 | class Visualizer: 14 | def __init__(self, extri, intri): 15 | if torch.is_tensor(extri): 16 | extri = extri.cpu().numpy() 17 | if torch.is_tensor(intri): 18 | intri = intri.cpu().numpy() 19 | self.w2c = extri 20 | self.k = intri 21 | 22 | def draw_trajectory(self, trajectory): 23 | if torch.is_tensor(trajectory): 24 | trajectory = trajectory.cpu().numpy() 25 | if hasattr(self, 'pcd'): 26 | pass 27 | else: 28 | self.pcd = o3d.geometry.PointCloud() 29 | self.pcd.points = o3d.utility.Vector3dVector(trajectory) 30 | self.pcd.transform(self.w2c) 31 | 32 | colors = np.random.uniform(0, 1, size=(trajectory.shape[0], 3)) 33 | self.pcd.colors = o3d.utility.Vector3dVector(colors) 34 | 35 | def cluster(features, method="KMeans", k=4): 36 | if method == "KMeans": 37 | kmeans = KMeans(n_clusters=k) 38 | kmeans.fit(features) 39 | labels = kmeans.labels_ 40 | else: 41 | meanShift = MeanShift() 42 | meanShift.fit(features) 43 | labels = meanShift.labels_ 44 | return labels 45 | 46 | def draw_motion_feature(points, features, output_path="vis_feature.txt", down_dim=1): 47 | pca = PCA(n_components=down_dim) 48 | features_reduced = pca.fit_transform(features) 49 | print(features_reduced.shape) 50 | cmap = plt.cm.get_cmap('jet') 51 | # features_color = cmap(features_reduced[:, 0] / np.max(features_reduced[:, 0])) 52 | features_color = cmap(features_reduced[:, 0]) 53 | points_vis = np.concatenate([points, features_color], axis=-1) 54 | np.savetxt(output_path, points_vis) 55 | 56 | 57 | def PCA_vis(xyzs, features, output_path="default.ply", dim=3, return_only=False, finish_return=False): 58 | xyzs = xyzs.cpu().numpy() 59 | pca = PCA(dim, random_state=42) 60 | f_samples = features.cpu().numpy() if dim == 3 else np.concatenate([xyzs, features.cpu().numpy()], axis=-1) 61 | transformed = pca.fit_transform(f_samples) 62 | feature_pca_mean = f_samples.mean(0) 63 | feature_pca_components = pca.components_ 64 | q1, q99 = np.percentile(transformed, [1, 99]) 65 | feature_pca_postprocess_sub = q1 66 | feature_pca_postprocess_div = (q99 - q1) 67 | 68 | vis_feature = (f_samples - feature_pca_mean[None, :]) @ feature_pca_components.T 69 | vis_feature = (vis_feature - feature_pca_postprocess_sub) / feature_pca_postprocess_div 70 | if dim > 3: 71 | xyzs = vis_feature[..., :2] 72 | vis_feature = vis_feature[..., 2:] 73 | z = np.ones_like(xyzs) 74 | xyzs = np.concatenate([xyzs, z[..., 0:1]],axis=-1) 75 | vis_feature = vis_feature.clip(0.0, 1.0).astype(np.float32) 76 | 77 | if return_only: 78 | return xyzs, vis_feature 79 | pcd = trimesh.PointCloud(xyzs, vis_feature) 80 | pcd.export(output_path) 81 | if finish_return: 82 | return xyzs, vis_feature 83 | 84 | def feature_kmeans(xyzs, features, K=10, vis=False, return_idx=False): 85 | cluster_ids_x, cluster_centers = kmeans(X=features, num_clusters=K, device=features.device) 86 | xyzs_means = scatter(xyzs, cluster_ids_x.to(xyzs.device), dim=0, reduce="mean") 87 | 88 | if vis: 89 | PCA_vis(xyzs_means, cluster_centers, "keypoints.ply") 90 | # PCA_vis(self._xyz, cluster_centers, "keypoints.ply") 91 | if return_idx: 92 | return xyzs_means, cluster_centers.to(xyzs.device), cluster_ids_x 93 | return xyzs_means, cluster_centers.to(xyzs.device) 94 | 95 | def draw_weights(weights_xyz, pcd, super_gaussians, idx): 96 | weights_xyz = weights_xyz[:, idx].cpu().numpy() 97 | pcd_ = pcd.cpu().numpy() 98 | super_gaussians = super_gaussians.cpu().numpy()[..., :3][idx].reshape([1, 3]).repeat(3, axis=0) 99 | color = plt.cm.jet(weights_xyz[None]).reshape([pcd_.shape[0], -1])[..., :3] 100 | pcd_ = np.concatenate([pcd_, color], axis=-1) 101 | np.savetxt(f"visualize/weights_xyz_{idx}.txt", pcd_) 102 | np.savetxt(f"visualize/kpts{idx}.txt", super_gaussians) 103 | 104 | def draw_trajectory(pcd, start_xyz, final_xyz): 105 | pcd = pcd.cpu().numpy() 106 | start_xyz = start_xyz.cpu().numpy() 107 | final_xyz = final_xyz.cpu().numpy() 108 | point_cloud = o3d.geometry.PointCloud() 109 | num_points = pcd.shape[1] 110 | num_frames = pcd.shape[0] 111 | trajectory_colors = np.random.uniform(0, 1, (num_points, 3)) 112 | black_color = np.array([0., 0., 0.]).reshape([1, 3]).repeat(start_xyz.shape[0], axis=0) 113 | shape_sizes = np.array([0.5]).reshape([1]).repeat(2 * start_xyz.shape[0], axis=0) 114 | trajectory_sizes = np.array([1.5]).reshape([1]).repeat(num_points * num_frames, axis=0) 115 | points = [] 116 | colors = [] 117 | 118 | for point_index in range(num_points): 119 | for frame in range(num_frames): 120 | # 在每个帧上生成点的位置 121 | x = pcd[frame, point_index, 0] 122 | y = pcd[frame, point_index, 1] 123 | z = pcd[frame, point_index, 2] 124 | 125 | points.append([x, y, z]) 126 | colors.append(trajectory_colors[point_index]) 127 | 128 | # 创建可视化窗口并添加点云 129 | points = np.array(points) 130 | colors = np.array(colors) 131 | points = np.concatenate([points, start_xyz, final_xyz], axis=0) 132 | colors = np.concatenate([colors, black_color, black_color], axis=0) 133 | sizes = np.concatenate([trajectory_sizes, shape_sizes], axis=0) 134 | point_cloud.points = o3d.utility.Vector3dVector(points) 135 | point_cloud.colors = o3d.utility.Vector3dVector(colors) 136 | o3d.visualization.draw_geometries([point_cloud]) 137 | --------------------------------------------------------------------------------