├── .gitignore ├── .gitmodules ├── README.md ├── Trim2DGS ├── LICENSE.md ├── arguments │ └── __init__.py ├── assets │ ├── teaser.jpg │ └── unbounded.gif ├── convert.py ├── cull_pcd.py ├── gaussian_renderer │ ├── __init__.py │ └── network_gui.py ├── lpipsPyTorch │ ├── __init__.py │ └── modules │ │ ├── lpips.py │ │ ├── networks.py │ │ └── utils.py ├── metrics.py ├── print_results.py ├── render.py ├── scene │ ├── __init__.py │ ├── cameras.py │ ├── colmap_loader.py │ ├── dataset_readers.py │ └── gaussian_model.py ├── scripts │ ├── dtu_eval.py │ ├── eval_TNT.py │ ├── eval_dtu │ │ ├── eval.py │ │ ├── evaluate_single_scene.py │ │ └── render_utils.py │ ├── eval_dtu_pcd │ │ ├── eval.py │ │ ├── evaluate_single_scene.py │ │ └── render_utils.py │ ├── eval_tnt │ │ ├── README.md │ │ ├── compute_bbox_for_mesh.py │ │ ├── config.py │ │ ├── cull_mesh.py │ │ ├── evaluate_single_scene.py │ │ ├── evaluation.py │ │ ├── help_func.py │ │ ├── plot.py │ │ ├── registration.py │ │ ├── requirements.txt │ │ ├── run.py │ │ ├── trajectory_io.py │ │ └── util.py │ ├── m360_eval.py │ ├── run_Mipnerf360.py │ ├── run_TNT.py │ ├── run_dtu.py │ └── train_from_scratch │ │ ├── run_Mipnerf360.py │ │ └── run_dtu.py ├── train.py ├── train_TrimGS.py ├── tune.py └── utils │ ├── camera_utils.py │ ├── general_utils.py │ ├── graphics_utils.py │ ├── image_utils.py │ ├── loss_utils.py │ ├── mcube_utils.py │ ├── mesh_utils.py │ ├── point_utils.py │ ├── render_utils.py │ ├── sh_utils.py │ └── system_utils.py ├── Trim3DGS ├── LICENSE.md ├── arguments │ └── __init__.py ├── assets │ └── logo_mpi.svg ├── convert.py ├── eval_dtu │ ├── eval.py │ ├── evaluate_single_scene.py │ └── render_utils.py ├── eval_dtu_pcd │ ├── eval.py │ ├── evaluate_single_scene.py │ └── render_utils.py ├── extract_mesh_tsdf.py ├── full_eval.py ├── gaussian_renderer │ ├── __init__.py │ └── network_gui.py ├── lpipsPyTorch │ ├── __init__.py │ └── modules │ │ ├── lpips.py │ │ ├── networks.py │ │ └── utils.py ├── metrics.py ├── print_results.py ├── render.py ├── scene │ ├── __init__.py │ ├── cameras.py │ ├── colmap_loader.py │ ├── dataset_readers.py │ └── gaussian_model.py ├── scripts │ ├── run_Mipnerf360.py │ └── run_dtu.py ├── train.py ├── tune.py └── utils │ ├── camera_utils.py │ ├── general_utils.py │ ├── graphics_utils.py │ ├── image_utils.py │ ├── loss_utils.py │ ├── sh_utils.py │ └── system_utils.py ├── assets └── teaser.jpg ├── blender ├── .gitignore ├── README.md ├── cam_pose_utils │ ├── cam_reader.py │ ├── colmap_loader.py │ └── graphic_utils.py ├── generate_video.py ├── render.py ├── render_cfgs │ ├── dtu │ │ ├── background.json │ │ ├── background.py │ │ ├── light.json │ │ └── light.py │ └── mip │ │ ├── background.json │ │ ├── background.py │ │ ├── light.json │ │ └── light.py └── render_utils │ ├── background_generator.py │ └── texture_allocator.py └── environment.yml /.gitignore: -------------------------------------------------------------------------------- 1 | *.pyc 2 | .vscode 3 | output 4 | build 5 | data 6 | *.ply 7 | **/PKG-INFO 8 | # submodules 9 | **__pycache__** 10 | *.png 11 | *.out 12 | eval 13 | *.npz 14 | **/tmp -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "Trim3DGS/SIBR_viewers"] 2 | path = Trim3DGS/SIBR_viewers 3 | url = https://gitlab.inria.fr/sibr/sibr_core.git 4 | [submodule "Trim3DGS/submodules/simple-knn"] 5 | path = Trim3DGS/submodules/simple-knn 6 | url = https://gitlab.inria.fr/bkerbl/simple-knn.git 7 | [submodule "Trim3DGS/submodules/diff-gaussian-rasterization"] 8 | path = Trim3DGS/submodules/diff-gaussian-rasterization 9 | url = https://github.com/Abyssaledge/diff-gaussian-rasterization.git 10 | [submodule "Trim2DGS/submodules/diff-surfel-rasterization"] 11 | path = Trim2DGS/submodules/diff-surfel-rasterization 12 | url = https://github.com/YuxueYang1204/diff-surfel-rasterization.git 13 | [submodule "Trim2DGS/submodules/simple-knn"] 14 | path = Trim2DGS/submodules/simple-knn 15 | url = https://gitlab.inria.fr/bkerbl/simple-knn.git 16 | -------------------------------------------------------------------------------- /Trim2DGS/LICENSE.md: -------------------------------------------------------------------------------- 1 | Gaussian-Splatting License 2 | =========================== 3 | 4 | **Inria** and **the Max Planck Institut for Informatik (MPII)** hold all the ownership rights on the *Software* named **gaussian-splatting**. 5 | The *Software* is in the process of being registered with the Agence pour la Protection des 6 | Programmes (APP). 7 | 8 | The *Software* is still being developed by the *Licensor*. 9 | 10 | *Licensor*'s goal is to allow the research community to use, test and evaluate 11 | the *Software*. 12 | 13 | ## 1. Definitions 14 | 15 | *Licensee* means any person or entity that uses the *Software* and distributes 16 | its *Work*. 17 | 18 | *Licensor* means the owners of the *Software*, i.e Inria and MPII 19 | 20 | *Software* means the original work of authorship made available under this 21 | License ie gaussian-splatting. 22 | 23 | *Work* means the *Software* and any additions to or derivative works of the 24 | *Software* that are made available under this License. 25 | 26 | 27 | ## 2. Purpose 28 | This license is intended to define the rights granted to the *Licensee* by 29 | Licensors under the *Software*. 30 | 31 | ## 3. Rights granted 32 | 33 | For the above reasons Licensors have decided to distribute the *Software*. 34 | Licensors grant non-exclusive rights to use the *Software* for research purposes 35 | to research users (both academic and industrial), free of charge, without right 36 | to sublicense.. The *Software* may be used "non-commercially", i.e., for research 37 | and/or evaluation purposes only. 38 | 39 | Subject to the terms and conditions of this License, you are granted a 40 | non-exclusive, royalty-free, license to reproduce, prepare derivative works of, 41 | publicly display, publicly perform and distribute its *Work* and any resulting 42 | derivative works in any form. 43 | 44 | ## 4. Limitations 45 | 46 | **4.1 Redistribution.** You may reproduce or distribute the *Work* only if (a) you do 47 | so under this License, (b) you include a complete copy of this License with 48 | your distribution, and (c) you retain without modification any copyright, 49 | patent, trademark, or attribution notices that are present in the *Work*. 50 | 51 | **4.2 Derivative Works.** You may specify that additional or different terms apply 52 | to the use, reproduction, and distribution of your derivative works of the *Work* 53 | ("Your Terms") only if (a) Your Terms provide that the use limitation in 54 | Section 2 applies to your derivative works, and (b) you identify the specific 55 | derivative works that are subject to Your Terms. Notwithstanding Your Terms, 56 | this License (including the redistribution requirements in Section 3.1) will 57 | continue to apply to the *Work* itself. 58 | 59 | **4.3** Any other use without of prior consent of Licensors is prohibited. Research 60 | users explicitly acknowledge having received from Licensors all information 61 | allowing to appreciate the adequacy between of the *Software* and their needs and 62 | to undertake all necessary precautions for its execution and use. 63 | 64 | **4.4** The *Software* is provided both as a compiled library file and as source 65 | code. In case of using the *Software* for a publication or other results obtained 66 | through the use of the *Software*, users are strongly encouraged to cite the 67 | corresponding publications as explained in the documentation of the *Software*. 68 | 69 | ## 5. Disclaimer 70 | 71 | THE USER CANNOT USE, EXPLOIT OR DISTRIBUTE THE *SOFTWARE* FOR COMMERCIAL PURPOSES 72 | WITHOUT PRIOR AND EXPLICIT CONSENT OF LICENSORS. YOU MUST CONTACT INRIA FOR ANY 73 | UNAUTHORIZED USE: stip-sophia.transfert@inria.fr . ANY SUCH ACTION WILL 74 | CONSTITUTE A FORGERY. THIS *SOFTWARE* IS PROVIDED "AS IS" WITHOUT ANY WARRANTIES 75 | OF ANY NATURE AND ANY EXPRESS OR IMPLIED WARRANTIES, WITH REGARDS TO COMMERCIAL 76 | USE, PROFESSIONNAL USE, LEGAL OR NOT, OR OTHER, OR COMMERCIALISATION OR 77 | ADAPTATION. UNLESS EXPLICITLY PROVIDED BY LAW, IN NO EVENT, SHALL INRIA OR THE 78 | AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 79 | CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE 80 | GOODS OR SERVICES, LOSS OF USE, DATA, OR PROFITS OR BUSINESS INTERRUPTION) 81 | HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT 82 | LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING FROM, OUT OF OR 83 | IN CONNECTION WITH THE *SOFTWARE* OR THE USE OR OTHER DEALINGS IN THE *SOFTWARE*. 84 | -------------------------------------------------------------------------------- /Trim2DGS/assets/teaser.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuxueYang1204/TrimGS/6da2f21271b79676fc9205b770d4d390b8739b89/Trim2DGS/assets/teaser.jpg -------------------------------------------------------------------------------- /Trim2DGS/assets/unbounded.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuxueYang1204/TrimGS/6da2f21271b79676fc9205b770d4d390b8739b89/Trim2DGS/assets/unbounded.gif -------------------------------------------------------------------------------- /Trim2DGS/cull_pcd.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | from scene import Scene 4 | import os 5 | from os import makedirs 6 | from gaussian_renderer import render 7 | import random 8 | from tqdm import tqdm 9 | from argparse import ArgumentParser 10 | from arguments import ModelParams, PipelineParams, get_combined_args 11 | from gaussian_renderer import GaussianModel 12 | import numpy as np 13 | import math 14 | 15 | def cull_pcd(dataset : ModelParams, iteration : int, pipeline : PipelineParams): 16 | with torch.no_grad(): 17 | gaussians = GaussianModel(dataset.sh_degree) 18 | scene = Scene(dataset, gaussians, load_iteration=iteration, shuffle=False) 19 | 20 | train_cameras = scene.getTrainCameras() 21 | 22 | bg_color = [1,1,1] if dataset.white_background else [0, 0, 0] 23 | background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") 24 | 25 | with torch.no_grad(): 26 | prune_mask = torch.zeros(gaussians.get_xyz.shape[0], dtype=torch.bool, device="cuda") 27 | for view in train_cameras: 28 | p_homo = gaussians.get_xyz @ view.full_proj_transform[:3] + view.full_proj_transform[3] 29 | p_ndc = p_homo[:, :2] / (p_homo[:, 3:4] + 1e-6) 30 | view_mask = (p_ndc[:, 0] >= -1) & (p_ndc[:, 0] <= 1) & (p_ndc[:, 1] >= -1) & (p_ndc[:, 1] <= 1) 31 | alpha_mask = F.grid_sample(view.gt_alpha_mask[None], p_ndc[view_mask][None, None], mode="nearest", align_corners=True).squeeze() < 0.5 32 | prune_mask[view_mask.nonzero()] |= alpha_mask.unsqueeze(-1) 33 | gaussians.prune_points_without_optimizer(prune_mask) 34 | output = os.path.join(dataset.model_path, "point_cloud/iteration_{}".format(iteration), "point_cloud_culled.ply") 35 | gaussians.save_ply(output) 36 | 37 | if __name__ == "__main__": 38 | # Set up command line argument parser 39 | parser = ArgumentParser(description="Testing script parameters") 40 | model = ModelParams(parser, sentinel=True) 41 | pipeline = PipelineParams(parser) 42 | parser.add_argument("--iteration", default=30000, type=int) 43 | parser.add_argument("--quiet", action="store_true") 44 | args = get_combined_args(parser) 45 | print("Rendering " + args.model_path) 46 | 47 | random.seed(0) 48 | np.random.seed(0) 49 | torch.manual_seed(0) 50 | torch.cuda.set_device(torch.device("cuda:0")) 51 | 52 | cull_pcd(model.extract(args), args.iteration, pipeline.extract(args)) -------------------------------------------------------------------------------- /Trim2DGS/gaussian_renderer/network_gui.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 traceback 14 | import socket 15 | import json 16 | from scene.cameras import MiniCam 17 | 18 | host = "127.0.0.1" 19 | port = 6009 20 | 21 | conn = None 22 | addr = None 23 | 24 | listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 25 | 26 | def init(wish_host, wish_port): 27 | global host, port, listener 28 | host = wish_host 29 | port = wish_port 30 | listener.bind((host, port)) 31 | listener.listen() 32 | listener.settimeout(0) 33 | 34 | def try_connect(): 35 | global conn, addr, listener 36 | try: 37 | conn, addr = listener.accept() 38 | print(f"\nConnected by {addr}") 39 | conn.settimeout(None) 40 | except Exception as inst: 41 | pass 42 | 43 | def read(): 44 | global conn 45 | messageLength = conn.recv(4) 46 | messageLength = int.from_bytes(messageLength, 'little') 47 | message = conn.recv(messageLength) 48 | return json.loads(message.decode("utf-8")) 49 | 50 | def send(message_bytes, verify): 51 | global conn 52 | if message_bytes != None: 53 | conn.sendall(message_bytes) 54 | conn.sendall(len(verify).to_bytes(4, 'little')) 55 | conn.sendall(bytes(verify, 'ascii')) 56 | 57 | def receive(): 58 | message = read() 59 | 60 | width = message["resolution_x"] 61 | height = message["resolution_y"] 62 | 63 | if width != 0 and height != 0: 64 | try: 65 | do_training = bool(message["train"]) 66 | fovy = message["fov_y"] 67 | fovx = message["fov_x"] 68 | znear = message["z_near"] 69 | zfar = message["z_far"] 70 | do_shs_python = bool(message["shs_python"]) 71 | do_rot_scale_python = bool(message["rot_scale_python"]) 72 | keep_alive = bool(message["keep_alive"]) 73 | scaling_modifier = message["scaling_modifier"] 74 | world_view_transform = torch.reshape(torch.tensor(message["view_matrix"]), (4, 4)).cuda() 75 | world_view_transform[:,1] = -world_view_transform[:,1] 76 | world_view_transform[:,2] = -world_view_transform[:,2] 77 | full_proj_transform = torch.reshape(torch.tensor(message["view_projection_matrix"]), (4, 4)).cuda() 78 | full_proj_transform[:,1] = -full_proj_transform[:,1] 79 | custom_cam = MiniCam(width, height, fovy, fovx, znear, zfar, world_view_transform, full_proj_transform) 80 | except Exception as e: 81 | print("") 82 | traceback.print_exc() 83 | raise e 84 | return custom_cam, do_training, do_shs_python, do_rot_scale_python, keep_alive, scaling_modifier 85 | else: 86 | return None, None, None, None, None, None -------------------------------------------------------------------------------- /Trim2DGS/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 | -------------------------------------------------------------------------------- /Trim2DGS/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 | -------------------------------------------------------------------------------- /Trim2DGS/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 | -------------------------------------------------------------------------------- /Trim2DGS/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 | -------------------------------------------------------------------------------- /Trim2DGS/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 torch 16 | import torchvision.transforms.functional as tf 17 | from utils.loss_utils import ssim 18 | from lpipsPyTorch import lpips 19 | import json 20 | from tqdm import tqdm 21 | from utils.image_utils import psnr 22 | from argparse import ArgumentParser 23 | 24 | def readImages(renders_dir, gt_dir): 25 | renders = [] 26 | gts = [] 27 | image_names = [] 28 | for fname in os.listdir(renders_dir): 29 | render = Image.open(renders_dir / fname) 30 | gt = Image.open(gt_dir / fname) 31 | renders.append(tf.to_tensor(render).unsqueeze(0)[:, :3, :, :].cuda()) 32 | gts.append(tf.to_tensor(gt).unsqueeze(0)[:, :3, :, :].cuda()) 33 | image_names.append(fname) 34 | return renders, gts, image_names 35 | 36 | def evaluate(model_paths): 37 | 38 | full_dict = {} 39 | per_view_dict = {} 40 | full_dict_polytopeonly = {} 41 | per_view_dict_polytopeonly = {} 42 | 43 | for scene_dir in model_paths: 44 | try: 45 | print("Scene:", scene_dir) 46 | full_dict[scene_dir] = {} 47 | per_view_dict[scene_dir] = {} 48 | full_dict_polytopeonly[scene_dir] = {} 49 | per_view_dict_polytopeonly[scene_dir] = {} 50 | 51 | test_dir = Path(scene_dir) / "test" 52 | 53 | for method in os.listdir(test_dir): 54 | print("Method:", method) 55 | 56 | full_dict[scene_dir][method] = {} 57 | per_view_dict[scene_dir][method] = {} 58 | full_dict_polytopeonly[scene_dir][method] = {} 59 | per_view_dict_polytopeonly[scene_dir][method] = {} 60 | 61 | method_dir = test_dir / method 62 | gt_dir = method_dir/ "gt" 63 | renders_dir = method_dir / "renders" 64 | renders, gts, image_names = readImages(renders_dir, gt_dir) 65 | 66 | ssims = [] 67 | psnrs = [] 68 | lpipss = [] 69 | 70 | for idx in tqdm(range(len(renders)), desc="Metric evaluation progress"): 71 | ssims.append(ssim(renders[idx], gts[idx])) 72 | psnrs.append(psnr(renders[idx], gts[idx])) 73 | lpipss.append(lpips(renders[idx], gts[idx], net_type='vgg')) 74 | 75 | print(" SSIM : {:>12.7f}".format(torch.tensor(ssims).mean(), ".5")) 76 | print(" PSNR : {:>12.7f}".format(torch.tensor(psnrs).mean(), ".5")) 77 | print(" LPIPS: {:>12.7f}".format(torch.tensor(lpipss).mean(), ".5")) 78 | print("") 79 | 80 | full_dict[scene_dir][method].update({"SSIM": torch.tensor(ssims).mean().item(), 81 | "PSNR": torch.tensor(psnrs).mean().item(), 82 | "LPIPS": torch.tensor(lpipss).mean().item()}) 83 | per_view_dict[scene_dir][method].update({"SSIM": {name: ssim for ssim, name in zip(torch.tensor(ssims).tolist(), image_names)}, 84 | "PSNR": {name: psnr for psnr, name in zip(torch.tensor(psnrs).tolist(), image_names)}, 85 | "LPIPS": {name: lp for lp, name in zip(torch.tensor(lpipss).tolist(), image_names)}}) 86 | 87 | with open(scene_dir + "/results.json", 'w') as fp: 88 | json.dump(full_dict[scene_dir], fp, indent=True) 89 | with open(scene_dir + "/per_view.json", 'w') as fp: 90 | json.dump(per_view_dict[scene_dir], fp, indent=True) 91 | except: 92 | print("Unable to compute metrics for model", scene_dir) 93 | 94 | if __name__ == "__main__": 95 | device = torch.device("cuda:0") 96 | torch.cuda.set_device(device) 97 | 98 | # Set up command line argument parser 99 | parser = ArgumentParser(description="Training script parameters") 100 | parser.add_argument('--model_paths', '-m', required=True, nargs="+", type=str, default=[]) 101 | args = parser.parse_args() 102 | evaluate(args.model_paths) 103 | -------------------------------------------------------------------------------- /Trim2DGS/print_results.py: -------------------------------------------------------------------------------- 1 | import json 2 | import os 3 | from argparse import ArgumentParser 4 | 5 | 6 | def report_dtu(path, iteration): 7 | print(f'Results of {path}') 8 | scans = os.listdir(path) 9 | 10 | print("***************** mesh *****************") 11 | sum_overall = 0 12 | n = 0 13 | for scan in sorted(scans, key=lambda x: int(x.replace('scan', ''))): 14 | p = os.path.join(path, scan, f'train/ours_{iteration}/results.json') 15 | if not os.path.exists(p): 16 | continue 17 | with open(p, 'r') as f: 18 | data = json.load(f) 19 | print(scan, data, scan) 20 | sum_overall += data['overall'] 21 | n += 1 22 | print(f"Overall: {sum_overall / n}") 23 | 24 | sum_overall = 0 25 | n = 0 26 | print("***************** pcd *****************") 27 | for scan in sorted(scans, key=lambda x: int(x.replace('scan', ''))): 28 | p = os.path.join(path, scan, f'train/ours_{iteration}/results_pcd.json') 29 | if not os.path.exists(p): 30 | continue 31 | with open(p, 'r') as f: 32 | data = json.load(f) 33 | print(scan, data, scan) 34 | sum_overall += data['overall'] 35 | n += 1 36 | print(f"Overall: {sum_overall / n}") 37 | 38 | def report_mipnerf360(path, iteration): 39 | print(f'Results of {path}') 40 | scans = os.listdir(path) 41 | sum_overall = 0 42 | n = 0 43 | for scan in sorted(scans): 44 | p = os.path.join(path, scan, f'point_cloud/iteration_{iteration}/point_cloud.ply') 45 | if not os.path.exists(p): 46 | print(f"Missing {p}") 47 | continue 48 | # check the storage size of the point cloud 49 | size = os.path.getsize(p) 50 | mb_size = size / 1024 / 1024 51 | print(scan, f"{mb_size:.2f} MB") 52 | sum_overall += mb_size 53 | n += 1 54 | print(f"Overall: {sum_overall / n:.2f} MB") 55 | 56 | indoor = ['room', 'counter', 'kitchen', 'bonsai'] 57 | outdoor = ['bicycle', 'flowers', 'garden', 'stump', 'treehill'] 58 | sum_overall_indoor = dict() 59 | sum_overall_outdoor = dict() 60 | n_indoor = 0 61 | n_outdoor = 0 62 | for scan in sorted(scans): 63 | p = os.path.join(path, scan, 'results.json') 64 | if not os.path.exists(p): 65 | continue 66 | with open(p, 'r') as f: 67 | data = json.load(f) 68 | print(scan, data, scan) 69 | if scan in indoor: 70 | for k, v in data[f'ours_{iteration}'].items(): 71 | if k not in sum_overall_indoor: 72 | sum_overall_indoor[k] = 0.0 73 | sum_overall_indoor[k] += v 74 | n_indoor += 1 75 | if scan in outdoor: 76 | for k, v in data[f'ours_{iteration}'].items(): 77 | if k not in sum_overall_outdoor: 78 | sum_overall_outdoor[k] = 0.0 79 | sum_overall_outdoor[k] += v 80 | n_outdoor += 1 81 | 82 | print("Outdoor") 83 | for k, v in sum_overall_outdoor.items(): 84 | print(f"{k}: {v / n_outdoor:.3f}") 85 | 86 | print("Indoor") 87 | for k, v in sum_overall_indoor.items(): 88 | print(f"{k}: {v / n_indoor:.3f}") 89 | 90 | def report_tnt(path, iteration): 91 | print(f'Results of {path}') 92 | scans = os.listdir(path) 93 | 94 | sum_overall = 0 95 | n = 0 96 | print("***************** mesh *****************") 97 | for scan in sorted(scans): 98 | p = os.path.join(path, scan, f'train/ours_{iteration}/evaluation/results.json') 99 | if not os.path.exists(p): 100 | continue 101 | with open(p, 'r') as f: 102 | data = json.load(f) 103 | print(scan, f"f-score: {data['f-score']:.2f}", scan) 104 | sum_overall += data['f-score'] 105 | n += 1 106 | 107 | print(f"Overall f-score: {sum_overall / n:.2f}") 108 | 109 | sum_overall = dict() 110 | n = 0 111 | print("***************** render *****************") 112 | for scan in sorted(scans): 113 | p = os.path.join(path, scan, "results.json") 114 | if not os.path.exists(p): 115 | continue 116 | with open(p, 'r') as f: 117 | data = json.load(f) 118 | print(scan, data, scan) 119 | for k, v in data[f'ours_{iteration}'].items(): 120 | if k not in sum_overall: 121 | sum_overall[k] = 0.0 122 | sum_overall[k] += v 123 | n += 1 124 | print("Mean") 125 | for k, v in sum_overall.items(): 126 | print(f"{k}: {v / n:.3f}") 127 | 128 | if __name__ == '__main__': 129 | parser = ArgumentParser() 130 | parser.add_argument('--output_path', '-o', type=str) 131 | parser.add_argument('--iteration', type=int, default=7000) 132 | parser.add_argument('--dataset', type=str, choices=['dtu', 'mipnerf360', 'tnt']) 133 | args = parser.parse_args() 134 | 135 | if args.dataset == 'dtu': 136 | report_dtu(args.output_path, args.iteration) 137 | elif args.dataset == 'mipnerf360': 138 | report_mipnerf360(args.output_path, args.iteration) 139 | elif args.dataset == 'tnt': 140 | report_tnt(args.output_path, args.iteration) 141 | else: 142 | raise ValueError(f"Unknown dataset {args.dataset}") -------------------------------------------------------------------------------- /Trim2DGS/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], pretrained_ply_path=None): 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 | self.pretrained_ply_path = pretrained_ply_path 33 | 34 | if load_iteration: 35 | if load_iteration == -1: 36 | self.loaded_iter = searchForMaxIteration(os.path.join(self.model_path, "point_cloud")) 37 | else: 38 | self.loaded_iter = load_iteration 39 | print("Loading trained model at iteration {}".format(self.loaded_iter)) 40 | 41 | self.train_cameras = {} 42 | self.test_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) 49 | else: 50 | assert False, "Could not recognize scene type!" 51 | 52 | if not self.loaded_iter: 53 | with open(scene_info.ply_path, 'rb') as src_file, open(os.path.join(self.model_path, "input.ply") , 'wb') as dest_file: 54 | dest_file.write(src_file.read()) 55 | json_cams = [] 56 | camlist = [] 57 | if scene_info.test_cameras: 58 | camlist.extend(scene_info.test_cameras) 59 | if scene_info.train_cameras: 60 | camlist.extend(scene_info.train_cameras) 61 | for id, cam in enumerate(camlist): 62 | json_cams.append(camera_to_JSON(id, cam)) 63 | with open(os.path.join(self.model_path, "cameras.json"), 'w') as file: 64 | json.dump(json_cams, file) 65 | 66 | if shuffle: 67 | random.shuffle(scene_info.train_cameras) # Multi-res consistent random shuffling 68 | random.shuffle(scene_info.test_cameras) # Multi-res consistent random shuffling 69 | 70 | self.cameras_extent = scene_info.nerf_normalization["radius"] 71 | 72 | for resolution_scale in resolution_scales: 73 | print("Loading Training Cameras") 74 | self.train_cameras[resolution_scale] = cameraList_from_camInfos(scene_info.train_cameras, resolution_scale, args) 75 | print("Loading Test Cameras") 76 | self.test_cameras[resolution_scale] = cameraList_from_camInfos(scene_info.test_cameras, resolution_scale, args) 77 | 78 | if self.loaded_iter: 79 | self.gaussians.load_ply(os.path.join(self.model_path, 80 | "point_cloud", 81 | "iteration_" + str(self.loaded_iter), 82 | "point_cloud.ply")) 83 | elif self.pretrained_ply_path: 84 | self.gaussians.load_ply(self.pretrained_ply_path, self.cameras_extent) 85 | else: 86 | self.gaussians.create_from_pcd(scene_info.point_cloud, self.cameras_extent) 87 | 88 | def save(self, iteration): 89 | point_cloud_path = os.path.join(self.model_path, "point_cloud/iteration_{}".format(iteration)) 90 | self.gaussians.save_ply(os.path.join(point_cloud_path, "point_cloud.ply")) 91 | 92 | def getTrainCameras(self, scale=1.0): 93 | return self.train_cameras[scale] 94 | 95 | def getTestCameras(self, scale=1.0): 96 | return self.test_cameras[scale] -------------------------------------------------------------------------------- /Trim2DGS/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" 21 | ): 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 | 32 | try: 33 | self.data_device = torch.device(data_device) 34 | except Exception as e: 35 | print(e) 36 | print(f"[Warning] Custom device {data_device} failed, fallback to default cuda device" ) 37 | self.data_device = torch.device("cuda") 38 | 39 | self.original_image = image.clamp(0.0, 1.0).to(self.data_device) 40 | self.image_width = self.original_image.shape[2] 41 | self.image_height = self.original_image.shape[1] 42 | 43 | if gt_alpha_mask is not None: 44 | # self.original_image *= gt_alpha_mask.to(self.data_device) 45 | self.gt_alpha_mask = gt_alpha_mask.to(self.data_device) 46 | else: 47 | self.original_image *= torch.ones((1, self.image_height, self.image_width), device=self.data_device) 48 | self.gt_alpha_mask = None 49 | 50 | self.zfar = 100.0 51 | self.znear = 0.01 52 | 53 | self.trans = trans 54 | self.scale = scale 55 | 56 | self.world_view_transform = torch.tensor(getWorld2View2(R, T, trans, scale)).transpose(0, 1).cuda() 57 | self.projection_matrix = getProjectionMatrix(znear=self.znear, zfar=self.zfar, fovX=self.FoVx, fovY=self.FoVy).transpose(0,1).cuda() 58 | self.full_proj_transform = (self.world_view_transform.unsqueeze(0).bmm(self.projection_matrix.unsqueeze(0))).squeeze(0) 59 | self.camera_center = self.world_view_transform.inverse()[3, :3] 60 | 61 | class MiniCam: 62 | def __init__(self, width, height, fovy, fovx, znear, zfar, world_view_transform, full_proj_transform): 63 | self.image_width = width 64 | self.image_height = height 65 | self.FoVy = fovy 66 | self.FoVx = fovx 67 | self.znear = znear 68 | self.zfar = zfar 69 | self.world_view_transform = world_view_transform 70 | self.full_proj_transform = full_proj_transform 71 | view_inv = torch.inverse(self.world_view_transform) 72 | self.camera_center = view_inv[3][:3] 73 | 74 | -------------------------------------------------------------------------------- /Trim2DGS/scripts/dtu_eval.py: -------------------------------------------------------------------------------- 1 | import os 2 | from argparse import ArgumentParser 3 | 4 | dtu_scenes = ['scan24', 'scan37', 'scan40', 'scan55', 'scan63', 'scan65', 'scan69', 'scan83', 'scan97', 'scan105', 'scan106', 'scan110', 'scan114', 'scan118', 'scan122'] 5 | 6 | parser = ArgumentParser(description="Full evaluation script parameters") 7 | parser.add_argument("--skip_training", action="store_true") 8 | parser.add_argument("--skip_rendering", action="store_true") 9 | parser.add_argument("--skip_metrics", action="store_true") 10 | parser.add_argument("--output_path", default="./eval/dtu") 11 | parser.add_argument('--dtu', "-dtu", required=True, type=str) 12 | args, _ = parser.parse_known_args() 13 | 14 | all_scenes = [] 15 | all_scenes.extend(dtu_scenes) 16 | 17 | if not args.skip_metrics: 18 | parser.add_argument('--DTU_Official', "-DTU", required=True, type=str) 19 | parser.add_argument('--eval_path', required=True, type=str) 20 | args = parser.parse_args() 21 | 22 | 23 | if not args.skip_training: 24 | common_args = " --quiet --test_iterations -1 --depth_ratio 1.0 -r 2 --lambda_dist 1000" 25 | for scene in dtu_scenes: 26 | source = args.dtu + "/" + scene 27 | print("python train.py -s " + source + " -m " + args.output_path + "/" + scene + common_args) 28 | os.system("python train.py -s " + source + " -m " + args.output_path + "/" + scene + common_args) 29 | 30 | 31 | if not args.skip_rendering: 32 | all_sources = [] 33 | common_args = " --quiet --skip_train --depth_ratio 1.0 --num_cluster 1 --voxel_size 0.004 --sdf_trunc 0.016 --depth_trunc 3.0" 34 | for scene in dtu_scenes: 35 | source = args.dtu + "/" + scene 36 | print("python render.py --iteration 30000 -s " + source + " -m" + args.output_path + "/" + scene + common_args) 37 | os.system("python render.py --iteration 30000 -s " + source + " -m" + args.output_path + "/" + scene + common_args) 38 | 39 | 40 | if not args.skip_metrics: 41 | script_dir = os.path.dirname(os.path.abspath(__file__)) 42 | for scene in dtu_scenes: 43 | scan_id = scene[4:] 44 | ply_file = f"{args.output_path}/{scene}/train/ours_30000/" 45 | iteration = 30000 46 | string = f"python {script_dir}/eval_dtu/evaluate_single_scene.py " + \ 47 | f"--input_mesh {args.output_path}/{scene}/train/ours_30000/fuse_post.ply " + \ 48 | f"--scan_id {scan_id} --output_dir {script_dir}/tmp/scan{scan_id} " + \ 49 | f"--mask_dir {args.dtu} " + \ 50 | f"--DTU {args.DTU_Official}" 51 | 52 | os.system(string) -------------------------------------------------------------------------------- /Trim2DGS/scripts/eval_TNT.py: -------------------------------------------------------------------------------- 1 | # evaluation script for TnT dataset 2 | 3 | from concurrent.futures import ProcessPoolExecutor 4 | import subprocess 5 | 6 | scenes = ['Barn', 'Caterpillar', 'Ignatius', 'Truck', 'Meetingroom', 'Courthouse'] 7 | 8 | excluded_gpus = set([]) 9 | 10 | output_dir = "output/TNT_Trim2DGS" 11 | iteration = 7000 12 | 13 | jobs = scenes 14 | 15 | def eval(scene): 16 | cmds = [ 17 | f"python scripts/eval_tnt/run.py --dataset-dir data/TNT_GOF/ground_truth/{scene} --traj-path data/TNT_GOF/TrainingSet/{scene}/{scene}_COLMAP_SfM.log --ply-path {output_dir}/{scene}/train/ours_{iteration}/fuse_post.ply", 18 | ] 19 | 20 | for cmd in cmds: 21 | print(cmd) 22 | subprocess.run(cmd, shell=True, check=True) 23 | return True 24 | 25 | def main(): 26 | with ProcessPoolExecutor(max_workers=len(jobs)) as executor: 27 | futures = [executor.submit(eval, scene) for scene in jobs] 28 | 29 | for future in futures: 30 | try: 31 | result = future.result() 32 | print(f"Finished job with result: {result}\n") 33 | except Exception as e: 34 | print(f"Failed job with exception: {e}\n") 35 | 36 | if __name__ == "__main__": 37 | main() -------------------------------------------------------------------------------- /Trim2DGS/scripts/eval_dtu/evaluate_single_scene.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | import cv2 5 | import numpy as np 6 | import os 7 | import glob 8 | from skimage.morphology import binary_dilation, disk 9 | import argparse 10 | 11 | import trimesh 12 | from pathlib import Path 13 | import subprocess 14 | 15 | import sys 16 | import render_utils as rend_util 17 | from tqdm import tqdm 18 | 19 | def cull_scan(scan, mesh_path, result_mesh_file, instance_dir): 20 | 21 | # load poses 22 | image_dir = '{0}/images'.format(instance_dir) 23 | image_paths = sorted(glob.glob(os.path.join(image_dir, "*.png"))) 24 | n_images = len(image_paths) 25 | cam_file = '{0}/cameras.npz'.format(instance_dir) 26 | camera_dict = np.load(cam_file) 27 | scale_mats = [camera_dict['scale_mat_%d' % idx].astype(np.float32) for idx in range(n_images)] 28 | world_mats = [camera_dict['world_mat_%d' % idx].astype(np.float32) for idx in range(n_images)] 29 | 30 | intrinsics_all = [] 31 | pose_all = [] 32 | for scale_mat, world_mat in zip(scale_mats, world_mats): 33 | P = world_mat @ scale_mat 34 | P = P[:3, :4] 35 | intrinsics, pose = rend_util.load_K_Rt_from_P(None, P) 36 | intrinsics_all.append(torch.from_numpy(intrinsics).float()) 37 | pose_all.append(torch.from_numpy(pose).float()) 38 | 39 | # load mask 40 | mask_dir = '{0}/mask'.format(instance_dir) 41 | mask_paths = sorted(glob.glob(os.path.join(mask_dir, "*.png"))) 42 | masks = [] 43 | for p in mask_paths: 44 | mask = cv2.imread(p) 45 | masks.append(mask) 46 | 47 | # hard-coded image shape 48 | W, H = 1600, 1200 49 | 50 | # load mesh 51 | mesh = trimesh.load(mesh_path) 52 | 53 | # load transformation matrix 54 | 55 | vertices = mesh.vertices 56 | 57 | # project and filter 58 | vertices = torch.from_numpy(vertices).cuda() 59 | vertices = torch.cat((vertices, torch.ones_like(vertices[:, :1])), dim=-1) 60 | vertices = vertices.permute(1, 0) 61 | vertices = vertices.float() 62 | 63 | sampled_masks = [] 64 | for i in tqdm(range(n_images), desc="Culling mesh given masks"): 65 | pose = pose_all[i] 66 | w2c = torch.inverse(pose).cuda() 67 | intrinsic = intrinsics_all[i].cuda() 68 | 69 | with torch.no_grad(): 70 | # transform and project 71 | cam_points = intrinsic @ w2c @ vertices 72 | pix_coords = cam_points[:2, :] / (cam_points[2, :].unsqueeze(0) + 1e-6) 73 | pix_coords = pix_coords.permute(1, 0) 74 | pix_coords[..., 0] /= W - 1 75 | pix_coords[..., 1] /= H - 1 76 | pix_coords = (pix_coords - 0.5) * 2 77 | valid = ((pix_coords > -1. ) & (pix_coords < 1.)).all(dim=-1).float() 78 | 79 | # dialate mask similar to unisurf 80 | maski = masks[i][:, :, 0].astype(np.float32) / 256. 81 | maski = torch.from_numpy(binary_dilation(maski, disk(24))).float()[None, None].cuda() 82 | 83 | sampled_mask = F.grid_sample(maski, pix_coords[None, None], mode='nearest', padding_mode='zeros', align_corners=True)[0, -1, 0] 84 | 85 | sampled_mask = sampled_mask + (1. - valid) 86 | sampled_masks.append(sampled_mask) 87 | 88 | sampled_masks = torch.stack(sampled_masks, -1) 89 | # filter 90 | 91 | mask = (sampled_masks > 0.).all(dim=-1).cpu().numpy() 92 | face_mask = mask[mesh.faces].all(axis=1) 93 | 94 | mesh.update_vertices(mask) 95 | mesh.update_faces(face_mask) 96 | 97 | # transform vertices to world 98 | scale_mat = scale_mats[0] 99 | mesh.vertices = mesh.vertices * scale_mat[0, 0] + scale_mat[:3, 3][None] 100 | mesh.export(result_mesh_file) 101 | del mesh 102 | 103 | 104 | if __name__ == "__main__": 105 | 106 | parser = argparse.ArgumentParser( 107 | description='Arguments to evaluate the mesh.' 108 | ) 109 | 110 | parser.add_argument('--input_mesh', type=str, help='path to the mesh to be evaluated') 111 | parser.add_argument('--scan_id', type=str, help='scan id of the input mesh') 112 | parser.add_argument('--output_dir', type=str, default='evaluation_results_single', help='path to the output folder') 113 | parser.add_argument('--mask_dir', type=str, default='mask', help='path to uncropped mask') 114 | parser.add_argument('--DTU', type=str, default='Offical_DTU_Dataset', help='path to the GT DTU point clouds') 115 | args = parser.parse_args() 116 | 117 | Offical_DTU_Dataset = args.DTU 118 | out_dir = args.output_dir 119 | Path(out_dir).mkdir(parents=True, exist_ok=True) 120 | 121 | scan = args.scan_id 122 | ply_file = args.input_mesh 123 | print("cull mesh ....") 124 | result_mesh_file = os.path.join(out_dir, "culled_mesh.ply") 125 | cull_scan(scan, ply_file, result_mesh_file, instance_dir=os.path.join(args.mask_dir, f'scan{args.scan_id}')) 126 | 127 | script_dir = os.path.dirname(os.path.abspath(__file__)) 128 | cmd = f"python {script_dir}/eval.py --data {result_mesh_file} --scan {scan} --mode mesh --dataset_dir {Offical_DTU_Dataset} --vis_out_dir {out_dir}" 129 | os.system(cmd) -------------------------------------------------------------------------------- /Trim2DGS/scripts/eval_dtu/render_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import imageio 3 | import skimage 4 | import cv2 5 | import torch 6 | from torch.nn import functional as F 7 | 8 | 9 | def get_psnr(img1, img2, normalize_rgb=False): 10 | if normalize_rgb: # [-1,1] --> [0,1] 11 | img1 = (img1 + 1.) / 2. 12 | img2 = (img2 + 1. ) / 2. 13 | 14 | mse = torch.mean((img1 - img2) ** 2) 15 | psnr = -10. * torch.log(mse) / torch.log(torch.Tensor([10.]).cuda()) 16 | 17 | return psnr 18 | 19 | 20 | def load_rgb(path, normalize_rgb = False): 21 | img = imageio.imread(path) 22 | img = skimage.img_as_float32(img) 23 | 24 | if normalize_rgb: # [-1,1] --> [0,1] 25 | img -= 0.5 26 | img *= 2. 27 | img = img.transpose(2, 0, 1) 28 | return img 29 | 30 | 31 | def load_K_Rt_from_P(filename, P=None): 32 | if P is None: 33 | lines = open(filename).read().splitlines() 34 | if len(lines) == 4: 35 | lines = lines[1:] 36 | lines = [[x[0], x[1], x[2], x[3]] for x in (x.split(" ") for x in lines)] 37 | P = np.asarray(lines).astype(np.float32).squeeze() 38 | 39 | out = cv2.decomposeProjectionMatrix(P) 40 | K = out[0] 41 | R = out[1] 42 | t = out[2] 43 | 44 | K = K/K[2,2] 45 | intrinsics = np.eye(4) 46 | intrinsics[:3, :3] = K 47 | 48 | pose = np.eye(4, dtype=np.float32) 49 | pose[:3, :3] = R.transpose() 50 | pose[:3,3] = (t[:3] / t[3])[:,0] 51 | 52 | return intrinsics, pose 53 | 54 | 55 | def get_camera_params(uv, pose, intrinsics): 56 | if pose.shape[1] == 7: #In case of quaternion vector representation 57 | cam_loc = pose[:, 4:] 58 | R = quat_to_rot(pose[:,:4]) 59 | p = torch.eye(4).repeat(pose.shape[0],1,1).cuda().float() 60 | p[:, :3, :3] = R 61 | p[:, :3, 3] = cam_loc 62 | else: # In case of pose matrix representation 63 | cam_loc = pose[:, :3, 3] 64 | p = pose 65 | 66 | batch_size, num_samples, _ = uv.shape 67 | 68 | depth = torch.ones((batch_size, num_samples)).cuda() 69 | x_cam = uv[:, :, 0].view(batch_size, -1) 70 | y_cam = uv[:, :, 1].view(batch_size, -1) 71 | z_cam = depth.view(batch_size, -1) 72 | 73 | pixel_points_cam = lift(x_cam, y_cam, z_cam, intrinsics=intrinsics) 74 | 75 | # permute for batch matrix product 76 | pixel_points_cam = pixel_points_cam.permute(0, 2, 1) 77 | 78 | world_coords = torch.bmm(p, pixel_points_cam).permute(0, 2, 1)[:, :, :3] 79 | ray_dirs = world_coords - cam_loc[:, None, :] 80 | ray_dirs = F.normalize(ray_dirs, dim=2) 81 | 82 | return ray_dirs, cam_loc 83 | 84 | 85 | def get_camera_for_plot(pose): 86 | if pose.shape[1] == 7: #In case of quaternion vector representation 87 | cam_loc = pose[:, 4:].detach() 88 | R = quat_to_rot(pose[:,:4].detach()) 89 | else: # In case of pose matrix representation 90 | cam_loc = pose[:, :3, 3] 91 | R = pose[:, :3, :3] 92 | cam_dir = R[:, :3, 2] 93 | return cam_loc, cam_dir 94 | 95 | 96 | def lift(x, y, z, intrinsics): 97 | # parse intrinsics 98 | intrinsics = intrinsics.cuda() 99 | fx = intrinsics[:, 0, 0] 100 | fy = intrinsics[:, 1, 1] 101 | cx = intrinsics[:, 0, 2] 102 | cy = intrinsics[:, 1, 2] 103 | sk = intrinsics[:, 0, 1] 104 | 105 | x_lift = (x - cx.unsqueeze(-1) + cy.unsqueeze(-1)*sk.unsqueeze(-1)/fy.unsqueeze(-1) - sk.unsqueeze(-1)*y/fy.unsqueeze(-1)) / fx.unsqueeze(-1) * z 106 | y_lift = (y - cy.unsqueeze(-1)) / fy.unsqueeze(-1) * z 107 | 108 | # homogeneous 109 | return torch.stack((x_lift, y_lift, z, torch.ones_like(z).cuda()), dim=-1) 110 | 111 | 112 | def quat_to_rot(q): 113 | batch_size, _ = q.shape 114 | q = F.normalize(q, dim=1) 115 | R = torch.ones((batch_size, 3,3)).cuda() 116 | qr=q[:,0] 117 | qi = q[:, 1] 118 | qj = q[:, 2] 119 | qk = q[:, 3] 120 | R[:, 0, 0]=1-2 * (qj**2 + qk**2) 121 | R[:, 0, 1] = 2 * (qj *qi -qk*qr) 122 | R[:, 0, 2] = 2 * (qi * qk + qr * qj) 123 | R[:, 1, 0] = 2 * (qj * qi + qk * qr) 124 | R[:, 1, 1] = 1-2 * (qi**2 + qk**2) 125 | R[:, 1, 2] = 2*(qj*qk - qi*qr) 126 | R[:, 2, 0] = 2 * (qk * qi-qj * qr) 127 | R[:, 2, 1] = 2 * (qj*qk + qi*qr) 128 | R[:, 2, 2] = 1-2 * (qi**2 + qj**2) 129 | return R 130 | 131 | 132 | def rot_to_quat(R): 133 | batch_size, _,_ = R.shape 134 | q = torch.ones((batch_size, 4)).cuda() 135 | 136 | R00 = R[:, 0,0] 137 | R01 = R[:, 0, 1] 138 | R02 = R[:, 0, 2] 139 | R10 = R[:, 1, 0] 140 | R11 = R[:, 1, 1] 141 | R12 = R[:, 1, 2] 142 | R20 = R[:, 2, 0] 143 | R21 = R[:, 2, 1] 144 | R22 = R[:, 2, 2] 145 | 146 | q[:,0]=torch.sqrt(1.0+R00+R11+R22)/2 147 | q[:, 1]=(R21-R12)/(4*q[:,0]) 148 | q[:, 2] = (R02 - R20) / (4 * q[:, 0]) 149 | q[:, 3] = (R10 - R01) / (4 * q[:, 0]) 150 | return q 151 | 152 | 153 | def get_sphere_intersections(cam_loc, ray_directions, r = 1.0): 154 | # Input: n_rays x 3 ; n_rays x 3 155 | # Output: n_rays x 1, n_rays x 1 (close and far) 156 | 157 | ray_cam_dot = torch.bmm(ray_directions.view(-1, 1, 3), 158 | cam_loc.view(-1, 3, 1)).squeeze(-1) 159 | under_sqrt = ray_cam_dot ** 2 - (cam_loc.norm(2, 1, keepdim=True) ** 2 - r ** 2) 160 | 161 | # sanity check 162 | if (under_sqrt <= 0).sum() > 0: 163 | print('BOUNDING SPHERE PROBLEM!') 164 | exit() 165 | 166 | sphere_intersections = torch.sqrt(under_sqrt) * torch.Tensor([-1, 1]).cuda().float() - ray_cam_dot 167 | sphere_intersections = sphere_intersections.clamp_min(0.0) 168 | 169 | return sphere_intersections -------------------------------------------------------------------------------- /Trim2DGS/scripts/eval_dtu_pcd/evaluate_single_scene.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | import cv2 5 | import numpy as np 6 | import os 7 | import glob 8 | from skimage.morphology import binary_dilation, disk 9 | import argparse 10 | 11 | import open3d as o3d 12 | from pathlib import Path 13 | import subprocess 14 | 15 | import sys 16 | import render_utils as rend_util 17 | from tqdm import tqdm 18 | 19 | def cull_scan(scan, ply_path, result_ply_file, instance_dir): 20 | 21 | # load poses 22 | image_dir = '{0}/images'.format(instance_dir) 23 | image_paths = sorted(glob.glob(os.path.join(image_dir, "*.png"))) 24 | n_images = len(image_paths) 25 | cam_file = '{0}/cameras.npz'.format(instance_dir) 26 | camera_dict = np.load(cam_file) 27 | scale_mats = [camera_dict['scale_mat_%d' % idx].astype(np.float32) for idx in range(n_images)] 28 | world_mats = [camera_dict['world_mat_%d' % idx].astype(np.float32) for idx in range(n_images)] 29 | 30 | intrinsics_all = [] 31 | pose_all = [] 32 | for scale_mat, world_mat in zip(scale_mats, world_mats): 33 | P = world_mat @ scale_mat 34 | P = P[:3, :4] 35 | intrinsics, pose = rend_util.load_K_Rt_from_P(None, P) 36 | intrinsics_all.append(torch.from_numpy(intrinsics).float()) 37 | pose_all.append(torch.from_numpy(pose).float()) 38 | 39 | # load mask 40 | mask_dir = '{0}/mask'.format(instance_dir) 41 | mask_paths = sorted(glob.glob(os.path.join(mask_dir, "*.png"))) 42 | masks = [] 43 | for p in mask_paths: 44 | mask = cv2.imread(p) 45 | masks.append(mask) 46 | 47 | # hard-coded image shape 48 | W, H = 1600, 1200 49 | 50 | # load pcd 51 | pcd = o3d.io.read_point_cloud(ply_path) 52 | # load transformation matrix 53 | 54 | vertices = np.asarray(pcd.points) 55 | 56 | # project and filter 57 | vertices = torch.from_numpy(vertices).cuda() 58 | vertices = torch.cat((vertices, torch.ones_like(vertices[:, :1])), dim=-1) 59 | vertices = vertices.permute(1, 0) 60 | vertices = vertices.float() 61 | 62 | sampled_masks = [] 63 | for i in tqdm(range(n_images), desc="Culling point clouds given masks"): 64 | pose = pose_all[i] 65 | w2c = torch.inverse(pose).cuda() 66 | intrinsic = intrinsics_all[i].cuda() 67 | 68 | with torch.no_grad(): 69 | # transform and project 70 | cam_points = intrinsic @ w2c @ vertices 71 | pix_coords = cam_points[:2, :] / (cam_points[2, :].unsqueeze(0) + 1e-6) 72 | pix_coords = pix_coords.permute(1, 0) 73 | pix_coords[..., 0] /= W - 1 74 | pix_coords[..., 1] /= H - 1 75 | pix_coords = (pix_coords - 0.5) * 2 76 | valid = ((pix_coords > -1. ) & (pix_coords < 1.)).all(dim=-1).float() 77 | 78 | # dialate mask similar to unisurf 79 | maski = masks[i][:, :, 0].astype(np.float32) / 256. 80 | maski = torch.from_numpy(binary_dilation(maski, disk(24))).float()[None, None].cuda() 81 | 82 | sampled_mask = F.grid_sample(maski, pix_coords[None, None], mode='nearest', padding_mode='zeros', align_corners=True)[0, -1, 0] 83 | 84 | sampled_mask = sampled_mask + (1. - valid) 85 | sampled_masks.append(sampled_mask) 86 | 87 | sampled_masks = torch.stack(sampled_masks, -1) 88 | # filter 89 | 90 | mask = (sampled_masks > 0.).all(dim=-1).cpu().numpy() 91 | 92 | points = np.asarray(pcd.points) 93 | points = points[mask] 94 | # transform vertices to world 95 | scale_mat = scale_mats[0] 96 | points = points * scale_mat[0, 0] + scale_mat[:3, 3][None] 97 | 98 | pcd = o3d.geometry.PointCloud() 99 | pcd.points = o3d.utility.Vector3dVector(points) 100 | o3d.io.write_point_cloud(result_ply_file, pcd) 101 | 102 | 103 | if __name__ == "__main__": 104 | 105 | parser = argparse.ArgumentParser( 106 | description='Arguments to evaluate the point clouds.' 107 | ) 108 | 109 | parser.add_argument('--input_pcd', type=str, help='path to the point clouds to be evaluated') 110 | parser.add_argument('--scan_id', type=str, help='scan id of the input point clouds') 111 | parser.add_argument('--output_dir', type=str, default='evaluation_results_single', help='path to the output folder') 112 | parser.add_argument('--mask_dir', type=str, default='mask', help='path to uncropped mask') 113 | parser.add_argument('--DTU', type=str, default='Offical_DTU_Dataset', help='path to the GT DTU point clouds') 114 | args = parser.parse_args() 115 | 116 | Offical_DTU_Dataset = args.DTU 117 | out_dir = args.output_dir 118 | Path(out_dir).mkdir(parents=True, exist_ok=True) 119 | 120 | scan = args.scan_id 121 | ply_file = args.input_pcd 122 | print("cull pcd ....") 123 | result_ply_file = os.path.join(out_dir, "culled_pcd.ply") 124 | cull_scan(scan, ply_file, result_ply_file, instance_dir=os.path.join(args.mask_dir, f'scan{args.scan_id}')) 125 | 126 | script_dir = os.path.dirname(os.path.abspath(__file__)) 127 | cmd = f"python {script_dir}/eval.py --data {result_ply_file} --scan {scan} --dataset_dir {Offical_DTU_Dataset} --vis_out_dir {out_dir}" 128 | os.system(cmd) -------------------------------------------------------------------------------- /Trim2DGS/scripts/eval_dtu_pcd/render_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import imageio 3 | import skimage 4 | import cv2 5 | import torch 6 | from torch.nn import functional as F 7 | 8 | 9 | def get_psnr(img1, img2, normalize_rgb=False): 10 | if normalize_rgb: # [-1,1] --> [0,1] 11 | img1 = (img1 + 1.) / 2. 12 | img2 = (img2 + 1. ) / 2. 13 | 14 | mse = torch.mean((img1 - img2) ** 2) 15 | psnr = -10. * torch.log(mse) / torch.log(torch.Tensor([10.]).cuda()) 16 | 17 | return psnr 18 | 19 | 20 | def load_rgb(path, normalize_rgb = False): 21 | img = imageio.imread(path) 22 | img = skimage.img_as_float32(img) 23 | 24 | if normalize_rgb: # [-1,1] --> [0,1] 25 | img -= 0.5 26 | img *= 2. 27 | img = img.transpose(2, 0, 1) 28 | return img 29 | 30 | 31 | def load_K_Rt_from_P(filename, P=None): 32 | if P is None: 33 | lines = open(filename).read().splitlines() 34 | if len(lines) == 4: 35 | lines = lines[1:] 36 | lines = [[x[0], x[1], x[2], x[3]] for x in (x.split(" ") for x in lines)] 37 | P = np.asarray(lines).astype(np.float32).squeeze() 38 | 39 | out = cv2.decomposeProjectionMatrix(P) 40 | K = out[0] 41 | R = out[1] 42 | t = out[2] 43 | 44 | K = K/K[2,2] 45 | intrinsics = np.eye(4) 46 | intrinsics[:3, :3] = K 47 | 48 | pose = np.eye(4, dtype=np.float32) 49 | pose[:3, :3] = R.transpose() 50 | pose[:3,3] = (t[:3] / t[3])[:,0] 51 | 52 | return intrinsics, pose 53 | 54 | 55 | def get_camera_params(uv, pose, intrinsics): 56 | if pose.shape[1] == 7: #In case of quaternion vector representation 57 | cam_loc = pose[:, 4:] 58 | R = quat_to_rot(pose[:,:4]) 59 | p = torch.eye(4).repeat(pose.shape[0],1,1).cuda().float() 60 | p[:, :3, :3] = R 61 | p[:, :3, 3] = cam_loc 62 | else: # In case of pose matrix representation 63 | cam_loc = pose[:, :3, 3] 64 | p = pose 65 | 66 | batch_size, num_samples, _ = uv.shape 67 | 68 | depth = torch.ones((batch_size, num_samples)).cuda() 69 | x_cam = uv[:, :, 0].view(batch_size, -1) 70 | y_cam = uv[:, :, 1].view(batch_size, -1) 71 | z_cam = depth.view(batch_size, -1) 72 | 73 | pixel_points_cam = lift(x_cam, y_cam, z_cam, intrinsics=intrinsics) 74 | 75 | # permute for batch matrix product 76 | pixel_points_cam = pixel_points_cam.permute(0, 2, 1) 77 | 78 | world_coords = torch.bmm(p, pixel_points_cam).permute(0, 2, 1)[:, :, :3] 79 | ray_dirs = world_coords - cam_loc[:, None, :] 80 | ray_dirs = F.normalize(ray_dirs, dim=2) 81 | 82 | return ray_dirs, cam_loc 83 | 84 | 85 | def get_camera_for_plot(pose): 86 | if pose.shape[1] == 7: #In case of quaternion vector representation 87 | cam_loc = pose[:, 4:].detach() 88 | R = quat_to_rot(pose[:,:4].detach()) 89 | else: # In case of pose matrix representation 90 | cam_loc = pose[:, :3, 3] 91 | R = pose[:, :3, :3] 92 | cam_dir = R[:, :3, 2] 93 | return cam_loc, cam_dir 94 | 95 | 96 | def lift(x, y, z, intrinsics): 97 | # parse intrinsics 98 | intrinsics = intrinsics.cuda() 99 | fx = intrinsics[:, 0, 0] 100 | fy = intrinsics[:, 1, 1] 101 | cx = intrinsics[:, 0, 2] 102 | cy = intrinsics[:, 1, 2] 103 | sk = intrinsics[:, 0, 1] 104 | 105 | x_lift = (x - cx.unsqueeze(-1) + cy.unsqueeze(-1)*sk.unsqueeze(-1)/fy.unsqueeze(-1) - sk.unsqueeze(-1)*y/fy.unsqueeze(-1)) / fx.unsqueeze(-1) * z 106 | y_lift = (y - cy.unsqueeze(-1)) / fy.unsqueeze(-1) * z 107 | 108 | # homogeneous 109 | return torch.stack((x_lift, y_lift, z, torch.ones_like(z).cuda()), dim=-1) 110 | 111 | 112 | def quat_to_rot(q): 113 | batch_size, _ = q.shape 114 | q = F.normalize(q, dim=1) 115 | R = torch.ones((batch_size, 3,3)).cuda() 116 | qr=q[:,0] 117 | qi = q[:, 1] 118 | qj = q[:, 2] 119 | qk = q[:, 3] 120 | R[:, 0, 0]=1-2 * (qj**2 + qk**2) 121 | R[:, 0, 1] = 2 * (qj *qi -qk*qr) 122 | R[:, 0, 2] = 2 * (qi * qk + qr * qj) 123 | R[:, 1, 0] = 2 * (qj * qi + qk * qr) 124 | R[:, 1, 1] = 1-2 * (qi**2 + qk**2) 125 | R[:, 1, 2] = 2*(qj*qk - qi*qr) 126 | R[:, 2, 0] = 2 * (qk * qi-qj * qr) 127 | R[:, 2, 1] = 2 * (qj*qk + qi*qr) 128 | R[:, 2, 2] = 1-2 * (qi**2 + qj**2) 129 | return R 130 | 131 | 132 | def rot_to_quat(R): 133 | batch_size, _,_ = R.shape 134 | q = torch.ones((batch_size, 4)).cuda() 135 | 136 | R00 = R[:, 0,0] 137 | R01 = R[:, 0, 1] 138 | R02 = R[:, 0, 2] 139 | R10 = R[:, 1, 0] 140 | R11 = R[:, 1, 1] 141 | R12 = R[:, 1, 2] 142 | R20 = R[:, 2, 0] 143 | R21 = R[:, 2, 1] 144 | R22 = R[:, 2, 2] 145 | 146 | q[:,0]=torch.sqrt(1.0+R00+R11+R22)/2 147 | q[:, 1]=(R21-R12)/(4*q[:,0]) 148 | q[:, 2] = (R02 - R20) / (4 * q[:, 0]) 149 | q[:, 3] = (R10 - R01) / (4 * q[:, 0]) 150 | return q 151 | 152 | 153 | def get_sphere_intersections(cam_loc, ray_directions, r = 1.0): 154 | # Input: n_rays x 3 ; n_rays x 3 155 | # Output: n_rays x 1, n_rays x 1 (close and far) 156 | 157 | ray_cam_dot = torch.bmm(ray_directions.view(-1, 1, 3), 158 | cam_loc.view(-1, 3, 1)).squeeze(-1) 159 | under_sqrt = ray_cam_dot ** 2 - (cam_loc.norm(2, 1, keepdim=True) ** 2 - r ** 2) 160 | 161 | # sanity check 162 | if (under_sqrt <= 0).sum() > 0: 163 | print('BOUNDING SPHERE PROBLEM!') 164 | exit() 165 | 166 | sphere_intersections = torch.sqrt(under_sqrt) * torch.Tensor([-1, 1]).cuda().float() - ray_cam_dot 167 | sphere_intersections = sphere_intersections.clamp_min(0.0) 168 | 169 | return sphere_intersections -------------------------------------------------------------------------------- /Trim2DGS/scripts/eval_tnt/README.md: -------------------------------------------------------------------------------- 1 | # Python Toolbox for Evaluation 2 | 3 | This Python script evaluates **training** dataset of TanksAndTemples benchmark. 4 | The script requires ``Open3D`` and a few Python packages such as ``matplotlib``, ``json``, and ``numpy``. 5 | 6 | ## How to use: 7 | **Step 0**. Reconstruct 3D models and recover camera poses from the training dataset. 8 | The raw videos of the training dataset can be found from: 9 | https://tanksandtemples.org/download/ 10 | 11 | **Step 1**. Download evaluation data (ground truth geometry + reference reconstruction) using 12 | [this link](https://drive.google.com/open?id=1UoKPiUUsKa0AVHFOrnMRhc5hFngjkE-t). In this example, we regard ``TanksAndTemples/evaluation/data/`` as a dataset folder. 13 | 14 | **Step 2**. Install Open3D. Follow instructions in http://open3d.org/docs/getting_started.html 15 | 16 | **Step 3**. Run the evaluation script and grab some coffee. 17 | ``` 18 | # firstly, run cull_mesh.py to cull mesh and then 19 | ./run.sh Barn 20 | ``` 21 | Output (evaluation of Ignatius): 22 | ``` 23 | =========================== 24 | Evaluating Ignatius 25 | =========================== 26 | path/to/TanksAndTemples/evaluation/data/Ignatius/Ignatius_COLMAP.ply 27 | Reading PLY: [========================================] 100% 28 | Read PointCloud: 6929586 vertices. 29 | path/to/TanksAndTemples/evaluation/data/Ignatius/Ignatius.ply 30 | Reading PLY: [========================================] 100% 31 | : 32 | ICP Iteration #0: Fitness 0.9980, RMSE 0.0044 33 | ICP Iteration #1: Fitness 0.9980, RMSE 0.0043 34 | ICP Iteration #2: Fitness 0.9980, RMSE 0.0043 35 | ICP Iteration #3: Fitness 0.9980, RMSE 0.0043 36 | ICP Iteration #4: Fitness 0.9980, RMSE 0.0042 37 | ICP Iteration #5: Fitness 0.9980, RMSE 0.0042 38 | ICP Iteration #6: Fitness 0.9979, RMSE 0.0042 39 | ICP Iteration #7: Fitness 0.9979, RMSE 0.0042 40 | ICP Iteration #8: Fitness 0.9979, RMSE 0.0042 41 | ICP Iteration #9: Fitness 0.9979, RMSE 0.0042 42 | ICP Iteration #10: Fitness 0.9979, RMSE 0.0042 43 | [EvaluateHisto] 44 | Cropping geometry: [========================================] 100% 45 | Pointcloud down sampled from 6929586 points to 1449840 points. 46 | Pointcloud down sampled from 1449840 points to 1365628 points. 47 | path/to/TanksAndTemples/evaluation/data/Ignatius/evaluation//Ignatius.precision.ply 48 | Cropping geometry: [========================================] 100% 49 | Pointcloud down sampled from 5016769 points to 4957123 points. 50 | Pointcloud down sampled from 4957123 points to 4181506 points. 51 | [compute_point_cloud_to_point_cloud_distance] 52 | [compute_point_cloud_to_point_cloud_distance] 53 | : 54 | [ViewDistances] Add color coding to visualize error 55 | [ViewDistances] Add color coding to visualize error 56 | : 57 | [get_f1_score_histo2] 58 | ============================== 59 | evaluation result : Ignatius 60 | ============================== 61 | distance tau : 0.003 62 | precision : 0.7679 63 | recall : 0.7937 64 | f-score : 0.7806 65 | ============================== 66 | ``` 67 | 68 | **Step 5**. Go to the evaluation folder. ``TanksAndTemples/evaluation/data/Ignatius/evaluation/`` will have the following outputs. 69 | 70 | 71 | 72 | ``PR_Ignatius_@d_th_0_0030.pdf`` (Precision and recall curves with a F-score) 73 | 74 | | | | 75 | |--|--| 76 | | ``Ignatius.precision.ply`` | ``Ignatius.recall.ply`` | 77 | 78 | (3D visualization of precision and recall. Each mesh is color coded using hot colormap) 79 | 80 | # Requirements 81 | 82 | - Python 3 83 | - open3d v0.9.0 84 | - matplotlib 85 | -------------------------------------------------------------------------------- /Trim2DGS/scripts/eval_tnt/config.py: -------------------------------------------------------------------------------- 1 | # ---------------------------------------------------------------------------- 2 | # - TanksAndTemples Website Toolbox - 3 | # - http://www.tanksandtemples.org - 4 | # ---------------------------------------------------------------------------- 5 | # The MIT License (MIT) 6 | # 7 | # Copyright (c) 2017 8 | # Arno Knapitsch 9 | # Jaesik Park 10 | # Qian-Yi Zhou 11 | # Vladlen Koltun 12 | # 13 | # Permission is hereby granted, free of charge, to any person obtaining a copy 14 | # of this software and associated documentation files (the "Software"), to deal 15 | # in the Software without restriction, including without limitation the rights 16 | # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 17 | # copies of the Software, and to permit persons to whom the Software is 18 | # furnished to do so, subject to the following conditions: 19 | # 20 | # The above copyright notice and this permission notice shall be included in 21 | # all copies or substantial portions of the Software. 22 | # 23 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 24 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 25 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 26 | # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 27 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 28 | # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 29 | # THE SOFTWARE. 30 | # ---------------------------------------------------------------------------- 31 | 32 | # some global parameters - do not modify 33 | scenes_tau_dict = { 34 | "Barn": 0.01, 35 | "Caterpillar": 0.005, 36 | "Church": 0.025, 37 | "Courthouse": 0.025, 38 | "Ignatius": 0.003, 39 | "Meetingroom": 0.01, 40 | "Truck": 0.005, 41 | } 42 | -------------------------------------------------------------------------------- /Trim2DGS/scripts/eval_tnt/evaluate_single_scene.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | import cv2 5 | import numpy as np 6 | import os 7 | import glob 8 | from skimage.morphology import binary_dilation, disk 9 | import argparse 10 | 11 | import trimesh 12 | from pathlib import Path 13 | import subprocess 14 | import sys 15 | import json 16 | 17 | 18 | if __name__ == "__main__": 19 | 20 | parser = argparse.ArgumentParser( 21 | description='Arguments to evaluate the mesh.' 22 | ) 23 | 24 | parser.add_argument('--input_mesh', type=str, help='path to the mesh to be evaluated') 25 | parser.add_argument('--scene', type=str, help='scan id of the input mesh') 26 | parser.add_argument('--output_dir', type=str, default='evaluation_results_single', help='path to the output folder') 27 | parser.add_argument('--TNT', type=str, default='Offical_DTU_Dataset', help='path to the GT DTU point clouds') 28 | args = parser.parse_args() 29 | 30 | 31 | TNT_Dataset = args.TNT 32 | out_dir = args.output_dir 33 | Path(out_dir).mkdir(parents=True, exist_ok=True) 34 | scene = args.scene 35 | ply_file = args.input_mesh 36 | result_mesh_file = os.path.join(out_dir, "culled_mesh.ply") 37 | # read scene.json 38 | f"python run.py --dataset-dir {ply_file} --traj-path {TNT_Dataset}/{scene}/{scene}_COLMAP_SfM.log --ply-path {TNT_Dataset}/{scene}/{scene}_COLMAP.ply" -------------------------------------------------------------------------------- /Trim2DGS/scripts/eval_tnt/help_func.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # coding=utf-8 3 | import torch 4 | 5 | def rotation_matrix(a, b): 6 | """Compute the rotation matrix that rotates vector a to vector b. 7 | 8 | Args: 9 | a: The vector to rotate. 10 | b: The vector to rotate to. 11 | Returns: 12 | The rotation matrix. 13 | """ 14 | a = a / torch.linalg.norm(a) 15 | b = b / torch.linalg.norm(b) 16 | v = torch.cross(a, b) 17 | c = torch.dot(a, b) 18 | # If vectors are exactly opposite, we add a little noise to one of them 19 | if c < -1 + 1e-8: 20 | eps = (torch.rand(3) - 0.5) * 0.01 21 | return rotation_matrix(a + eps, b) 22 | s = torch.linalg.norm(v) 23 | skew_sym_mat = torch.Tensor( 24 | [ 25 | [0, -v[2], v[1]], 26 | [v[2], 0, -v[0]], 27 | [-v[1], v[0], 0], 28 | ] 29 | ) 30 | return torch.eye(3) + skew_sym_mat + skew_sym_mat @ skew_sym_mat * ((1 - c) / (s**2 + 1e-8)) 31 | 32 | 33 | def auto_orient_and_center_poses( 34 | poses, method="up", center_poses=True 35 | ): 36 | """Orients and centers the poses. We provide two methods for orientation: pca and up. 37 | 38 | pca: Orient the poses so that the principal component of the points is aligned with the axes. 39 | This method works well when all of the cameras are in the same plane. 40 | up: Orient the poses so that the average up vector is aligned with the z axis. 41 | This method works well when images are not at arbitrary angles. 42 | 43 | 44 | Args: 45 | poses: The poses to orient. 46 | method: The method to use for orientation. 47 | center_poses: If True, the poses are centered around the origin. 48 | 49 | Returns: 50 | The oriented poses. 51 | """ 52 | 53 | translation = poses[..., :3, 3] 54 | 55 | mean_translation = torch.mean(translation, dim=0) 56 | translation_diff = translation - mean_translation 57 | 58 | if center_poses: 59 | translation = mean_translation 60 | else: 61 | translation = torch.zeros_like(mean_translation) 62 | 63 | if method == "pca": 64 | _, eigvec = torch.linalg.eigh(translation_diff.T @ translation_diff) 65 | eigvec = torch.flip(eigvec, dims=(-1,)) 66 | 67 | if torch.linalg.det(eigvec) < 0: 68 | eigvec[:, 2] = -eigvec[:, 2] 69 | 70 | transform = torch.cat([eigvec, eigvec @ -translation[..., None]], dim=-1) 71 | oriented_poses = transform @ poses 72 | 73 | if oriented_poses.mean(axis=0)[2, 1] < 0: 74 | oriented_poses[:, 1:3] = -1 * oriented_poses[:, 1:3] 75 | elif method == "up": 76 | up = torch.mean(poses[:, :3, 1], dim=0) 77 | up = up / torch.linalg.norm(up) 78 | 79 | rotation = rotation_matrix(up, torch.Tensor([0, 0, 1])) 80 | transform = torch.cat([rotation, rotation @ -translation[..., None]], dim=-1) 81 | oriented_poses = transform @ poses 82 | elif method == "none": 83 | transform = torch.eye(4) 84 | transform[:3, 3] = -translation 85 | transform = transform[:3, :] 86 | oriented_poses = transform @ poses 87 | 88 | return oriented_poses, transform 89 | 90 | 91 | -------------------------------------------------------------------------------- /Trim2DGS/scripts/eval_tnt/plot.py: -------------------------------------------------------------------------------- 1 | # ---------------------------------------------------------------------------- 2 | # - TanksAndTemples Website Toolbox - 3 | # - http://www.tanksandtemples.org - 4 | # ---------------------------------------------------------------------------- 5 | # The MIT License (MIT) 6 | # 7 | # Copyright (c) 2017 8 | # Arno Knapitsch 9 | # Jaesik Park 10 | # Qian-Yi Zhou 11 | # Vladlen Koltun 12 | # 13 | # Permission is hereby granted, free of charge, to any person obtaining a copy 14 | # of this software and associated documentation files (the "Software"), to deal 15 | # in the Software without restriction, including without limitation the rights 16 | # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 17 | # copies of the Software, and to permit persons to whom the Software is 18 | # furnished to do so, subject to the following conditions: 19 | # 20 | # The above copyright notice and this permission notice shall be included in 21 | # all copies or substantial portions of the Software. 22 | # 23 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 24 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 25 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 26 | # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 27 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 28 | # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 29 | # THE SOFTWARE. 30 | # ---------------------------------------------------------------------------- 31 | # 32 | # This python script is for downloading dataset from www.tanksandtemples.org 33 | # The dataset has a different license, please refer to 34 | # https://tanksandtemples.org/license/ 35 | 36 | import matplotlib.pyplot as plt 37 | from cycler import cycler 38 | 39 | 40 | def plot_graph( 41 | scene, 42 | fscore, 43 | dist_threshold, 44 | edges_source, 45 | cum_source, 46 | edges_target, 47 | cum_target, 48 | plot_stretch, 49 | mvs_outpath, 50 | show_figure=False, 51 | ): 52 | f = plt.figure() 53 | plt_size = [14, 7] 54 | pfontsize = "medium" 55 | 56 | ax = plt.subplot(111) 57 | label_str = "precision" 58 | ax.plot( 59 | edges_source[1::], 60 | cum_source * 100, 61 | c="red", 62 | label=label_str, 63 | linewidth=2.0, 64 | ) 65 | 66 | label_str = "recall" 67 | ax.plot( 68 | edges_target[1::], 69 | cum_target * 100, 70 | c="blue", 71 | label=label_str, 72 | linewidth=2.0, 73 | ) 74 | 75 | ax.grid(True) 76 | plt.rcParams["figure.figsize"] = plt_size 77 | plt.rc("axes", prop_cycle=cycler("color", ["r", "g", "b", "y"])) 78 | plt.title("Precision and Recall: " + scene + ", " + "%02.2f f-score" % 79 | (fscore * 100)) 80 | plt.axvline(x=dist_threshold, c="black", ls="dashed", linewidth=2.0) 81 | 82 | plt.ylabel("# of points (%)", fontsize=15) 83 | plt.xlabel("Meters", fontsize=15) 84 | plt.axis([0, dist_threshold * plot_stretch, 0, 100]) 85 | ax.legend(shadow=True, fancybox=True, fontsize=pfontsize) 86 | # plt.axis([0, dist_threshold*plot_stretch, 0, 100]) 87 | 88 | plt.setp(ax.get_legend().get_texts(), fontsize=pfontsize) 89 | 90 | plt.legend(loc=2, borderaxespad=0.0, fontsize=pfontsize) 91 | plt.legend(loc=4) 92 | leg = plt.legend(loc="lower right") 93 | 94 | box = ax.get_position() 95 | ax.set_position([box.x0, box.y0, box.width * 0.8, box.height]) 96 | 97 | # Put a legend to the right of the current axis 98 | ax.legend(loc="center left", bbox_to_anchor=(1, 0.5)) 99 | plt.setp(ax.get_legend().get_texts(), fontsize=pfontsize) 100 | png_name = mvs_outpath + "/PR_{0}_@d_th_0_{1}.png".format( 101 | scene, "%04d" % (dist_threshold * 10000)) 102 | pdf_name = mvs_outpath + "/PR_{0}_@d_th_0_{1}.pdf".format( 103 | scene, "%04d" % (dist_threshold * 10000)) 104 | 105 | # save figure and display 106 | f.savefig(png_name, format="png", bbox_inches="tight") 107 | f.savefig(pdf_name, format="pdf", bbox_inches="tight") 108 | if show_figure: 109 | plt.show() 110 | -------------------------------------------------------------------------------- /Trim2DGS/scripts/eval_tnt/requirements.txt: -------------------------------------------------------------------------------- 1 | matplotlib>=1.3 2 | open3d==0.10 3 | -------------------------------------------------------------------------------- /Trim2DGS/scripts/eval_tnt/trajectory_io.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import open3d as o3d 3 | 4 | 5 | class CameraPose: 6 | 7 | def __init__(self, meta, mat): 8 | self.metadata = meta 9 | self.pose = mat 10 | 11 | def __str__(self): 12 | return ("Metadata : " + " ".join(map(str, self.metadata)) + "\n" + 13 | "Pose : " + "\n" + np.array_str(self.pose)) 14 | 15 | 16 | def convert_trajectory_to_pointcloud(traj): 17 | pcd = o3d.geometry.PointCloud() 18 | for t in traj: 19 | pcd.points.append(t.pose[:3, 3]) 20 | return pcd 21 | 22 | 23 | def read_trajectory(filename): 24 | traj = [] 25 | with open(filename, "r") as f: 26 | metastr = f.readline() 27 | while metastr: 28 | metadata = map(int, metastr.split()) 29 | mat = np.zeros(shape=(4, 4)) 30 | for i in range(4): 31 | matstr = f.readline() 32 | mat[i, :] = np.fromstring(matstr, dtype=float, sep=" \t") 33 | traj.append(CameraPose(metadata, mat)) 34 | metastr = f.readline() 35 | return traj 36 | 37 | 38 | def write_trajectory(traj, filename): 39 | with open(filename, "w") as f: 40 | for x in traj: 41 | p = x.pose.tolist() 42 | f.write(" ".join(map(str, x.metadata)) + "\n") 43 | f.write("\n".join( 44 | " ".join(map("{0:.12f}".format, p[i])) for i in range(4))) 45 | f.write("\n") 46 | -------------------------------------------------------------------------------- /Trim2DGS/scripts/eval_tnt/util.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | 4 | def make_dir(path): 5 | if not os.path.exists(path): 6 | os.makedirs(path) 7 | -------------------------------------------------------------------------------- /Trim2DGS/scripts/m360_eval.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 | from argparse import ArgumentParser 14 | 15 | mipnerf360_outdoor_scenes = ["bicycle", "flowers", "garden", "stump", "treehill"] 16 | mipnerf360_indoor_scenes = ["room", "counter", "kitchen", "bonsai"] 17 | # tanks_and_temples_scenes = ["truck", "train"] 18 | # deep_blending_scenes = ["drjohnson", "playroom"] 19 | 20 | parser = ArgumentParser(description="Full evaluation script parameters") 21 | parser.add_argument("--skip_training", action="store_true") 22 | parser.add_argument("--skip_rendering", action="store_true") 23 | parser.add_argument("--skip_metrics", action="store_true") 24 | parser.add_argument("--output_path", default="eval/mipnerf360") 25 | args, _ = parser.parse_known_args() 26 | 27 | all_scenes = [] 28 | all_scenes.extend(mipnerf360_outdoor_scenes) 29 | all_scenes.extend(mipnerf360_indoor_scenes) 30 | # all_scenes.extend(tanks_and_temples_scenes) 31 | # all_scenes.extend(deep_blending_scenes) 32 | 33 | if not args.skip_training or not args.skip_rendering: 34 | parser.add_argument('--mipnerf360', "-m360", required=True, type=str) 35 | # parser.add_argument("--tanksandtemples", "-tat", required=True, type=str) 36 | # parser.add_argument("--deepblending", "-db", required=True, type=str) 37 | args = parser.parse_args() 38 | 39 | if not args.skip_training: 40 | common_args = " --quiet --eval --test_iterations -1" 41 | for scene in mipnerf360_outdoor_scenes: 42 | source = args.mipnerf360 + "/" + scene 43 | os.system("python train.py -s " + source + " -i images_4 -m " + args.output_path + "/" + scene + common_args) 44 | for scene in mipnerf360_indoor_scenes: 45 | source = args.mipnerf360 + "/" + scene 46 | os.system("python train.py -s " + source + " -i images_2 -m " + args.output_path + "/" + scene + common_args) 47 | # for scene in tanks_and_temples_scenes: 48 | # source = args.tanksandtemples + "/" + scene 49 | # os.system("python train.py -s " + source + " -m " + args.output_path + "/" + scene + common_args) 50 | # for scene in deep_blending_scenes: 51 | # source = args.deepblending + "/" + scene 52 | # os.system("python train.py -s " + source + " -m " + args.output_path + "/" + scene + common_args) 53 | 54 | if not args.skip_rendering: 55 | all_sources = [] 56 | for scene in mipnerf360_outdoor_scenes: 57 | all_sources.append(args.mipnerf360 + "/" + scene) 58 | for scene in mipnerf360_indoor_scenes: 59 | all_sources.append(args.mipnerf360 + "/" + scene) 60 | # for scene in tanks_and_temples_scenes: 61 | # all_sources.append(args.tanksandtemples + "/" + scene) 62 | # for scene in deep_blending_scenes: 63 | # all_sources.append(args.deepblending + "/" + scene) 64 | 65 | common_args = " --quiet --eval --skip_train" 66 | for scene, source in zip(all_scenes, all_sources): 67 | os.system("python render.py --iteration 30000 -s " + source + " -m " + args.output_path + "/" + scene + common_args) 68 | 69 | if not args.skip_metrics: 70 | scenes_string = "" 71 | for scene in all_scenes: 72 | scenes_string += "\"" + args.output_path + "/" + scene + "\" " 73 | 74 | os.system("python metrics.py -m " + scenes_string) -------------------------------------------------------------------------------- /Trim2DGS/scripts/run_Mipnerf360.py: -------------------------------------------------------------------------------- 1 | # training script for MipNeRF360 dataset 2 | # adapted from https://github.com/autonomousvision/gaussian-opacity-fields/blob/main/scripts/run_mipnerf360.py 3 | 4 | import os 5 | from concurrent.futures import ThreadPoolExecutor 6 | import subprocess 7 | import time 8 | import torch 9 | 10 | scenes = ["bicycle", "bonsai", "counter", "flowers", "garden", "stump", "treehill", "kitchen", "room"] 11 | 12 | factors = [4, 2, 2, 4, 4, 4, 4, 2, 2] 13 | 14 | excluded_gpus = set([]) 15 | 16 | output_dir = "output/MipNeRF360_2DGS" 17 | tune_output_dir = f"output/MipNeRF360_Trim2DGS" 18 | iteration = 7000 19 | split = "scale" 20 | 21 | jobs = list(zip(scenes, factors)) 22 | 23 | def train_scene(gpu, scene, factor): 24 | cmds = [ 25 | f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python train.py -s data/MipNeRF360/{scene} -m {output_dir}/{scene} --eval -i images_{factor} --test_iterations -1 --quiet", 26 | f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python tune.py -s data/MipNeRF360/{scene} -m {tune_output_dir}/{scene} --eval -i images_{factor} --pretrained_ply {output_dir}/{scene}/point_cloud/iteration_30000/point_cloud.ply --test_iterations -1 --quiet --split {split} --max_screen_size 100", 27 | f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python render.py --iteration {iteration} -m {tune_output_dir}/{scene} --quiet --eval --skip_train", 28 | f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python metrics.py -m {tune_output_dir}/{scene}", 29 | ] 30 | 31 | for cmd in cmds: 32 | print(cmd) 33 | subprocess.run(cmd, shell=True, check=True) 34 | return True 35 | 36 | 37 | def worker(gpu, scene, factor): 38 | print(f"Starting job on GPU {gpu} with scene {scene}\n") 39 | train_scene(gpu, scene, factor) 40 | print(f"Finished job on GPU {gpu} with scene {scene}\n") 41 | # This worker function starts a job and returns when it's done. 42 | 43 | def dispatch_jobs(jobs, executor): 44 | future_to_job = {} 45 | reserved_gpus = set() # GPUs that are slated for work but may not be active yet 46 | 47 | while jobs or future_to_job: 48 | # Get the list of available GPUs, not including those that are reserved. 49 | all_available_gpus = set(range(torch.cuda.device_count())) 50 | available_gpus = list(all_available_gpus - reserved_gpus - excluded_gpus) 51 | 52 | # Launch new jobs on available GPUs 53 | while available_gpus and jobs: 54 | gpu = available_gpus.pop(0) 55 | job = jobs.pop(0) 56 | future = executor.submit(worker, gpu, *job) # Unpacking job as arguments to worker 57 | future_to_job[future] = (gpu, job) 58 | 59 | reserved_gpus.add(gpu) # Reserve this GPU until the job starts processing 60 | 61 | # Check for completed jobs and remove them from the list of running jobs. 62 | # Also, release the GPUs they were using. 63 | done_futures = [future for future in future_to_job if future.done()] 64 | for future in done_futures: 65 | job = future_to_job.pop(future) # Remove the job associated with the completed future 66 | gpu = job[0] # The GPU is the first element in each job tuple 67 | reserved_gpus.discard(gpu) # Release this GPU 68 | print(f"Job {job} has finished., rellasing GPU {gpu}") 69 | # (Optional) You might want to introduce a small delay here to prevent this loop from spinning very fast 70 | # when there are no GPUs available. 71 | time.sleep(5) 72 | 73 | print("All jobs have been processed.") 74 | 75 | 76 | # Using ThreadPoolExecutor to manage the thread pool 77 | with ThreadPoolExecutor(max_workers=8) as executor: 78 | dispatch_jobs(jobs, executor) 79 | 80 | -------------------------------------------------------------------------------- /Trim2DGS/scripts/run_TNT.py: -------------------------------------------------------------------------------- 1 | # training script for TnT dataset 2 | # adapted from https://github.com/autonomousvision/gaussian-opacity-fields/blob/main/scripts/run_tnt.py 3 | 4 | import os 5 | from concurrent.futures import ThreadPoolExecutor 6 | import subprocess 7 | import time 8 | import torch 9 | 10 | scenes = ['Barn', 'Caterpillar', 'Ignatius', 'Truck', 'Meetingroom', 'Courthouse'] 11 | 12 | lambda_dist = { 13 | 'Barn': 100, 'Caterpillar': 100, 'Ignatius': 100, 'Truck': 100, 14 | 'Meetingroom': 10, 'Courthouse': 10 15 | } 16 | 17 | voxel_size = { 18 | 'Barn': 0.004, 'Caterpillar': 0.004, 'Ignatius': 0.004, 'Truck': 0.004, 19 | 'Meetingroom': 0.006, 'Courthouse': 0.006 20 | } 21 | 22 | sdf_trunc = { 23 | 'Barn': 0.016, 'Caterpillar': 0.016, 'Ignatius': 0.016, 'Truck': 0.016, 24 | 'Meetingroom': 0.024, 'Courthouse': 0.024 25 | } 26 | 27 | depth_trunc = { 28 | 'Barn': 3.0, 'Caterpillar': 3.0, 'Ignatius': 3.0, 'Truck': 3.0, 29 | 'Meetingroom': 4.5, 'Courthouse': 4.5 30 | } 31 | 32 | excluded_gpus = set([]) 33 | 34 | output_dir = "output/TNT_2DGS" 35 | tune_output_dir = f"output/TNT_Trim2DGS" 36 | iteration = 7000 37 | split = "scale" 38 | extra_cmd = "--position_lr_init 0.0000016 --contribution_prune_interval 300 --opacity_reset_interval 99999 --depth_grad_thresh 0.0" 39 | jobs = scenes 40 | 41 | def train_scene(gpu, scene): 42 | cmds = [ 43 | f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python train.py -s data/TNT_GOF/TrainingSet/{scene} -m {output_dir}/{scene} --eval --test_iterations -1 --quiet --depth_ratio 1.0 -r 2 --lambda_dist {lambda_dist[scene]}", 44 | f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python tune.py -s data/TNT_GOF/TrainingSet/{scene} -m {tune_output_dir}/{scene} --eval --pretrained_ply {output_dir}/{scene}/point_cloud/iteration_30000/point_cloud.ply --test_iterations -1 --quiet --depth_ratio 1.0 -r 2 --lambda_dist {lambda_dist[scene]} --split {split} {extra_cmd}", 45 | f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python render.py --iteration {iteration} -m {tune_output_dir}/{scene} --quiet --depth_ratio 1.0 --num_cluster 1 --voxel_size {voxel_size[scene]} --sdf_trunc {sdf_trunc[scene]} --depth_trunc {depth_trunc[scene]}", 46 | f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python metrics.py -m {tune_output_dir}/{scene}", 47 | ] 48 | 49 | for cmd in cmds: 50 | print(cmd) 51 | subprocess.run(cmd, shell=True, check=True) 52 | return True 53 | 54 | 55 | def worker(gpu, scene): 56 | print(f"Starting job on GPU {gpu} with scene {scene}\n") 57 | train_scene(gpu, scene) 58 | print(f"Finished job on GPU {gpu} with scene {scene}\n") 59 | # This worker function starts a job and returns when it's done. 60 | 61 | def dispatch_jobs(jobs, executor): 62 | future_to_job = {} 63 | reserved_gpus = set() # GPUs that are slated for work but may not be active yet 64 | 65 | while jobs or future_to_job: 66 | # Get the list of available GPUs, not including those that are reserved. 67 | all_available_gpus = set(range(torch.cuda.device_count())) 68 | available_gpus = list(all_available_gpus - reserved_gpus - excluded_gpus) 69 | 70 | # Launch new jobs on available GPUs 71 | while available_gpus and jobs: 72 | gpu = available_gpus.pop(0) 73 | job = jobs.pop(0) 74 | future = executor.submit(worker, gpu, job) # Unpacking job as arguments to worker 75 | future_to_job[future] = (gpu, job) 76 | 77 | reserved_gpus.add(gpu) # Reserve this GPU until the job starts processing 78 | 79 | # Check for completed jobs and remove them from the list of running jobs. 80 | # Also, release the GPUs they were using. 81 | done_futures = [future for future in future_to_job if future.done()] 82 | for future in done_futures: 83 | job = future_to_job.pop(future) # Remove the job associated with the completed future 84 | gpu = job[0] # The GPU is the first element in each job tuple 85 | reserved_gpus.discard(gpu) # Release this GPU 86 | print(f"Job {job} has finished., rellasing GPU {gpu}") 87 | # (Optional) You might want to introduce a small delay here to prevent this loop from spinning very fast 88 | # when there are no GPUs available. 89 | time.sleep(5) 90 | 91 | print("All jobs have been processed.") 92 | 93 | 94 | # Using ThreadPoolExecutor to manage the thread pool 95 | with ThreadPoolExecutor(max_workers=8) as executor: 96 | dispatch_jobs(jobs, executor) 97 | 98 | -------------------------------------------------------------------------------- /Trim2DGS/scripts/run_dtu.py: -------------------------------------------------------------------------------- 1 | # training script for DTU dataset 2 | # adapted from https://github.com/autonomousvision/gaussian-opacity-fields/blob/main/scripts/run_dtu.py 3 | 4 | import os 5 | from concurrent.futures import ThreadPoolExecutor 6 | import subprocess 7 | import time 8 | import torch 9 | 10 | scenes = ['scan24', 'scan37', 'scan40', 'scan55', 'scan63', 'scan65', 'scan69', 'scan83', 'scan97', 'scan105', 'scan106', 'scan110', 'scan114', 'scan118', 'scan122'] 11 | 12 | factors = [2] * len(scenes) 13 | 14 | excluded_gpus = set([]) 15 | 16 | output_dir = "output/DTU_2DGS" 17 | tune_output_dir = "output/DTU_Trim2DGS" 18 | iteration = 7000 19 | position_lr_init = {"scan63": 0.0000016} 20 | 21 | jobs = list(zip(scenes, factors)) 22 | 23 | def train_scene(gpu, scene, factor): 24 | scan_id = scene[4:] 25 | cmds = [ 26 | f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python train.py -s data/dtu_dataset/DTU/{scene} -m {output_dir}/{scene} --quiet --test_iterations -1 --depth_ratio 1.0 -r {factor} --lambda_dist 1000", 27 | # cull the points that are not visible in the training set 28 | f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python cull_pcd.py -m {output_dir}/{scene} --iteration 30000", 29 | f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python tune.py -s data/dtu_dataset/DTU/{scene} -m {tune_output_dir}/{scene} --pretrained_ply {output_dir}/{scene}/point_cloud/iteration_30000/point_cloud_culled.ply --quiet --test_iterations -1 --depth_ratio 1.0 -r {factor} --lambda_dist 1000 --densification_interval 100 --contribution_prune_interval 100 --position_lr_init {position_lr_init.get(scene, 0.00016)} --split mix", 30 | f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python render.py -s data/dtu_dataset/DTU/{scene} -m {tune_output_dir}/{scene} --quiet --skip_train --depth_ratio 1.0 --num_cluster 1 --iteration {iteration} --voxel_size 0.004 --voxel_size 0.004 --sdf_trunc 0.016 --depth_trunc 3.0", 31 | f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python scripts/eval_dtu/evaluate_single_scene.py --input_mesh {tune_output_dir}/{scene}/train/ours_{iteration}/fuse_post.ply --scan_id {scan_id} --output_dir {tune_output_dir}/{scene}/train/ours_{iteration} --mask_dir data/dtu_dataset/DTU --DTU data/dtu_dataset/Official_DTU_Dataset", 32 | f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python scripts/eval_dtu_pcd/evaluate_single_scene.py --input_pcd {tune_output_dir}/{scene}/point_cloud/iteration_{iteration}/point_cloud.ply --scan_id {scan_id} --output_dir {tune_output_dir}/{scene}/train/ours_{iteration} --mask_dir data/dtu_dataset/DTU --DTU data/dtu_dataset/Official_DTU_Dataset", 33 | ] 34 | 35 | for cmd in cmds: 36 | print(cmd) 37 | subprocess.run(cmd, shell=True, check=True) 38 | return True 39 | 40 | 41 | def worker(gpu, scene, factor): 42 | print(f"Starting job on GPU {gpu} with scene {scene}\n") 43 | train_scene(gpu, scene, factor) 44 | print(f"Finished job on GPU {gpu} with scene {scene}\n") 45 | # This worker function starts a job and returns when it's done. 46 | 47 | def dispatch_jobs(jobs, executor): 48 | future_to_job = {} 49 | reserved_gpus = set() # GPUs that are slated for work but may not be active yet 50 | 51 | while jobs or future_to_job: 52 | # Get the list of available GPUs, not including those that are reserved. 53 | all_available_gpus = set(range(torch.cuda.device_count())) 54 | available_gpus = list(all_available_gpus - reserved_gpus - excluded_gpus) 55 | 56 | # Launch new jobs on available GPUs 57 | while available_gpus and jobs: 58 | gpu = available_gpus.pop(0) 59 | job = jobs.pop(0) 60 | future = executor.submit(worker, gpu, *job) # Unpacking job as arguments to worker 61 | future_to_job[future] = (gpu, job) 62 | 63 | reserved_gpus.add(gpu) # Reserve this GPU until the job starts processing 64 | 65 | # Check for completed jobs and remove them from the list of running jobs. 66 | # Also, release the GPUs they were using. 67 | done_futures = [future for future in future_to_job if future.done()] 68 | for future in done_futures: 69 | job = future_to_job.pop(future) # Remove the job associated with the completed future 70 | gpu = job[0] # The GPU is the first element in each job tuple 71 | reserved_gpus.discard(gpu) # Release this GPU 72 | print(f"Job {job} has finished., rellasing GPU {gpu}") 73 | # (Optional) You might want to introduce a small delay here to prevent this loop from spinning very fast 74 | # when there are no GPUs available. 75 | time.sleep(5) 76 | 77 | print("All jobs have been processed.") 78 | 79 | 80 | # Using ThreadPoolExecutor to manage the thread pool 81 | with ThreadPoolExecutor(max_workers=8) as executor: 82 | dispatch_jobs(jobs, executor) 83 | 84 | -------------------------------------------------------------------------------- /Trim2DGS/scripts/train_from_scratch/run_Mipnerf360.py: -------------------------------------------------------------------------------- 1 | # training script for MipNeRF360 dataset 2 | # adapted from https://github.com/autonomousvision/gaussian-opacity-fields/blob/main/scripts/run_mipnerf360.py 3 | 4 | import os 5 | from concurrent.futures import ThreadPoolExecutor 6 | import subprocess 7 | import time 8 | import torch 9 | 10 | scenes = ["bicycle", "bonsai", "counter", "flowers", "garden", "stump", "treehill", "kitchen", "room"] 11 | 12 | factors = [4, 2, 2, 4, 4, 4, 4, 2, 2] 13 | 14 | excluded_gpus = set([]) 15 | 16 | output_dir = "output/MipNeRF360_Trim2DGS" 17 | iteration = 30000 18 | 19 | jobs = list(zip(scenes, factors)) 20 | 21 | def train_scene(gpu, scene, factor): 22 | cmds = [ 23 | f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python train_TrimGS.py -s data/MipNeRF360/{scene} -m {output_dir}/{scene} --eval -i images_{factor} --test_iterations -1 --quiet --split mix", 24 | f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python render.py --iteration {iteration} -m {output_dir}/{scene} --quiet --eval --skip_train --skip_mesh", 25 | f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python metrics.py -m {output_dir}/{scene}", 26 | ] 27 | 28 | for cmd in cmds: 29 | print(cmd) 30 | subprocess.run(cmd, shell=True, check=True) 31 | return True 32 | 33 | 34 | def worker(gpu, scene, factor): 35 | print(f"Starting job on GPU {gpu} with scene {scene}\n") 36 | train_scene(gpu, scene, factor) 37 | print(f"Finished job on GPU {gpu} with scene {scene}\n") 38 | # This worker function starts a job and returns when it's done. 39 | 40 | def dispatch_jobs(jobs, executor): 41 | future_to_job = {} 42 | reserved_gpus = set() # GPUs that are slated for work but may not be active yet 43 | 44 | while jobs or future_to_job: 45 | # Get the list of available GPUs, not including those that are reserved. 46 | all_available_gpus = set(range(torch.cuda.device_count())) 47 | available_gpus = list(all_available_gpus - reserved_gpus - excluded_gpus) 48 | 49 | # Launch new jobs on available GPUs 50 | while available_gpus and jobs: 51 | gpu = available_gpus.pop(0) 52 | job = jobs.pop(0) 53 | future = executor.submit(worker, gpu, *job) # Unpacking job as arguments to worker 54 | future_to_job[future] = (gpu, job) 55 | 56 | reserved_gpus.add(gpu) # Reserve this GPU until the job starts processing 57 | 58 | # Check for completed jobs and remove them from the list of running jobs. 59 | # Also, release the GPUs they were using. 60 | done_futures = [future for future in future_to_job if future.done()] 61 | for future in done_futures: 62 | job = future_to_job.pop(future) # Remove the job associated with the completed future 63 | gpu = job[0] # The GPU is the first element in each job tuple 64 | reserved_gpus.discard(gpu) # Release this GPU 65 | print(f"Job {job} has finished., rellasing GPU {gpu}") 66 | # (Optional) You might want to introduce a small delay here to prevent this loop from spinning very fast 67 | # when there are no GPUs available. 68 | time.sleep(5) 69 | 70 | print("All jobs have been processed.") 71 | 72 | 73 | # Using ThreadPoolExecutor to manage the thread pool 74 | with ThreadPoolExecutor(max_workers=8) as executor: 75 | dispatch_jobs(jobs, executor) 76 | 77 | -------------------------------------------------------------------------------- /Trim2DGS/scripts/train_from_scratch/run_dtu.py: -------------------------------------------------------------------------------- 1 | # training script for DTU dataset 2 | # adapted from https://github.com/autonomousvision/gaussian-opacity-fields/blob/main/scripts/run_dtu.py 3 | 4 | import os 5 | from concurrent.futures import ThreadPoolExecutor 6 | import subprocess 7 | import time 8 | import torch 9 | 10 | scenes = ['scan24', 'scan37', 'scan40', 'scan55', 'scan63', 'scan65', 'scan69', 'scan83', 'scan97', 'scan105', 'scan106', 'scan110', 'scan114', 'scan118', 'scan122'] 11 | 12 | factors = [2] * len(scenes) 13 | 14 | excluded_gpus = set([]) 15 | 16 | output_dir = "output/DTU_Trim2DGS" 17 | iteration = 30000 18 | 19 | jobs = list(zip(scenes, factors)) 20 | 21 | def train_scene(gpu, scene, factor): 22 | scan_id = scene[4:] 23 | cmds = [ 24 | f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python train_TrimGS.py -s data/dtu_dataset/DTU/{scene} -m {output_dir}/{scene} --quiet --test_iterations -1 --depth_ratio 1.0 -r {factor} --lambda_dist 1000 --split mix --contribution_prune_interval 100", 25 | f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python render.py -s data/dtu_dataset/DTU/{scene} -m {output_dir}/{scene} --quiet --skip_train --depth_ratio 1.0 --num_cluster 1 --iteration {iteration} --voxel_size 0.004 --voxel_size 0.004 --sdf_trunc 0.016 --depth_trunc 3.0", 26 | f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python scripts/eval_dtu/evaluate_single_scene.py --input_mesh {output_dir}/{scene}/train/ours_{iteration}/fuse_post.ply --scan_id {scan_id} --output_dir {output_dir}/{scene}/train/ours_{iteration} --mask_dir data/dtu_dataset/DTU --DTU data/dtu_dataset/Official_DTU_Dataset", 27 | f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python scripts/eval_dtu_pcd/evaluate_single_scene.py --input_pcd {output_dir}/{scene}/point_cloud/iteration_{iteration}/point_cloud.ply --scan_id {scan_id} --output_dir {output_dir}/{scene}/train/ours_{iteration} --mask_dir data/dtu_dataset/DTU --DTU data/dtu_dataset/Official_DTU_Dataset", 28 | ] 29 | 30 | for cmd in cmds: 31 | print(cmd) 32 | subprocess.run(cmd, shell=True, check=True) 33 | return True 34 | 35 | 36 | def worker(gpu, scene, factor): 37 | print(f"Starting job on GPU {gpu} with scene {scene}\n") 38 | train_scene(gpu, scene, factor) 39 | print(f"Finished job on GPU {gpu} with scene {scene}\n") 40 | # This worker function starts a job and returns when it's done. 41 | 42 | def dispatch_jobs(jobs, executor): 43 | future_to_job = {} 44 | reserved_gpus = set() # GPUs that are slated for work but may not be active yet 45 | 46 | while jobs or future_to_job: 47 | # Get the list of available GPUs, not including those that are reserved. 48 | all_available_gpus = set(range(torch.cuda.device_count())) 49 | available_gpus = list(all_available_gpus - reserved_gpus - excluded_gpus) 50 | 51 | # Launch new jobs on available GPUs 52 | while available_gpus and jobs: 53 | gpu = available_gpus.pop(0) 54 | job = jobs.pop(0) 55 | future = executor.submit(worker, gpu, *job) # Unpacking job as arguments to worker 56 | future_to_job[future] = (gpu, job) 57 | 58 | reserved_gpus.add(gpu) # Reserve this GPU until the job starts processing 59 | 60 | # Check for completed jobs and remove them from the list of running jobs. 61 | # Also, release the GPUs they were using. 62 | done_futures = [future for future in future_to_job if future.done()] 63 | for future in done_futures: 64 | job = future_to_job.pop(future) # Remove the job associated with the completed future 65 | gpu = job[0] # The GPU is the first element in each job tuple 66 | reserved_gpus.discard(gpu) # Release this GPU 67 | print(f"Job {job} has finished., rellasing GPU {gpu}") 68 | # (Optional) You might want to introduce a small delay here to prevent this loop from spinning very fast 69 | # when there are no GPUs available. 70 | time.sleep(5) 71 | 72 | print("All jobs have been processed.") 73 | 74 | 75 | # Using ThreadPoolExecutor to manage the thread pool 76 | with ThreadPoolExecutor(max_workers=8) as executor: 77 | dispatch_jobs(jobs, executor) 78 | 79 | -------------------------------------------------------------------------------- /Trim2DGS/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 | from utils.general_utils import PILtoTorch 15 | from utils.graphics_utils import fov2focal 16 | 17 | WARNED = False 18 | 19 | def loadCam(args, id, cam_info, resolution_scale): 20 | orig_w, orig_h = cam_info.image.size 21 | 22 | if args.resolution in [1, 2, 4, 8]: 23 | resolution = round(orig_w/(resolution_scale * args.resolution)), round(orig_h/(resolution_scale * args.resolution)) 24 | else: # should be a type that converts to float 25 | if args.resolution == -1: 26 | if orig_w > 1600: 27 | global WARNED 28 | if not WARNED: 29 | print("[ INFO ] Encountered quite large input images (>1.6K pixels width), rescaling to 1.6K.\n " 30 | "If this is not desired, please explicitly specify '--resolution/-r' as 1") 31 | WARNED = True 32 | global_down = orig_w / 1600 33 | else: 34 | global_down = 1 35 | else: 36 | global_down = orig_w / args.resolution 37 | 38 | scale = float(global_down) * float(resolution_scale) 39 | resolution = (int(orig_w / scale), int(orig_h / scale)) 40 | 41 | if len(cam_info.image.split()) > 3: 42 | import torch 43 | resized_image_rgb = torch.cat([PILtoTorch(im, resolution) for im in cam_info.image.split()[:3]], dim=0) 44 | loaded_mask = PILtoTorch(cam_info.image.split()[3], resolution) 45 | gt_image = resized_image_rgb 46 | else: 47 | resized_image_rgb = PILtoTorch(cam_info.image, resolution) 48 | loaded_mask = None 49 | gt_image = resized_image_rgb 50 | 51 | return Camera(colmap_id=cam_info.uid, R=cam_info.R, T=cam_info.T, 52 | FoVx=cam_info.FovX, FoVy=cam_info.FovY, 53 | image=gt_image, gt_alpha_mask=loaded_mask, 54 | image_name=cam_info.image_name, uid=id, data_device=args.data_device) 55 | 56 | def cameraList_from_camInfos(cam_infos, resolution_scale, args): 57 | camera_list = [] 58 | 59 | for id, c in enumerate(cam_infos): 60 | camera_list.append(loadCam(args, id, c, resolution_scale)) 61 | 62 | return camera_list 63 | 64 | def camera_to_JSON(id, camera : Camera): 65 | Rt = np.zeros((4, 4)) 66 | Rt[:3, :3] = camera.R.transpose() 67 | Rt[:3, 3] = camera.T 68 | Rt[3, 3] = 1.0 69 | 70 | W2C = np.linalg.inv(Rt) 71 | pos = W2C[:3, 3] 72 | rot = W2C[:3, :3] 73 | serializable_array_2d = [x.tolist() for x in rot] 74 | camera_entry = { 75 | 'id' : id, 76 | 'img_name' : camera.image_name, 77 | 'width' : camera.width, 78 | 'height' : camera.height, 79 | 'position': pos.tolist(), 80 | 'rotation': serializable_array_2d, 81 | 'fy' : fov2focal(camera.FovY, camera.height), 82 | 'fx' : fov2focal(camera.FovX, camera.width) 83 | } 84 | return camera_entry -------------------------------------------------------------------------------- /Trim2DGS/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)) -------------------------------------------------------------------------------- /Trim2DGS/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 | 14 | def mse(img1, img2): 15 | return (((img1 - img2)) ** 2).view(img1.shape[0], -1).mean(1, keepdim=True) 16 | 17 | def psnr(img1, img2): 18 | mse = (((img1 - img2)) ** 2).view(img1.shape[0], -1).mean(1, keepdim=True) 19 | return 20 * torch.log10(1.0 / torch.sqrt(mse)) 20 | -------------------------------------------------------------------------------- /Trim2DGS/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 l1_loss(network_output, gt): 18 | return torch.abs((network_output - gt)).mean() 19 | 20 | def l2_loss(network_output, gt): 21 | return ((network_output - gt) ** 2).mean() 22 | 23 | def gaussian(window_size, sigma): 24 | gauss = torch.Tensor([exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) 25 | return gauss / gauss.sum() 26 | 27 | 28 | def smooth_loss(disp, img): 29 | grad_disp_x = torch.abs(disp[:,1:-1, :-2] + disp[:,1:-1,2:] - 2 * disp[:,1:-1,1:-1]) 30 | grad_disp_y = torch.abs(disp[:,:-2, 1:-1] + disp[:,2:,1:-1] - 2 * disp[:,1:-1,1:-1]) 31 | grad_img_x = torch.mean(torch.abs(img[:, 1:-1, :-2] - img[:, 1:-1, 2:]), 0, keepdim=True) * 0.5 32 | grad_img_y = torch.mean(torch.abs(img[:, :-2, 1:-1] - img[:, 2:, 1:-1]), 0, keepdim=True) * 0.5 33 | grad_disp_x *= torch.exp(-grad_img_x) 34 | grad_disp_y *= torch.exp(-grad_img_y) 35 | return grad_disp_x.mean() + grad_disp_y.mean() 36 | 37 | def create_window(window_size, channel): 38 | _1D_window = gaussian(window_size, 1.5).unsqueeze(1) 39 | _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) 40 | window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) 41 | return window 42 | 43 | def ssim(img1, img2, window_size=11, size_average=True): 44 | channel = img1.size(-3) 45 | window = create_window(window_size, channel) 46 | 47 | if img1.is_cuda: 48 | window = window.cuda(img1.get_device()) 49 | window = window.type_as(img1) 50 | 51 | return _ssim(img1, img2, window, window_size, channel, size_average) 52 | 53 | def _ssim(img1, img2, window, window_size, channel, size_average=True): 54 | mu1 = F.conv2d(img1, window, padding=window_size // 2, groups=channel) 55 | mu2 = F.conv2d(img2, window, padding=window_size // 2, groups=channel) 56 | 57 | mu1_sq = mu1.pow(2) 58 | mu2_sq = mu2.pow(2) 59 | mu1_mu2 = mu1 * mu2 60 | 61 | sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size // 2, groups=channel) - mu1_sq 62 | sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size // 2, groups=channel) - mu2_sq 63 | sigma12 = F.conv2d(img1 * img2, window, padding=window_size // 2, groups=channel) - mu1_mu2 64 | 65 | C1 = 0.01 ** 2 66 | C2 = 0.03 ** 2 67 | 68 | ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) 69 | 70 | if size_average: 71 | return ssim_map.mean() 72 | else: 73 | return ssim_map.mean(1).mean(1).mean(1) 74 | 75 | -------------------------------------------------------------------------------- /Trim2DGS/utils/mcube_utils.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (C) 2024, ShanghaiTech 3 | # SVIP research group, https://github.com/svip-lab 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 huangbb@shanghaitech.edu.cn 10 | # 11 | 12 | import numpy as np 13 | import torch 14 | import trimesh 15 | from skimage import measure 16 | # modified from here https://github.com/autonomousvision/sdfstudio/blob/370902a10dbef08cb3fe4391bd3ed1e227b5c165/nerfstudio/utils/marching_cubes.py#L201 17 | def marching_cubes_with_contraction( 18 | sdf, 19 | resolution=512, 20 | bounding_box_min=(-1.0, -1.0, -1.0), 21 | bounding_box_max=(1.0, 1.0, 1.0), 22 | return_mesh=False, 23 | level=0, 24 | simplify_mesh=True, 25 | inv_contraction=None, 26 | max_range=32.0, 27 | ): 28 | assert resolution % 512 == 0 29 | 30 | resN = resolution 31 | cropN = 512 32 | level = 0 33 | N = resN // cropN 34 | 35 | grid_min = bounding_box_min 36 | grid_max = bounding_box_max 37 | xs = np.linspace(grid_min[0], grid_max[0], N + 1) 38 | ys = np.linspace(grid_min[1], grid_max[1], N + 1) 39 | zs = np.linspace(grid_min[2], grid_max[2], N + 1) 40 | 41 | meshes = [] 42 | for i in range(N): 43 | for j in range(N): 44 | for k in range(N): 45 | print(i, j, k) 46 | x_min, x_max = xs[i], xs[i + 1] 47 | y_min, y_max = ys[j], ys[j + 1] 48 | z_min, z_max = zs[k], zs[k + 1] 49 | 50 | x = torch.linspace(x_min, x_max, cropN).cuda() 51 | y = torch.linspace(y_min, y_max, cropN).cuda() 52 | z = torch.linspace(z_min, z_max, cropN).cuda() 53 | 54 | xx, yy, zz = torch.meshgrid(x, y, z, indexing="ij") 55 | points = torch.tensor(torch.vstack([xx.ravel(), yy.ravel(), zz.ravel()]).T, dtype=torch.float).cuda() 56 | 57 | @torch.no_grad() 58 | def evaluate(points): 59 | z = [] 60 | for _, pnts in enumerate(torch.split(points, 256**3, dim=0)): 61 | z.append(sdf(pnts)) 62 | z = torch.cat(z, axis=0) 63 | return z 64 | 65 | # construct point pyramids 66 | points = points.reshape(cropN, cropN, cropN, 3) 67 | points = points.reshape(-1, 3) 68 | pts_sdf = evaluate(points.contiguous()) 69 | z = pts_sdf.detach().cpu().numpy() 70 | if not (np.min(z) > level or np.max(z) < level): 71 | z = z.astype(np.float32) 72 | verts, faces, normals, _ = measure.marching_cubes( 73 | volume=z.reshape(cropN, cropN, cropN), 74 | level=level, 75 | spacing=( 76 | (x_max - x_min) / (cropN - 1), 77 | (y_max - y_min) / (cropN - 1), 78 | (z_max - z_min) / (cropN - 1), 79 | ), 80 | ) 81 | verts = verts + np.array([x_min, y_min, z_min]) 82 | meshcrop = trimesh.Trimesh(verts, faces, normals) 83 | meshes.append(meshcrop) 84 | 85 | print("finished one block") 86 | 87 | combined = trimesh.util.concatenate(meshes) 88 | combined.merge_vertices(digits_vertex=6) 89 | 90 | # inverse contraction and clipping the points range 91 | if inv_contraction is not None: 92 | combined.vertices = inv_contraction(torch.from_numpy(combined.vertices).float().cuda()).cpu().numpy() 93 | combined.vertices = np.clip(combined.vertices, -max_range, max_range) 94 | 95 | return combined -------------------------------------------------------------------------------- /Trim2DGS/utils/point_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | import numpy as np 5 | import os, cv2 6 | import matplotlib.pyplot as plt 7 | import math 8 | 9 | def depths_to_points(view, depthmap): 10 | c2w = (view.world_view_transform.T).inverse() 11 | W, H = view.image_width, view.image_height 12 | fx = W / (2 * math.tan(view.FoVx / 2.)) 13 | fy = H / (2 * math.tan(view.FoVy / 2.)) 14 | intrins = torch.tensor( 15 | [[fx, 0., W/2.], 16 | [0., fy, H/2.], 17 | [0., 0., 1.0]] 18 | ).float().cuda() 19 | grid_x, grid_y = torch.meshgrid(torch.arange(W, device='cuda').float(), torch.arange(H, device='cuda').float(), indexing='xy') 20 | points = torch.stack([grid_x, grid_y, torch.ones_like(grid_x)], dim=-1).reshape(-1, 3) 21 | rays_d = points @ intrins.inverse().T @ c2w[:3,:3].T 22 | rays_o = c2w[:3,3] 23 | points = depthmap.reshape(-1, 1) * rays_d + rays_o 24 | return points 25 | 26 | def depth_to_normal(view, depth): 27 | """ 28 | view: view camera 29 | depth: depthmap 30 | """ 31 | points = depths_to_points(view, depth).reshape(*depth.shape[1:], 3) 32 | output = torch.zeros_like(points) 33 | dx = torch.cat([points[2:, 1:-1] - points[:-2, 1:-1]], dim=0) 34 | dy = torch.cat([points[1:-1, 2:] - points[1:-1, :-2]], dim=1) 35 | normal_map = torch.nn.functional.normalize(torch.cross(dx, dy, dim=-1), dim=-1) 36 | output[1:-1, 1:-1, :] = normal_map 37 | return output -------------------------------------------------------------------------------- /Trim2DGS/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 -------------------------------------------------------------------------------- /Trim2DGS/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 | -------------------------------------------------------------------------------- /Trim3DGS/LICENSE.md: -------------------------------------------------------------------------------- 1 | Gaussian-Splatting License 2 | =========================== 3 | 4 | **Inria** and **the Max Planck Institut for Informatik (MPII)** hold all the ownership rights on the *Software* named **gaussian-splatting**. 5 | The *Software* is in the process of being registered with the Agence pour la Protection des 6 | Programmes (APP). 7 | 8 | The *Software* is still being developed by the *Licensor*. 9 | 10 | *Licensor*'s goal is to allow the research community to use, test and evaluate 11 | the *Software*. 12 | 13 | ## 1. Definitions 14 | 15 | *Licensee* means any person or entity that uses the *Software* and distributes 16 | its *Work*. 17 | 18 | *Licensor* means the owners of the *Software*, i.e Inria and MPII 19 | 20 | *Software* means the original work of authorship made available under this 21 | License ie gaussian-splatting. 22 | 23 | *Work* means the *Software* and any additions to or derivative works of the 24 | *Software* that are made available under this License. 25 | 26 | 27 | ## 2. Purpose 28 | This license is intended to define the rights granted to the *Licensee* by 29 | Licensors under the *Software*. 30 | 31 | ## 3. Rights granted 32 | 33 | For the above reasons Licensors have decided to distribute the *Software*. 34 | Licensors grant non-exclusive rights to use the *Software* for research purposes 35 | to research users (both academic and industrial), free of charge, without right 36 | to sublicense.. The *Software* may be used "non-commercially", i.e., for research 37 | and/or evaluation purposes only. 38 | 39 | Subject to the terms and conditions of this License, you are granted a 40 | non-exclusive, royalty-free, license to reproduce, prepare derivative works of, 41 | publicly display, publicly perform and distribute its *Work* and any resulting 42 | derivative works in any form. 43 | 44 | ## 4. Limitations 45 | 46 | **4.1 Redistribution.** You may reproduce or distribute the *Work* only if (a) you do 47 | so under this License, (b) you include a complete copy of this License with 48 | your distribution, and (c) you retain without modification any copyright, 49 | patent, trademark, or attribution notices that are present in the *Work*. 50 | 51 | **4.2 Derivative Works.** You may specify that additional or different terms apply 52 | to the use, reproduction, and distribution of your derivative works of the *Work* 53 | ("Your Terms") only if (a) Your Terms provide that the use limitation in 54 | Section 2 applies to your derivative works, and (b) you identify the specific 55 | derivative works that are subject to Your Terms. Notwithstanding Your Terms, 56 | this License (including the redistribution requirements in Section 3.1) will 57 | continue to apply to the *Work* itself. 58 | 59 | **4.3** Any other use without of prior consent of Licensors is prohibited. Research 60 | users explicitly acknowledge having received from Licensors all information 61 | allowing to appreciate the adequacy between of the *Software* and their needs and 62 | to undertake all necessary precautions for its execution and use. 63 | 64 | **4.4** The *Software* is provided both as a compiled library file and as source 65 | code. In case of using the *Software* for a publication or other results obtained 66 | through the use of the *Software*, users are strongly encouraged to cite the 67 | corresponding publications as explained in the documentation of the *Software*. 68 | 69 | ## 5. Disclaimer 70 | 71 | THE USER CANNOT USE, EXPLOIT OR DISTRIBUTE THE *SOFTWARE* FOR COMMERCIAL PURPOSES 72 | WITHOUT PRIOR AND EXPLICIT CONSENT OF LICENSORS. YOU MUST CONTACT INRIA FOR ANY 73 | UNAUTHORIZED USE: stip-sophia.transfert@inria.fr . ANY SUCH ACTION WILL 74 | CONSTITUTE A FORGERY. THIS *SOFTWARE* IS PROVIDED "AS IS" WITHOUT ANY WARRANTIES 75 | OF ANY NATURE AND ANY EXPRESS OR IMPLIED WARRANTIES, WITH REGARDS TO COMMERCIAL 76 | USE, PROFESSIONNAL USE, LEGAL OR NOT, OR OTHER, OR COMMERCIALISATION OR 77 | ADAPTATION. UNLESS EXPLICITLY PROVIDED BY LAW, IN NO EVENT, SHALL INRIA OR THE 78 | AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 79 | CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE 80 | GOODS OR SERVICES, LOSS OF USE, DATA, OR PROFITS OR BUSINESS INTERRUPTION) 81 | HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT 82 | LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING FROM, OUT OF OR 83 | IN CONNECTION WITH THE *SOFTWARE* OR THE USE OR OTHER DEALINGS IN THE *SOFTWARE*. 84 | 85 | ## 6. Files subject to permissive licenses 86 | The contents of the file ```utils/loss_utils.py``` are based on publicly available code authored by Evan Su, which falls under the permissive MIT license. 87 | 88 | Title: pytorch-ssim\ 89 | Project code: https://github.com/Po-Hsun-Su/pytorch-ssim\ 90 | Copyright Evan Su, 2017\ 91 | License: https://github.com/Po-Hsun-Su/pytorch-ssim/blob/master/LICENSE.txt (MIT) -------------------------------------------------------------------------------- /Trim3DGS/eval_dtu/evaluate_single_scene.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | import cv2 5 | import numpy as np 6 | import os 7 | import glob 8 | from skimage.morphology import binary_dilation, disk 9 | import argparse 10 | 11 | import trimesh 12 | from pathlib import Path 13 | import subprocess 14 | 15 | import sys 16 | import render_utils as rend_util 17 | from tqdm import tqdm 18 | 19 | def cull_scan(scan, mesh_path, result_mesh_file, instance_dir): 20 | 21 | # load poses 22 | image_dir = '{0}/images'.format(instance_dir) 23 | image_paths = sorted(glob.glob(os.path.join(image_dir, "*.png"))) 24 | n_images = len(image_paths) 25 | cam_file = '{0}/cameras.npz'.format(instance_dir) 26 | camera_dict = np.load(cam_file) 27 | scale_mats = [camera_dict['scale_mat_%d' % idx].astype(np.float32) for idx in range(n_images)] 28 | world_mats = [camera_dict['world_mat_%d' % idx].astype(np.float32) for idx in range(n_images)] 29 | 30 | intrinsics_all = [] 31 | pose_all = [] 32 | for scale_mat, world_mat in zip(scale_mats, world_mats): 33 | P = world_mat @ scale_mat 34 | P = P[:3, :4] 35 | intrinsics, pose = rend_util.load_K_Rt_from_P(None, P) 36 | intrinsics_all.append(torch.from_numpy(intrinsics).float()) 37 | pose_all.append(torch.from_numpy(pose).float()) 38 | 39 | # load mask 40 | mask_dir = '{0}/mask'.format(instance_dir) 41 | mask_paths = sorted(glob.glob(os.path.join(mask_dir, "*.png"))) 42 | masks = [] 43 | for p in mask_paths: 44 | mask = cv2.imread(p) 45 | masks.append(mask) 46 | 47 | # hard-coded image shape 48 | W, H = 1600, 1200 49 | 50 | # load mesh 51 | mesh = trimesh.load(mesh_path) 52 | 53 | # load transformation matrix 54 | 55 | vertices = mesh.vertices 56 | 57 | # project and filter 58 | vertices = torch.from_numpy(vertices).cuda() 59 | vertices = torch.cat((vertices, torch.ones_like(vertices[:, :1])), dim=-1) 60 | vertices = vertices.permute(1, 0) 61 | vertices = vertices.float() 62 | 63 | sampled_masks = [] 64 | for i in tqdm(range(n_images), desc="Culling mesh given masks"): 65 | pose = pose_all[i] 66 | w2c = torch.inverse(pose).cuda() 67 | intrinsic = intrinsics_all[i].cuda() 68 | 69 | with torch.no_grad(): 70 | # transform and project 71 | cam_points = intrinsic @ w2c @ vertices 72 | pix_coords = cam_points[:2, :] / (cam_points[2, :].unsqueeze(0) + 1e-6) 73 | pix_coords = pix_coords.permute(1, 0) 74 | pix_coords[..., 0] /= W - 1 75 | pix_coords[..., 1] /= H - 1 76 | pix_coords = (pix_coords - 0.5) * 2 77 | valid = ((pix_coords > -1. ) & (pix_coords < 1.)).all(dim=-1).float() 78 | 79 | # dialate mask similar to unisurf 80 | maski = masks[i][:, :, 0].astype(np.float32) / 256. 81 | maski = torch.from_numpy(binary_dilation(maski, disk(24))).float()[None, None].cuda() 82 | 83 | sampled_mask = F.grid_sample(maski, pix_coords[None, None], mode='nearest', padding_mode='zeros', align_corners=True)[0, -1, 0] 84 | 85 | sampled_mask = sampled_mask + (1. - valid) 86 | sampled_masks.append(sampled_mask) 87 | 88 | sampled_masks = torch.stack(sampled_masks, -1) 89 | # filter 90 | 91 | mask = (sampled_masks > 0.).all(dim=-1).cpu().numpy() 92 | face_mask = mask[mesh.faces].all(axis=1) 93 | 94 | mesh.update_vertices(mask) 95 | mesh.update_faces(face_mask) 96 | 97 | # transform vertices to world 98 | scale_mat = scale_mats[0] 99 | mesh.vertices = mesh.vertices * scale_mat[0, 0] + scale_mat[:3, 3][None] 100 | mesh.export(result_mesh_file) 101 | del mesh 102 | 103 | 104 | if __name__ == "__main__": 105 | 106 | parser = argparse.ArgumentParser( 107 | description='Arguments to evaluate the mesh.' 108 | ) 109 | 110 | parser.add_argument('--input_mesh', type=str, help='path to the mesh to be evaluated') 111 | parser.add_argument('--scan_id', type=str, help='scan id of the input mesh') 112 | parser.add_argument('--output_dir', type=str, default='evaluation_results_single', help='path to the output folder') 113 | parser.add_argument('--mask_dir', type=str, default='mask', help='path to uncropped mask') 114 | parser.add_argument('--DTU', type=str, default='Official_DTU_Dataset', help='path to the GT DTU point clouds') 115 | args = parser.parse_args() 116 | 117 | Official_DTU_Dataset = args.DTU 118 | out_dir = args.output_dir 119 | Path(out_dir).mkdir(parents=True, exist_ok=True) 120 | 121 | scan = args.scan_id 122 | ply_file = args.input_mesh 123 | print("cull mesh ....") 124 | result_mesh_file = os.path.join(out_dir, "culled_mesh.ply") 125 | cull_scan(scan, ply_file, result_mesh_file, instance_dir=os.path.join(args.mask_dir, f'scan{args.scan_id}')) 126 | 127 | script_dir = os.path.dirname(os.path.abspath(__file__)) 128 | cmd = f"python {script_dir}/eval.py --data {result_mesh_file} --scan {scan} --mode mesh --dataset_dir {Official_DTU_Dataset} --vis_out_dir {out_dir}" 129 | os.system(cmd) -------------------------------------------------------------------------------- /Trim3DGS/eval_dtu/render_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import imageio 3 | import skimage 4 | import cv2 5 | import torch 6 | from torch.nn import functional as F 7 | 8 | 9 | def get_psnr(img1, img2, normalize_rgb=False): 10 | if normalize_rgb: # [-1,1] --> [0,1] 11 | img1 = (img1 + 1.) / 2. 12 | img2 = (img2 + 1. ) / 2. 13 | 14 | mse = torch.mean((img1 - img2) ** 2) 15 | psnr = -10. * torch.log(mse) / torch.log(torch.Tensor([10.]).cuda()) 16 | 17 | return psnr 18 | 19 | 20 | def load_rgb(path, normalize_rgb = False): 21 | img = imageio.imread(path) 22 | img = skimage.img_as_float32(img) 23 | 24 | if normalize_rgb: # [-1,1] --> [0,1] 25 | img -= 0.5 26 | img *= 2. 27 | img = img.transpose(2, 0, 1) 28 | return img 29 | 30 | 31 | def load_K_Rt_from_P(filename, P=None): 32 | if P is None: 33 | lines = open(filename).read().splitlines() 34 | if len(lines) == 4: 35 | lines = lines[1:] 36 | lines = [[x[0], x[1], x[2], x[3]] for x in (x.split(" ") for x in lines)] 37 | P = np.asarray(lines).astype(np.float32).squeeze() 38 | 39 | out = cv2.decomposeProjectionMatrix(P) 40 | K = out[0] 41 | R = out[1] 42 | t = out[2] 43 | 44 | K = K/K[2,2] 45 | intrinsics = np.eye(4) 46 | intrinsics[:3, :3] = K 47 | 48 | pose = np.eye(4, dtype=np.float32) 49 | pose[:3, :3] = R.transpose() 50 | pose[:3,3] = (t[:3] / t[3])[:,0] 51 | 52 | return intrinsics, pose 53 | 54 | 55 | def get_camera_params(uv, pose, intrinsics): 56 | if pose.shape[1] == 7: #In case of quaternion vector representation 57 | cam_loc = pose[:, 4:] 58 | R = quat_to_rot(pose[:,:4]) 59 | p = torch.eye(4).repeat(pose.shape[0],1,1).cuda().float() 60 | p[:, :3, :3] = R 61 | p[:, :3, 3] = cam_loc 62 | else: # In case of pose matrix representation 63 | cam_loc = pose[:, :3, 3] 64 | p = pose 65 | 66 | batch_size, num_samples, _ = uv.shape 67 | 68 | depth = torch.ones((batch_size, num_samples)).cuda() 69 | x_cam = uv[:, :, 0].view(batch_size, -1) 70 | y_cam = uv[:, :, 1].view(batch_size, -1) 71 | z_cam = depth.view(batch_size, -1) 72 | 73 | pixel_points_cam = lift(x_cam, y_cam, z_cam, intrinsics=intrinsics) 74 | 75 | # permute for batch matrix product 76 | pixel_points_cam = pixel_points_cam.permute(0, 2, 1) 77 | 78 | world_coords = torch.bmm(p, pixel_points_cam).permute(0, 2, 1)[:, :, :3] 79 | ray_dirs = world_coords - cam_loc[:, None, :] 80 | ray_dirs = F.normalize(ray_dirs, dim=2) 81 | 82 | return ray_dirs, cam_loc 83 | 84 | 85 | def get_camera_for_plot(pose): 86 | if pose.shape[1] == 7: #In case of quaternion vector representation 87 | cam_loc = pose[:, 4:].detach() 88 | R = quat_to_rot(pose[:,:4].detach()) 89 | else: # In case of pose matrix representation 90 | cam_loc = pose[:, :3, 3] 91 | R = pose[:, :3, :3] 92 | cam_dir = R[:, :3, 2] 93 | return cam_loc, cam_dir 94 | 95 | 96 | def lift(x, y, z, intrinsics): 97 | # parse intrinsics 98 | intrinsics = intrinsics.cuda() 99 | fx = intrinsics[:, 0, 0] 100 | fy = intrinsics[:, 1, 1] 101 | cx = intrinsics[:, 0, 2] 102 | cy = intrinsics[:, 1, 2] 103 | sk = intrinsics[:, 0, 1] 104 | 105 | x_lift = (x - cx.unsqueeze(-1) + cy.unsqueeze(-1)*sk.unsqueeze(-1)/fy.unsqueeze(-1) - sk.unsqueeze(-1)*y/fy.unsqueeze(-1)) / fx.unsqueeze(-1) * z 106 | y_lift = (y - cy.unsqueeze(-1)) / fy.unsqueeze(-1) * z 107 | 108 | # homogeneous 109 | return torch.stack((x_lift, y_lift, z, torch.ones_like(z).cuda()), dim=-1) 110 | 111 | 112 | def quat_to_rot(q): 113 | batch_size, _ = q.shape 114 | q = F.normalize(q, dim=1) 115 | R = torch.ones((batch_size, 3,3)).cuda() 116 | qr=q[:,0] 117 | qi = q[:, 1] 118 | qj = q[:, 2] 119 | qk = q[:, 3] 120 | R[:, 0, 0]=1-2 * (qj**2 + qk**2) 121 | R[:, 0, 1] = 2 * (qj *qi -qk*qr) 122 | R[:, 0, 2] = 2 * (qi * qk + qr * qj) 123 | R[:, 1, 0] = 2 * (qj * qi + qk * qr) 124 | R[:, 1, 1] = 1-2 * (qi**2 + qk**2) 125 | R[:, 1, 2] = 2*(qj*qk - qi*qr) 126 | R[:, 2, 0] = 2 * (qk * qi-qj * qr) 127 | R[:, 2, 1] = 2 * (qj*qk + qi*qr) 128 | R[:, 2, 2] = 1-2 * (qi**2 + qj**2) 129 | return R 130 | 131 | 132 | def rot_to_quat(R): 133 | batch_size, _,_ = R.shape 134 | q = torch.ones((batch_size, 4)).cuda() 135 | 136 | R00 = R[:, 0,0] 137 | R01 = R[:, 0, 1] 138 | R02 = R[:, 0, 2] 139 | R10 = R[:, 1, 0] 140 | R11 = R[:, 1, 1] 141 | R12 = R[:, 1, 2] 142 | R20 = R[:, 2, 0] 143 | R21 = R[:, 2, 1] 144 | R22 = R[:, 2, 2] 145 | 146 | q[:,0]=torch.sqrt(1.0+R00+R11+R22)/2 147 | q[:, 1]=(R21-R12)/(4*q[:,0]) 148 | q[:, 2] = (R02 - R20) / (4 * q[:, 0]) 149 | q[:, 3] = (R10 - R01) / (4 * q[:, 0]) 150 | return q 151 | 152 | 153 | def get_sphere_intersections(cam_loc, ray_directions, r = 1.0): 154 | # Input: n_rays x 3 ; n_rays x 3 155 | # Output: n_rays x 1, n_rays x 1 (close and far) 156 | 157 | ray_cam_dot = torch.bmm(ray_directions.view(-1, 1, 3), 158 | cam_loc.view(-1, 3, 1)).squeeze(-1) 159 | under_sqrt = ray_cam_dot ** 2 - (cam_loc.norm(2, 1, keepdim=True) ** 2 - r ** 2) 160 | 161 | # sanity check 162 | if (under_sqrt <= 0).sum() > 0: 163 | print('BOUNDING SPHERE PROBLEM!') 164 | exit() 165 | 166 | sphere_intersections = torch.sqrt(under_sqrt) * torch.Tensor([-1, 1]).cuda().float() - ray_cam_dot 167 | sphere_intersections = sphere_intersections.clamp_min(0.0) 168 | 169 | return sphere_intersections -------------------------------------------------------------------------------- /Trim3DGS/eval_dtu_pcd/evaluate_single_scene.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | import cv2 5 | import numpy as np 6 | import os 7 | import glob 8 | from skimage.morphology import binary_dilation, disk 9 | import argparse 10 | 11 | import open3d as o3d 12 | from pathlib import Path 13 | import subprocess 14 | 15 | import sys 16 | import render_utils as rend_util 17 | from tqdm import tqdm 18 | 19 | def cull_scan(scan, ply_path, result_ply_file, instance_dir): 20 | 21 | # load poses 22 | image_dir = '{0}/images'.format(instance_dir) 23 | image_paths = sorted(glob.glob(os.path.join(image_dir, "*.png"))) 24 | n_images = len(image_paths) 25 | cam_file = '{0}/cameras.npz'.format(instance_dir) 26 | camera_dict = np.load(cam_file) 27 | scale_mats = [camera_dict['scale_mat_%d' % idx].astype(np.float32) for idx in range(n_images)] 28 | world_mats = [camera_dict['world_mat_%d' % idx].astype(np.float32) for idx in range(n_images)] 29 | 30 | intrinsics_all = [] 31 | pose_all = [] 32 | for scale_mat, world_mat in zip(scale_mats, world_mats): 33 | P = world_mat @ scale_mat 34 | P = P[:3, :4] 35 | intrinsics, pose = rend_util.load_K_Rt_from_P(None, P) 36 | intrinsics_all.append(torch.from_numpy(intrinsics).float()) 37 | pose_all.append(torch.from_numpy(pose).float()) 38 | 39 | # load mask 40 | mask_dir = '{0}/mask'.format(instance_dir) 41 | mask_paths = sorted(glob.glob(os.path.join(mask_dir, "*.png"))) 42 | masks = [] 43 | for p in mask_paths: 44 | mask = cv2.imread(p) 45 | masks.append(mask) 46 | 47 | # hard-coded image shape 48 | W, H = 1600, 1200 49 | 50 | # load pcd 51 | pcd = o3d.io.read_point_cloud(ply_path) 52 | # load transformation matrix 53 | 54 | vertices = np.asarray(pcd.points) 55 | 56 | # project and filter 57 | vertices = torch.from_numpy(vertices).cuda() 58 | vertices = torch.cat((vertices, torch.ones_like(vertices[:, :1])), dim=-1) 59 | vertices = vertices.permute(1, 0) 60 | vertices = vertices.float() 61 | 62 | sampled_masks = [] 63 | for i in tqdm(range(n_images), desc="Culling point clouds given masks"): 64 | pose = pose_all[i] 65 | w2c = torch.inverse(pose).cuda() 66 | intrinsic = intrinsics_all[i].cuda() 67 | 68 | with torch.no_grad(): 69 | # transform and project 70 | cam_points = intrinsic @ w2c @ vertices 71 | pix_coords = cam_points[:2, :] / (cam_points[2, :].unsqueeze(0) + 1e-6) 72 | pix_coords = pix_coords.permute(1, 0) 73 | pix_coords[..., 0] /= W - 1 74 | pix_coords[..., 1] /= H - 1 75 | pix_coords = (pix_coords - 0.5) * 2 76 | valid = ((pix_coords > -1. ) & (pix_coords < 1.)).all(dim=-1).float() 77 | 78 | # dialate mask similar to unisurf 79 | maski = masks[i][:, :, 0].astype(np.float32) / 256. 80 | maski = torch.from_numpy(binary_dilation(maski, disk(24))).float()[None, None].cuda() 81 | 82 | sampled_mask = F.grid_sample(maski, pix_coords[None, None], mode='nearest', padding_mode='zeros', align_corners=True)[0, -1, 0] 83 | 84 | sampled_mask = sampled_mask + (1. - valid) 85 | sampled_masks.append(sampled_mask) 86 | 87 | sampled_masks = torch.stack(sampled_masks, -1) 88 | # filter 89 | 90 | mask = (sampled_masks > 0.).all(dim=-1).cpu().numpy() 91 | 92 | points = np.asarray(pcd.points) 93 | points = points[mask] 94 | # transform vertices to world 95 | scale_mat = scale_mats[0] 96 | points = points * scale_mat[0, 0] + scale_mat[:3, 3][None] 97 | 98 | pcd = o3d.geometry.PointCloud() 99 | pcd.points = o3d.utility.Vector3dVector(points) 100 | o3d.io.write_point_cloud(result_ply_file, pcd) 101 | 102 | 103 | if __name__ == "__main__": 104 | 105 | parser = argparse.ArgumentParser( 106 | description='Arguments to evaluate the point clouds.' 107 | ) 108 | 109 | parser.add_argument('--input_pcd', type=str, help='path to the point clouds to be evaluated') 110 | parser.add_argument('--scan_id', type=str, help='scan id of the input point clouds') 111 | parser.add_argument('--output_dir', type=str, default='evaluation_results_single', help='path to the output folder') 112 | parser.add_argument('--mask_dir', type=str, default='mask', help='path to uncropped mask') 113 | parser.add_argument('--DTU', type=str, default='Offical_DTU_Dataset', help='path to the GT DTU point clouds') 114 | args = parser.parse_args() 115 | 116 | Offical_DTU_Dataset = args.DTU 117 | out_dir = args.output_dir 118 | Path(out_dir).mkdir(parents=True, exist_ok=True) 119 | 120 | scan = args.scan_id 121 | ply_file = args.input_pcd 122 | print("cull pcd ....") 123 | result_ply_file = os.path.join(out_dir, "culled_pcd.ply") 124 | cull_scan(scan, ply_file, result_ply_file, instance_dir=os.path.join(args.mask_dir, f'scan{args.scan_id}')) 125 | 126 | script_dir = os.path.dirname(os.path.abspath(__file__)) 127 | cmd = f"python {script_dir}/eval.py --data {result_ply_file} --scan {scan} --dataset_dir {Offical_DTU_Dataset} --vis_out_dir {out_dir}" 128 | os.system(cmd) -------------------------------------------------------------------------------- /Trim3DGS/eval_dtu_pcd/render_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import imageio 3 | import skimage 4 | import cv2 5 | import torch 6 | from torch.nn import functional as F 7 | 8 | 9 | def get_psnr(img1, img2, normalize_rgb=False): 10 | if normalize_rgb: # [-1,1] --> [0,1] 11 | img1 = (img1 + 1.) / 2. 12 | img2 = (img2 + 1. ) / 2. 13 | 14 | mse = torch.mean((img1 - img2) ** 2) 15 | psnr = -10. * torch.log(mse) / torch.log(torch.Tensor([10.]).cuda()) 16 | 17 | return psnr 18 | 19 | 20 | def load_rgb(path, normalize_rgb = False): 21 | img = imageio.imread(path) 22 | img = skimage.img_as_float32(img) 23 | 24 | if normalize_rgb: # [-1,1] --> [0,1] 25 | img -= 0.5 26 | img *= 2. 27 | img = img.transpose(2, 0, 1) 28 | return img 29 | 30 | 31 | def load_K_Rt_from_P(filename, P=None): 32 | if P is None: 33 | lines = open(filename).read().splitlines() 34 | if len(lines) == 4: 35 | lines = lines[1:] 36 | lines = [[x[0], x[1], x[2], x[3]] for x in (x.split(" ") for x in lines)] 37 | P = np.asarray(lines).astype(np.float32).squeeze() 38 | 39 | out = cv2.decomposeProjectionMatrix(P) 40 | K = out[0] 41 | R = out[1] 42 | t = out[2] 43 | 44 | K = K/K[2,2] 45 | intrinsics = np.eye(4) 46 | intrinsics[:3, :3] = K 47 | 48 | pose = np.eye(4, dtype=np.float32) 49 | pose[:3, :3] = R.transpose() 50 | pose[:3,3] = (t[:3] / t[3])[:,0] 51 | 52 | return intrinsics, pose 53 | 54 | 55 | def get_camera_params(uv, pose, intrinsics): 56 | if pose.shape[1] == 7: #In case of quaternion vector representation 57 | cam_loc = pose[:, 4:] 58 | R = quat_to_rot(pose[:,:4]) 59 | p = torch.eye(4).repeat(pose.shape[0],1,1).cuda().float() 60 | p[:, :3, :3] = R 61 | p[:, :3, 3] = cam_loc 62 | else: # In case of pose matrix representation 63 | cam_loc = pose[:, :3, 3] 64 | p = pose 65 | 66 | batch_size, num_samples, _ = uv.shape 67 | 68 | depth = torch.ones((batch_size, num_samples)).cuda() 69 | x_cam = uv[:, :, 0].view(batch_size, -1) 70 | y_cam = uv[:, :, 1].view(batch_size, -1) 71 | z_cam = depth.view(batch_size, -1) 72 | 73 | pixel_points_cam = lift(x_cam, y_cam, z_cam, intrinsics=intrinsics) 74 | 75 | # permute for batch matrix product 76 | pixel_points_cam = pixel_points_cam.permute(0, 2, 1) 77 | 78 | world_coords = torch.bmm(p, pixel_points_cam).permute(0, 2, 1)[:, :, :3] 79 | ray_dirs = world_coords - cam_loc[:, None, :] 80 | ray_dirs = F.normalize(ray_dirs, dim=2) 81 | 82 | return ray_dirs, cam_loc 83 | 84 | 85 | def get_camera_for_plot(pose): 86 | if pose.shape[1] == 7: #In case of quaternion vector representation 87 | cam_loc = pose[:, 4:].detach() 88 | R = quat_to_rot(pose[:,:4].detach()) 89 | else: # In case of pose matrix representation 90 | cam_loc = pose[:, :3, 3] 91 | R = pose[:, :3, :3] 92 | cam_dir = R[:, :3, 2] 93 | return cam_loc, cam_dir 94 | 95 | 96 | def lift(x, y, z, intrinsics): 97 | # parse intrinsics 98 | intrinsics = intrinsics.cuda() 99 | fx = intrinsics[:, 0, 0] 100 | fy = intrinsics[:, 1, 1] 101 | cx = intrinsics[:, 0, 2] 102 | cy = intrinsics[:, 1, 2] 103 | sk = intrinsics[:, 0, 1] 104 | 105 | x_lift = (x - cx.unsqueeze(-1) + cy.unsqueeze(-1)*sk.unsqueeze(-1)/fy.unsqueeze(-1) - sk.unsqueeze(-1)*y/fy.unsqueeze(-1)) / fx.unsqueeze(-1) * z 106 | y_lift = (y - cy.unsqueeze(-1)) / fy.unsqueeze(-1) * z 107 | 108 | # homogeneous 109 | return torch.stack((x_lift, y_lift, z, torch.ones_like(z).cuda()), dim=-1) 110 | 111 | 112 | def quat_to_rot(q): 113 | batch_size, _ = q.shape 114 | q = F.normalize(q, dim=1) 115 | R = torch.ones((batch_size, 3,3)).cuda() 116 | qr=q[:,0] 117 | qi = q[:, 1] 118 | qj = q[:, 2] 119 | qk = q[:, 3] 120 | R[:, 0, 0]=1-2 * (qj**2 + qk**2) 121 | R[:, 0, 1] = 2 * (qj *qi -qk*qr) 122 | R[:, 0, 2] = 2 * (qi * qk + qr * qj) 123 | R[:, 1, 0] = 2 * (qj * qi + qk * qr) 124 | R[:, 1, 1] = 1-2 * (qi**2 + qk**2) 125 | R[:, 1, 2] = 2*(qj*qk - qi*qr) 126 | R[:, 2, 0] = 2 * (qk * qi-qj * qr) 127 | R[:, 2, 1] = 2 * (qj*qk + qi*qr) 128 | R[:, 2, 2] = 1-2 * (qi**2 + qj**2) 129 | return R 130 | 131 | 132 | def rot_to_quat(R): 133 | batch_size, _,_ = R.shape 134 | q = torch.ones((batch_size, 4)).cuda() 135 | 136 | R00 = R[:, 0,0] 137 | R01 = R[:, 0, 1] 138 | R02 = R[:, 0, 2] 139 | R10 = R[:, 1, 0] 140 | R11 = R[:, 1, 1] 141 | R12 = R[:, 1, 2] 142 | R20 = R[:, 2, 0] 143 | R21 = R[:, 2, 1] 144 | R22 = R[:, 2, 2] 145 | 146 | q[:,0]=torch.sqrt(1.0+R00+R11+R22)/2 147 | q[:, 1]=(R21-R12)/(4*q[:,0]) 148 | q[:, 2] = (R02 - R20) / (4 * q[:, 0]) 149 | q[:, 3] = (R10 - R01) / (4 * q[:, 0]) 150 | return q 151 | 152 | 153 | def get_sphere_intersections(cam_loc, ray_directions, r = 1.0): 154 | # Input: n_rays x 3 ; n_rays x 3 155 | # Output: n_rays x 1, n_rays x 1 (close and far) 156 | 157 | ray_cam_dot = torch.bmm(ray_directions.view(-1, 1, 3), 158 | cam_loc.view(-1, 3, 1)).squeeze(-1) 159 | under_sqrt = ray_cam_dot ** 2 - (cam_loc.norm(2, 1, keepdim=True) ** 2 - r ** 2) 160 | 161 | # sanity check 162 | if (under_sqrt <= 0).sum() > 0: 163 | print('BOUNDING SPHERE PROBLEM!') 164 | exit() 165 | 166 | sphere_intersections = torch.sqrt(under_sqrt) * torch.Tensor([-1, 1]).cuda().float() - ray_cam_dot 167 | sphere_intersections = sphere_intersections.clamp_min(0.0) 168 | 169 | return sphere_intersections -------------------------------------------------------------------------------- /Trim3DGS/full_eval.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 | from argparse import ArgumentParser 14 | 15 | mipnerf360_outdoor_scenes = ["bicycle", "flowers", "garden", "stump", "treehill"] 16 | mipnerf360_indoor_scenes = ["room", "counter", "kitchen", "bonsai"] 17 | tanks_and_temples_scenes = ["truck", "train"] 18 | deep_blending_scenes = ["drjohnson", "playroom"] 19 | 20 | parser = ArgumentParser(description="Full evaluation script parameters") 21 | parser.add_argument("--skip_training", action="store_true") 22 | parser.add_argument("--skip_rendering", action="store_true") 23 | parser.add_argument("--skip_metrics", action="store_true") 24 | parser.add_argument("--output_path", default="./eval") 25 | args, _ = parser.parse_known_args() 26 | 27 | all_scenes = [] 28 | all_scenes.extend(mipnerf360_outdoor_scenes) 29 | all_scenes.extend(mipnerf360_indoor_scenes) 30 | all_scenes.extend(tanks_and_temples_scenes) 31 | all_scenes.extend(deep_blending_scenes) 32 | 33 | if not args.skip_training or not args.skip_rendering: 34 | parser.add_argument('--mipnerf360', "-m360", required=True, type=str) 35 | parser.add_argument("--tanksandtemples", "-tat", required=True, type=str) 36 | parser.add_argument("--deepblending", "-db", required=True, type=str) 37 | args = parser.parse_args() 38 | 39 | if not args.skip_training: 40 | common_args = " --quiet --eval --test_iterations -1 " 41 | for scene in mipnerf360_outdoor_scenes: 42 | source = args.mipnerf360 + "/" + scene 43 | os.system("python train.py -s " + source + " -i images_4 -m " + args.output_path + "/" + scene + common_args) 44 | for scene in mipnerf360_indoor_scenes: 45 | source = args.mipnerf360 + "/" + scene 46 | os.system("python train.py -s " + source + " -i images_2 -m " + args.output_path + "/" + scene + common_args) 47 | for scene in tanks_and_temples_scenes: 48 | source = args.tanksandtemples + "/" + scene 49 | os.system("python train.py -s " + source + " -m " + args.output_path + "/" + scene + common_args) 50 | for scene in deep_blending_scenes: 51 | source = args.deepblending + "/" + scene 52 | os.system("python train.py -s " + source + " -m " + args.output_path + "/" + scene + common_args) 53 | 54 | if not args.skip_rendering: 55 | all_sources = [] 56 | for scene in mipnerf360_outdoor_scenes: 57 | all_sources.append(args.mipnerf360 + "/" + scene) 58 | for scene in mipnerf360_indoor_scenes: 59 | all_sources.append(args.mipnerf360 + "/" + scene) 60 | for scene in tanks_and_temples_scenes: 61 | all_sources.append(args.tanksandtemples + "/" + scene) 62 | for scene in deep_blending_scenes: 63 | all_sources.append(args.deepblending + "/" + scene) 64 | 65 | common_args = " --quiet --eval --skip_train" 66 | for scene, source in zip(all_scenes, all_sources): 67 | os.system("python render.py --iteration 7000 -s " + source + " -m " + args.output_path + "/" + scene + common_args) 68 | os.system("python render.py --iteration 30000 -s " + source + " -m " + args.output_path + "/" + scene + common_args) 69 | 70 | if not args.skip_metrics: 71 | scenes_string = "" 72 | for scene in all_scenes: 73 | scenes_string += "\"" + args.output_path + "/" + scene + "\" " 74 | 75 | os.system("python metrics.py -m " + scenes_string) -------------------------------------------------------------------------------- /Trim3DGS/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, record_transmittance=False, extra_feats=None): 19 | """ 20 | Render the scene. 21 | 22 | Background tensor (bg_color) must be on GPU! 23 | """ 24 | 25 | # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means 26 | screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0 27 | try: 28 | screenspace_points.retain_grad() 29 | except: 30 | pass 31 | 32 | # Set up rasterization configuration 33 | tanfovx = math.tan(viewpoint_camera.FoVx * 0.5) 34 | tanfovy = math.tan(viewpoint_camera.FoVy * 0.5) 35 | 36 | raster_settings = GaussianRasterizationSettings( 37 | image_height=int(viewpoint_camera.image_height), 38 | image_width=int(viewpoint_camera.image_width), 39 | tanfovx=tanfovx, 40 | tanfovy=tanfovy, 41 | bg=bg_color, 42 | scale_modifier=scaling_modifier, 43 | viewmatrix=viewpoint_camera.world_view_transform, 44 | projmatrix=viewpoint_camera.full_proj_transform, 45 | sh_degree=pc.active_sh_degree, 46 | campos=viewpoint_camera.camera_center, 47 | prefiltered=False, 48 | record_transmittance=record_transmittance, 49 | debug=pipe.debug 50 | ) 51 | 52 | rasterizer = GaussianRasterizer(raster_settings=raster_settings) 53 | 54 | means3D = pc.get_xyz 55 | means2D = screenspace_points 56 | opacity = pc.get_opacity 57 | 58 | # If precomputed 3d covariance is provided, use it. If not, then it will be computed from 59 | # scaling / rotation by the rasterizer. 60 | scales = None 61 | rotations = None 62 | cov3D_precomp = None 63 | if pipe.compute_cov3D_python: 64 | cov3D_precomp = pc.get_covariance(scaling_modifier) 65 | else: 66 | scales = pc.get_scaling 67 | rotations = pc.get_rotation 68 | 69 | # If precomputed colors are provided, use them. Otherwise, if it is desired to precompute colors 70 | # from SHs in Python, do it. If not, then SH -> RGB conversion will be done by rasterizer. 71 | shs = None 72 | colors_precomp = None 73 | if override_color is None: 74 | if pipe.convert_SHs_python: 75 | shs_view = pc.get_features.transpose(1, 2).view(-1, 3, (pc.max_sh_degree+1)**2) 76 | dir_pp = (pc.get_xyz - viewpoint_camera.camera_center.repeat(pc.get_features.shape[0], 1)) 77 | dir_pp_normalized = dir_pp/dir_pp.norm(dim=1, keepdim=True) 78 | sh2rgb = eval_sh(pc.active_sh_degree, shs_view, dir_pp_normalized) 79 | colors_precomp = torch.clamp_min(sh2rgb + 0.5, 0.0) 80 | else: 81 | shs = pc.get_features 82 | else: 83 | colors_precomp = override_color 84 | 85 | # Rasterize visible Gaussians to image, obtain their radii (on screen). 86 | if extra_feats is not None and extra_feats.ndim == 1: 87 | extra_feats = extra_feats.unsqueeze(1) 88 | output = rasterizer( 89 | means3D = means3D, 90 | means2D = means2D, 91 | shs = shs, 92 | colors_precomp = colors_precomp, 93 | extra_feats = extra_feats, 94 | opacities = opacity, 95 | scales = scales, 96 | rotations = rotations, 97 | cov3D_precomp = cov3D_precomp) 98 | 99 | if record_transmittance: 100 | transmittance_sum, num_covered_pixels, radii = output 101 | transmittance = transmittance_sum / (num_covered_pixels + 1e-6) 102 | return transmittance 103 | else: 104 | rendered_image, rendered_extra_feats, median_depth, radii = output 105 | # Those Gaussians that were frustum culled or had a radius of 0 were not visible. 106 | # They will be excluded from value updates used in the splitting criteria. 107 | return {"render": rendered_image, 108 | "extra_feats": rendered_extra_feats, 109 | "viewspace_points": screenspace_points, 110 | "visibility_filter" : radii > 0, 111 | "radii": radii, 112 | "median_depth": median_depth} 113 | -------------------------------------------------------------------------------- /Trim3DGS/gaussian_renderer/network_gui.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 traceback 14 | import socket 15 | import json 16 | from scene.cameras import MiniCam 17 | 18 | host = "127.0.0.1" 19 | port = 6009 20 | 21 | conn = None 22 | addr = None 23 | 24 | listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 25 | 26 | def init(wish_host, wish_port): 27 | global host, port, listener 28 | host = wish_host 29 | port = wish_port 30 | listener.bind((host, port)) 31 | listener.listen() 32 | listener.settimeout(0) 33 | 34 | def try_connect(): 35 | global conn, addr, listener 36 | try: 37 | conn, addr = listener.accept() 38 | print(f"\nConnected by {addr}") 39 | conn.settimeout(None) 40 | except Exception as inst: 41 | pass 42 | 43 | def read(): 44 | global conn 45 | messageLength = conn.recv(4) 46 | messageLength = int.from_bytes(messageLength, 'little') 47 | message = conn.recv(messageLength) 48 | return json.loads(message.decode("utf-8")) 49 | 50 | def send(message_bytes, verify): 51 | global conn 52 | if message_bytes != None: 53 | conn.sendall(message_bytes) 54 | conn.sendall(len(verify).to_bytes(4, 'little')) 55 | conn.sendall(bytes(verify, 'ascii')) 56 | 57 | def receive(): 58 | message = read() 59 | 60 | width = message["resolution_x"] 61 | height = message["resolution_y"] 62 | 63 | if width != 0 and height != 0: 64 | try: 65 | do_training = bool(message["train"]) 66 | fovy = message["fov_y"] 67 | fovx = message["fov_x"] 68 | znear = message["z_near"] 69 | zfar = message["z_far"] 70 | do_shs_python = bool(message["shs_python"]) 71 | do_rot_scale_python = bool(message["rot_scale_python"]) 72 | keep_alive = bool(message["keep_alive"]) 73 | scaling_modifier = message["scaling_modifier"] 74 | world_view_transform = torch.reshape(torch.tensor(message["view_matrix"]), (4, 4)).cuda() 75 | world_view_transform[:,1] = -world_view_transform[:,1] 76 | world_view_transform[:,2] = -world_view_transform[:,2] 77 | full_proj_transform = torch.reshape(torch.tensor(message["view_projection_matrix"]), (4, 4)).cuda() 78 | full_proj_transform[:,1] = -full_proj_transform[:,1] 79 | custom_cam = MiniCam(width, height, fovy, fovx, znear, zfar, world_view_transform, full_proj_transform) 80 | except Exception as e: 81 | print("") 82 | traceback.print_exc() 83 | raise e 84 | return custom_cam, do_training, do_shs_python, do_rot_scale_python, keep_alive, scaling_modifier 85 | else: 86 | return None, None, None, None, None, None -------------------------------------------------------------------------------- /Trim3DGS/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 | -------------------------------------------------------------------------------- /Trim3DGS/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 | -------------------------------------------------------------------------------- /Trim3DGS/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 | -------------------------------------------------------------------------------- /Trim3DGS/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 | -------------------------------------------------------------------------------- /Trim3DGS/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 torch 16 | import torchvision.transforms.functional as tf 17 | from utils.loss_utils import ssim 18 | from lpipsPyTorch import lpips 19 | import json 20 | from tqdm import tqdm 21 | from utils.image_utils import psnr 22 | from argparse import ArgumentParser 23 | 24 | def readImages(renders_dir, gt_dir): 25 | renders = [] 26 | gts = [] 27 | image_names = [] 28 | for fname in os.listdir(renders_dir): 29 | render = Image.open(renders_dir / fname) 30 | gt = Image.open(gt_dir / fname) 31 | renders.append(tf.to_tensor(render).unsqueeze(0)[:, :3, :, :].cuda()) 32 | gts.append(tf.to_tensor(gt).unsqueeze(0)[:, :3, :, :].cuda()) 33 | image_names.append(fname) 34 | return renders, gts, image_names 35 | 36 | def evaluate(model_paths): 37 | 38 | full_dict = {} 39 | per_view_dict = {} 40 | full_dict_polytopeonly = {} 41 | per_view_dict_polytopeonly = {} 42 | print("") 43 | 44 | for scene_dir in model_paths: 45 | try: 46 | print("Scene:", scene_dir) 47 | full_dict[scene_dir] = {} 48 | per_view_dict[scene_dir] = {} 49 | full_dict_polytopeonly[scene_dir] = {} 50 | per_view_dict_polytopeonly[scene_dir] = {} 51 | 52 | test_dir = Path(scene_dir) / "test" 53 | 54 | for method in os.listdir(test_dir): 55 | print("Method:", method) 56 | 57 | full_dict[scene_dir][method] = {} 58 | per_view_dict[scene_dir][method] = {} 59 | full_dict_polytopeonly[scene_dir][method] = {} 60 | per_view_dict_polytopeonly[scene_dir][method] = {} 61 | 62 | method_dir = test_dir / method 63 | gt_dir = method_dir/ "gt" 64 | renders_dir = method_dir / "renders" 65 | renders, gts, image_names = readImages(renders_dir, gt_dir) 66 | 67 | ssims = [] 68 | psnrs = [] 69 | lpipss = [] 70 | 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 | 76 | print(" SSIM : {:>12.7f}".format(torch.tensor(ssims).mean(), ".5")) 77 | print(" PSNR : {:>12.7f}".format(torch.tensor(psnrs).mean(), ".5")) 78 | print(" LPIPS: {:>12.7f}".format(torch.tensor(lpipss).mean(), ".5")) 79 | print("") 80 | 81 | full_dict[scene_dir][method].update({"SSIM": torch.tensor(ssims).mean().item(), 82 | "PSNR": torch.tensor(psnrs).mean().item(), 83 | "LPIPS": torch.tensor(lpipss).mean().item()}) 84 | per_view_dict[scene_dir][method].update({"SSIM": {name: ssim for ssim, name in zip(torch.tensor(ssims).tolist(), image_names)}, 85 | "PSNR": {name: psnr for psnr, name in zip(torch.tensor(psnrs).tolist(), image_names)}, 86 | "LPIPS": {name: lp for lp, name in zip(torch.tensor(lpipss).tolist(), image_names)}}) 87 | 88 | with open(scene_dir + "/results.json", 'w') as fp: 89 | json.dump(full_dict[scene_dir], fp, indent=True) 90 | with open(scene_dir + "/per_view.json", 'w') as fp: 91 | json.dump(per_view_dict[scene_dir], fp, indent=True) 92 | except: 93 | print("Unable to compute metrics for model", scene_dir) 94 | 95 | if __name__ == "__main__": 96 | device = torch.device("cuda:0") 97 | torch.cuda.set_device(device) 98 | 99 | # Set up command line argument parser 100 | parser = ArgumentParser(description="Training script parameters") 101 | parser.add_argument('--model_paths', '-m', required=True, nargs="+", type=str, default=[]) 102 | args = parser.parse_args() 103 | evaluate(args.model_paths) 104 | -------------------------------------------------------------------------------- /Trim3DGS/print_results.py: -------------------------------------------------------------------------------- 1 | import json 2 | import os 3 | from argparse import ArgumentParser 4 | 5 | 6 | def report_dtu(path, iteration): 7 | print(f'Results of {path}') 8 | scans = os.listdir(path) 9 | 10 | print("***************** mesh *****************") 11 | sum_overall = 0 12 | n = 0 13 | for scan in sorted(scans, key=lambda x: int(x.replace('scan', ''))): 14 | p = os.path.join(path, scan, f'tsdf/ours_{iteration}/results.json') 15 | if not os.path.exists(p): 16 | continue 17 | with open(p, 'r') as f: 18 | data = json.load(f) 19 | print(scan, data, scan) 20 | sum_overall += data['overall'] 21 | n += 1 22 | print(f"Overall: {sum_overall / n}") 23 | 24 | sum_overall = 0 25 | n = 0 26 | print("***************** pcd *****************") 27 | for scan in sorted(scans, key=lambda x: int(x.replace('scan', ''))): 28 | p = os.path.join(path, scan, f'train/ours_{iteration}/results_pcd.json') 29 | if not os.path.exists(p): 30 | continue 31 | with open(p, 'r') as f: 32 | data = json.load(f) 33 | print(scan, data, scan) 34 | sum_overall += data['overall'] 35 | n += 1 36 | print(f"Overall: {sum_overall / n}") 37 | 38 | def report_mipnerf360(path, iteration): 39 | print(f'Results of {path}') 40 | scans = os.listdir(path) 41 | sum_overall = 0 42 | n = 0 43 | for scan in sorted(scans): 44 | p = os.path.join(path, scan, f'point_cloud/iteration_{iteration}/point_cloud.ply') 45 | if not os.path.exists(p): 46 | print(f"Missing {p}") 47 | continue 48 | # check the storage size of the point cloud 49 | size = os.path.getsize(p) 50 | mb_size = size / 1024 / 1024 51 | print(scan, f"{mb_size:.2f} MB") 52 | sum_overall += mb_size 53 | n += 1 54 | print(f"Overall: {sum_overall / n:.2f} MB") 55 | 56 | indoor = ['room', 'counter', 'kitchen', 'bonsai'] 57 | outdoor = ['bicycle', 'flowers', 'garden', 'stump', 'treehill'] 58 | sum_overall_indoor = dict() 59 | sum_overall_outdoor = dict() 60 | n_indoor = 0 61 | n_outdoor = 0 62 | for scan in sorted(scans): 63 | p = os.path.join(path, scan, f'results.json') 64 | if not os.path.exists(p): 65 | continue 66 | with open(p, 'r') as f: 67 | data = json.load(f) 68 | print(scan, data, scan) 69 | if scan in indoor: 70 | for k, v in data[f'ours_{iteration}'].items(): 71 | if k not in sum_overall_indoor: 72 | sum_overall_indoor[k] = 0.0 73 | sum_overall_indoor[k] += v 74 | n_indoor += 1 75 | if scan in outdoor: 76 | for k, v in data[f'ours_{iteration}'].items(): 77 | if k not in sum_overall_outdoor: 78 | sum_overall_outdoor[k] = 0.0 79 | sum_overall_outdoor[k] += v 80 | n_outdoor += 1 81 | 82 | print("Outdoor") 83 | for k, v in sum_overall_outdoor.items(): 84 | print(f"{k}: {v / n_outdoor:.3f}") 85 | 86 | print("Indoor") 87 | for k, v in sum_overall_indoor.items(): 88 | print(f"{k}: {v / n_indoor:.3f}") 89 | 90 | 91 | if __name__ == '__main__': 92 | parser = ArgumentParser() 93 | parser.add_argument('--output_path', '-o', type=str) 94 | parser.add_argument('--iteration', type=int, default=7000) 95 | parser.add_argument('--dataset', type=str, choices=['dtu', 'mipnerf360']) 96 | args = parser.parse_args() 97 | 98 | if args.dataset == 'dtu': 99 | report_dtu(args.output_path, args.iteration) 100 | elif args.dataset == 'mipnerf360': 101 | report_mipnerf360(args.output_path, args.iteration) 102 | else: 103 | raise ValueError(f"Unknown dataset {args.dataset}") -------------------------------------------------------------------------------- /Trim3DGS/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, FlattenGaussianModel 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], pretrained_ply_path=None): 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 | self.pretrained_ply_path = pretrained_ply_path 33 | 34 | if load_iteration: 35 | if load_iteration == -1: 36 | self.loaded_iter = searchForMaxIteration(os.path.join(self.model_path, "point_cloud")) 37 | else: 38 | self.loaded_iter = load_iteration 39 | print("Loading trained model at iteration {}".format(self.loaded_iter)) 40 | 41 | self.train_cameras = {} 42 | self.test_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) 49 | else: 50 | assert False, "Could not recognize scene type!" 51 | 52 | if not self.loaded_iter: 53 | with open(scene_info.ply_path, 'rb') as src_file, open(os.path.join(self.model_path, "input.ply") , 'wb') as dest_file: 54 | dest_file.write(src_file.read()) 55 | json_cams = [] 56 | camlist = [] 57 | if scene_info.test_cameras: 58 | camlist.extend(scene_info.test_cameras) 59 | if scene_info.train_cameras: 60 | camlist.extend(scene_info.train_cameras) 61 | for id, cam in enumerate(camlist): 62 | json_cams.append(camera_to_JSON(id, cam)) 63 | with open(os.path.join(self.model_path, "cameras.json"), 'w') as file: 64 | json.dump(json_cams, file) 65 | 66 | if shuffle: 67 | random.shuffle(scene_info.train_cameras) # Multi-res consistent random shuffling 68 | random.shuffle(scene_info.test_cameras) # Multi-res consistent random shuffling 69 | 70 | self.cameras_extent = scene_info.nerf_normalization["radius"] 71 | 72 | for resolution_scale in resolution_scales: 73 | print("Loading Training Cameras") 74 | self.train_cameras[resolution_scale] = cameraList_from_camInfos(scene_info.train_cameras, resolution_scale, args) 75 | print("Loading Test Cameras") 76 | self.test_cameras[resolution_scale] = cameraList_from_camInfos(scene_info.test_cameras, resolution_scale, args) 77 | 78 | if self.loaded_iter: 79 | self.gaussians.load_ply(os.path.join(self.model_path, 80 | "point_cloud", 81 | "iteration_" + str(self.loaded_iter), 82 | "point_cloud.ply")) 83 | elif self.pretrained_ply_path: 84 | self.gaussians.load_ply(self.pretrained_ply_path, self.cameras_extent) 85 | else: 86 | self.gaussians.create_from_pcd(scene_info.point_cloud, self.cameras_extent) 87 | 88 | def save(self, iteration): 89 | point_cloud_path = os.path.join(self.model_path, "point_cloud/iteration_{}".format(iteration)) 90 | self.gaussians.save_ply(os.path.join(point_cloud_path, "point_cloud.ply")) 91 | 92 | def getTrainCameras(self, scale=1.0): 93 | return self.train_cameras[scale] 94 | 95 | def getTestCameras(self, scale=1.0): 96 | return self.test_cameras[scale] -------------------------------------------------------------------------------- /Trim3DGS/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" 21 | ): 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 | 32 | try: 33 | self.data_device = torch.device(data_device) 34 | except Exception as e: 35 | print(e) 36 | print(f"[Warning] Custom device {data_device} failed, fallback to default cuda device" ) 37 | self.data_device = torch.device("cuda") 38 | 39 | self.original_image = image.clamp(0.0, 1.0).to(self.data_device) 40 | self.image_width = self.original_image.shape[2] 41 | self.image_height = self.original_image.shape[1] 42 | 43 | if gt_alpha_mask is not None: 44 | self.original_image *= gt_alpha_mask.to(self.data_device) 45 | self.gt_alpha_mask = gt_alpha_mask.to(self.data_device) 46 | else: 47 | self.original_image *= torch.ones((1, self.image_height, self.image_width), device=self.data_device) 48 | self.gt_alpha_mask = None 49 | 50 | self.zfar = 100.0 51 | self.znear = 0.01 52 | 53 | self.trans = trans 54 | self.scale = scale 55 | 56 | self.world_view_transform = torch.tensor(getWorld2View2(R, T, trans, scale)).transpose(0, 1).cuda() 57 | self.projection_matrix = getProjectionMatrix(znear=self.znear, zfar=self.zfar, fovX=self.FoVx, fovY=self.FoVy).transpose(0,1).cuda() 58 | self.full_proj_transform = (self.world_view_transform.unsqueeze(0).bmm(self.projection_matrix.unsqueeze(0))).squeeze(0) 59 | self.camera_center = self.world_view_transform.inverse()[3, :3] 60 | 61 | tan_fovx = np.tan(self.FoVx / 2.0) 62 | tan_fovy = np.tan(self.FoVy / 2.0) 63 | self.focal_y = self.image_height / (2.0 * tan_fovy) 64 | self.focal_x = self.image_width / (2.0 * tan_fovx) 65 | 66 | class MiniCam: 67 | def __init__(self, width, height, fovy, fovx, znear, zfar, world_view_transform, full_proj_transform): 68 | self.image_width = width 69 | self.image_height = height 70 | self.FoVy = fovy 71 | self.FoVx = fovx 72 | self.znear = znear 73 | self.zfar = zfar 74 | self.world_view_transform = world_view_transform 75 | self.full_proj_transform = full_proj_transform 76 | view_inv = torch.inverse(self.world_view_transform) 77 | self.camera_center = view_inv[3][:3] 78 | 79 | -------------------------------------------------------------------------------- /Trim3DGS/scripts/run_Mipnerf360.py: -------------------------------------------------------------------------------- 1 | # training script for MipNeRF360 dataset 2 | # adapted from https://github.com/autonomousvision/gaussian-opacity-fields/blob/main/scripts/run_mipnerf360.py 3 | 4 | import os 5 | from concurrent.futures import ThreadPoolExecutor 6 | import subprocess 7 | import time 8 | import torch 9 | 10 | scenes = ["bicycle", "bonsai", "counter", "flowers", "garden", "stump", "treehill", "kitchen", "room"] 11 | 12 | factors = [4, 2, 2, 4, 4, 4, 4, 2, 2] 13 | 14 | normal_weights = [0.1, 0.01, 0.01, 0.1, 0.1, 0.1, 0.1, 0.01, 0.01] 15 | 16 | excluded_gpus = set([]) 17 | 18 | split = "scale" 19 | output_dir = "output/MipNeRF360_3DGS" 20 | tune_output_dir = f"output/MipNeRF360_Trim3DGS" 21 | iteration = 7000 22 | 23 | jobs = list(zip(scenes, factors, normal_weights)) 24 | 25 | 26 | def train_scene(gpu, scene, factor, weight): 27 | cmds = [ 28 | f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python train.py -s data/MipNeRF360/{scene} -m {output_dir}/{scene} --eval -i images_{factor} --test_iterations -1 --quiet", 29 | f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python tune.py -s data/MipNeRF360/{scene} -m {tune_output_dir}/{scene} --eval -i images_{factor} --pretrained_ply {output_dir}/{scene}/point_cloud/iteration_30000/point_cloud.ply --test_iterations -1 --quiet --split {split} --position_lr_init 0.0000016 --densification_interval 1000 --opacity_reset_interval 999999 --normal_regularity_param {weight} --contribution_prune_from_iter 0 --contribution_prune_interval 1000", 30 | f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python render.py --iteration {iteration} -m {tune_output_dir}/{scene} --eval --skip_train --render_other", 31 | f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python metrics.py -m {tune_output_dir}/{scene}", 32 | # f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python extract_mesh_tsdf.py -s data/MipNeRF360/{scene} -m {tune_output_dir}/{scene} --eval -i images_{factor} --iteration {iteration} --voxel_size 0.004 --sdf_trunc 0.04", 33 | 34 | ] 35 | 36 | for cmd in cmds: 37 | print(cmd) 38 | subprocess.run(cmd, shell=True, check=True) 39 | return True 40 | 41 | 42 | def worker(gpu, scene, factor, weight): 43 | print(f"Starting job on GPU {gpu} with scene {scene}\n") 44 | train_scene(gpu, scene, factor, weight) 45 | print(f"Finished job on GPU {gpu} with scene {scene}\n") 46 | # This worker function starts a job and returns when it's done. 47 | 48 | def dispatch_jobs(jobs, executor): 49 | future_to_job = {} 50 | reserved_gpus = set() # GPUs that are slated for work but may not be active yet 51 | 52 | while jobs or future_to_job: 53 | # Get the list of available GPUs, not including those that are reserved. 54 | all_available_gpus = set(range(torch.cuda.device_count())) 55 | available_gpus = list(all_available_gpus - reserved_gpus - excluded_gpus) 56 | 57 | # Launch new jobs on available GPUs 58 | while available_gpus and jobs: 59 | gpu = available_gpus.pop(0) 60 | job = jobs.pop(0) 61 | future = executor.submit(worker, gpu, *job) # Unpacking job as arguments to worker 62 | future_to_job[future] = (gpu, job) 63 | 64 | reserved_gpus.add(gpu) # Reserve this GPU until the job starts processing 65 | 66 | # Check for completed jobs and remove them from the list of running jobs. 67 | # Also, release the GPUs they were using. 68 | done_futures = [future for future in future_to_job if future.done()] 69 | for future in done_futures: 70 | job = future_to_job.pop(future) # Remove the job associated with the completed future 71 | gpu = job[0] # The GPU is the first element in each job tuple 72 | reserved_gpus.discard(gpu) # Release this GPU 73 | print(f"Job {job} has finished., rellasing GPU {gpu}") 74 | # (Optional) You might want to introduce a small delay here to prevent this loop from spinning very fast 75 | # when there are no GPUs available. 76 | time.sleep(5) 77 | 78 | print("All jobs have been processed.") 79 | 80 | 81 | # Using ThreadPoolExecutor to manage the thread pool 82 | with ThreadPoolExecutor(max_workers=8) as executor: 83 | dispatch_jobs(jobs, executor) 84 | 85 | -------------------------------------------------------------------------------- /Trim3DGS/scripts/run_dtu.py: -------------------------------------------------------------------------------- 1 | # training script for DTU dataset 2 | # adapted from https://github.com/autonomousvision/gaussian-opacity-fields/blob/main/scripts/run_dtu.py 3 | 4 | import os 5 | from concurrent.futures import ThreadPoolExecutor 6 | import subprocess 7 | import time 8 | import torch 9 | 10 | scenes = [24, 37, 40, 55, 63, 65, 69, 83, 97, 105, 106, 110, 114, 118, 122] 11 | 12 | factors = [2] * len(scenes) 13 | 14 | excluded_gpus = set([]) 15 | 16 | output_dir = "output/DTU_3DGS" 17 | tune_output_dir = "output/DTU_Trim3DGS" 18 | iteration = 7000 19 | 20 | jobs = list(zip(scenes, factors)) 21 | 22 | normal_weight = {24:0.1, 63:0.1, 65:0.1, 97:0.01, 110:0.01, 118:0.02} 23 | prune_ratio = {55:0.2, 114:0.2} 24 | 25 | def train_scene(gpu, scene, factor): 26 | cmds = [ 27 | f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python train.py -s data/dtu_dataset/DTU/scan{scene} -m {output_dir}/scan{scene} -r {factor}", 28 | f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python tune.py -s data/dtu_dataset/DTU/scan{scene} -m {tune_output_dir}/scan{scene} --pretrained_ply {output_dir}/scan{scene}/point_cloud/iteration_30000/point_cloud.ply --densify_scale_factor 2.0 --normal_regularity_from_iter 5000 --normal_dilation 3 --contribution_prune_ratio {prune_ratio.get(scene, 0.1)} --normal_regularity_param {normal_weight.get(scene, 0.05)} --split mix", 29 | f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python extract_mesh_tsdf.py -s data/dtu_dataset/DTU/scan{scene} -m {tune_output_dir}/scan{scene} --iteration {iteration} --num_cluster 1 --voxel_size 0.004 --sdf_trunc 0.016 --depth_trunc 3.0", 30 | f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python eval_dtu/evaluate_single_scene.py --input_mesh {tune_output_dir}/scan{scene}/tsdf/ours_{iteration}/mesh_post.ply --scan_id {scene} --output_dir {tune_output_dir}/scan{scene}/tsdf/ours_{iteration} --mask_dir data/dtu_dataset/DTU --DTU data/dtu_dataset/Official_DTU_Dataset", 31 | f"OMP_NUM_THREADS=4 CUDA_VISIBLE_DEVICES={gpu} python eval_dtu_pcd/evaluate_single_scene.py --input_pcd {tune_output_dir}/scan{scene}/point_cloud/iteration_{iteration}/point_cloud.ply --scan_id {scene} --output_dir {tune_output_dir}/scan{scene}/train/ours_{iteration} --mask_dir data/dtu_dataset/DTU --DTU data/dtu_dataset/Official_DTU_Dataset" 32 | ] 33 | 34 | for cmd in cmds: 35 | print(cmd) 36 | subprocess.run(cmd, shell=True, check=True) 37 | return True 38 | 39 | 40 | def worker(gpu, scene, factor): 41 | print(f"Starting job on GPU {gpu} with scene {scene}\n") 42 | train_scene(gpu, scene, factor) 43 | print(f"Finished job on GPU {gpu} with scene {scene}\n") 44 | # This worker function starts a job and returns when it's done. 45 | 46 | def dispatch_jobs(jobs, executor): 47 | future_to_job = {} 48 | reserved_gpus = set() # GPUs that are slated for work but may not be active yet 49 | 50 | while jobs or future_to_job: 51 | # Get the list of available GPUs, not including those that are reserved. 52 | all_available_gpus = set(range(torch.cuda.device_count())) 53 | available_gpus = list(all_available_gpus - reserved_gpus - excluded_gpus) 54 | 55 | # Launch new jobs on available GPUs 56 | while available_gpus and jobs: 57 | gpu = available_gpus.pop(0) 58 | job = jobs.pop(0) 59 | future = executor.submit(worker, gpu, *job) # Unpacking job as arguments to worker 60 | future_to_job[future] = (gpu, job) 61 | 62 | reserved_gpus.add(gpu) # Reserve this GPU until the job starts processing 63 | 64 | # Check for completed jobs and remove them from the list of running jobs. 65 | # Also, release the GPUs they were using. 66 | done_futures = [future for future in future_to_job if future.done()] 67 | for future in done_futures: 68 | job = future_to_job.pop(future) # Remove the job associated with the completed future 69 | gpu = job[0] # The GPU is the first element in each job tuple 70 | reserved_gpus.discard(gpu) # Release this GPU 71 | print(f"Job {job} has finished., rellasing GPU {gpu}") 72 | # (Optional) You might want to introduce a small delay here to prevent this loop from spinning very fast 73 | # when there are no GPUs available. 74 | time.sleep(5) 75 | 76 | print("All jobs have been processed.") 77 | 78 | 79 | # Using ThreadPoolExecutor to manage the thread pool 80 | with ThreadPoolExecutor(max_workers=8) as executor: 81 | dispatch_jobs(jobs, executor) 82 | 83 | -------------------------------------------------------------------------------- /Trim3DGS/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 | from utils.general_utils import PILtoTorch 15 | from utils.graphics_utils import fov2focal 16 | 17 | WARNED = False 18 | 19 | def loadCam(args, id, cam_info, resolution_scale): 20 | orig_w, orig_h = cam_info.image.size 21 | 22 | if args.resolution in [1, 2, 4, 8]: 23 | resolution = round(orig_w/(resolution_scale * args.resolution)), round(orig_h/(resolution_scale * args.resolution)) 24 | else: # should be a type that converts to float 25 | if args.resolution == -1: 26 | if orig_w > 1600: 27 | global WARNED 28 | if not WARNED: 29 | print("[ INFO ] Encountered quite large input images (>1.6K pixels width), rescaling to 1.6K.\n " 30 | "If this is not desired, please explicitly specify '--resolution/-r' as 1") 31 | WARNED = True 32 | global_down = orig_w / 1600 33 | else: 34 | global_down = 1 35 | else: 36 | global_down = orig_w / args.resolution 37 | 38 | scale = float(global_down) * float(resolution_scale) 39 | resolution = (int(orig_w / scale), int(orig_h / scale)) 40 | 41 | resized_image_rgb = PILtoTorch(cam_info.image, resolution) 42 | 43 | gt_image = resized_image_rgb[:3, ...] 44 | loaded_mask = None 45 | 46 | if resized_image_rgb.shape[0] == 4: 47 | loaded_mask = resized_image_rgb[3:4, ...] 48 | 49 | return Camera(colmap_id=cam_info.uid, R=cam_info.R, T=cam_info.T, 50 | FoVx=cam_info.FovX, FoVy=cam_info.FovY, 51 | image=gt_image, gt_alpha_mask=loaded_mask, 52 | image_name=cam_info.image_name, uid=id, data_device=args.data_device) 53 | 54 | def cameraList_from_camInfos(cam_infos, resolution_scale, args): 55 | camera_list = [] 56 | 57 | for id, c in enumerate(cam_infos): 58 | camera_list.append(loadCam(args, id, c, resolution_scale)) 59 | 60 | return camera_list 61 | 62 | def camera_to_JSON(id, camera : Camera): 63 | Rt = np.zeros((4, 4)) 64 | Rt[:3, :3] = camera.R.transpose() 65 | Rt[:3, 3] = camera.T 66 | Rt[3, 3] = 1.0 67 | 68 | W2C = np.linalg.inv(Rt) 69 | pos = W2C[:3, 3] 70 | rot = W2C[:3, :3] 71 | serializable_array_2d = [x.tolist() for x in rot] 72 | camera_entry = { 73 | 'id' : id, 74 | 'img_name' : camera.image_name, 75 | 'width' : camera.width, 76 | 'height' : camera.height, 77 | 'position': pos.tolist(), 78 | 'rotation': serializable_array_2d, 79 | 'fy' : fov2focal(camera.FovY, camera.height), 80 | 'fx' : fov2focal(camera.FovX, camera.width) 81 | } 82 | return camera_entry 83 | -------------------------------------------------------------------------------- /Trim3DGS/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.resize(resolution) 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 | -------------------------------------------------------------------------------- /Trim3DGS/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)) -------------------------------------------------------------------------------- /Trim3DGS/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 torch.nn.functional as F 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 generate_grid(x_low, x_high, x_num, y_low, y_high, y_num, device): 23 | xs = torch.linspace(x_low, x_high, x_num, device=device) 24 | ys = torch.linspace(y_low, y_high, y_num, device=device) 25 | xv, yv = torch.meshgrid([xs, ys], indexing='xy') 26 | grid = torch.stack((xv.flatten(), yv.flatten())).T 27 | return grid 28 | 29 | def interpolate(data, xy): 30 | """ 31 | Interpolates values from a grid of data based on given coordinates. 32 | 33 | Args: 34 | data (torch.Tensor): The input data grid of shape (..., H, W). 35 | xy (torch.Tensor): The coordinates to interpolate the values from, of shape (N, 2). The coordinates are 36 | expected to be in the range [0, 1]. 37 | 38 | Returns: 39 | torch.Tensor: The interpolated values of shape (..., N). 40 | """ 41 | 42 | pos = xy * torch.tensor([data.shape[-1], data.shape[-2]], dtype=torch.float32, device=xy.device) 43 | indices = pos.long() 44 | lerp_weights = pos - indices.float() 45 | x0 = indices[:, 0].clamp(min=0, max=data.shape[-1]-1) 46 | y0 = indices[:, 1].clamp(min=0, max=data.shape[-2]-1) 47 | x1 = (x0 + 1).clamp(max=data.shape[-1]-1) 48 | y1 = (y0 + 1).clamp(max=data.shape[-2]-1) 49 | 50 | return ( 51 | data[..., y0, x0] * (1.0 - lerp_weights[:,0]) * (1.0 - lerp_weights[:,1]) + 52 | data[..., y0, x1] * lerp_weights[:,0] * (1.0 - lerp_weights[:,1]) + 53 | data[..., y1, x0] * (1.0 - lerp_weights[:,0]) * lerp_weights[:,1] + 54 | data[..., y1, x1] * lerp_weights[:,0] * lerp_weights[:,1] 55 | ) 56 | 57 | def compute_gradient(image, RGB2GRAY=False): 58 | assert image.ndim == 4, "image must have 4 dimensions" 59 | assert image.shape[1] == 1 or image.shape[1] == 3, "image must have 1 or 3 channels" 60 | if image.shape[1] == 3: 61 | assert RGB2GRAY == True, "RGB image must be converted to grayscale first" 62 | image = rgb_to_gray(image) 63 | sobel_kernel_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=image.dtype, device=image.device).view(1, 1, 3, 3) 64 | sobel_kernel_y = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=image.dtype, device=image.device).view(1, 1, 3, 3) 65 | 66 | image_for_pad = F.pad(image, pad=(1, 1, 1, 1), mode="replicate") 67 | gradient_x = F.conv2d(image_for_pad, sobel_kernel_x) / 3 68 | gradient_y = F.conv2d(image_for_pad, sobel_kernel_y) / 3 69 | 70 | return gradient_x, gradient_y 71 | 72 | def rgb_to_gray(image): 73 | gray_image = (0.299 * image[:, 0, :, :] + 0.587 * image[:, 1, :, :] + 74 | 0.114 * image[:, 2, :, :]) 75 | gray_image = gray_image.unsqueeze(1) 76 | 77 | return gray_image 78 | 79 | def blur(image): 80 | if image.ndim == 2: 81 | image = image[None, None, ...] 82 | channel = image.shape[1] 83 | guassian_kernel = torch.tensor([[1/16, 1/8, 1/16], [1/8, 1/4, 1/8], [1/16, 1/8, 1/16]], dtype=image.dtype, device=image.device).view(1, 1, 3, 3).repeat(channel, 1, 1, 1) 84 | output = F.conv2d(image, guassian_kernel, padding=1, groups=channel) 85 | return output 86 | 87 | import matplotlib.pyplot as plt 88 | import numpy as np 89 | 90 | def draw_hist(tensor, path, density=False): 91 | tensor = tensor.reshape(-1).detach().cpu().numpy() 92 | _ = plt.hist(tensor, bins=50, density=density) 93 | plt.savefig(path) -------------------------------------------------------------------------------- /Trim3DGS/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 l1_loss(network_output, gt): 18 | return torch.abs((network_output - gt)).mean() 19 | 20 | def l2_loss(network_output, gt): 21 | return ((network_output - gt) ** 2).mean() 22 | 23 | def gaussian(window_size, sigma): 24 | gauss = torch.Tensor([exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) 25 | return gauss / gauss.sum() 26 | 27 | def create_window(window_size, channel): 28 | _1D_window = gaussian(window_size, 1.5).unsqueeze(1) 29 | _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) 30 | window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) 31 | return window 32 | 33 | def ssim(img1, img2, window_size=11, size_average=True): 34 | channel = img1.size(-3) 35 | window = create_window(window_size, channel) 36 | 37 | if img1.is_cuda: 38 | window = window.cuda(img1.get_device()) 39 | window = window.type_as(img1) 40 | 41 | return _ssim(img1, img2, window, window_size, channel, size_average) 42 | 43 | def _ssim(img1, img2, window, window_size, channel, size_average=True): 44 | mu1 = F.conv2d(img1, window, padding=window_size // 2, groups=channel) 45 | mu2 = F.conv2d(img2, window, padding=window_size // 2, groups=channel) 46 | 47 | mu1_sq = mu1.pow(2) 48 | mu2_sq = mu2.pow(2) 49 | mu1_mu2 = mu1 * mu2 50 | 51 | sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size // 2, groups=channel) - mu1_sq 52 | sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size // 2, groups=channel) - mu2_sq 53 | sigma12 = F.conv2d(img1 * img2, window, padding=window_size // 2, groups=channel) - mu1_mu2 54 | 55 | C1 = 0.01 ** 2 56 | C2 = 0.03 ** 2 57 | 58 | ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) 59 | 60 | if size_average: 61 | return ssim_map.mean() 62 | else: 63 | return ssim_map.mean(1).mean(1).mean(1) 64 | 65 | -------------------------------------------------------------------------------- /Trim3DGS/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 -------------------------------------------------------------------------------- /Trim3DGS/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 | -------------------------------------------------------------------------------- /assets/teaser.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/YuxueYang1204/TrimGS/6da2f21271b79676fc9205b770d4d390b8739b89/assets/teaser.jpg -------------------------------------------------------------------------------- /blender/.gitignore: -------------------------------------------------------------------------------- 1 | data 2 | render_res 3 | __pycache__ -------------------------------------------------------------------------------- /blender/README.md: -------------------------------------------------------------------------------- 1 | # Instruction for Rendering using bpy 2 | 3 | ## Installation 4 | 5 | ```bash 6 | cd blender 7 | # python 3.11 is necessary for blender 4.1 on which our scripts based 8 | conda create -n blender python=3.11 -y 9 | conda activate blender 10 | pip install bpy==4.1 11 | pip install imageio 12 | pip install tqdm 13 | pip install imageio[ffmpeg] 14 | ``` 15 | 16 | ## Trajectory Preparation 17 | 18 | To render a video for a targeted mesh `mesh.ply`, please first generate a trajectory containing a series of cameras around the mesh. After that, pack the intrinsics and extrinsics of each camera on the desired trajectory into dictionaries with following 4 keys: 19 | "image_height", 20 | "image_width", 21 | "world_view_transform" and 22 | "FoVx", 23 | and dump them into pickle files naming `f{camera_index}.pkl`. Then organize the data, making the folder look as follows: 24 | 25 | ```bash 26 | data_path 27 | ├── mesh.ply 28 | └── traj 29 | ├── 00000.pkl 30 | ├── 00001.pkl 31 | ├── 00002.pkl 32 | └── ... 33 | ``` 34 | 35 | ## Rendering and Video Generation 36 | 37 | Run 38 | ```bash 39 | python render.py --load_dir --config_dir 40 | ``` 41 | to render images and generate a video for the mesh and trajectory located at the `load_dir`. To render textured mesh, add the command line argument `--is_texture`. 42 | 43 | ### Command Line Arguments for render.py 44 | 45 | #### --load_dir 46 | Path to the mesh and trajectory data (`data_path` in the above section). 47 | #### --save_dir 48 | Path where the rendered images and video should be stored (```output/``` by default). 49 | #### --config_dir 50 | Path where the render config file is stored. Set `render_cfgs/dtu` for `DTU` and `render_cfgs/mip` for `mip`. 51 | #### --is_texture 52 | The mesh file(.ply) generated by our method already contains texture. Use this argument to render textured images and video 53 | . 54 | ### Video Generation Only 55 | With images already rendered, run 56 | ```bash 57 | python generate_video.py --load_dir 58 | ``` 59 | to concatenate all of them into desired videos. 60 | 61 | ## Useful Command Line Arguments 62 | 63 | ### Debug Mode 64 | 65 | You can specify command line arguments `--debug_mode` and `--debug_video_step` in `render.py` to render some images of a scene only, for quickly debugging. 66 | -------------------------------------------------------------------------------- /blender/cam_pose_utils/graphic_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import math 3 | 4 | def getWorld2View(R, t): 5 | Rt = np.zeros((4, 4)) 6 | Rt[:3, :3] = R 7 | Rt[:3, 3] = t 8 | Rt[3, 3] = 1.0 9 | return np.float32(Rt) 10 | 11 | def fov2focal(fov, pixels): 12 | return pixels / (2 * math.tan(fov / 2)) 13 | 14 | def focal2fov(focal, pixels): 15 | return 2*math.atan(pixels/(2*focal)) 16 | 17 | 18 | class Virtual(): 19 | pass -------------------------------------------------------------------------------- /blender/generate_video.py: -------------------------------------------------------------------------------- 1 | import imageio 2 | import os 3 | from tqdm import tqdm 4 | from argparse import ArgumentParser 5 | 6 | def generate_video(path, fps): 7 | # if not is_texture: 8 | # write_dir = os.path.join(path, 'videos', 'mesh') 9 | # load_dir = os.path.join(path, 'mesh') 10 | # else: 11 | # write_dir = os.path.join(path, 'videos', 'texture') 12 | # load_dir = os.path.join(path, 'texture') 13 | load_dir = path 14 | if not os.path.isdir(load_dir): 15 | assert False 16 | write_dir = os.path.join(path, 'videos') 17 | if not os.path.exists(write_dir): 18 | os.makedirs(write_dir) 19 | video = imageio.get_writer(os.path.join(write_dir, 'video.mp4'), fps=fps) 20 | image_list = sorted(os.listdir(load_dir)) 21 | for i in tqdm(range(len(image_list)), desc=f"Creating video"): 22 | path = os.path.join(load_dir, image_list[i]) 23 | if os.path.isdir(path): 24 | continue 25 | image = imageio.imread(path) 26 | video.append_data(image) 27 | video.close() 28 | 29 | if __name__ == "__main__": 30 | 31 | parser = ArgumentParser(description='video generator arg parser') 32 | parser.add_argument('--load_dir', type=str, default="render_res") 33 | parser.add_argument("--is_texture", action="store_true") 34 | parser.add_argument("--fps", type=int, default=60) 35 | args = parser.parse_args() 36 | generate_video(path=args.load_dir, fps=args.fps) -------------------------------------------------------------------------------- /blender/render_cfgs/dtu/background.json: -------------------------------------------------------------------------------- 1 | 3.8 -------------------------------------------------------------------------------- /blender/render_cfgs/dtu/background.py: -------------------------------------------------------------------------------- 1 | import json 2 | 3 | back_sphere_radius = 3.8 4 | 5 | with open("background.json", "w") as f: 6 | json.dump(back_sphere_radius, f) -------------------------------------------------------------------------------- /blender/render_cfgs/dtu/light.json: -------------------------------------------------------------------------------- 1 | {"mesh": {"pose": [[1, 0, -2], [0, 1, -2], [0, -1, -2]], "energy": [12.0, 8.0, 7.0]}, "texture": {"pose": [[1, 0, -2], [0, 1, -2], [0, -1, -2]], "energy": [70.0, 45.0, 45.0]}} -------------------------------------------------------------------------------- /blender/render_cfgs/dtu/light.py: -------------------------------------------------------------------------------- 1 | import json 2 | light_cfg = {'mesh':{'pose':((1, 0, -2), (0, 1, -2), (0, -1, -2)), 'energy':(12.0, 8.0, 7.0)}, 'texture':{'pose':((1, 0, -2), (0, 1, -2), (0, -1, -2)), 'energy':(70.0, 45.0, 45.0)}} 3 | 4 | with open("light.json", "w") as f: 5 | json.dump(light_cfg, f) -------------------------------------------------------------------------------- /blender/render_cfgs/mip/background.json: -------------------------------------------------------------------------------- 1 | 7.5 -------------------------------------------------------------------------------- /blender/render_cfgs/mip/background.py: -------------------------------------------------------------------------------- 1 | import json 2 | 3 | back_sphere_radius = 7.5 4 | 5 | with open("background.json", "w") as f: 6 | json.dump(back_sphere_radius, f) -------------------------------------------------------------------------------- /blender/render_cfgs/mip/light.json: -------------------------------------------------------------------------------- 1 | {"mesh": {"pose": [[2.0, -3.5, 0], [-1, -3, 1], [0, -3, -1]], "energy": [60.0, 50.0, 30.0]}, "texture": {"pose": [[2.0, -3.5, 0], [-1, -3, 1], [0, -3, -1]], "energy": [200.0, 150.0, 150.0]}} -------------------------------------------------------------------------------- /blender/render_cfgs/mip/light.py: -------------------------------------------------------------------------------- 1 | import json 2 | light_cfg = {'mesh':{'pose':((2., -3.5, 0), (-1, -3, 1), (0, -3, -1)), 'energy':(60.0, 50.0, 30.0)}, 'texture':{'pose':((2., -3.5, 0), (-1, -3, 1), (0, -3, -1)), 'energy':(200.0, 150.0, 150.0)}} 3 | 4 | with open("light.json", "w") as f: 5 | json.dump(light_cfg, f) -------------------------------------------------------------------------------- /blender/render_utils/background_generator.py: -------------------------------------------------------------------------------- 1 | import bpy 2 | 3 | color_mesh = (1, 0.878, 0.949) 4 | color_texture = (0.5, 0.5, 0.5) 5 | 6 | def newMaterial(id): 7 | 8 | mat = bpy.data.materials.get(id) 9 | if mat is None: 10 | mat = bpy.data.materials.new(name=id) 11 | 12 | mat.use_nodes = True 13 | if mat.node_tree: 14 | mat.node_tree.links.clear() 15 | mat.node_tree.nodes.clear() 16 | 17 | return mat 18 | 19 | 20 | def newShader(id, type, r, g, b): 21 | 22 | mat = newMaterial(id) 23 | nodes = mat.node_tree.nodes 24 | links = mat.node_tree.links 25 | output = nodes.new(type='ShaderNodeOutputMaterial') 26 | 27 | if type == "diffuse": 28 | shader = nodes.new(type='ShaderNodeBsdfDiffuse') 29 | nodes["Diffuse BSDF"].inputs[0].default_value = (r, g, b, 1) 30 | else: 31 | assert False 32 | links.new(shader.outputs[0], output.inputs[0]) 33 | 34 | return mat 35 | 36 | 37 | def draw_background(is_texture): 38 | 39 | if is_texture: 40 | mat = newShader("Texture", "diffuse", *color_texture) 41 | else: 42 | mat = newShader("Mesh", "diffuse", *color_mesh) 43 | bpy.ops.surface.primitive_nurbs_surface_sphere_add() 44 | bpy.context.active_object.data.materials.append(mat) 45 | -------------------------------------------------------------------------------- /blender/render_utils/texture_allocator.py: -------------------------------------------------------------------------------- 1 | 2 | class TextureAllocator: 3 | 4 | def __init__(self, bpy, texture_name='texture_material'): 5 | self.bpy = bpy 6 | self.texture_name = texture_name 7 | # self.init_texture() 8 | 9 | def init_texture(self): 10 | bpy = self.bpy 11 | texture_name = self.texture_name 12 | mat = bpy.data.materials.new(name=texture_name) 13 | mat.use_nodes = True 14 | if mat.node_tree: 15 | mat.node_tree.links.clear() 16 | mat.node_tree.nodes.clear() 17 | 18 | nodes = mat.node_tree.nodes 19 | links = mat.node_tree.links 20 | output = nodes.new(type='ShaderNodeOutputMaterial') 21 | # shader = nodes.new(type='ShaderNodeBsdfDiffuse') 22 | shader = nodes.new(type='ShaderNodeBsdfPrincipled') 23 | links.new(shader.outputs[0], output.inputs[0]) 24 | 25 | input_attribute = nodes.new(type='ShaderNodeAttribute') 26 | input_attribute.attribute_name = 'Col' 27 | links.new(input_attribute.outputs[0], shader.inputs[0]) 28 | # return mat 29 | 30 | def set_texture(self): 31 | bpy = self.bpy 32 | bpy.context.active_object.data.materials.append(bpy.data.materials[self.texture_name]) 33 | -------------------------------------------------------------------------------- /environment.yml: -------------------------------------------------------------------------------- 1 | name: trimgs 2 | channels: 3 | - pytorch3d 4 | - pytorch 5 | - nvidia 6 | - conda-forge 7 | - pyg 8 | - defaults 9 | dependencies: 10 | - setuptools==69.0.3 11 | - ffmpeg=4.2.2 12 | - pillow=10.2.0 13 | - pip=23.3.1 14 | - python=3.8.18 15 | - pytorch=2.1.0 16 | - pytorch3d=0.7.5 17 | - pytorch-scatter 18 | - torchvision=0.16.0 19 | - typing_extensions=4.9.0 20 | - pip: 21 | - open3d==0.18.0 22 | - mediapy==1.1.2 23 | - lpips==0.1.4 24 | - scikit-image==0.21.0 25 | - tqdm==4.66.2 26 | - trimesh==4.3.2 27 | - Trim3DGS/submodules/diff-gaussian-rasterization 28 | - Trim3DGS/submodules/simple-knn 29 | - Trim2DGS/submodules/diff-surfel-rasterization 30 | # the submodule simple-knn in Trim2DGS is the same as the one in Trim3DGS 31 | # - Trim2DGS/submodules/simple-knn 32 | - plyfile 33 | - opencv-python 34 | --------------------------------------------------------------------------------