├── instant-nsr-pl ├── datasets │ ├── utils.py │ ├── __init__.py │ ├── openmaterial_womask.py │ ├── blender.py │ ├── openmaterial_wmask.py │ └── colmap_utils.py ├── utils │ ├── __init__.py │ ├── loggers.py │ ├── misc.py │ ├── obj.py │ ├── callbacks.py │ └── mixins.py ├── requirements.txt ├── .pylintrc ├── models │ ├── __init__.py │ ├── base.py │ ├── ray_utils.py │ ├── texture.py │ ├── utils.py │ ├── nerf.py │ ├── network_utils.py │ └── geometry.py ├── systems │ ├── __init__.py │ ├── loss_utils.py │ ├── base.py │ ├── criterions.py │ ├── nerf.py │ └── neus.py ├── LICENSE ├── configs │ ├── nerf-openmaterial-wmask.yaml │ ├── nerf-openmaterial-womask.yaml │ ├── neus-openmaterial-wmask.yaml │ └── neus-openmaterial-womask.yaml ├── run_openmaterial.sh ├── run_batch_ablation.sh ├── .gitignore ├── scripts │ └── imgs2poses.py ├── launch.py └── README.md ├── .gitignore ├── assets └── teaser.png ├── use_depth.py ├── Openmaterial-main ├── .gitignore └── eval │ ├── eval.sh │ ├── eval.py │ └── clean_mesh.py ├── sum_metrics.py ├── download.py ├── README.md ├── gen_colmap.py └── sum_metrics-ablation.py /instant-nsr-pl/datasets/utils.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /instant-nsr-pl/utils/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | datasets/ 2 | output/ 3 | Mesh/ -------------------------------------------------------------------------------- /assets/teaser.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Christy61/OpenMaterial/HEAD/assets/teaser.png -------------------------------------------------------------------------------- /instant-nsr-pl/requirements.txt: -------------------------------------------------------------------------------- 1 | pytorch-lightning<2 2 | omegaconf==2.2.3 3 | nerfacc==0.3.3 4 | matplotlib 5 | opencv-python 6 | imageio 7 | imageio-ffmpeg 8 | scipy 9 | PyMCubes 10 | pyransac3d 11 | torch_efficient_distloss 12 | tensorboard 13 | -------------------------------------------------------------------------------- /instant-nsr-pl/.pylintrc: -------------------------------------------------------------------------------- 1 | disable=R,C 2 | 3 | [TYPECHECK] 4 | # List of members which are set dynamically and missed by pylint inference 5 | # system, and so shouldn't trigger E1101 when accessed. Python regular 6 | # expressions are accepted. 7 | generated-members=numpy.*,torch.*,cv2.* 8 | -------------------------------------------------------------------------------- /use_depth.py: -------------------------------------------------------------------------------- 1 | import h5py 2 | 3 | # This is an example of using depth data 4 | filename = 'check/a99731685e4a44c7b3ce80e41633f486/driving+school_4k-diffuse/train/depth/000.h5' 5 | with h5py.File(filename, 'r') as hdf: 6 | dataset = hdf['depth'] 7 | data = dataset[:] 8 | print(data.shape) 9 | print(data.max()) -------------------------------------------------------------------------------- /instant-nsr-pl/models/__init__.py: -------------------------------------------------------------------------------- 1 | models = {} 2 | 3 | 4 | def register(name): 5 | def decorator(cls): 6 | models[name] = cls 7 | return cls 8 | return decorator 9 | 10 | 11 | def make(name, config): 12 | model = models[name](config) 13 | return model 14 | 15 | 16 | from . import nerf, neus, geometry, texture 17 | -------------------------------------------------------------------------------- /instant-nsr-pl/datasets/__init__.py: -------------------------------------------------------------------------------- 1 | datasets = {} 2 | 3 | 4 | def register(name): 5 | def decorator(cls): 6 | datasets[name] = cls 7 | return cls 8 | return decorator 9 | 10 | 11 | def make(name, config): 12 | dataset = datasets[name](config) 13 | return dataset 14 | 15 | 16 | from . import blender, colmap, openmaterial_wmask, openmaterial_womask 17 | -------------------------------------------------------------------------------- /instant-nsr-pl/systems/__init__.py: -------------------------------------------------------------------------------- 1 | systems = {} 2 | 3 | 4 | def register(name): 5 | def decorator(cls): 6 | systems[name] = cls 7 | return cls 8 | return decorator 9 | 10 | 11 | def make(name, config, load_from_checkpoint=None): 12 | if load_from_checkpoint is None: 13 | system = systems[name](config) 14 | else: 15 | system = systems[name].load_from_checkpoint(load_from_checkpoint, strict=False, config=config) 16 | return system 17 | 18 | 19 | from . import nerf, neus 20 | -------------------------------------------------------------------------------- /Openmaterial-main/.gitignore: -------------------------------------------------------------------------------- 1 | render/ 2 | render.reproduce/ 3 | main.py 4 | test_diffuse/ 5 | test_metallic/ 6 | test_pricipled_diffuse/ 7 | test_sample/ 8 | test_material.py 9 | __pycache__/ 10 | integrators/__pycache__/ 11 | scene_teapot/pic/ 12 | scene_teapot/preview/ 13 | teapot/scene_material.xml 14 | glass-of-water/ 15 | banner_06/ 16 | datasets/ 17 | banner_06.zip 18 | test_trans.py 19 | datasets_test/ 20 | emitters/ 21 | groundtruth/ 22 | scene_xml/ 23 | objaverse/ 24 | scene 25 | textures/ 26 | assert/ 27 | gen_sphere_box.sh 28 | cam_render_sphere_box.py 29 | output.log 30 | resource/ 31 | eval/visualization.py 32 | eval/vis.py -------------------------------------------------------------------------------- /instant-nsr-pl/models/base.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | from utils.misc import get_rank 5 | 6 | class BaseModel(nn.Module): 7 | def __init__(self, config): 8 | super().__init__() 9 | self.config = config 10 | self.rank = get_rank() 11 | self.setup() 12 | if self.config.get('weights', None): 13 | self.load_state_dict(torch.load(self.config.weights)) 14 | 15 | def setup(self): 16 | raise NotImplementedError 17 | 18 | def update_step(self, epoch, global_step): 19 | pass 20 | 21 | def train(self, mode=True): 22 | return super().train(mode=mode) 23 | 24 | def eval(self): 25 | return super().eval() 26 | 27 | def regularizations(self, out): 28 | return {} 29 | 30 | @torch.no_grad() 31 | def export(self, export_config): 32 | return {} 33 | -------------------------------------------------------------------------------- /instant-nsr-pl/LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 Yuanchen Guo 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /instant-nsr-pl/utils/loggers.py: -------------------------------------------------------------------------------- 1 | import re 2 | import pprint 3 | import logging 4 | 5 | from pytorch_lightning.loggers.base import LightningLoggerBase, rank_zero_experiment 6 | from pytorch_lightning.utilities.rank_zero import rank_zero_only 7 | 8 | 9 | class ConsoleLogger(LightningLoggerBase): 10 | def __init__(self, log_keys=[]): 11 | super().__init__() 12 | self.log_keys = [re.compile(k) for k in log_keys] 13 | self.dict_printer = pprint.PrettyPrinter(indent=2, compact=False).pformat 14 | 15 | def match_log_keys(self, s): 16 | return True if not self.log_keys else any(r.search(s) for r in self.log_keys) 17 | 18 | @property 19 | def name(self): 20 | return 'console' 21 | 22 | @property 23 | def version(self): 24 | return '0' 25 | 26 | @property 27 | @rank_zero_experiment 28 | def experiment(self): 29 | return logging.getLogger('pytorch_lightning') 30 | 31 | @rank_zero_only 32 | def log_hyperparams(self, params): 33 | pass 34 | 35 | @rank_zero_only 36 | def log_metrics(self, metrics, step): 37 | metrics_ = {k: v for k, v in metrics.items() if self.match_log_keys(k)} 38 | if not metrics_: 39 | return 40 | self.experiment.info(f"\nEpoch{metrics['epoch']} Step{step}\n{self.dict_printer(metrics_)}") 41 | 42 | -------------------------------------------------------------------------------- /Openmaterial-main/eval/eval.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | directory=$1 4 | output_dir=$2 5 | method=$3 6 | ablation=$4 7 | if [ -d "$directory" ]; then 8 | subdir_name1=$method 9 | # for subdir1 in "$directory"/*; do 10 | # subdir_name1=$(basename "$subdir1") 11 | for subdir2 in "${directory}/${subdir_name1}/meshes/"*; do 12 | subdir_name2=$(basename "$subdir2") 13 | echo "clean_mesh start: ${subdir_name1}" 14 | if [ "$ablation" = "true" ]; then 15 | python eval/clean_mesh.py --dataset_dir ../datasets/ablation \ 16 | --groundtruth_dir ../datasets/groundtruth_ablation \ 17 | --method ${subdir_name1} \ 18 | --directory ${directory}/${subdir_name1} \ 19 | --object_name ${subdir_name2} 20 | else 21 | python eval/clean_mesh.py --method ${subdir_name1} --directory ${directory}/${subdir_name1} --object_name ${subdir_name2} 22 | fi 23 | echo "evaluation start:" 24 | for subdir3 in "${directory}/${subdir_name1}/CleanedMesh/${subdir_name2}/"*; do 25 | subdir_name3=$(basename "$subdir3") 26 | python eval/eval.py \ 27 | --pr ${directory}/${subdir_name1}/CleanedMesh/${subdir_name2}/${subdir_name3} \ 28 | --gt ../datasets/groundtruth/${subdir_name2}/clean_${subdir_name2}.ply \ 29 | --object ${subdir_name2} \ 30 | --method ${subdir_name1} \ 31 | --output ${output_dir} 32 | done 33 | done 34 | # done 35 | else 36 | echo "no dataset, please check again." 37 | fi -------------------------------------------------------------------------------- /instant-nsr-pl/models/ray_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | 4 | 5 | def cast_rays(ori, dir, z_vals): 6 | return ori[..., None, :] + z_vals[..., None] * dir[..., None, :] 7 | 8 | 9 | def get_ray_directions(W, H, fx, fy, cx, cy, use_pixel_centers=True): 10 | pixel_center = 0.5 if use_pixel_centers else 0 11 | i, j = np.meshgrid( 12 | np.arange(W, dtype=np.float32) + pixel_center, 13 | np.arange(H, dtype=np.float32) + pixel_center, 14 | indexing='xy' 15 | ) 16 | i, j = torch.from_numpy(i), torch.from_numpy(j) 17 | 18 | directions = torch.stack([(i - cx) / fx, -(j - cy) / fy, -torch.ones_like(i)], -1) # (H, W, 3) 19 | 20 | return directions 21 | 22 | 23 | def get_rays(directions, c2w, keepdim=False): 24 | # Rotate ray directions from camera coordinate to the world coordinate 25 | # rays_d = directions @ c2w[:, :3].T # (H, W, 3) # slow? 26 | assert directions.shape[-1] == 3 27 | 28 | if directions.ndim == 2: # (N_rays, 3) 29 | assert c2w.ndim == 3 # (N_rays, 4, 4) / (1, 4, 4) 30 | rays_d = (directions[:,None,:] * c2w[:,:3,:3]).sum(-1) # (N_rays, 3) 31 | rays_o = c2w[:,:,3].expand(rays_d.shape) 32 | elif directions.ndim == 3: # (H, W, 3) 33 | if c2w.ndim == 2: # (4, 4) 34 | rays_d = (directions[:,:,None,:] * c2w[None,None,:3,:3]).sum(-1) # (H, W, 3) 35 | rays_o = c2w[None,None,:,3].expand(rays_d.shape) 36 | elif c2w.ndim == 3: # (B, 4, 4) 37 | rays_d = (directions[None,:,:,None,:] * c2w[:,None,None,:3,:3]).sum(-1) # (B, H, W, 3) 38 | rays_o = c2w[:,None,None,:,3].expand(rays_d.shape) 39 | 40 | if not keepdim: 41 | rays_o, rays_d = rays_o.reshape(-1, 3), rays_d.reshape(-1, 3) 42 | 43 | return rays_o, rays_d 44 | -------------------------------------------------------------------------------- /instant-nsr-pl/utils/misc.py: -------------------------------------------------------------------------------- 1 | import os 2 | from omegaconf import OmegaConf 3 | from packaging import version 4 | 5 | 6 | # ============ Register OmegaConf Recolvers ============= # 7 | OmegaConf.register_new_resolver('calc_exp_lr_decay_rate', lambda factor, n: factor**(1./n)) 8 | OmegaConf.register_new_resolver('add', lambda a, b: a + b) 9 | OmegaConf.register_new_resolver('sub', lambda a, b: a - b) 10 | OmegaConf.register_new_resolver('mul', lambda a, b: a * b) 11 | OmegaConf.register_new_resolver('div', lambda a, b: a / b) 12 | OmegaConf.register_new_resolver('idiv', lambda a, b: a // b) 13 | OmegaConf.register_new_resolver('basename', lambda p: os.path.basename(p)) 14 | # ======================================================= # 15 | 16 | 17 | def prompt(question): 18 | inp = input(f"{question} (y/n)").lower().strip() 19 | if inp and inp == 'y': 20 | return True 21 | if inp and inp == 'n': 22 | return False 23 | return prompt(question) 24 | 25 | 26 | def load_config(*yaml_files, cli_args=[]): 27 | yaml_confs = [OmegaConf.load(f) for f in yaml_files] 28 | cli_conf = OmegaConf.from_cli(cli_args) 29 | conf = OmegaConf.merge(*yaml_confs, cli_conf) 30 | OmegaConf.resolve(conf) 31 | return conf 32 | 33 | 34 | def config_to_primitive(config, resolve=True): 35 | return OmegaConf.to_container(config, resolve=resolve) 36 | 37 | 38 | def dump_config(path, config): 39 | with open(path, 'w') as fp: 40 | OmegaConf.save(config=config, f=fp) 41 | 42 | def get_rank(): 43 | # SLURM_PROCID can be set even if SLURM is not managing the multiprocessing, 44 | # therefore LOCAL_RANK needs to be checked first 45 | rank_keys = ("RANK", "LOCAL_RANK", "SLURM_PROCID", "JSM_NAMESPACE_RANK") 46 | for key in rank_keys: 47 | rank = os.environ.get(key) 48 | if rank is not None: 49 | return int(rank) 50 | return 0 51 | 52 | 53 | def parse_version(ver): 54 | return version.parse(ver) 55 | -------------------------------------------------------------------------------- /Openmaterial-main/eval/eval.py: -------------------------------------------------------------------------------- 1 | # copy from: https://github.com/liuyuan-pal/NeRO/blob/main/eval_real_shape.py 2 | import torch 3 | import argparse 4 | from pathlib import Path 5 | import trimesh 6 | from pytorch3d.structures import Meshes 7 | from pytorch3d.ops import sample_points_from_meshes 8 | from tqdm import tqdm 9 | import os 10 | 11 | def load_mesh(file_path): 12 | mesh = trimesh.load(file_path, process=False) 13 | verts = torch.tensor(mesh.vertices, dtype=torch.float32).cuda() 14 | faces = torch.tensor(mesh.faces, dtype=torch.int64).cuda() 15 | return Meshes(verts=[verts], faces=[faces]) , len(verts) 16 | 17 | 18 | def nearest_dist(pts0, pts1, batch_size=512): 19 | pn0, pn1 = pts0.shape[0], pts1.shape[0] 20 | dists = [] 21 | for i in tqdm(range(0, pn0, batch_size), desc='evaluting...'): 22 | dist = torch.norm(pts0[i:i+batch_size,None,:] - pts1[None,:,:], dim=-1) 23 | dists.append(torch.min(dist,1)[0]) 24 | dists = torch.cat(dists,0).cpu().numpy() 25 | return dists 26 | 27 | 28 | if __name__=="__main__": 29 | parser = argparse.ArgumentParser() 30 | parser.add_argument('--pr',type=str,) 31 | parser.add_argument('--gt',type=str,) 32 | parser.add_argument('--object',type=str,) 33 | parser.add_argument('--method',type=str,) 34 | parser.add_argument('--output',type=str,) 35 | args = parser.parse_args() 36 | max_dist = 0.15 37 | 38 | mesh_pr, len_pr = load_mesh(f'{args.pr}') 39 | mesh_gt, len_gt = load_mesh(f'{args.gt}') 40 | 41 | num_samples = 1000000 42 | 43 | pts_pr = sample_points_from_meshes(mesh_pr, num_samples=num_samples).squeeze() 44 | pts_gt = sample_points_from_meshes(mesh_gt, num_samples=num_samples).squeeze() 45 | 46 | bn = 512 47 | dist_gt = nearest_dist(pts_gt, pts_pr, bn) 48 | mean_gt = dist_gt[dist_gt < max_dist].mean() 49 | dist_pr = nearest_dist(pts_pr, pts_gt, bn) 50 | mean_pr = dist_pr[dist_pr < max_dist].mean() 51 | 52 | stem = Path(args.pr).stem 53 | print(stem) 54 | chamfer = (mean_gt + mean_pr) / 2 * 100 55 | results = f'{args.object}:{args.method}:{stem}:{chamfer:.5f}' 56 | os.makedirs(f'{args.output}/{args.object}', exist_ok=True) 57 | with open(os.path.join(args.output, args.object, f"{args.method}-mesh-output.txt"), "a") as file: 58 | file.write(results + '\n') 59 | print(results) 60 | -------------------------------------------------------------------------------- /instant-nsr-pl/systems/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 | -------------------------------------------------------------------------------- /instant-nsr-pl/models/texture.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | import models 5 | from models.utils import get_activation 6 | from models.network_utils import get_encoding, get_mlp 7 | from systems.utils import update_module_step 8 | 9 | 10 | @models.register('volume-radiance') 11 | class VolumeRadiance(nn.Module): 12 | def __init__(self, config): 13 | super(VolumeRadiance, self).__init__() 14 | self.config = config 15 | self.n_dir_dims = self.config.get('n_dir_dims', 3) 16 | self.n_output_dims = 3 17 | encoding = get_encoding(self.n_dir_dims, self.config.dir_encoding_config) 18 | self.n_input_dims = self.config.input_feature_dim + encoding.n_output_dims 19 | network = get_mlp(self.n_input_dims, self.n_output_dims, self.config.mlp_network_config) 20 | self.encoding = encoding 21 | self.network = network 22 | 23 | def forward(self, features, dirs, *args): 24 | dirs = (dirs + 1.) / 2. # (-1, 1) => (0, 1) 25 | dirs_embd = self.encoding(dirs.view(-1, self.n_dir_dims)) 26 | network_inp = torch.cat([features.view(-1, features.shape[-1]), dirs_embd] + [arg.view(-1, arg.shape[-1]) for arg in args], dim=-1) 27 | color = self.network(network_inp).view(*features.shape[:-1], self.n_output_dims).float() 28 | if 'color_activation' in self.config: 29 | color = get_activation(self.config.color_activation)(color) 30 | return color 31 | 32 | def update_step(self, epoch, global_step): 33 | update_module_step(self.encoding, epoch, global_step) 34 | 35 | def regularizations(self, out): 36 | return {} 37 | 38 | 39 | @models.register('volume-color') 40 | class VolumeColor(nn.Module): 41 | def __init__(self, config): 42 | super(VolumeColor, self).__init__() 43 | self.config = config 44 | self.n_output_dims = 3 45 | self.n_input_dims = self.config.input_feature_dim 46 | network = get_mlp(self.n_input_dims, self.n_output_dims, self.config.mlp_network_config) 47 | self.network = network 48 | 49 | def forward(self, features, *args): 50 | network_inp = features.view(-1, features.shape[-1]) 51 | color = self.network(network_inp).view(*features.shape[:-1], self.n_output_dims).float() 52 | if 'color_activation' in self.config: 53 | color = get_activation(self.config.color_activation)(color) 54 | return color 55 | 56 | def regularizations(self, out): 57 | return {} 58 | -------------------------------------------------------------------------------- /instant-nsr-pl/configs/nerf-openmaterial-wmask.yaml: -------------------------------------------------------------------------------- 1 | name: nerf-openmaterial-wmask-${dataset.scene} 2 | tag: "" 3 | seed: 42 4 | 5 | dataset: 6 | name: openmaterial-wmask 7 | scene: abandoned_bakery_4k-conductor-Ag 8 | root_dir: ../datasets/openmaterial/Pineapple/${dataset.scene} 9 | img_wh: 10 | - 1600 11 | - 1200 12 | near_plane: 1.3 13 | far_plane: 4.6 14 | train_split: "train" 15 | test_split: "test" 16 | bsdf_name: "" 17 | object: "" 18 | 19 | model: 20 | name: nerf 21 | radius: 1.0 22 | num_samples_per_ray: 1024 23 | train_num_rays: 256 24 | max_train_num_rays: 8192 25 | grid_prune: true 26 | dynamic_ray_sampling: true 27 | batch_image_sampling: true 28 | randomized: true 29 | ray_chunk: 32768 30 | learned_background: false 31 | background_color: black 32 | geometry: 33 | name: volume-density 34 | radius: ${model.radius} 35 | feature_dim: 16 36 | density_activation: trunc_exp 37 | density_bias: -1 38 | isosurface: 39 | method: mc 40 | resolution: 256 41 | chunk: 2097152 42 | threshold: 2.5 43 | xyz_encoding_config: 44 | otype: HashGrid 45 | n_levels: 16 46 | n_features_per_level: 2 47 | log2_hashmap_size: 19 48 | base_resolution: 16 49 | per_level_scale: 1.447269237440378 50 | mlp_network_config: 51 | otype: FullyFusedMLP 52 | activation: ReLU 53 | output_activation: none 54 | n_neurons: 64 55 | n_hidden_layers: 1 56 | texture: 57 | name: volume-radiance 58 | input_feature_dim: ${model.geometry.feature_dim} 59 | dir_encoding_config: 60 | otype: SphericalHarmonics 61 | degree: 4 62 | mlp_network_config: 63 | otype: FullyFusedMLP 64 | activation: ReLU 65 | output_activation: Sigmoid 66 | n_neurons: 64 67 | n_hidden_layers: 2 68 | 69 | system: 70 | name: nerf-system 71 | loss: 72 | lambda_rgb: 1.0 73 | lambda_distortion: 0.0 74 | optimizer: 75 | name: AdamW 76 | args: 77 | lr: 0.01 78 | betas: [0.9, 0.99] 79 | eps: 1.e-15 80 | scheduler: 81 | name: MultiStepLR 82 | interval: step 83 | args: 84 | milestones: [10000, 15000, 18000] 85 | gamma: 0.33 86 | 87 | checkpoint: 88 | save_top_k: -1 89 | every_n_train_steps: ${trainer.max_steps} 90 | 91 | export: 92 | chunk_size: 2097152 93 | export_vertex_color: False 94 | 95 | trainer: 96 | max_steps: 15000 97 | log_every_n_steps: 200 98 | num_sanity_val_steps: 0 99 | val_check_interval: 10000 100 | limit_train_batches: 1.0 101 | limit_val_batches: 2 102 | enable_progress_bar: true 103 | precision: 16 104 | -------------------------------------------------------------------------------- /instant-nsr-pl/configs/nerf-openmaterial-womask.yaml: -------------------------------------------------------------------------------- 1 | name: nerf-openmaterial-womask-${dataset.scene} 2 | tag: "" 3 | seed: 42 4 | 5 | dataset: 6 | name: openmaterial-womask 7 | scene: abandoned_bakery_4k-conductor-Ag 8 | root_dir: ../datasets/openmaterial/Pineapple/${dataset.scene} 9 | img_wh: 10 | - 1600 11 | - 1200 12 | near_plane: 1.3 13 | far_plane: 4.6 14 | train_split: "train" 15 | test_split: "test" 16 | bsdf_name: "" 17 | object: "" 18 | 19 | model: 20 | name: nerf 21 | radius: 1.0 22 | num_samples_per_ray: 1024 23 | train_num_rays: 256 24 | max_train_num_rays: 8192 25 | grid_prune: true 26 | dynamic_ray_sampling: true 27 | batch_image_sampling: true 28 | randomized: true 29 | ray_chunk: 32768 30 | learned_background: true 31 | background_color: random 32 | geometry: 33 | name: volume-density 34 | radius: ${model.radius} 35 | feature_dim: 16 36 | density_activation: trunc_exp 37 | density_bias: -1 38 | isosurface: 39 | method: mc 40 | resolution: 256 41 | chunk: 2097152 42 | threshold: 2.5 43 | xyz_encoding_config: 44 | otype: HashGrid 45 | n_levels: 16 46 | n_features_per_level: 2 47 | log2_hashmap_size: 19 48 | base_resolution: 16 49 | per_level_scale: 1.447269237440378 50 | mlp_network_config: 51 | otype: FullyFusedMLP 52 | activation: ReLU 53 | output_activation: none 54 | n_neurons: 64 55 | n_hidden_layers: 1 56 | texture: 57 | name: volume-radiance 58 | input_feature_dim: ${model.geometry.feature_dim} 59 | dir_encoding_config: 60 | otype: SphericalHarmonics 61 | degree: 4 62 | mlp_network_config: 63 | otype: FullyFusedMLP 64 | activation: ReLU 65 | output_activation: Sigmoid 66 | n_neurons: 64 67 | n_hidden_layers: 2 68 | 69 | system: 70 | name: nerf-system 71 | loss: 72 | lambda_rgb: 1.0 73 | lambda_distortion: 0.0 74 | optimizer: 75 | name: AdamW 76 | args: 77 | lr: 0.01 78 | betas: [0.9, 0.99] 79 | eps: 1.e-15 80 | scheduler: 81 | name: MultiStepLR 82 | interval: step 83 | args: 84 | milestones: [10000, 15000, 18000] 85 | gamma: 0.33 86 | 87 | checkpoint: 88 | save_top_k: -1 89 | every_n_train_steps: ${trainer.max_steps} 90 | 91 | export: 92 | chunk_size: 2097152 93 | export_vertex_color: False 94 | 95 | trainer: 96 | max_steps: 15000 97 | log_every_n_steps: 200 98 | num_sanity_val_steps: 0 99 | val_check_interval: 10000 100 | limit_train_batches: 1.0 101 | limit_val_batches: 2 102 | enable_progress_bar: true 103 | precision: 16 104 | -------------------------------------------------------------------------------- /instant-nsr-pl/run_openmaterial.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | start=${1:-0} 3 | end=${2:-10} 4 | 5 | directory="../datasets/openmaterial" 6 | 7 | # rm -rf ../Mesh/instant-nsr-pl-wmask/ 8 | 9 | if [ -d "$root_dir" ]; then 10 | root_dir="${root_dir%/}" 11 | base_dir=$(basename "$root_dir") 12 | 13 | IFS=$'\n' read -d '' -r -a directories < <(find "$root_dir" -maxdepth 1 -type d | while IFS= read -r directory; do basename "$directory"; done | sort | grep -v "^$base_dir$" && printf '\0') 14 | # cut into groups 15 | echo "[+] Number of directories: ${#directories[@]}" 16 | declare -a group 17 | count=0 18 | for dir in "${directories[@]}"; do 19 | let count+=1 20 | if [ $count -gt $start ]; then 21 | group+=("$dir") 22 | if [ $count -eq $end ]; then 23 | echo "[+] from $start to $end, total count: ${#group[@]} in group" 24 | for g in "${group[@]}"; do 25 | echo "group: $g" 26 | done 27 | break 28 | fi 29 | fi 30 | done 31 | 32 | for subdir1 in "${group[@]}"; do 33 | subdir_name1=$(basename "$subdir1") 34 | for subdir2 in "$root_dir/$subdir_name1"/*; do 35 | subdir_name2=$(basename "$subdir2") 36 | echo "[+] Case $subdir_name1" 37 | bsdf_name="${subdir_name2#*-}" 38 | 39 | python launch.py \ 40 | --config configs/neus-openmaterial-wmask.yaml \ 41 | --output_dir ../Mesh/instant-nsr-pl-wmask/meshes/${subdir_name1} \ 42 | --gpu $3 \ 43 | --train \ 44 | dataset.bsdf_name=${bsdf_name} \ 45 | dataset.object=${subdir_name1} \ 46 | dataset.scene=${subdir_name2} \ 47 | dataset.root_dir=${root_dir}/${subdir_name1}/${subdir_name2} \ 48 | trial_name=${subdir_name1} 49 | rm -r exp/neus-openmaterial-wmask-${subdir_name2}/${subdir_name1} 50 | 51 | new_line=" with open(os.path.join(f'../output', self.config.dataset.object, \"instant-nsr-pl-wmask.txt\"), \"w\") as file:" 52 | sed -i "219s|.*|$new_line|" systems/nerf.py 53 | python launch.py \ 54 | --config configs/nerf-openmaterial-wmask.yaml \ 55 | --gpu $3 \ 56 | --train \ 57 | dataset.bsdf_name=${bsdf_name} \ 58 | dataset.object=${subdir_name1} \ 59 | dataset.scene=${subdir_name2} \ 60 | dataset.root_dir=${root_dir}/${subdir_name1}/${subdir_name2} \ 61 | trial_name=${subdir_name1} \ 62 | render_save_dir=../output 63 | rm -r exp/nerf-openmaterial-wmask-${subdir_name2}/${subdir_name1} 64 | 65 | done 66 | done 67 | else 68 | echo "no openmaterial dataset, please generate before running." 69 | fi 70 | # rm -rf exp/ -------------------------------------------------------------------------------- /instant-nsr-pl/utils/obj.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | def load_obj(filename): 5 | # Read entire file 6 | with open(filename, 'r') as f: 7 | lines = f.readlines() 8 | 9 | # load vertices 10 | vertices, texcoords = [], [] 11 | for line in lines: 12 | if len(line.split()) == 0: 13 | continue 14 | 15 | prefix = line.split()[0].lower() 16 | if prefix == 'v': 17 | vertices.append([float(v) for v in line.split()[1:]]) 18 | elif prefix == 'vt': 19 | val = [float(v) for v in line.split()[1:]] 20 | texcoords.append([val[0], 1.0 - val[1]]) 21 | 22 | uv = len(texcoords) > 0 23 | faces, tfaces = [], [] 24 | for line in lines: 25 | if len(line.split()) == 0: 26 | continue 27 | prefix = line.split()[0].lower() 28 | if prefix == 'usemtl': # Track used materials 29 | pass 30 | elif prefix == 'f': # Parse face 31 | vs = line.split()[1:] 32 | nv = len(vs) 33 | vv = vs[0].split('/') 34 | v0 = int(vv[0]) - 1 35 | if uv: 36 | t0 = int(vv[1]) - 1 if vv[1] != "" else -1 37 | for i in range(nv - 2): # Triangulate polygons 38 | vv1 = vs[i + 1].split('/') 39 | v1 = int(vv1[0]) - 1 40 | vv2 = vs[i + 2].split('/') 41 | v2 = int(vv2[0]) - 1 42 | faces.append([v0, v1, v2]) 43 | if uv: 44 | t1 = int(vv1[1]) - 1 if vv1[1] != "" else -1 45 | t2 = int(vv2[1]) - 1 if vv2[1] != "" else -1 46 | tfaces.append([t0, t1, t2]) 47 | vertices = np.array(vertices, dtype=np.float32) 48 | faces = np.array(faces, dtype=np.int64) 49 | if uv: 50 | assert len(tfaces) == len(faces) 51 | texcoords = np.array(texcoords, dtype=np.float32) 52 | tfaces = np.array(tfaces, dtype=np.int64) 53 | else: 54 | texcoords, tfaces = None, None 55 | 56 | return vertices, faces, texcoords, tfaces 57 | 58 | 59 | def write_obj(filename, v_pos, t_pos_idx, v_tex, t_tex_idx): 60 | with open(filename, "w") as f: 61 | for v in v_pos: 62 | f.write('v {} {} {} \n'.format(v[0], v[1], v[2])) 63 | 64 | if v_tex is not None: 65 | assert(len(t_pos_idx) == len(t_tex_idx)) 66 | for v in v_tex: 67 | f.write('vt {} {} \n'.format(v[0], 1.0 - v[1])) 68 | 69 | # Write faces 70 | for i in range(len(t_pos_idx)): 71 | f.write("f ") 72 | for j in range(3): 73 | f.write(' %s/%s' % (str(t_pos_idx[i][j]+1), '' if v_tex is None else str(t_tex_idx[i][j]+1))) 74 | f.write("\n") 75 | -------------------------------------------------------------------------------- /instant-nsr-pl/run_batch_ablation.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | start=${1:-0} 3 | end=${2:-10} 4 | 5 | root_dir="../datasets/ablation" 6 | 7 | # rm -rf ../Mesh/instant-nsr-pl-wmask/ 8 | 9 | if [ -d "$root_dir" ]; then 10 | root_dir="${root_dir%/}" 11 | base_dir=$(basename "$root_dir") 12 | 13 | IFS=$'\n' read -d '' -r -a directories < <(find "$root_dir" -maxdepth 1 -type d | while IFS= read -r directory; do basename "$directory"; done | sort | grep -v "^$base_dir$" && printf '\0') 14 | # cut into groups 15 | echo "[+] Number of directories: ${#directories[@]}" 16 | declare -a group 17 | count=0 18 | for dir in "${directories[@]}"; do 19 | let count+=1 20 | if [ $count -gt $start ]; then 21 | group+=("$dir") 22 | if [ $count -eq $end ]; then 23 | echo "[+] from $start to $end, total count: ${#group[@]} in group" 24 | for g in "${group[@]}"; do 25 | echo "group: $g" 26 | done 27 | break 28 | fi 29 | fi 30 | done 31 | 32 | for subdir1 in "${group[@]}"; do 33 | subdir_name1=$(basename "$subdir1") 34 | for subdir2 in "$root_dir/$subdir_name1"/*; do 35 | subdir_name2=$(basename "$subdir2") 36 | echo "[+] Case $subdir_name1" 37 | bsdf_name="${subdir_name2#*-}" 38 | 39 | python launch.py \ 40 | --config configs/neus-openmaterial-wmask.yaml \ 41 | --output_dir ../Mesh-ablation/instant-nsr-pl-wmask/meshes/${subdir_name1} \ 42 | --gpu 0 \ 43 | --train \ 44 | dataset.bsdf_name=${bsdf_name} \ 45 | dataset.object=${subdir_name1} \ 46 | dataset.scene=${subdir_name2} \ 47 | dataset.root_dir=${root_dir}/${subdir_name1}/${subdir_name2} \ 48 | trial_name=${subdir_name1} 49 | rm -r exp/neus-openmaterial-wmask-${subdir_name2}/${subdir_name1} 50 | 51 | new_line=" with open(os.path.join(f'../output-ablation', self.config.dataset.object, f'{self.config.dataset.scene}+insr', \"instant-nsr-pl-wmask.txt\"), \"a\") as file:" 52 | sed -i "219s|.*|$new_line|" systems/nerf.py 53 | CUDA_VISIBLE_DEVICES=1 python launch.py \ 54 | --config configs/nerf-openmaterial-wmask.yaml \ 55 | --gpu 0 \ 56 | --train \ 57 | dataset.bsdf_name=${bsdf_name} \ 58 | dataset.object=${subdir_name1} \ 59 | dataset.scene=${subdir_name2} \ 60 | dataset.root_dir=${root_dir}/${subdir_name1}/${subdir_name2} \ 61 | trial_name=${subdir_name1} \ 62 | render_save_dir=../output-ablation 63 | rm -r exp/nerf-openmaterial-wmask-${subdir_name2}/${subdir_name1} 64 | 65 | done 66 | done 67 | else 68 | echo "no openmaterial dataset, please generate before running." 69 | fi 70 | # rm -rf exp/ -------------------------------------------------------------------------------- /instant-nsr-pl/.gitignore: -------------------------------------------------------------------------------- 1 | 2 | # Created by https://www.toptal.com/developers/gitignore/api/python 3 | # Edit at https://www.toptal.com/developers/gitignore?templates=python 4 | 5 | ### Python ### 6 | # Byte-compiled / optimized / DLL files 7 | __pycache__/ 8 | *.py[cod] 9 | *$py.class 10 | 11 | # C extensions 12 | *.so 13 | 14 | # Distribution / packaging 15 | .Python 16 | build/ 17 | develop-eggs/ 18 | dist/ 19 | downloads/ 20 | eggs/ 21 | .eggs/ 22 | lib/ 23 | lib64/ 24 | parts/ 25 | sdist/ 26 | var/ 27 | wheels/ 28 | share/python-wheels/ 29 | *.egg-info/ 30 | .installed.cfg 31 | *.egg 32 | MANIFEST 33 | 34 | # PyInstaller 35 | # Usually these files are written by a python script from a template 36 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 37 | *.manifest 38 | *.spec 39 | 40 | # Installer logs 41 | pip-log.txt 42 | pip-delete-this-directory.txt 43 | 44 | # Unit test / coverage reports 45 | htmlcov/ 46 | .tox/ 47 | .nox/ 48 | .coverage 49 | .coverage.* 50 | .cache 51 | nosetests.xml 52 | coverage.xml 53 | *.cover 54 | *.py,cover 55 | .hypothesis/ 56 | .pytest_cache/ 57 | cover/ 58 | 59 | # Translations 60 | *.mo 61 | *.pot 62 | 63 | # Django stuff: 64 | *.log 65 | local_settings.py 66 | db.sqlite3 67 | db.sqlite3-journal 68 | 69 | # Flask stuff: 70 | instance/ 71 | .webassets-cache 72 | 73 | # Scrapy stuff: 74 | .scrapy 75 | 76 | # Sphinx documentation 77 | docs/_build/ 78 | 79 | # PyBuilder 80 | .pybuilder/ 81 | target/ 82 | 83 | # Jupyter Notebook 84 | .ipynb_checkpoints 85 | 86 | # IPython 87 | profile_default/ 88 | ipython_config.py 89 | 90 | # pyenv 91 | # For a library or package, you might want to ignore these files since the code is 92 | # intended to run in multiple environments; otherwise, check them in: 93 | # .python-version 94 | 95 | # pipenv 96 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 97 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 98 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 99 | # install all needed dependencies. 100 | #Pipfile.lock 101 | 102 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 103 | __pypackages__/ 104 | 105 | # Celery stuff 106 | celerybeat-schedule 107 | celerybeat.pid 108 | 109 | # SageMath parsed files 110 | *.sage.py 111 | 112 | # Environments 113 | .env 114 | .venv 115 | env/ 116 | venv/ 117 | ENV/ 118 | env.bak/ 119 | venv.bak/ 120 | 121 | # Spyder project settings 122 | .spyderproject 123 | .spyproject 124 | 125 | # Rope project settings 126 | .ropeproject 127 | 128 | # mkdocs documentation 129 | /site 130 | 131 | # mypy 132 | .mypy_cache/ 133 | .dmypy.json 134 | dmypy.json 135 | 136 | # Pyre type checker 137 | .pyre/ 138 | 139 | # pytype static type analyzer 140 | .pytype/ 141 | 142 | # Cython debug symbols 143 | cython_debug/ 144 | 145 | # End of https://www.toptal.com/developers/gitignore/api/python 146 | 147 | .DS_Store 148 | .vscode/ 149 | exp/ 150 | runs/ 151 | load/ 152 | extern/ 153 | -------------------------------------------------------------------------------- /instant-nsr-pl/scripts/imgs2poses.py: -------------------------------------------------------------------------------- 1 | 2 | """ 3 | This file is adapted from https://github.com/Fyusion/LLFF. 4 | """ 5 | 6 | import os 7 | import sys 8 | import argparse 9 | import subprocess 10 | 11 | 12 | def run_colmap(basedir, match_type): 13 | logfile_name = os.path.join(basedir, 'colmap_output.txt') 14 | logfile = open(logfile_name, 'w') 15 | 16 | feature_extractor_args = [ 17 | 'colmap', 'feature_extractor', 18 | '--database_path', os.path.join(basedir, 'database.db'), 19 | '--image_path', os.path.join(basedir, 'images'), 20 | '--ImageReader.single_camera', '1' 21 | ] 22 | feat_output = ( subprocess.check_output(feature_extractor_args, universal_newlines=True) ) 23 | logfile.write(feat_output) 24 | print('Features extracted') 25 | 26 | exhaustive_matcher_args = [ 27 | 'colmap', match_type, 28 | '--database_path', os.path.join(basedir, 'database.db'), 29 | ] 30 | 31 | match_output = ( subprocess.check_output(exhaustive_matcher_args, universal_newlines=True) ) 32 | logfile.write(match_output) 33 | print('Features matched') 34 | 35 | p = os.path.join(basedir, 'sparse') 36 | if not os.path.exists(p): 37 | os.makedirs(p) 38 | 39 | mapper_args = [ 40 | 'colmap', 'mapper', 41 | '--database_path', os.path.join(basedir, 'database.db'), 42 | '--image_path', os.path.join(basedir, 'images'), 43 | '--output_path', os.path.join(basedir, 'sparse'), # --export_path changed to --output_path in colmap 3.6 44 | '--Mapper.num_threads', '16', 45 | '--Mapper.init_min_tri_angle', '4', 46 | '--Mapper.multiple_models', '0', 47 | '--Mapper.extract_colors', '0', 48 | ] 49 | 50 | map_output = ( subprocess.check_output(mapper_args, universal_newlines=True) ) 51 | logfile.write(map_output) 52 | logfile.close() 53 | print('Sparse map created') 54 | 55 | print( 'Finished running COLMAP, see {} for logs'.format(logfile_name) ) 56 | 57 | 58 | def gen_poses(basedir, match_type): 59 | files_needed = ['{}.bin'.format(f) for f in ['cameras', 'images', 'points3D']] 60 | if os.path.exists(os.path.join(basedir, 'sparse/0')): 61 | files_had = os.listdir(os.path.join(basedir, 'sparse/0')) 62 | else: 63 | files_had = [] 64 | if not all([f in files_had for f in files_needed]): 65 | print( 'Need to run COLMAP' ) 66 | run_colmap(basedir, match_type) 67 | else: 68 | print('Don\'t need to run COLMAP') 69 | 70 | return True 71 | 72 | 73 | if __name__=='__main__': 74 | parser = argparse.ArgumentParser() 75 | parser.add_argument('--match_type', type=str, 76 | default='exhaustive_matcher', help='type of matcher used. Valid options: \ 77 | exhaustive_matcher sequential_matcher. Other matchers not supported at this time') 78 | parser.add_argument('scenedir', type=str, 79 | help='input scene directory') 80 | args = parser.parse_args() 81 | 82 | if args.match_type != 'exhaustive_matcher' and args.match_type != 'sequential_matcher': 83 | print('ERROR: matcher type ' + args.match_type + ' is not valid. Aborting') 84 | sys.exit() 85 | gen_poses(args.scenedir, args.match_type) 86 | -------------------------------------------------------------------------------- /instant-nsr-pl/configs/neus-openmaterial-wmask.yaml: -------------------------------------------------------------------------------- 1 | name: neus-openmaterial-wmask-${dataset.scene} 2 | tag: "" 3 | seed: 42 4 | 5 | dataset: 6 | name: openmaterial-wmask 7 | scene: abandoned_bakery_4k-conductor-Ag 8 | root_dir: ../datasets/openmaterial/Pineapple/${dataset.scene} 9 | img_wh: 10 | - 1600 11 | - 1200 12 | # img_downscale: 1 # specify training image size by either img_wh or img_downscale 13 | near_plane: 1.3 14 | far_plane: 4.6 15 | train_split: "train" 16 | test_split: "test" 17 | bsdf_name: "" 18 | object: "" 19 | 20 | model: 21 | name: neus 22 | radius: 1.0 23 | num_samples_per_ray: 1024 24 | train_num_rays: 256 25 | max_train_num_rays: 8192 26 | grid_prune: true 27 | grid_prune_occ_thre: 0.001 28 | dynamic_ray_sampling: true 29 | batch_image_sampling: true 30 | randomized: true 31 | ray_chunk: 4096 32 | cos_anneal_end: 20000 33 | learned_background: false 34 | background_color: white 35 | variance: 36 | init_val: 0.3 37 | modulate: false 38 | geometry: 39 | name: volume-sdf 40 | radius: ${model.radius} 41 | feature_dim: 13 42 | grad_type: analytic 43 | isosurface: 44 | method: mc 45 | resolution: 512 46 | chunk: 2097152 47 | threshold: 0.0 48 | xyz_encoding_config: 49 | otype: HashGrid 50 | n_levels: 16 51 | n_features_per_level: 2 52 | log2_hashmap_size: 19 53 | base_resolution: 32 54 | per_level_scale: 1.3195079107728942 55 | include_xyz: true 56 | mlp_network_config: 57 | otype: VanillaMLP 58 | activation: ReLU 59 | output_activation: none 60 | n_neurons: 64 61 | n_hidden_layers: 1 62 | sphere_init: true 63 | sphere_init_radius: 0.5 64 | weight_norm: true 65 | texture: 66 | name: volume-radiance 67 | input_feature_dim: ${add:${model.geometry.feature_dim},3} # surface normal as additional input 68 | dir_encoding_config: 69 | otype: SphericalHarmonics 70 | degree: 4 71 | mlp_network_config: 72 | otype: FullyFusedMLP 73 | activation: ReLU 74 | output_activation: none 75 | n_neurons: 64 76 | n_hidden_layers: 2 77 | color_activation: sigmoid 78 | 79 | system: 80 | name: neus-system 81 | loss: 82 | lambda_rgb_mse: 0.0 83 | lambda_rgb_l1: 1.0 84 | lambda_mask: 0.1 85 | lambda_eikonal: 0.1 86 | lambda_curvature: 0.0 87 | lambda_sparsity: 0.0 88 | lambda_distortion: 0.0 89 | lambda_distortion_bg: 0.0 90 | lambda_opaque: 0.0 91 | sparsity_scale: 1.0 92 | optimizer: 93 | name: AdamW 94 | args: 95 | lr: 0.01 96 | betas: [0.9, 0.99] 97 | eps: 1.e-15 98 | params: 99 | geometry: 100 | lr: 0.01 101 | texture: 102 | lr: 0.01 103 | variance: 104 | lr: 0.001 105 | warmup_steps: 1000 106 | scheduler: 107 | name: SequentialLR 108 | interval: step 109 | milestones: 110 | - ${system.warmup_steps} 111 | schedulers: 112 | - name: LinearLR # linear warm-up in the first system.warmup_steps steps 113 | args: 114 | start_factor: 0.01 115 | end_factor: 1.0 116 | total_iters: ${system.warmup_steps} 117 | - name: ExponentialLR 118 | args: 119 | gamma: ${calc_exp_lr_decay_rate:0.1,${sub:${trainer.max_steps},${system.warmup_steps}}} 120 | 121 | checkpoint: 122 | save_top_k: -1 123 | every_n_train_steps: ${trainer.max_steps} 124 | 125 | export: 126 | chunk_size: 2097152 127 | export_vertex_color: True 128 | 129 | trainer: 130 | max_steps: 15000 131 | log_every_n_steps: 100 132 | num_sanity_val_steps: 0 133 | val_check_interval: 5000 134 | limit_train_batches: 1.0 135 | limit_val_batches: 2 136 | enable_progress_bar: true 137 | precision: 16 138 | -------------------------------------------------------------------------------- /instant-nsr-pl/utils/callbacks.py: -------------------------------------------------------------------------------- 1 | import os 2 | import subprocess 3 | import shutil 4 | from utils.misc import dump_config, parse_version 5 | 6 | 7 | import pytorch_lightning 8 | if parse_version(pytorch_lightning.__version__) > parse_version('1.8'): 9 | from pytorch_lightning.callbacks import Callback 10 | else: 11 | from pytorch_lightning.callbacks.base import Callback 12 | from pytorch_lightning.utilities.rank_zero import rank_zero_only, rank_zero_warn 13 | from pytorch_lightning.callbacks.progress import TQDMProgressBar 14 | 15 | 16 | class VersionedCallback(Callback): 17 | def __init__(self, save_root, version=None, use_version=True): 18 | self.save_root = save_root 19 | self._version = version 20 | self.use_version = use_version 21 | 22 | @property 23 | def version(self) -> int: 24 | """Get the experiment version. 25 | 26 | Returns: 27 | The experiment version if specified else the next version. 28 | """ 29 | if self._version is None: 30 | self._version = self._get_next_version() 31 | return self._version 32 | 33 | def _get_next_version(self): 34 | existing_versions = [] 35 | if os.path.isdir(self.save_root): 36 | for f in os.listdir(self.save_root): 37 | bn = os.path.basename(f) 38 | if bn.startswith("version_"): 39 | dir_ver = os.path.splitext(bn)[0].split("_")[1].replace("/", "") 40 | existing_versions.append(int(dir_ver)) 41 | if len(existing_versions) == 0: 42 | return 0 43 | return max(existing_versions) + 1 44 | 45 | @property 46 | def savedir(self): 47 | if not self.use_version: 48 | return self.save_root 49 | return os.path.join(self.save_root, self.version if isinstance(self.version, str) else f"version_{self.version}") 50 | 51 | 52 | class CodeSnapshotCallback(VersionedCallback): 53 | def __init__(self, save_root, version=None, use_version=True): 54 | super().__init__(save_root, version, use_version) 55 | 56 | def get_file_list(self): 57 | return [ 58 | b.decode() for b in 59 | set(subprocess.check_output('git ls-files', shell=True).splitlines()) | 60 | set(subprocess.check_output('git ls-files --others --exclude-standard', shell=True).splitlines()) 61 | ] 62 | 63 | @rank_zero_only 64 | def save_code_snapshot(self): 65 | os.makedirs(self.savedir, exist_ok=True) 66 | for f in self.get_file_list(): 67 | if not os.path.exists(f) or os.path.isdir(f): 68 | continue 69 | os.makedirs(os.path.join(self.savedir, os.path.dirname(f)), exist_ok=True) 70 | shutil.copyfile(f, os.path.join(self.savedir, f)) 71 | 72 | def on_fit_start(self, trainer, pl_module): 73 | try: 74 | self.save_code_snapshot() 75 | except: 76 | rank_zero_warn("Code snapshot is not saved. Please make sure you have git installed and are in a git repository.") 77 | 78 | 79 | class ConfigSnapshotCallback(VersionedCallback): 80 | def __init__(self, config, save_root, version=None, use_version=True): 81 | super().__init__(save_root, version, use_version) 82 | self.config = config 83 | 84 | @rank_zero_only 85 | def save_config_snapshot(self): 86 | os.makedirs(self.savedir, exist_ok=True) 87 | dump_config(os.path.join(self.savedir, 'parsed.yaml'), self.config) 88 | shutil.copyfile(self.config.cmd_args['config'], os.path.join(self.savedir, 'raw.yaml')) 89 | 90 | def on_fit_start(self, trainer, pl_module): 91 | self.save_config_snapshot() 92 | 93 | 94 | class CustomProgressBar(TQDMProgressBar): 95 | def get_metrics(self, *args, **kwargs): 96 | # don't show the version number 97 | items = super().get_metrics(*args, **kwargs) 98 | items.pop("v_num", None) 99 | return items 100 | -------------------------------------------------------------------------------- /instant-nsr-pl/models/utils.py: -------------------------------------------------------------------------------- 1 | import gc 2 | from collections import defaultdict 3 | 4 | import torch 5 | import torch.nn as nn 6 | import torch.nn.functional as F 7 | from torch.autograd import Function 8 | from torch.cuda.amp import custom_bwd, custom_fwd 9 | 10 | import tinycudann as tcnn 11 | 12 | 13 | def chunk_batch(func, chunk_size, move_to_cpu, *args, **kwargs): 14 | B = None 15 | for arg in args: 16 | if isinstance(arg, torch.Tensor): 17 | B = arg.shape[0] 18 | break 19 | out = defaultdict(list) 20 | out_type = None 21 | for i in range(0, B, chunk_size): 22 | out_chunk = func(*[arg[i:i+chunk_size] if isinstance(arg, torch.Tensor) else arg for arg in args], **kwargs) 23 | if out_chunk is None: 24 | continue 25 | out_type = type(out_chunk) 26 | if isinstance(out_chunk, torch.Tensor): 27 | out_chunk = {0: out_chunk} 28 | elif isinstance(out_chunk, tuple) or isinstance(out_chunk, list): 29 | chunk_length = len(out_chunk) 30 | out_chunk = {i: chunk for i, chunk in enumerate(out_chunk)} 31 | elif isinstance(out_chunk, dict): 32 | pass 33 | else: 34 | print(f'Return value of func must be in type [torch.Tensor, list, tuple, dict], get {type(out_chunk)}.') 35 | exit(1) 36 | for k, v in out_chunk.items(): 37 | v = v if torch.is_grad_enabled() else v.detach() 38 | v = v.cpu() if move_to_cpu else v 39 | out[k].append(v) 40 | 41 | if out_type is None: 42 | return 43 | 44 | out = {k: torch.cat(v, dim=0) for k, v in out.items()} 45 | if out_type is torch.Tensor: 46 | return out[0] 47 | elif out_type in [tuple, list]: 48 | return out_type([out[i] for i in range(chunk_length)]) 49 | elif out_type is dict: 50 | return out 51 | 52 | 53 | class _TruncExp(Function): # pylint: disable=abstract-method 54 | # Implementation from torch-ngp: 55 | # https://github.com/ashawkey/torch-ngp/blob/93b08a0d4ec1cc6e69d85df7f0acdfb99603b628/activation.py 56 | @staticmethod 57 | @custom_fwd(cast_inputs=torch.float32) 58 | def forward(ctx, x): # pylint: disable=arguments-differ 59 | ctx.save_for_backward(x) 60 | return torch.exp(x) 61 | 62 | @staticmethod 63 | @custom_bwd 64 | def backward(ctx, g): # pylint: disable=arguments-differ 65 | x = ctx.saved_tensors[0] 66 | return g * torch.exp(torch.clamp(x, max=15)) 67 | 68 | trunc_exp = _TruncExp.apply 69 | 70 | 71 | def get_activation(name): 72 | if name is None: 73 | return lambda x: x 74 | name = name.lower() 75 | if name == 'none': 76 | return lambda x: x 77 | elif name.startswith('scale'): 78 | scale_factor = float(name[5:]) 79 | return lambda x: x.clamp(0., scale_factor) / scale_factor 80 | elif name.startswith('clamp'): 81 | clamp_max = float(name[5:]) 82 | return lambda x: x.clamp(0., clamp_max) 83 | elif name.startswith('mul'): 84 | mul_factor = float(name[3:]) 85 | return lambda x: x * mul_factor 86 | elif name == 'lin2srgb': 87 | return lambda x: torch.where(x > 0.0031308, torch.pow(torch.clamp(x, min=0.0031308), 1.0/2.4)*1.055 - 0.055, 12.92*x).clamp(0., 1.) 88 | elif name == 'trunc_exp': 89 | return trunc_exp 90 | elif name.startswith('+') or name.startswith('-'): 91 | return lambda x: x + float(name) 92 | elif name == 'sigmoid': 93 | return lambda x: torch.sigmoid(x) 94 | elif name == 'tanh': 95 | return lambda x: torch.tanh(x) 96 | else: 97 | return getattr(F, name) 98 | 99 | 100 | def dot(x, y): 101 | return torch.sum(x*y, -1, keepdim=True) 102 | 103 | 104 | def reflect(x, n): 105 | return 2 * dot(x, n) * n - x 106 | 107 | 108 | def scale_anything(dat, inp_scale, tgt_scale): 109 | if inp_scale is None: 110 | inp_scale = [dat.min(), dat.max()] 111 | dat = (dat - inp_scale[0]) / (inp_scale[1] - inp_scale[0]) 112 | dat = dat * (tgt_scale[1] - tgt_scale[0]) + tgt_scale[0] 113 | return dat 114 | 115 | 116 | def cleanup(): 117 | gc.collect() 118 | torch.cuda.empty_cache() 119 | tcnn.free_temporary_memory() 120 | -------------------------------------------------------------------------------- /sum_metrics.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | import pandas as pd 4 | import re 5 | import argparse 6 | import glob 7 | 8 | bsdf_names = [ 9 | 'diffuse', 10 | 'dielectric', 11 | 'roughdielectric', 12 | 'conductor', 13 | 'roughconductor', 14 | 'plastic', 15 | 'roughplastic' 16 | ] 17 | 18 | parser = argparse.ArgumentParser() 19 | parser.add_argument('--output_dir', type=str, default='output') 20 | args = parser.parse_args() 21 | 22 | def sum_instant_nsr_pl(): 23 | psnr_dir, ssim_dir, lpips_dir, count_dir = {}, {}, {}, {} 24 | folder_path = args.output_dir 25 | object_names = os.listdir(folder_path) 26 | for name in bsdf_names: 27 | count_dir[name] = 0 28 | psnr_dir[name] = 0.0 29 | ssim_dir[name] = 0.0 30 | lpips_dir[name] = 0.0 31 | 32 | for file_name in object_names: 33 | txt_path = os.path.join(folder_path, file_name, 'instant-nsr-pl-wmask.txt') 34 | 35 | with open(txt_path, 'r') as f: 36 | txt_data = f.read() 37 | psnr_ = re.search(r'PSNR=(\d+\.\d+)', txt_data).group(1) 38 | ssim_ = re.search(r'SSIM=(\d+\.\d+)', txt_data).group(1) 39 | lpips_ = re.search(r'lpips=(\d+\.\d+)', txt_data).group(1) 40 | method = txt_data.split(':')[1] 41 | bsdf_name = txt_data.split(':')[2] 42 | 43 | for name in bsdf_names: 44 | if bsdf_name.startswith(name): 45 | psnr_dir[name] += float(psnr_) 46 | ssim_dir[name] += float(ssim_) 47 | lpips_dir[name] += float(lpips_) 48 | count_dir[name] += 1 49 | 50 | # print(psnr_dir) 51 | # print(ssim_dir) 52 | # print(lpips_dir) 53 | 54 | for name in bsdf_names: 55 | if count_dir[name] > 0: 56 | print(f"[+] {name} result: {count_dir[name]}") 57 | psnr_dir[name] = psnr_dir[name] / count_dir[name] 58 | ssim_dir[name] = ssim_dir[name] / count_dir[name] 59 | lpips_dir[name] = lpips_dir[name] / count_dir[name] 60 | return psnr_dir, ssim_dir, lpips_dir 61 | 62 | def mesh_cds(): 63 | cds_dir, count_dir = {}, {} 64 | folder_path = args.output_dir 65 | object_names = os.listdir(folder_path) 66 | count_dir['instant-nsr-pl-wmask'] = {} 67 | cds_dir['instant-nsr-pl-wmask'] = {} 68 | for name in bsdf_names: 69 | count_dir['instant-nsr-pl-wmask'][name] = 0 70 | cds_dir['instant-nsr-pl-wmask'][name] = 0.0 71 | 72 | for file_name in object_names: 73 | txt_paths = glob.glob(os.path.join(folder_path, file_name, '*mesh-output.txt')) 74 | print(txt_paths) 75 | for txt_path in txt_paths: 76 | with open(txt_path, 'r') as f: 77 | for line in f: 78 | txt_data = line.strip() 79 | method = txt_data.split(':')[1] 80 | bsdf_name = txt_data.split(':')[2] 81 | cds = txt_data.split(':')[-1] 82 | 83 | for name in bsdf_names: 84 | if bsdf_name.startswith(name): 85 | cds_dir[method][name] += float(cds) 86 | count_dir[method][name] += 1 87 | 88 | for name in bsdf_names: 89 | cds_dir['instant-nsr-pl-wmask'][name] = cds_dir['instant-nsr-pl-wmask'][name] / count_dir['instant-nsr-pl-wmask'][name] 90 | return cds_dir 91 | 92 | 93 | if __name__ == '__main__': 94 | psnr, ssim, lpips = {}, {}, {} 95 | psnr['instant-nsr-pl(nerf-wmask)'], ssim['instant-nsr-pl(nerf-wmask)'], lpips['instant-nsr-pl(nerf-wmask)'] = sum_instant_nsr_pl() 96 | cds = mesh_cds() 97 | rows = ['instant-nsr-pl(nerf-wmask)'] 98 | columns = [name for name in bsdf_names] 99 | PSNR_df = pd.DataFrame(index=rows, columns=columns) 100 | for r in rows: 101 | for c in columns: 102 | PSNR_df.at[r, c] = psnr[r][c] 103 | print("PSNR:") 104 | print(PSNR_df) 105 | print('\n') 106 | SSIM_df = pd.DataFrame(index=rows, columns=columns) 107 | for r in rows: 108 | for c in columns: 109 | SSIM_df.at[r, c] = ssim[r][c] 110 | print("SSIM:") 111 | print(SSIM_df) 112 | print('\n') 113 | LPIPS_df = pd.DataFrame(index=rows, columns=columns) 114 | for r in rows: 115 | for c in columns: 116 | LPIPS_df.at[r, c] = lpips[r][c] 117 | print("LPIPS:") 118 | print(LPIPS_df) 119 | print('\n') 120 | 121 | rows = ['instant-nsr-pl-wmask'] 122 | columns = [name for name in bsdf_names] 123 | CDs_df = pd.DataFrame(index=rows, columns=columns) 124 | for r in rows: 125 | for c in columns: 126 | CDs_df.at[r, c] = cds[r][c] 127 | print("Chamfer Distance:") 128 | print(CDs_df) 129 | -------------------------------------------------------------------------------- /instant-nsr-pl/datasets/openmaterial_womask.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | import math 4 | import numpy as np 5 | from PIL import Image 6 | 7 | import torch 8 | from torch.utils.data import Dataset, DataLoader, IterableDataset 9 | import torchvision.transforms.functional as TF 10 | 11 | import pytorch_lightning as pl 12 | 13 | import datasets 14 | from models.ray_utils import get_ray_directions 15 | from utils.misc import get_rank 16 | 17 | 18 | class BlenderDatasetBase(): 19 | def setup(self, config, split): 20 | self.config = config 21 | self.split = split 22 | self.rank = get_rank() 23 | 24 | self.has_mask = True 25 | self.apply_mask = False 26 | 27 | with open(os.path.join(self.config.root_dir, f"transforms_{self.split}.json"), 'r') as f: 28 | meta = json.load(f) 29 | 30 | if 'w' in meta and 'h' in meta: 31 | W, H = int(meta['w']), int(meta['h']) 32 | else: 33 | W, H = 1600, 1200 34 | 35 | if 'img_wh' in self.config: 36 | w, h = self.config.img_wh 37 | assert round(W / w * h) == H 38 | elif 'img_downscale' in self.config: 39 | w, h = W // self.config.img_downscale, H // self.config.img_downscale 40 | else: 41 | raise KeyError("Either img_wh or img_downscale should be specified.") 42 | 43 | self.w, self.h = w, h 44 | self.img_wh = (self.w, self.h) 45 | 46 | self.near, self.far = self.config.near_plane, self.config.far_plane 47 | 48 | self.focal_x = meta['fl_x'] # scaled focal length 49 | self.focal_y = meta['fl_y'] 50 | 51 | # ray directions for all pixels, same for all images (same H, W, focal) 52 | self.directions = \ 53 | get_ray_directions(self.w, self.h, self.focal_x, self.focal_y, self.w//2, self.h//2).to(self.rank) # (h, w, 3) 54 | 55 | self.all_c2w, self.all_images, self.all_fg_masks = [], [], [] 56 | 57 | for i, frame in enumerate(meta['frames']): 58 | c2w = torch.from_numpy(np.array(frame['transform_matrix'])[:3, :4]) 59 | self.all_c2w.append(c2w) 60 | 61 | img_path = os.path.join(self.config.root_dir, f"{frame['file_path']}") 62 | img = Image.open(img_path) 63 | img = img.resize(self.img_wh, Image.BICUBIC) 64 | img = TF.to_tensor(img).permute(1, 2, 0) # (4, h, w) => (h, w, 4) 65 | 66 | self.all_fg_masks.append(torch.ones_like(img[...,0], device=img.device)) # (h, w) 67 | self.all_images.append(img[...,:3]) 68 | 69 | self.all_c2w, self.all_images, self.all_fg_masks = \ 70 | torch.stack(self.all_c2w, dim=0).float().to(self.rank), \ 71 | torch.stack(self.all_images, dim=0).float().to(self.rank), \ 72 | torch.stack(self.all_fg_masks, dim=0).float().to(self.rank) 73 | 74 | 75 | class BlenderDataset(Dataset, BlenderDatasetBase): 76 | def __init__(self, config, split): 77 | self.setup(config, split) 78 | 79 | def __len__(self): 80 | return len(self.all_images) 81 | 82 | def __getitem__(self, index): 83 | return { 84 | 'index': index 85 | } 86 | 87 | 88 | class BlenderIterableDataset(IterableDataset, BlenderDatasetBase): 89 | def __init__(self, config, split): 90 | self.setup(config, split) 91 | 92 | def __iter__(self): 93 | while True: 94 | yield {} 95 | 96 | 97 | @datasets.register('openmaterial-womask') 98 | class BlenderDataModule(pl.LightningDataModule): 99 | def __init__(self, config): 100 | super().__init__() 101 | self.config = config 102 | 103 | def setup(self, stage=None): 104 | if stage in [None, 'fit']: 105 | self.train_dataset = BlenderIterableDataset(self.config, self.config.train_split) 106 | if stage in [None, 'test']: 107 | self.test_dataset = BlenderDataset(self.config, self.config.test_split) 108 | if stage in [None, 'predict']: 109 | self.predict_dataset = BlenderDataset(self.config, self.config.train_split) 110 | 111 | def prepare_data(self): 112 | pass 113 | 114 | def general_loader(self, dataset, batch_size): 115 | sampler = None 116 | return DataLoader( 117 | dataset, 118 | num_workers=os.cpu_count(), 119 | batch_size=batch_size, 120 | pin_memory=True, 121 | sampler=sampler 122 | ) 123 | 124 | def train_dataloader(self): 125 | return self.general_loader(self.train_dataset, batch_size=1) 126 | 127 | def test_dataloader(self): 128 | return self.general_loader(self.test_dataset, batch_size=1) 129 | 130 | def predict_dataloader(self): 131 | return self.general_loader(self.predict_dataset, batch_size=1) 132 | -------------------------------------------------------------------------------- /instant-nsr-pl/configs/neus-openmaterial-womask.yaml: -------------------------------------------------------------------------------- 1 | name: neus-openmaterial-womask-${dataset.scene} 2 | tag: "" 3 | seed: 42 4 | 5 | dataset: 6 | name: openmaterial-womask 7 | scene: abandoned_bakery_4k-conductor-Ag 8 | root_dir: ../datasets/openmaterial/Pineapple/${dataset.scene} 9 | img_wh: 10 | - 1600 11 | - 1200 12 | # img_downscale: 1 # specify training image size by either img_wh or img_downscale 13 | near_plane: 1.3 14 | far_plane: 4.6 15 | train_split: "train" 16 | test_split: "test" 17 | bsdf_name: "" 18 | object: "" 19 | 20 | model: 21 | name: neus 22 | radius: 1.0 23 | num_samples_per_ray: 1024 24 | train_num_rays: 256 25 | max_train_num_rays: 8192 26 | grid_prune: true 27 | grid_prune_occ_thre: 0.001 28 | dynamic_ray_sampling: true 29 | batch_image_sampling: true 30 | randomized: true 31 | ray_chunk: 4096 32 | cos_anneal_end: 20000 33 | learned_background: true 34 | background_color: random 35 | variance: 36 | init_val: 0.3 37 | modulate: false 38 | geometry: 39 | name: volume-sdf 40 | radius: ${model.radius} 41 | feature_dim: 13 42 | grad_type: analytic 43 | isosurface: 44 | method: mc 45 | resolution: 512 46 | chunk: 2097152 47 | threshold: 0.0 48 | xyz_encoding_config: 49 | otype: HashGrid 50 | n_levels: 16 51 | n_features_per_level: 2 52 | log2_hashmap_size: 19 53 | base_resolution: 32 54 | per_level_scale: 1.3195079107728942 55 | include_xyz: true 56 | mlp_network_config: 57 | otype: VanillaMLP 58 | activation: ReLU 59 | output_activation: none 60 | n_neurons: 64 61 | n_hidden_layers: 1 62 | sphere_init: true 63 | sphere_init_radius: 0.5 64 | weight_norm: true 65 | texture: 66 | name: volume-radiance 67 | input_feature_dim: ${add:${model.geometry.feature_dim},3} # surface normal as additional input 68 | dir_encoding_config: 69 | otype: SphericalHarmonics 70 | degree: 4 71 | mlp_network_config: 72 | otype: FullyFusedMLP 73 | activation: ReLU 74 | output_activation: none 75 | n_neurons: 64 76 | n_hidden_layers: 2 77 | color_activation: sigmoid 78 | num_samples_per_ray_bg: 64 79 | geometry_bg: 80 | name: volume-density 81 | radius: ${model.radius} 82 | feature_dim: 8 83 | density_activation: trunc_exp 84 | density_bias: -1 85 | isosurface: null 86 | xyz_encoding_config: 87 | otype: HashGrid 88 | n_levels: 16 89 | n_features_per_level: 2 90 | log2_hashmap_size: 19 91 | base_resolution: 32 92 | per_level_scale: 1.3195079107728942 93 | mlp_network_config: 94 | otype: VanillaMLP 95 | activation: ReLU 96 | output_activation: none 97 | n_neurons: 64 98 | n_hidden_layers: 1 99 | texture_bg: 100 | name: volume-radiance 101 | input_feature_dim: ${model.geometry_bg.feature_dim} 102 | dir_encoding_config: 103 | otype: SphericalHarmonics 104 | degree: 4 105 | mlp_network_config: 106 | otype: VanillaMLP 107 | activation: ReLU 108 | output_activation: none 109 | n_neurons: 64 110 | n_hidden_layers: 2 111 | color_activation: sigmoid 112 | 113 | system: 114 | name: neus-system 115 | loss: 116 | lambda_rgb_mse: 0.0 117 | lambda_rgb_l1: 1.0 118 | lambda_mask: 0.0 119 | lambda_eikonal: 0.1 120 | lambda_curvature: 0.0 121 | lambda_sparsity: 0.0 122 | lambda_distortion: 0.0 123 | lambda_distortion_bg: 0.0 124 | lambda_opaque: 0.0 125 | sparsity_scale: 1.0 126 | optimizer: 127 | name: AdamW 128 | args: 129 | lr: 0.01 130 | betas: [0.9, 0.99] 131 | eps: 1.e-15 132 | params: 133 | geometry: 134 | lr: 0.01 135 | texture: 136 | lr: 0.01 137 | geometry_bg: 138 | lr: 0.01 139 | texture_bg: 140 | lr: 0.01 141 | variance: 142 | lr: 0.001 143 | warmup_steps: 1000 144 | scheduler: 145 | name: SequentialLR 146 | interval: step 147 | milestones: 148 | - ${system.warmup_steps} 149 | schedulers: 150 | - name: LinearLR # linear warm-up in the first system.warmup_steps steps 151 | args: 152 | start_factor: 0.01 153 | end_factor: 1.0 154 | total_iters: ${system.warmup_steps} 155 | - name: ExponentialLR 156 | args: 157 | gamma: ${calc_exp_lr_decay_rate:0.1,${sub:${trainer.max_steps},${system.warmup_steps}}} 158 | 159 | checkpoint: 160 | save_top_k: -1 161 | every_n_train_steps: ${trainer.max_steps} 162 | 163 | export: 164 | chunk_size: 2097152 165 | export_vertex_color: True 166 | 167 | trainer: 168 | max_steps: 15000 169 | log_every_n_steps: 100 170 | num_sanity_val_steps: 0 171 | val_check_interval: 5000 172 | limit_train_batches: 1.0 173 | limit_val_batches: 2 174 | enable_progress_bar: true 175 | precision: 16 176 | -------------------------------------------------------------------------------- /instant-nsr-pl/systems/base.py: -------------------------------------------------------------------------------- 1 | import pytorch_lightning as pl 2 | 3 | import models 4 | from systems.utils import parse_optimizer, parse_scheduler, update_module_step 5 | from utils.mixins import SaverMixin 6 | from utils.misc import config_to_primitive, get_rank 7 | 8 | 9 | class BaseSystem(pl.LightningModule, SaverMixin): 10 | """ 11 | Two ways to print to console: 12 | 1. self.print: correctly handle progress bar 13 | 2. rank_zero_info: use the logging module 14 | """ 15 | def __init__(self, config): 16 | super().__init__() 17 | self.config = config 18 | self.rank = get_rank() 19 | self.prepare() 20 | self.model = models.make(self.config.model.name, self.config.model) 21 | 22 | def prepare(self): 23 | pass 24 | 25 | def forward(self, batch): 26 | raise NotImplementedError 27 | 28 | def C(self, value): 29 | if isinstance(value, int) or isinstance(value, float): 30 | pass 31 | else: 32 | value = config_to_primitive(value) 33 | if not isinstance(value, list): 34 | raise TypeError('Scalar specification only supports list, got', type(value)) 35 | if len(value) == 3: 36 | value = [0] + value 37 | assert len(value) == 4 38 | start_step, start_value, end_value, end_step = value 39 | if isinstance(end_step, int): 40 | current_step = self.global_step 41 | value = start_value + (end_value - start_value) * max(min(1.0, (current_step - start_step) / (end_step - start_step)), 0.0) 42 | elif isinstance(end_step, float): 43 | current_step = self.current_epoch 44 | value = start_value + (end_value - start_value) * max(min(1.0, (current_step - start_step) / (end_step - start_step)), 0.0) 45 | return value 46 | 47 | def preprocess_data(self, batch, stage): 48 | pass 49 | 50 | """ 51 | Implementing on_after_batch_transfer of DataModule does the same. 52 | But on_after_batch_transfer does not support DP. 53 | """ 54 | def on_train_batch_start(self, batch, batch_idx, unused=0): 55 | self.dataset = self.trainer.datamodule.train_dataloader().dataset 56 | self.preprocess_data(batch, 'train') 57 | update_module_step(self.model, self.current_epoch, self.global_step) 58 | 59 | def on_validation_batch_start(self, batch, batch_idx, dataloader_idx): 60 | self.dataset = self.trainer.datamodule.val_dataloader().dataset 61 | self.preprocess_data(batch, 'validation') 62 | update_module_step(self.model, self.current_epoch, self.global_step) 63 | 64 | def on_test_batch_start(self, batch, batch_idx, dataloader_idx): 65 | self.dataset = self.trainer.datamodule.test_dataloader().dataset 66 | self.preprocess_data(batch, 'test') 67 | update_module_step(self.model, self.current_epoch, self.global_step) 68 | 69 | def on_predict_batch_start(self, batch, batch_idx, dataloader_idx): 70 | self.dataset = self.trainer.datamodule.predict_dataloader().dataset 71 | self.preprocess_data(batch, 'predict') 72 | update_module_step(self.model, self.current_epoch, self.global_step) 73 | 74 | def training_step(self, batch, batch_idx): 75 | raise NotImplementedError 76 | 77 | """ 78 | # aggregate outputs from different devices (DP) 79 | def training_step_end(self, out): 80 | pass 81 | """ 82 | 83 | """ 84 | # aggregate outputs from different iterations 85 | def training_epoch_end(self, out): 86 | pass 87 | """ 88 | 89 | def validation_step(self, batch, batch_idx): 90 | raise NotImplementedError 91 | 92 | """ 93 | # aggregate outputs from different devices when using DP 94 | def validation_step_end(self, out): 95 | pass 96 | """ 97 | 98 | def validation_epoch_end(self, out): 99 | """ 100 | Gather metrics from all devices, compute mean. 101 | Purge repeated results using data index. 102 | """ 103 | raise NotImplementedError 104 | 105 | def test_step(self, batch, batch_idx): 106 | raise NotImplementedError 107 | 108 | def test_epoch_end(self, out): 109 | """ 110 | Gather metrics from all devices, compute mean. 111 | Purge repeated results using data index. 112 | """ 113 | raise NotImplementedError 114 | 115 | def export(self): 116 | raise NotImplementedError 117 | 118 | def configure_optimizers(self): 119 | optim = parse_optimizer(self.config.system.optimizer, self.model) 120 | ret = { 121 | 'optimizer': optim, 122 | } 123 | if 'scheduler' in self.config.system: 124 | ret.update({ 125 | 'lr_scheduler': parse_scheduler(self.config.system.scheduler, optim), 126 | }) 127 | return ret 128 | 129 | -------------------------------------------------------------------------------- /instant-nsr-pl/datasets/blender.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | import math 4 | import numpy as np 5 | from PIL import Image 6 | 7 | import torch 8 | from torch.utils.data import Dataset, DataLoader, IterableDataset 9 | import torchvision.transforms.functional as TF 10 | 11 | import pytorch_lightning as pl 12 | 13 | import datasets 14 | from models.ray_utils import get_ray_directions 15 | from utils.misc import get_rank 16 | 17 | 18 | class BlenderDatasetBase(): 19 | def setup(self, config, split): 20 | self.config = config 21 | self.split = split 22 | self.rank = get_rank() 23 | 24 | self.has_mask = True 25 | self.apply_mask = True 26 | 27 | with open(os.path.join(self.config.root_dir, f"transforms_{self.split}.json"), 'r') as f: 28 | meta = json.load(f) 29 | 30 | if 'w' in meta and 'h' in meta: 31 | W, H = int(meta['w']), int(meta['h']) 32 | else: 33 | W, H = 800, 800 34 | 35 | if 'img_wh' in self.config: 36 | w, h = self.config.img_wh 37 | assert round(W / w * h) == H 38 | elif 'img_downscale' in self.config: 39 | w, h = W // self.config.img_downscale, H // self.config.img_downscale 40 | else: 41 | raise KeyError("Either img_wh or img_downscale should be specified.") 42 | 43 | self.w, self.h = w, h 44 | self.img_wh = (self.w, self.h) 45 | 46 | self.near, self.far = self.config.near_plane, self.config.far_plane 47 | 48 | self.focal = 0.5 * w / math.tan(0.5 * meta['camera_angle_x']) # scaled focal length 49 | 50 | # ray directions for all pixels, same for all images (same H, W, focal) 51 | self.directions = \ 52 | get_ray_directions(self.w, self.h, self.focal, self.focal, self.w//2, self.h//2).to(self.rank) # (h, w, 3) 53 | 54 | self.all_c2w, self.all_images, self.all_fg_masks = [], [], [] 55 | 56 | for i, frame in enumerate(meta['frames']): 57 | c2w = torch.from_numpy(np.array(frame['transform_matrix'])[:3, :4]) 58 | self.all_c2w.append(c2w) 59 | 60 | img_path = os.path.join(self.config.root_dir, f"{frame['file_path']}") 61 | img = Image.open(img_path) 62 | img = img.resize(self.img_wh, Image.BICUBIC) 63 | img = TF.to_tensor(img).permute(1, 2, 0) # (4, h, w) => (h, w, 4) 64 | 65 | self.all_fg_masks.append(img[..., -1]) # (h, w) 66 | self.all_images.append(img[...,:3]) 67 | 68 | self.all_c2w, self.all_images, self.all_fg_masks = \ 69 | torch.stack(self.all_c2w, dim=0).float().to(self.rank), \ 70 | torch.stack(self.all_images, dim=0).float().to(self.rank), \ 71 | torch.stack(self.all_fg_masks, dim=0).float().to(self.rank) 72 | 73 | 74 | class BlenderDataset(Dataset, BlenderDatasetBase): 75 | def __init__(self, config, split): 76 | self.setup(config, split) 77 | 78 | def __len__(self): 79 | return len(self.all_images) 80 | 81 | def __getitem__(self, index): 82 | return { 83 | 'index': index 84 | } 85 | 86 | 87 | class BlenderIterableDataset(IterableDataset, BlenderDatasetBase): 88 | def __init__(self, config, split): 89 | self.setup(config, split) 90 | 91 | def __iter__(self): 92 | while True: 93 | yield {} 94 | 95 | 96 | @datasets.register('blender') 97 | class BlenderDataModule(pl.LightningDataModule): 98 | def __init__(self, config): 99 | super().__init__() 100 | self.config = config 101 | 102 | def setup(self, stage=None): 103 | if stage in [None, 'fit']: 104 | self.train_dataset = BlenderIterableDataset(self.config, self.config.train_split) 105 | if stage in [None, 'fit', 'validate']: 106 | self.val_dataset = BlenderDataset(self.config, self.config.val_split) 107 | if stage in [None, 'test']: 108 | self.test_dataset = BlenderDataset(self.config, self.config.test_split) 109 | if stage in [None, 'predict']: 110 | self.predict_dataset = BlenderDataset(self.config, self.config.train_split) 111 | 112 | def prepare_data(self): 113 | pass 114 | 115 | def general_loader(self, dataset, batch_size): 116 | sampler = None 117 | return DataLoader( 118 | dataset, 119 | num_workers=os.cpu_count(), 120 | batch_size=batch_size, 121 | pin_memory=True, 122 | sampler=sampler 123 | ) 124 | 125 | def train_dataloader(self): 126 | return self.general_loader(self.train_dataset, batch_size=1) 127 | 128 | def val_dataloader(self): 129 | return self.general_loader(self.val_dataset, batch_size=1) 130 | 131 | def test_dataloader(self): 132 | return self.general_loader(self.test_dataset, batch_size=1) 133 | 134 | def predict_dataloader(self): 135 | return self.general_loader(self.predict_dataset, batch_size=1) 136 | -------------------------------------------------------------------------------- /download.py: -------------------------------------------------------------------------------- 1 | from huggingface_hub import login, snapshot_download 2 | import os 3 | import argparse 4 | from glob import glob 5 | 6 | bsdf_names = [ 7 | 'diffuse', 8 | 'dielectric', 9 | 'roughdielectric', 10 | 'conductor', 11 | 'roughconductor', 12 | 'plastic', 13 | 'roughplastic' 14 | ] 15 | 16 | parser = argparse.ArgumentParser() 17 | parser.add_argument('--token', type=str, help='your own token', required=True) 18 | # Here you have to change to your own token 19 | # 1. Click on your avatar in the upper right corner and select "Settings". 20 | # 2. On the "Settings" page, click "Access Tokens" on the left side. 21 | # 3. Generate a new Token and copy it. 22 | 23 | parser.add_argument('--type', type=str, help='Types of materials to download, all for all downloads', required=True) 24 | # Either choose one from bsdf_names or "all" 25 | 26 | parser.add_argument('--depth', action='store_true', help='Whether depth data is required') 27 | # Whether depth data is required 28 | 29 | args = parser.parse_args() 30 | 31 | login(token=args.token) 32 | os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1" 33 | 34 | REPO_ID = "EPFL-CVLab/OpenMaterial" 35 | 36 | if __name__ == "__main__": 37 | BASE_DIR = f"datasets" 38 | LOCAL_DIR = f"datasets/openmaterial" 39 | LOCAL_DIR_ABLATION = f"datasets/ablation" 40 | if args.type in bsdf_names: 41 | material_type = args.type 42 | os.makedirs(LOCAL_DIR, exist_ok=True) 43 | snapshot_download(repo_id=REPO_ID, repo_type="dataset", allow_patterns=f"{material_type}-*.tar", ignore_patterns=["depth-*.tar", "groundtruth.tar", "ablation-*.tar"],local_dir=LOCAL_DIR, token=args.token) 44 | tar_paths = glob(os.path.join(LOCAL_DIR, f"{material_type}-*.tar")) 45 | for tar_path in tar_paths: 46 | cmd = f'tar -xvf {tar_path} -C {LOCAL_DIR}' 47 | print(cmd) 48 | os.system(cmd) 49 | if args.depth: 50 | snapshot_download(repo_id=REPO_ID, repo_type="dataset", allow_patterns=f"depth-{material_type}*.tar",ignore_patterns="groundtruth.tar", local_dir=LOCAL_DIR, token=args.token) 51 | snapshot_download(repo_id=REPO_ID, repo_type="dataset", allow_patterns=f"groundtruth.tar", local_dir="./datasets", token=args.token) 52 | cmd = f'tar -xvf ./datasets/groundtruth.tar -C ./datasets' 53 | print(cmd) 54 | os.system(cmd) 55 | elif args.type == 'ablation': 56 | snapshot_download(repo_id=REPO_ID, repo_type="dataset", allow_patterns=f"ablation-*.tar", ignore_patterns=["depth-*.tar", "groundtruth.tar"],local_dir=LOCAL_DIR, token=args.token) 57 | tar_paths = glob(os.path.join(LOCAL_DIR, f"ablation-*.tar")) 58 | for tar_path in tar_paths: 59 | cmd = f'tar -xvf {tar_path} -C {LOCAL_DIR_ABLATION}' 60 | print(cmd) 61 | os.system(cmd) 62 | if args.depth: 63 | snapshot_download(repo_id=REPO_ID, repo_type="dataset", allow_patterns=f"depth-ablation*.tar", local_dir=LOCAL_DIR, token=args.token) 64 | snapshot_download(repo_id=REPO_ID, repo_type="dataset", allow_patterns=f"groundtruth-ablation.tar", local_dir="./datasets", token=args.token) 65 | cmd = f'tar -xvf ./datasets/groundtruth-ablation.tar -C ./datasets' 66 | print(cmd) 67 | os.system(cmd) 68 | elif args.type == 'all': 69 | os.makedirs(LOCAL_DIR, exist_ok=True) 70 | snapshot_download(repo_id=REPO_ID, repo_type="dataset", allow_patterns="*.tar", ignore_patterns=["depth-*.tar", "groundtruth.tar", "ablation-*.tar", "groundtruth-ablation.tar"], local_dir=LOCAL_DIR, token=args.token) 71 | tar_paths = glob(os.path.join(LOCAL_DIR, "*.tar")) 72 | for tar_path in tar_paths: 73 | cmd = f'tar -xvf {tar_path} -C {LOCAL_DIR}' 74 | print(cmd) 75 | os.system(cmd) 76 | if args.depth: 77 | snapshot_download(repo_id=REPO_ID, repo_type="dataset", allow_patterns="depth*.tar", local_dir=LOCAL_DIR, token=args.token) 78 | snapshot_download(repo_id=REPO_ID, repo_type="dataset", allow_patterns=f"groundtruth.tar", local_dir="./datasets", token=args.token) 79 | cmd = f'tar -xvf ./datasets/groundtruth.tar -C ./datasets' 80 | print(cmd) 81 | os.system(cmd) 82 | else: 83 | raise ValueError("There's no such material.") 84 | 85 | cmd = f'rm -rf {BASE_DIR}/*.tar' 86 | print(cmd) 87 | os.system(cmd) 88 | 89 | cmd = f'rm -rf {LOCAL_DIR}/*.tar' 90 | print(cmd) 91 | os.system(cmd) 92 | 93 | cmd = f'rm -r {LOCAL_DIR}/.cache' 94 | if os.path.exists(f"{LOCAL_DIR}/.cache"): 95 | print(cmd) 96 | os.system(cmd) 97 | 98 | cmd = f'rm -r {LOCAL_DIR}/.huggingface' 99 | if os.path.exists(f"{LOCAL_DIR}/.huggingface"): 100 | print(cmd) 101 | os.system(cmd) 102 | 103 | cmd = f'rm -r {BASE_DIR}/.cache' 104 | if os.path.exists(f"{BASE_DIR}/.cache"): 105 | print(cmd) 106 | os.system(cmd) 107 | 108 | cmd = f'rm -r {BASE_DIR}/.huggingface' 109 | if os.path.exists(f"{BASE_DIR}/.huggingface"): 110 | print(cmd) 111 | os.system(cmd) -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Benchmarking_Everything 2 | 3 | The repository contains scripts for download and evaluation on the openmaterial dataset. 4 | 5 | 6 | 7 | We introduce the OpenMaterial dataset, comprising 1001 objects made of 295 distinct materials—including conductors, dielectrics, plastics, and their roughened variants— and captured under 723 diverse lighting conditions. 8 | 9 | For working with this dataset, you can refer to the following steps: 10 | 11 | ## 1. Download dataset 12 | 13 | First get your own huggingface token 14 | 15 | 1. Click on your avatar in the upper right corner and select "Settings". 16 | 2. On the "Settings" page, click "Access Tokens" on the left side. 17 | 3. Generate a new Token and copy it. 18 | 19 | To verify the validity of the method for different materials, shapes and light conditions, you can start with our ablation dataset 20 | 21 | ```shell 22 | python download.py --token --type ablation 23 | ``` 24 | 25 | If you need the complete dataset, you can run the following command: 26 | 27 | ```shell 28 | python download.py --token --type all 29 | ``` 30 | 31 | Or if you need to download a subset of a certain material type separately, such as "conductor", you can run the following command: 32 | 33 | ```shell 34 | python download.py --token --type conductor 35 | ``` 36 | 37 | after downlaod complete dataset, The following file structure is obtained 38 | 39 | ```shell 40 | datasets 41 | ├── groundtruth 42 | │ ├── 5c4ae9c4a3cb47a4b6273eb2839a7b8c 43 | │ └── clean_5c4ae9c4a3cb47a4b6273eb2839a7b8c.ply 44 | │ ├── 5c0514eae1f94f22bc5475fe0970cd28 45 | │ └── clean_5c0514eae1f94f22bc5475fe0970cd28.ply 46 | │ └── ... 47 | ├── openmaterial 48 | │ ├── 5c4ae9c4a3cb47a4b6273eb2839a7b8c 49 | │ ├── train 50 | │ ├── images 51 | │ ├── mask 52 | │ ├── test 53 | │ ├── transforms_train.json 54 | │ └── transforms_test.json 55 | │ ├── 5c0514eae1f94f22bc5475fe0970cd28 56 | │ └── ... 57 | ``` 58 | 59 | (optional) If you need to use depth, please use the following command: 60 | 61 | ```shell 62 | python download.py --token --type all --depth 63 | ``` 64 | 65 | Here is an example for using our depth data, which are real depth values, not normalised: 66 | 67 | ```python 68 | with h5py.File(filename, 'r') as hdf: 69 | dataset = hdf['depth'] 70 | depth = dataset[:] # size: (1200, 1600) 71 | ``` 72 | 73 | ## 2. Start training 74 | 75 | for ${method}: (method can be instant-nsr-pl, NeuS2, and so on...) 76 | 77 | ```shell 78 | cd ${method} 79 | chmod +x run_openmaterial.sh 80 | bash run_openmaterial.sh $start $end $gpu 81 | # The $start and $end parameters are used to run batches, e.g. you need to run the 0th-50th case on gpu:0 82 | # for example: bash run_openmaterial.sh 0 50 0 83 | cd ../ 84 | ``` 85 | 86 | the result of nerf are stored in the "instant-nsr-pl-output-womask/output.txt" in the following format: 87 | 88 | ```shell 89 | ${object}:${method}:${material}:${PSNR}-${SSIM} 90 | ``` 91 | 92 | ## 3. Eval 93 | 94 | ### Ablation dataset 95 | 96 | If you want to use our script to evaluate on ablation dataset, make sure you store the mesh in the following file format: 97 | 98 | ```shell 99 | Mesh-ablation 100 | ├── instant-nsr-pl-wmask 101 | │ ├── meshes 102 | │ ├── 5c4ae9c4a3cb47a4b6273eb2839a7b8c 103 | │ └── cobblestone_street_night_4k-conductor.obj 104 | │ └── ... 105 | │ ├── 5c0514eae1f94f22bc5475fe0970cd28 106 | │ └── cobblestone_street_night_4k-conductor.obj 107 | │ └── ... 108 | │ └── ... 109 | ``` 110 | Then calculate chamfer distance after training: 111 | 112 | ```shell 113 | bash eval/eval_mitsuba.sh ../Mesh-ablation ../output-ablation ${method} true 114 | # for example: 115 | # bash eval/eval_mitsuba.sh ../Mesh-ablation ../output-ablation instant-nsr-pl-wmask true 116 | ``` 117 | 118 | the result are stored in the "output-ablation/${method}-mesh-output.txt" in the following format: 119 | 120 | ```shell 121 | ${object}:${method}:${envmap}-${material}:${cds} 122 | ``` 123 | 124 | Run the following command to integrate the results: 125 | 126 | ```shell 127 | python sum_metrics-ablation.py --method insr 128 | ``` 129 | 130 | 131 | ### Complete dataset 132 | 133 | If you want to use our script to evaluate on complete dataset, make sure you store the mesh in the following file format: 134 | 135 | ```shell 136 | Mesh 137 | ├── instant-nsr-pl-wmask 138 | │ ├── meshes 139 | │ ├── 5c4ae9c4a3cb47a4b6273eb2839a7b8c 140 | │ └── diffuse.ply 141 | │ ├── 5c0514eae1f94f22bc5475fe0970cd28 142 | │ └── diffuse.ply 143 | │ └── ... 144 | ``` 145 | 146 | Then calculate chamfer distance after training: 147 | 148 | ```shell 149 | bash eval/eval_mitsuba.sh ../Mesh ../output ${method} 150 | ``` 151 | 152 | the result are stored in the "output/${method}-mesh-output.txt" in the following format: 153 | 154 | ```shell 155 | ${object}:${method}:${material}:${cds} 156 | ``` 157 | 158 | Run the following command to integrate the results: 159 | 160 | ```shell 161 | python sum_metrics.py 162 | ``` 163 | -------------------------------------------------------------------------------- /instant-nsr-pl/datasets/openmaterial_wmask.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | import math 4 | import numpy as np 5 | from PIL import Image 6 | 7 | import torch 8 | from torch.utils.data import Dataset, DataLoader, IterableDataset 9 | import torchvision.transforms.functional as TF 10 | 11 | import pytorch_lightning as pl 12 | 13 | import datasets 14 | from models.ray_utils import get_ray_directions 15 | from utils.misc import get_rank 16 | 17 | 18 | class BlenderDatasetBase(): 19 | def setup(self, config, split): 20 | self.config = config 21 | self.split = split 22 | self.rank = get_rank() 23 | 24 | self.has_mask = True 25 | self.apply_mask = True 26 | 27 | with open(os.path.join(self.config.root_dir, f"transforms_{self.split}.json"), 'r') as f: 28 | meta = json.load(f) 29 | 30 | if 'w' in meta and 'h' in meta: 31 | W, H = int(meta['w']), int(meta['h']) 32 | else: 33 | W, H = 1600, 1200 34 | 35 | if 'img_wh' in self.config: 36 | w, h = self.config.img_wh 37 | assert round(W / w * h) == H 38 | elif 'img_downscale' in self.config: 39 | w, h = W // self.config.img_downscale, H // self.config.img_downscale 40 | else: 41 | raise KeyError("Either img_wh or img_downscale should be specified.") 42 | 43 | self.w, self.h = w, h 44 | self.img_wh = (self.w, self.h) 45 | 46 | self.near, self.far = self.config.near_plane, self.config.far_plane 47 | 48 | self.focal_x = meta['fl_x'] # scaled focal length 49 | self.focal_y = meta['fl_y'] 50 | 51 | # ray directions for all pixels, same for all images (same H, W, focal) 52 | self.directions = \ 53 | get_ray_directions(self.w, self.h, self.focal_x, self.focal_y, self.w//2, self.h//2).to(self.rank) # (h, w, 3) 54 | 55 | self.all_c2w, self.all_images, self.all_fg_masks = [], [], [] 56 | 57 | for i, frame in enumerate(meta['frames']): 58 | c2w = torch.from_numpy(np.array(frame['transform_matrix'])[:3, :4]) 59 | self.all_c2w.append(c2w) 60 | 61 | img_path = os.path.join(self.config.root_dir, f"{frame['file_path']}") 62 | img = Image.open(img_path) 63 | img = img.resize(self.img_wh, Image.BICUBIC) 64 | img = TF.to_tensor(img).permute(1, 2, 0) 65 | 66 | mask_path_ = frame['file_path'].split('/')[0] + '/mask/' + os.path.split(frame['file_path'])[-1] 67 | mask_path = os.path.join(self.config.root_dir, mask_path_) 68 | mask = Image.open(mask_path).convert('L') # (H, W, 1) 69 | mask = mask.resize(self.img_wh, Image.BICUBIC) 70 | mask = TF.to_tensor(mask)[0] 71 | self.all_fg_masks.append(mask) 72 | self.all_images.append(img[...,:3]) 73 | 74 | self.all_c2w, self.all_images, self.all_fg_masks = \ 75 | torch.stack(self.all_c2w, dim=0).float().to(self.rank), \ 76 | torch.stack(self.all_images, dim=0).float().to(self.rank), \ 77 | torch.stack(self.all_fg_masks, dim=0).float().to(self.rank) 78 | 79 | 80 | class BlenderDataset(Dataset, BlenderDatasetBase): 81 | def __init__(self, config, split): 82 | self.setup(config, split) 83 | 84 | def __len__(self): 85 | return len(self.all_images) 86 | 87 | def __getitem__(self, index): 88 | return { 89 | 'index': index 90 | } 91 | 92 | 93 | class BlenderIterableDataset(IterableDataset, BlenderDatasetBase): 94 | def __init__(self, config, split): 95 | self.setup(config, split) 96 | 97 | def __iter__(self): 98 | while True: 99 | yield {} 100 | 101 | 102 | @datasets.register('openmaterial-wmask') 103 | class BlenderDataModule(pl.LightningDataModule): 104 | def __init__(self, config): 105 | super().__init__() 106 | self.config = config 107 | 108 | def setup(self, stage=None): 109 | if stage in [None, 'fit']: 110 | self.train_dataset = BlenderIterableDataset(self.config, self.config.train_split) 111 | if stage in [None, 'test']: 112 | self.test_dataset = BlenderDataset(self.config, self.config.test_split) 113 | if stage in [None, 'predict']: 114 | self.predict_dataset = BlenderDataset(self.config, self.config.train_split) 115 | 116 | def prepare_data(self): 117 | pass 118 | 119 | def general_loader(self, dataset, batch_size): 120 | sampler = None 121 | return DataLoader( 122 | dataset, 123 | num_workers=os.cpu_count(), 124 | batch_size=batch_size, 125 | pin_memory=True, 126 | sampler=sampler 127 | ) 128 | 129 | def train_dataloader(self): 130 | return self.general_loader(self.train_dataset, batch_size=1) 131 | 132 | def test_dataloader(self): 133 | return self.general_loader(self.test_dataset, batch_size=1) 134 | 135 | def predict_dataloader(self): 136 | return self.general_loader(self.predict_dataset, batch_size=1) 137 | -------------------------------------------------------------------------------- /instant-nsr-pl/launch.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import argparse 3 | import os 4 | import time 5 | import logging 6 | from datetime import datetime 7 | 8 | 9 | def main(): 10 | parser = argparse.ArgumentParser() 11 | parser.add_argument('--config', required=True, help='path to config file') 12 | parser.add_argument('--gpu', default='0', help='GPU(s) to be used') 13 | parser.add_argument('--resume', default=None, help='path to the weights to be resumed') 14 | parser.add_argument( 15 | '--resume_weights_only', 16 | action='store_true', 17 | help='specify this argument to restore only the weights (w/o training states), e.g. --resume path/to/resume --resume_weights_only' 18 | ) 19 | 20 | group = parser.add_mutually_exclusive_group(required=True) 21 | group.add_argument('--train', action='store_true') 22 | group.add_argument('--validate', action='store_true') 23 | group.add_argument('--test', action='store_true') 24 | group.add_argument('--predict', action='store_true') 25 | # group.add_argument('--export', action='store_true') # TODO: a separate export action 26 | 27 | parser.add_argument('--exp_dir', default='./exp') 28 | parser.add_argument('--runs_dir', default='./runs') 29 | parser.add_argument('--verbose', action='store_true', help='if true, set logging level to DEBUG') 30 | 31 | # add 32 | parser.add_argument('--output_dir', default='') 33 | args, extras = parser.parse_known_args() 34 | 35 | # set CUDA_VISIBLE_DEVICES then import pytorch-lightning 36 | os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID' 37 | os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu 38 | n_gpus = len(args.gpu.split(',')) 39 | 40 | import datasets 41 | import systems 42 | import pytorch_lightning as pl 43 | from pytorch_lightning import Trainer 44 | from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor 45 | from pytorch_lightning.loggers import TensorBoardLogger, CSVLogger 46 | from utils.callbacks import CodeSnapshotCallback, ConfigSnapshotCallback, CustomProgressBar 47 | from utils.misc import load_config 48 | 49 | # parse YAML config to OmegaConf 50 | config = load_config(args.config, cli_args=extras) 51 | config.cmd_args = vars(args) 52 | 53 | config.trial_name = config.get('trial_name') or (config.tag + datetime.now().strftime('@%Y%m%d-%H%M%S')) 54 | config.exp_dir = config.get('exp_dir') or os.path.join(args.exp_dir, config.name) 55 | if args.output_dir != '': 56 | config.save_dir = args.output_dir 57 | else: 58 | config.save_dir = config.get('save_dir') or os.path.join(config.render_save_dir, config.trial_name, 'instant-nsr-pl-save') 59 | config.ckpt_dir = config.get('ckpt_dir') or os.path.join(config.exp_dir, config.trial_name, 'ckpt') 60 | config.code_dir = config.get('code_dir') or os.path.join(config.exp_dir, config.trial_name, 'code') 61 | config.config_dir = config.get('config_dir') or os.path.join(config.exp_dir, config.trial_name, 'config') 62 | 63 | logger = logging.getLogger('pytorch_lightning') 64 | if args.verbose: 65 | logger.setLevel(logging.DEBUG) 66 | 67 | if 'seed' not in config: 68 | config.seed = int(time.time() * 1000) % 1000 69 | pl.seed_everything(config.seed) 70 | 71 | dm = datasets.make(config.dataset.name, config.dataset) 72 | system = systems.make(config.system.name, config, load_from_checkpoint=None if not args.resume_weights_only else args.resume) 73 | 74 | callbacks = [] 75 | if args.train: 76 | callbacks += [ 77 | ModelCheckpoint( 78 | dirpath=config.ckpt_dir, 79 | **config.checkpoint 80 | ), 81 | LearningRateMonitor(logging_interval='step'), 82 | # CodeSnapshotCallback( 83 | # config.code_dir, use_version=False 84 | # ), 85 | # ConfigSnapshotCallback( 86 | # config, config.config_dir, use_version=False 87 | # ), 88 | CustomProgressBar(refresh_rate=1), 89 | ] 90 | 91 | loggers = [] 92 | if args.train: 93 | loggers += [ 94 | # TensorBoardLogger(args.runs_dir, name=config.name, version=config.trial_name), 95 | CSVLogger(config.exp_dir, name=config.trial_name, version='csv_logs') 96 | ] 97 | 98 | if sys.platform == 'win32': 99 | # does not support multi-gpu on windows 100 | strategy = 'dp' 101 | assert n_gpus == 1 102 | else: 103 | strategy = 'ddp_find_unused_parameters_false' 104 | 105 | trainer = Trainer( 106 | devices=n_gpus, 107 | accelerator='gpu', 108 | callbacks=callbacks, 109 | logger=loggers, 110 | strategy=strategy, 111 | **config.trainer 112 | ) 113 | 114 | if args.train: 115 | if args.resume and not args.resume_weights_only: 116 | # FIXME: different behavior in pytorch-lighting>1.9 ? 117 | trainer.fit(system, datamodule=dm, ckpt_path=args.resume) 118 | else: 119 | trainer.fit(system, datamodule=dm) 120 | trainer.test(system, datamodule=dm) 121 | elif args.validate: 122 | trainer.validate(system, datamodule=dm, ckpt_path=args.resume) 123 | elif args.test: 124 | trainer.test(system, datamodule=dm, ckpt_path=args.resume) 125 | elif args.predict: 126 | trainer.predict(system, datamodule=dm, ckpt_path=args.resume) 127 | 128 | 129 | if __name__ == '__main__': 130 | main() 131 | -------------------------------------------------------------------------------- /gen_colmap.py: -------------------------------------------------------------------------------- 1 | import os 2 | import math 3 | import numpy as np 4 | from scipy.spatial.transform import Rotation as Rot 5 | from shutil import copytree 6 | import json 7 | import numpy as np 8 | import cv2 as cv 9 | 10 | """ 11 | Part of the code is adapted from: 12 | https://gist.github.com/Totoro97/05e0b6afef5e580464731ad4c69c7a41 13 | 14 | """ 15 | 16 | def gen_w2c(pose): 17 | 18 | pose[:3, :1] = -pose[:3, :1] 19 | pose[:3, 1:2] = -pose[:3, 1:2] # Flip the x+ and y+ to align coordinate system 20 | 21 | R = pose[:3, :3].transpose() 22 | T = -R @ pose[:3, 3:] 23 | return R, T 24 | 25 | # "cv.decomposeProjectionMatrix(P)" generates some errors in special cases 26 | # Refer to: https://stackoverflow.com/questions/55814640/decomposeprojectionmatrix-gives-unexpected-result 27 | # so neither pitch nor yaw can be zero. 28 | def load_K_Rt_from_P(filename, P=None): 29 | if P is None: 30 | lines = open(filename).read().splitlines() 31 | if len(lines) == 4: 32 | lines = lines[1:] 33 | lines = [[x[0], x[1], x[2], x[3]] for x in (x.split(" ") for x in lines)] 34 | P = np.asarray(lines).astype(np.float32).squeeze() 35 | 36 | out = cv.decomposeProjectionMatrix(P) 37 | K = out[0] 38 | R = out[1] 39 | t = out[2] 40 | 41 | K = K/K[2,2] 42 | intrinsics = np.eye(4) 43 | intrinsics[:3, :3] = K 44 | 45 | pose = np.eye(4, dtype=np.float32) 46 | pose[:3, :3] = R.transpose() 47 | pose[:3, 3] = (t[:3] / t[3])[:, 0] 48 | 49 | return intrinsics, pose 50 | 51 | 52 | def gen_camera_intrinsic(width, height, fov_x, fov_y): 53 | fx = width / 2.0 / math.tan(fov_x / 180 * math.pi / 2.0) 54 | fy = height / 2.0 / math.tan(fov_y / 180 * math.pi / 2.0) 55 | return fx, fy 56 | 57 | 58 | def main(): 59 | LOCAL_PATH = "datasets/openmaterial" 60 | 61 | for path in os.listdir(LOCAL_PATH): 62 | path2 = os.listdir(os.path.join(LOCAL_PATH, path)) 63 | output_dir = os.path.join(LOCAL_PATH, path) 64 | images_folder = os.path.join(os.path.dirname(__file__), f'./{output_dir}') 65 | os.makedirs(images_folder, exist_ok=True) 66 | 67 | file_train_path = os.path.join(LOCAL_PATH, path, path2[0], 'transforms_train.json') 68 | file_test_path = os.path.join(LOCAL_PATH, path, path2[0], 'transforms_test.json') 69 | train_paths = os.path.join(LOCAL_PATH, path, path2[0], 'train') 70 | write_colmap(train_paths, output_dir, file_train_path, flag=True, is_train=True) 71 | test_paths = os.path.join(LOCAL_PATH, path, path2[0], 'test') 72 | write_colmap(test_paths, output_dir, file_test_path, flag=True, is_train=False) 73 | 74 | def write_colmap(img_path, output_dir, json_path, flag, is_train): 75 | object_name = os.path.split(output_dir)[-1] 76 | width, height = 1600, 1200 77 | with open(json_path, 'r', encoding='utf-8') as f: 78 | data = json.load(f) 79 | fov_x = 37.8492 80 | fov_y = 28.8415 81 | 82 | scale_mat = np.diag([1.0, 1.0, 1.0, 1.0]) 83 | os.makedirs(os.path.join(img_path, 'colmap', 'manually', 'sparse'), exist_ok=True) 84 | os.makedirs(os.path.join(img_path, 'colmap', 'dense'), exist_ok=True) 85 | 86 | # transform to Colmap format 87 | fx, fy = gen_camera_intrinsic(width, height, fov_x, fov_y) 88 | 89 | # use float64 to avoid loss of precision 90 | intrinsic = np.diag([fx, fy, 1.0, 1.0]).astype(np.float64) 91 | # The origin is in the center and not in the upper left corner of the image 92 | intrinsic[0, 2] = width / 2.0 93 | intrinsic[1, 2] = height / 2.0 94 | 95 | # manually construct a sparse model by creating a cameras.txt, points3D.txt, and images.txt 96 | # Follow the steps at https://colmap.github.io/faq.html#reconstruct-sparse-dense-model-from-known-camera-poses 97 | with open(os.path.join(img_path, 'colmap', 'manually', 'sparse', 'points3D.txt'), 'w') as f: 98 | pass 99 | 100 | with open(os.path.join(img_path, 'colmap', 'manually', 'sparse', 'cameras.txt'), 'w') as f_camera: 101 | f_camera.write('1 PINHOLE 1600 1200 {} {} {} {}\n'.format(intrinsic[0, 0], intrinsic[1, 1], intrinsic[0, 2], intrinsic[1, 2])) 102 | 103 | camera_params = {} 104 | bottom = np.array([0, 0, 0, 1.]).reshape([1, 4]) 105 | copytree(os.path.join(img_path, 'images'), os.path.join(img_path, 'colmap', 'images'), dirs_exist_ok=True) 106 | with open(os.path.join(img_path, 'colmap', 'manually', 'sparse', 'images.txt'), 'w') as f: 107 | for i, frame in enumerate(data['frames']): 108 | flip_mat = np.array([ 109 | [-1, 0, 0, 0], 110 | [0, 1, 0, 0], 111 | [0, 0, -1, 0], 112 | [0, 0, 0, 1] 113 | ]) 114 | cam_pose_ = np.matmul(frame['transform_matrix'], flip_mat) 115 | # P = K[R|t] 116 | cam_pose = np.array(cam_pose_) 117 | R, T = gen_w2c(cam_pose) 118 | w2c = np.concatenate([np.concatenate([R, T], 1), bottom], 0) 119 | world_mat = intrinsic @ w2c 120 | 121 | camera_params['world_mat_%d' % i] = world_mat 122 | # Since the object has been normalised into the unit sphere, scale_mat is the unit array. 123 | camera_params['scale_mat_%d' % i] = scale_mat 124 | camera_params['intrinsic_mat_%d' % i] = intrinsic 125 | camera_params['extrinsic_mat_%d' % i] = w2c 126 | # P = world_mat @ scale_mat 127 | # P = P[:3, :4] 128 | # _, c2w_ = load_K_Rt_from_P(None,P) 129 | # R = c2w_[:3, :3].transpose() 130 | # T = -R @ c2w_[:3, 3:] 131 | 132 | rot = Rot.from_matrix(R) 133 | rot = rot.as_quat() 134 | image_name = '{:03d}.png'.format(i) 135 | f.write('{} {} {} {} {} {} {} {} {} {}\n\n'.format(i + 1, rot[3], rot[0], rot[1], rot[2], T[0, 0], T[1, 0], T[2, 0], 1, image_name)) 136 | # np.savez(f"{img_path}/cameras_sphere.npz", **camera_params) 137 | # if flag and is_train: 138 | # os.makedirs(f"../groundtruth/{object_name}", exist_ok=True) 139 | # np.savez(f"../groundtruth/{object_name}/cameras_sphere.npz", **camera_params) 140 | # flag = False 141 | # elif flag and not is_train: 142 | # np.savez(f"../groundtruth/{object_name}/cameras_sphere_test.npz", **camera_params) 143 | # flag = False 144 | 145 | 146 | if __name__ == '__main__': 147 | main() 148 | -------------------------------------------------------------------------------- /instant-nsr-pl/systems/criterions.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | 6 | class WeightedLoss(nn.Module): 7 | @property 8 | def func(self): 9 | raise NotImplementedError 10 | 11 | def forward(self, inputs, targets, weight=None, reduction='mean'): 12 | assert reduction in ['none', 'sum', 'mean', 'valid_mean'] 13 | loss = self.func(inputs, targets, reduction='none') 14 | if weight is not None: 15 | while weight.ndim < inputs.ndim: 16 | weight = weight[..., None] 17 | loss *= weight.float() 18 | if reduction == 'none': 19 | return loss 20 | elif reduction == 'sum': 21 | return loss.sum() 22 | elif reduction == 'mean': 23 | return loss.mean() 24 | elif reduction == 'valid_mean': 25 | return loss.sum() / weight.float().sum() 26 | 27 | 28 | class MSELoss(WeightedLoss): 29 | @property 30 | def func(self): 31 | return F.mse_loss 32 | 33 | 34 | class L1Loss(WeightedLoss): 35 | @property 36 | def func(self): 37 | return F.l1_loss 38 | 39 | 40 | class PSNR(nn.Module): 41 | def __init__(self): 42 | super().__init__() 43 | 44 | def forward(self, inputs, targets, valid_mask=None, reduction='mean'): 45 | assert reduction in ['mean', 'none'] 46 | value = (inputs - targets)**2 47 | if valid_mask is not None: 48 | value = value[valid_mask] 49 | if reduction == 'mean': 50 | return -10 * torch.log10(torch.mean(value)) 51 | elif reduction == 'none': 52 | return -10 * torch.log10(torch.mean(value, dim=tuple(range(value.ndim)[1:]))) 53 | 54 | 55 | class SSIM(): 56 | def __init__(self, data_range=(0, 1), kernel_size=(11, 11), sigma=(1.5, 1.5), k1=0.01, k2=0.03, gaussian=True, device='cpu'): 57 | self.kernel_size = kernel_size 58 | self.sigma = sigma 59 | self.gaussian = gaussian 60 | 61 | if any(x % 2 == 0 or x <= 0 for x in self.kernel_size): 62 | raise ValueError(f"Expected kernel_size to have odd positive number. Got {kernel_size}.") 63 | if any(y <= 0 for y in self.sigma): 64 | raise ValueError(f"Expected sigma to have positive number. Got {sigma}.") 65 | 66 | data_scale = data_range[1] - data_range[0] 67 | self.c1 = (k1 * data_scale)**2 68 | self.c2 = (k2 * data_scale)**2 69 | self.pad_h = (self.kernel_size[0] - 1) // 2 70 | self.pad_w = (self.kernel_size[1] - 1) // 2 71 | self._kernel = self._gaussian_or_uniform_kernel(kernel_size=self.kernel_size, sigma=self.sigma).to(device) 72 | 73 | def _uniform(self, kernel_size): 74 | max, min = 2.5, -2.5 75 | ksize_half = (kernel_size - 1) * 0.5 76 | kernel = torch.linspace(-ksize_half, ksize_half, steps=kernel_size) 77 | for i, j in enumerate(kernel): 78 | if min <= j <= max: 79 | kernel[i] = 1 / (max - min) 80 | else: 81 | kernel[i] = 0 82 | 83 | return kernel.unsqueeze(dim=0) # (1, kernel_size) 84 | 85 | def _gaussian(self, kernel_size, sigma): 86 | ksize_half = (kernel_size - 1) * 0.5 87 | kernel = torch.linspace(-ksize_half, ksize_half, steps=kernel_size) 88 | gauss = torch.exp(-0.5 * (kernel / sigma).pow(2)) 89 | return (gauss / gauss.sum()).unsqueeze(dim=0) # (1, kernel_size) 90 | 91 | def _gaussian_or_uniform_kernel(self, kernel_size, sigma): 92 | if self.gaussian: 93 | kernel_x = self._gaussian(kernel_size[0], sigma[0]) 94 | kernel_y = self._gaussian(kernel_size[1], sigma[1]) 95 | else: 96 | kernel_x = self._uniform(kernel_size[0]) 97 | kernel_y = self._uniform(kernel_size[1]) 98 | 99 | return torch.matmul(kernel_x.t(), kernel_y) # (kernel_size, 1) * (1, kernel_size) 100 | 101 | def __call__(self, output, target, reduction='mean'): 102 | if output.dtype != target.dtype: 103 | raise TypeError( 104 | f"Expected output and target to have the same data type. Got output: {output.dtype} and y: {target.dtype}." 105 | ) 106 | 107 | if output.shape != target.shape: 108 | raise ValueError( 109 | f"Expected output and target to have the same shape. Got output: {output.shape} and y: {target.shape}." 110 | ) 111 | 112 | if len(output.shape) != 4 or len(target.shape) != 4: 113 | raise ValueError( 114 | f"Expected output and target to have BxCxHxW shape. Got output: {output.shape} and y: {target.shape}." 115 | ) 116 | 117 | assert reduction in ['mean', 'sum', 'none'] 118 | 119 | channel = output.size(1) 120 | if len(self._kernel.shape) < 4: 121 | self._kernel = self._kernel.expand(channel, 1, -1, -1) 122 | 123 | output = F.pad(output, [self.pad_w, self.pad_w, self.pad_h, self.pad_h], mode="reflect") 124 | target = F.pad(target, [self.pad_w, self.pad_w, self.pad_h, self.pad_h], mode="reflect") 125 | 126 | input_list = torch.cat([output, target, output * output, target * target, output * target]) 127 | outputs = F.conv2d(input_list, self._kernel, groups=channel) 128 | 129 | output_list = [outputs[x * output.size(0) : (x + 1) * output.size(0)] for x in range(len(outputs))] 130 | 131 | mu_pred_sq = output_list[0].pow(2) 132 | mu_target_sq = output_list[1].pow(2) 133 | mu_pred_target = output_list[0] * output_list[1] 134 | 135 | sigma_pred_sq = output_list[2] - mu_pred_sq 136 | sigma_target_sq = output_list[3] - mu_target_sq 137 | sigma_pred_target = output_list[4] - mu_pred_target 138 | 139 | a1 = 2 * mu_pred_target + self.c1 140 | a2 = 2 * sigma_pred_target + self.c2 141 | b1 = mu_pred_sq + mu_target_sq + self.c1 142 | b2 = sigma_pred_sq + sigma_target_sq + self.c2 143 | 144 | ssim_idx = (a1 * a2) / (b1 * b2) 145 | _ssim = torch.mean(ssim_idx, (1, 2, 3)) 146 | 147 | if reduction == 'none': 148 | return _ssim 149 | elif reduction == 'sum': 150 | return _ssim.sum() 151 | elif reduction == 'mean': 152 | return _ssim.mean() 153 | 154 | 155 | def binary_cross_entropy(input, target): 156 | """ 157 | F.binary_cross_entropy is not numerically stable in mixed-precision training. 158 | """ 159 | return -(target * torch.log(input) + (1 - target) * torch.log(1 - input)).mean() 160 | -------------------------------------------------------------------------------- /instant-nsr-pl/models/nerf.py: -------------------------------------------------------------------------------- 1 | import math 2 | 3 | import torch 4 | import torch.nn as nn 5 | import torch.nn.functional as F 6 | 7 | import models 8 | from models.base import BaseModel 9 | from models.utils import chunk_batch 10 | from systems.utils import update_module_step 11 | from nerfacc import ContractionType, OccupancyGrid, ray_marching, render_weight_from_density, accumulate_along_rays 12 | 13 | 14 | @models.register('nerf') 15 | class NeRFModel(BaseModel): 16 | def setup(self): 17 | self.geometry = models.make(self.config.geometry.name, self.config.geometry) 18 | self.texture = models.make(self.config.texture.name, self.config.texture) 19 | self.register_buffer('scene_aabb', torch.as_tensor([-self.config.radius, -self.config.radius, -self.config.radius, self.config.radius, self.config.radius, self.config.radius], dtype=torch.float32)) 20 | 21 | if self.config.learned_background: 22 | self.occupancy_grid_res = 256 23 | self.near_plane, self.far_plane = 0.2, 1e4 24 | self.cone_angle = 10**(math.log10(self.far_plane) / self.config.num_samples_per_ray) - 1. # approximate 25 | self.render_step_size = 0.01 # render_step_size = max(distance_to_camera * self.cone_angle, self.render_step_size) 26 | self.contraction_type = ContractionType.UN_BOUNDED_SPHERE 27 | else: 28 | self.occupancy_grid_res = 128 29 | self.near_plane, self.far_plane = None, None 30 | self.cone_angle = 0.0 31 | self.render_step_size = 1.732 * 2 * self.config.radius / self.config.num_samples_per_ray 32 | self.contraction_type = ContractionType.AABB 33 | 34 | self.geometry.contraction_type = self.contraction_type 35 | 36 | if self.config.grid_prune: 37 | self.occupancy_grid = OccupancyGrid( 38 | roi_aabb=self.scene_aabb, 39 | resolution=self.occupancy_grid_res, 40 | contraction_type=self.contraction_type 41 | ) 42 | self.randomized = self.config.randomized 43 | self.background_color = None 44 | 45 | def update_step(self, epoch, global_step): 46 | update_module_step(self.geometry, epoch, global_step) 47 | update_module_step(self.texture, epoch, global_step) 48 | 49 | def occ_eval_fn(x): 50 | density, _ = self.geometry(x) 51 | # approximate for 1 - torch.exp(-density[...,None] * self.render_step_size) based on taylor series 52 | return density[...,None] * self.render_step_size 53 | 54 | if self.training and self.config.grid_prune: 55 | self.occupancy_grid.every_n_step(step=global_step, occ_eval_fn=occ_eval_fn) 56 | 57 | def isosurface(self): 58 | mesh = self.geometry.isosurface() 59 | return mesh 60 | 61 | def forward_(self, rays): 62 | n_rays = rays.shape[0] 63 | rays_o, rays_d = rays[:, 0:3], rays[:, 3:6] # both (N_rays, 3) 64 | 65 | def sigma_fn(t_starts, t_ends, ray_indices): 66 | ray_indices = ray_indices.long() 67 | t_origins = rays_o[ray_indices] 68 | t_dirs = rays_d[ray_indices] 69 | positions = t_origins + t_dirs * (t_starts + t_ends) / 2. 70 | density, _ = self.geometry(positions) 71 | return density[...,None] 72 | 73 | def rgb_sigma_fn(t_starts, t_ends, ray_indices): 74 | ray_indices = ray_indices.long() 75 | t_origins = rays_o[ray_indices] 76 | t_dirs = rays_d[ray_indices] 77 | positions = t_origins + t_dirs * (t_starts + t_ends) / 2. 78 | density, feature = self.geometry(positions) 79 | rgb = self.texture(feature, t_dirs) 80 | return rgb, density[...,None] 81 | 82 | with torch.no_grad(): 83 | ray_indices, t_starts, t_ends = ray_marching( 84 | rays_o, rays_d, 85 | scene_aabb=None if self.config.learned_background else self.scene_aabb, 86 | grid=self.occupancy_grid if self.config.grid_prune else None, 87 | sigma_fn=sigma_fn, 88 | near_plane=self.near_plane, far_plane=self.far_plane, 89 | render_step_size=self.render_step_size, 90 | stratified=self.randomized, 91 | cone_angle=self.cone_angle, 92 | alpha_thre=0.0 93 | ) 94 | 95 | ray_indices = ray_indices.long() 96 | t_origins = rays_o[ray_indices] 97 | t_dirs = rays_d[ray_indices] 98 | midpoints = (t_starts + t_ends) / 2. 99 | positions = t_origins + t_dirs * midpoints 100 | intervals = t_ends - t_starts 101 | 102 | density, feature = self.geometry(positions) 103 | rgb = self.texture(feature, t_dirs) 104 | 105 | weights = render_weight_from_density(t_starts, t_ends, density[...,None], ray_indices=ray_indices, n_rays=n_rays) 106 | opacity = accumulate_along_rays(weights, ray_indices, values=None, n_rays=n_rays) 107 | depth = accumulate_along_rays(weights, ray_indices, values=midpoints, n_rays=n_rays) 108 | comp_rgb = accumulate_along_rays(weights, ray_indices, values=rgb, n_rays=n_rays) 109 | comp_rgb = comp_rgb + self.background_color * (1.0 - opacity) 110 | 111 | out = { 112 | 'comp_rgb': comp_rgb, 113 | 'opacity': opacity, 114 | 'depth': depth, 115 | 'rays_valid': opacity > 0, 116 | 'num_samples': torch.as_tensor([len(t_starts)], dtype=torch.int32, device=rays.device) 117 | } 118 | 119 | if self.training: 120 | out.update({ 121 | 'weights': weights.view(-1), 122 | 'points': midpoints.view(-1), 123 | 'intervals': intervals.view(-1), 124 | 'ray_indices': ray_indices.view(-1) 125 | }) 126 | 127 | return out 128 | 129 | def forward(self, rays): 130 | if self.training: 131 | out = self.forward_(rays) 132 | else: 133 | out = chunk_batch(self.forward_, self.config.ray_chunk, True, rays) 134 | return { 135 | **out, 136 | } 137 | 138 | def train(self, mode=True): 139 | self.randomized = mode and self.config.randomized 140 | return super().train(mode=mode) 141 | 142 | def eval(self): 143 | self.randomized = False 144 | return super().eval() 145 | 146 | def regularizations(self, out): 147 | losses = {} 148 | losses.update(self.geometry.regularizations(out)) 149 | losses.update(self.texture.regularizations(out)) 150 | return losses 151 | 152 | @torch.no_grad() 153 | def export(self, export_config): 154 | mesh = self.isosurface() 155 | if export_config.export_vertex_color: 156 | _, feature = chunk_batch(self.geometry, export_config.chunk_size, False, mesh['v_pos'].to(self.rank)) 157 | viewdirs = torch.zeros(feature.shape[0], 3).to(feature) 158 | viewdirs[...,2] = -1. # set the viewing directions to be -z (looking down) 159 | rgb = self.texture(feature, viewdirs).clamp(0,1) 160 | mesh['v_rgb'] = rgb.cpu() 161 | return mesh 162 | -------------------------------------------------------------------------------- /instant-nsr-pl/README.md: -------------------------------------------------------------------------------- 1 | # Instant Neural Surface Reconstruction 2 | 3 | This repository contains a concise and extensible implementation of NeRF and NeuS for neural surface reconstruction based on Instant-NGP and the Pytorch-Lightning framework. **Training on a NeRF-Synthetic scene takes ~5min for NeRF and ~10min for NeuS on a single RTX3090.** 4 | 5 | ||NeRF in 5min|NeuS in 10 min| 6 | |---|---|---| 7 | |Rendering|![rendering-nerf](https://user-images.githubusercontent.com/19284678/199078178-b719676b-7e60-47f1-813b-c0b533f5480d.png)|![rendering-neus](https://user-images.githubusercontent.com/19284678/199078300-ebcf249d-b05e-431f-b035-da354705d8db.png)| 8 | |Mesh|![mesh-nerf](https://user-images.githubusercontent.com/19284678/199078661-b5cd569a-c22b-4220-9c11-d5fd13a52fb8.png)|![mesh-neus](https://user-images.githubusercontent.com/19284678/199078481-164e36a6-6d55-45cc-aaf3-795a114e4a38.png)| 9 | 10 | 11 | ## Features 12 | **This repository aims to provide a highly efficient while customizable boilerplate for research projects based on NeRF or NeuS.** 13 | 14 | - acceleration techniques from [Instant-NGP](https://github.com/NVlabs/instant-ngp): multiresolution hash encoding and fully fused networks by [tiny-cuda-nn](https://github.com/NVlabs/tiny-cuda-nn), occupancy grid pruning and rendering by [nerfacc](https://github.com/KAIR-BAIR/nerfacc) 15 | - out-of-the-box multi-GPU and mixed precision training by [PyTorch-Lightning](https://github.com/Lightning-AI/lightning) 16 | - hierarchical project layout that is designed to be easily customized and extended, flexible experiment configuration by [OmegaConf](https://github.com/omry/omegaconf) 17 | 18 | **Please subscribe to [#26](https://github.com/bennyguo/instant-nsr-pl/issues/26) for our latest findings on quality improvements!** 19 | 20 | ## News 21 | 22 | 🔥🔥🔥 Check out my new project on 3D content generation: https://github.com/threestudio-project/threestudio 🔥🔥🔥 23 | 24 | - 06/03/2023: Add an implementation of [Neuralangelo](https://research.nvidia.com/labs/dir/neuralangelo/). See [here](https://github.com/bennyguo/instant-nsr-pl#training-on-DTU) for details. 25 | - 03/31/2023: NeuS model now supports background modeling. You could try on the DTU dataset provided by [NeuS](https://drive.google.com/drive/folders/1Nlzejs4mfPuJYORLbDEUDWlc9IZIbU0C?usp=sharing) or [IDR](https://www.dropbox.com/sh/5tam07ai8ch90pf/AADniBT3dmAexvm_J1oL__uoa) following [the instruction here](https://github.com/bennyguo/instant-nsr-pl#training-on-DTU). 26 | - 02/11/2023: NeRF model now supports unbounded 360 scenes with learned background. You could try on [MipNeRF 360 data](http://storage.googleapis.com/gresearch/refraw360/360_v2.zip) following [the COLMAP configuration](https://github.com/bennyguo/instant-nsr-pl#training-on-custom-colmap-data). 27 | 28 | ## Requirements 29 | **Note:** 30 | - To utilize multiresolution hash encoding or fully fused networks provided by tiny-cuda-nn, you should have least an RTX 2080Ti, see [https://github.com/NVlabs/tiny-cuda-nn#requirements](https://github.com/NVlabs/tiny-cuda-nn#requirements) for more details. 31 | - Multi-GPU training is currently not supported on Windows (see [#4](https://github.com/bennyguo/instant-nsr-pl/issues/4)). 32 | ### Environments 33 | - Install PyTorch>=1.10 [here](https://pytorch.org/get-started/locally/) based the package management tool you used and your cuda version (older PyTorch versions may work but have not been tested) 34 | - Install tiny-cuda-nn PyTorch extension: `pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch` 35 | - `pip install -r requirements.txt` 36 | 37 | 38 | ## Run 39 | ### Training on NeRF-Synthetic 40 | Download the NeRF-Synthetic data [here](https://drive.google.com/drive/folders/128yBriW1IG_3NJ5Rp7APSTZsJqdJdfc1) and put it under `load/`. The file structure should be like `load/nerf_synthetic/lego`. 41 | 42 | Run the launch script with `--train`, specifying the config file, the GPU(s) to be used (GPU 0 will be used by default), and the scene name: 43 | ```bash 44 | # train NeRF 45 | python launch.py --config configs/nerf-blender.yaml --gpu 0 --train dataset.scene=lego tag=example 46 | 47 | # train NeuS with mask 48 | python launch.py --config configs/neus-blender.yaml --gpu 0 --train dataset.scene=lego tag=example 49 | # train NeuS without mask 50 | python launch.py --config configs/neus-blender.yaml --gpu 0 --train dataset.scene=lego tag=example system.loss.lambda_mask=0.0 51 | ``` 52 | The code snapshots, checkpoints and experiment outputs are saved to `exp/[name]/[tag]@[timestamp]`, and tensorboard logs can be found at `runs/[name]/[tag]@[timestamp]`. You can change any configuration in the YAML file by specifying arguments without `--`, for example: 53 | ```bash 54 | python launch.py --config configs/nerf-blender.yaml --gpu 0 --train dataset.scene=lego tag=iter50k seed=0 trainer.max_steps=50000 55 | ``` 56 | ### Training on DTU 57 | Download preprocessed DTU data provided by [NeuS](https://drive.google.com/drive/folders/1Nlzejs4mfPuJYORLbDEUDWlc9IZIbU0C?usp=sharing) or [IDR](https://www.dropbox.com/sh/5tam07ai8ch90pf/AADniBT3dmAexvm_J1oL__uoa). In the provided config files we assume using NeuS DTU data. If you are using IDR DTU data, please set `dataset.cameras_file=cameras.npz`. You may also need to adjust `dataset.root_dir` to point to your downloaded data location. 58 | ```bash 59 | # train NeuS on DTU without mask 60 | python launch.py --config configs/neus-dtu.yaml --gpu 0 --train 61 | # train NeuS on DTU with mask 62 | python launch.py --config configs/neus-dtu-wmask.yaml --gpu 0 --train 63 | # train NeuS on DTU with mask using tricks from Neuralangelo (experimental) 64 | python launch.py --config configs/neuralangelo-dtu-wmask.yaml --gpu 0 --train 65 | ``` 66 | Notes: 67 | - PSNR in the testing stage is meaningless, as we simply compare to pure white images in testing. 68 | - The results of Neuralangelo can't reach those in the original paper. Some potential improvements: more iterations; larger `system.geometry.xyz_encoding_config.update_steps`; larger `system.geometry.xyz_encoding_config.n_features_per_level`; larger `system.geometry.xyz_encoding_config.log2_hashmap_size`; adopting curvature loss. 69 | 70 | ### Training on Custom COLMAP Data 71 | To get COLMAP data from custom images, you should have COLMAP installed (see [here](https://colmap.github.io/install.html) for installation instructions). Then put your images in the `images/` folder, and run `scripts/imgs2poses.py` specifying the path containing the `images/` folder. For example: 72 | ```bash 73 | python scripts/imgs2poses.py ./load/bmvs_dog # images are in ./load/bmvs_dog/images 74 | ``` 75 | Existing data following this file structure also works as long as images are store in `images/` and there is a `sparse/` folder for the COLMAP output, for example [the data provided by MipNeRF 360](http://storage.googleapis.com/gresearch/refraw360/360_v2.zip). An optional `masks/` folder could be provided for object mask supervision. To train on COLMAP data, please refer to the example config files `config/*-colmap.yaml`. Some notes: 76 | - Adapt the `root_dir` and `img_wh` (or `img_downscale`) option in the config file to your data; 77 | - The scene is normalized so that cameras have a minimum distance `1.0` to the center of the scene. Setting `model.radius=1.0` works in most cases. If not, try setting a smaller radius that wraps tightly to your foreground object. 78 | - There are three choices to determine the scene center: `dataset.center_est_method=camera` uses the center of all camera positions as the scene center; `dataset.center_est_method=lookat` assumes the cameras are looking at the same point and calculates an approximate look-at point as the scene center; `dataset.center_est_method=point` uses the center of all points (reconstructed by COLMAP) that are bounded by cameras as the scene center. Please choose an appropriate method according to your capture. 79 | - PSNR in the testing stage is meaningless, as we simply compare to pure white images in testing. 80 | 81 | ### Testing 82 | The training procedure are by default followed by testing, which computes metrics on test data, generates animations and exports the geometry as triangular meshes. If you want to do testing alone, just resume the pretrained model and replace `--train` with `--test`, for example: 83 | ```bash 84 | python launch.py --config path/to/your/exp/config/parsed.yaml --resume path/to/your/exp/ckpt/epoch=0-step=20000.ckpt --gpu 0 --test 85 | ``` 86 | 87 | 88 | ## Benchmarks 89 | All experiments are conducted on a single NVIDIA RTX3090. 90 | 91 | |PSNR|Chair|Drums|Ficus|Hotdog|Lego|Materials|Mic|Ship|Avg.| 92 | |---|---|---|---|---|---|---|---|---|---| 93 | |NeRF Paper|33.00|25.01|30.13|36.18|32.54|29.62|32.91|28.65|31.01| 94 | |NeRF Ours (20k)|34.80|26.04|33.89|37.42|35.33|29.46|35.22|31.17|32.92| 95 | |NeuS Ours (20k, with masks)|34.04|25.26|32.47|35.94|33.78|27.67|33.43|29.50|31.51| 96 | 97 | |Training Time (mm:ss)|Chair|Drums|Ficus|Hotdog|Lego|Materials|Mic|Ship|Avg.| 98 | |---|---|---|---|---|---|---|---|---|---| 99 | |NeRF Ours (20k)|04:34|04:35|04:18|04:46|04:39|04:35|04:26|05:41|04:42| 100 | |NeuS Ours (20k, with masks)|11:25|10:34|09:51|12:11|11:37|11:46|09:59|16:25|11:44| 101 | 102 | 103 | ## TODO 104 | - [✅] Support more dataset formats, like COLMAP outputs and DTU 105 | - [✅] Support simple background model 106 | - [ ] Support GUI training and interaction 107 | - [ ] More illustrations about the framework 108 | 109 | ## Related Projects 110 | - [ngp_pl](https://github.com/kwea123/ngp_pl): Great Instant-NGP implementation in PyTorch-Lightning! Background model and GUI supported. 111 | - [Instant-NSR](https://github.com/zhaofuq/Instant-NSR): NeuS implementation using multiresolution hash encoding. 112 | 113 | ## Citation 114 | If you find this codebase useful, please consider citing: 115 | ``` 116 | @misc{instant-nsr-pl, 117 | Author = {Yuan-Chen Guo}, 118 | Year = {2022}, 119 | Note = {https://github.com/bennyguo/instant-nsr-pl}, 120 | Title = {Instant Neural Surface Reconstruction} 121 | } 122 | ``` 123 | -------------------------------------------------------------------------------- /instant-nsr-pl/models/network_utils.py: -------------------------------------------------------------------------------- 1 | import math 2 | import numpy as np 3 | 4 | import torch 5 | import torch.nn as nn 6 | import tinycudann as tcnn 7 | 8 | from pytorch_lightning.utilities.rank_zero import rank_zero_debug, rank_zero_info 9 | 10 | from utils.misc import config_to_primitive, get_rank 11 | from models.utils import get_activation 12 | from systems.utils import update_module_step 13 | 14 | class VanillaFrequency(nn.Module): 15 | def __init__(self, in_channels, config): 16 | super().__init__() 17 | self.N_freqs = config['n_frequencies'] 18 | self.in_channels, self.n_input_dims = in_channels, in_channels 19 | self.funcs = [torch.sin, torch.cos] 20 | self.freq_bands = 2**torch.linspace(0, self.N_freqs-1, self.N_freqs) 21 | self.n_output_dims = self.in_channels * (len(self.funcs) * self.N_freqs) 22 | self.n_masking_step = config.get('n_masking_step', 0) 23 | self.update_step(None, None) # mask should be updated at the beginning each step 24 | 25 | def forward(self, x): 26 | out = [] 27 | for freq, mask in zip(self.freq_bands, self.mask): 28 | for func in self.funcs: 29 | out += [func(freq*x) * mask] 30 | return torch.cat(out, -1) 31 | 32 | def update_step(self, epoch, global_step): 33 | if self.n_masking_step <= 0 or global_step is None: 34 | self.mask = torch.ones(self.N_freqs, dtype=torch.float32) 35 | else: 36 | self.mask = (1. - torch.cos(math.pi * (global_step / self.n_masking_step * self.N_freqs - torch.arange(0, self.N_freqs)).clamp(0, 1))) / 2. 37 | rank_zero_debug(f'Update mask: {global_step}/{self.n_masking_step} {self.mask}') 38 | 39 | 40 | class ProgressiveBandHashGrid(nn.Module): 41 | def __init__(self, in_channels, config): 42 | super().__init__() 43 | self.n_input_dims = in_channels 44 | encoding_config = config.copy() 45 | encoding_config['otype'] = 'HashGrid' 46 | with torch.cuda.device(get_rank()): 47 | self.encoding = tcnn.Encoding(in_channels, encoding_config) 48 | self.n_output_dims = self.encoding.n_output_dims 49 | self.n_level = config['n_levels'] 50 | self.n_features_per_level = config['n_features_per_level'] 51 | self.start_level, self.start_step, self.update_steps = config['start_level'], config['start_step'], config['update_steps'] 52 | self.current_level = self.start_level 53 | self.mask = torch.zeros(self.n_level * self.n_features_per_level, dtype=torch.float32, device=get_rank()) 54 | 55 | def forward(self, x): 56 | enc = self.encoding(x) 57 | enc = enc * self.mask 58 | return enc 59 | 60 | def update_step(self, epoch, global_step): 61 | current_level = min(self.start_level + max(global_step - self.start_step, 0) // self.update_steps, self.n_level) 62 | if current_level > self.current_level: 63 | rank_zero_info(f'Update grid level to {current_level}') 64 | self.current_level = current_level 65 | self.mask[:self.current_level * self.n_features_per_level] = 1. 66 | 67 | 68 | class CompositeEncoding(nn.Module): 69 | def __init__(self, encoding, include_xyz=False, xyz_scale=1., xyz_offset=0.): 70 | super(CompositeEncoding, self).__init__() 71 | self.encoding = encoding 72 | self.include_xyz, self.xyz_scale, self.xyz_offset = include_xyz, xyz_scale, xyz_offset 73 | self.n_output_dims = int(self.include_xyz) * self.encoding.n_input_dims + self.encoding.n_output_dims 74 | 75 | def forward(self, x, *args): 76 | return self.encoding(x, *args) if not self.include_xyz else torch.cat([x * self.xyz_scale + self.xyz_offset, self.encoding(x, *args)], dim=-1) 77 | 78 | def update_step(self, epoch, global_step): 79 | update_module_step(self.encoding, epoch, global_step) 80 | 81 | 82 | def get_encoding(n_input_dims, config): 83 | # input suppose to be range [0, 1] 84 | if config.otype == 'VanillaFrequency': 85 | encoding = VanillaFrequency(n_input_dims, config_to_primitive(config)) 86 | elif config.otype == 'ProgressiveBandHashGrid': 87 | encoding = ProgressiveBandHashGrid(n_input_dims, config_to_primitive(config)) 88 | else: 89 | with torch.cuda.device(get_rank()): 90 | encoding = tcnn.Encoding(n_input_dims, config_to_primitive(config)) 91 | encoding = CompositeEncoding(encoding, include_xyz=config.get('include_xyz', False), xyz_scale=2., xyz_offset=-1.) 92 | return encoding 93 | 94 | 95 | class VanillaMLP(nn.Module): 96 | def __init__(self, dim_in, dim_out, config): 97 | super().__init__() 98 | self.n_neurons, self.n_hidden_layers = config['n_neurons'], config['n_hidden_layers'] 99 | self.sphere_init, self.weight_norm = config.get('sphere_init', False), config.get('weight_norm', False) 100 | self.sphere_init_radius = config.get('sphere_init_radius', 0.5) 101 | self.layers = [self.make_linear(dim_in, self.n_neurons, is_first=True, is_last=False), self.make_activation()] 102 | for i in range(self.n_hidden_layers - 1): 103 | self.layers += [self.make_linear(self.n_neurons, self.n_neurons, is_first=False, is_last=False), self.make_activation()] 104 | self.layers += [self.make_linear(self.n_neurons, dim_out, is_first=False, is_last=True)] 105 | self.layers = nn.Sequential(*self.layers) 106 | self.output_activation = get_activation(config['output_activation']) 107 | 108 | @torch.cuda.amp.autocast(False) 109 | def forward(self, x): 110 | x = self.layers(x.float()) 111 | x = self.output_activation(x) 112 | return x 113 | 114 | def make_linear(self, dim_in, dim_out, is_first, is_last): 115 | layer = nn.Linear(dim_in, dim_out, bias=True) # network without bias will degrade quality 116 | if self.sphere_init: 117 | if is_last: 118 | torch.nn.init.constant_(layer.bias, -self.sphere_init_radius) 119 | torch.nn.init.normal_(layer.weight, mean=math.sqrt(math.pi) / math.sqrt(dim_in), std=0.0001) 120 | elif is_first: 121 | torch.nn.init.constant_(layer.bias, 0.0) 122 | torch.nn.init.constant_(layer.weight[:, 3:], 0.0) 123 | torch.nn.init.normal_(layer.weight[:, :3], 0.0, math.sqrt(2) / math.sqrt(dim_out)) 124 | else: 125 | torch.nn.init.constant_(layer.bias, 0.0) 126 | torch.nn.init.normal_(layer.weight, 0.0, math.sqrt(2) / math.sqrt(dim_out)) 127 | else: 128 | torch.nn.init.constant_(layer.bias, 0.0) 129 | torch.nn.init.kaiming_uniform_(layer.weight, nonlinearity='relu') 130 | 131 | if self.weight_norm: 132 | layer = nn.utils.weight_norm(layer) 133 | return layer 134 | 135 | def make_activation(self): 136 | if self.sphere_init: 137 | return nn.Softplus(beta=100) 138 | else: 139 | return nn.ReLU(inplace=True) 140 | 141 | 142 | def sphere_init_tcnn_network(n_input_dims, n_output_dims, config, network): 143 | rank_zero_debug('Initialize tcnn MLP to approximately represent a sphere.') 144 | """ 145 | from https://github.com/NVlabs/tiny-cuda-nn/issues/96 146 | It's the weight matrices of each layer laid out in row-major order and then concatenated. 147 | Notably: inputs and output dimensions are padded to multiples of 8 (CutlassMLP) or 16 (FullyFusedMLP). 148 | The padded input dimensions get a constant value of 1.0, 149 | whereas the padded output dimensions are simply ignored, 150 | so the weights pertaining to those can have any value. 151 | """ 152 | padto = 16 if config.otype == 'FullyFusedMLP' else 8 153 | n_input_dims = n_input_dims + (padto - n_input_dims % padto) % padto 154 | n_output_dims = n_output_dims + (padto - n_output_dims % padto) % padto 155 | data = list(network.parameters())[0].data 156 | assert data.shape[0] == (n_input_dims + n_output_dims) * config.n_neurons + (config.n_hidden_layers - 1) * config.n_neurons**2 157 | new_data = [] 158 | # first layer 159 | weight = torch.zeros((config.n_neurons, n_input_dims)).to(data) 160 | torch.nn.init.constant_(weight[:, 3:], 0.0) 161 | torch.nn.init.normal_(weight[:, :3], 0.0, math.sqrt(2) / math.sqrt(config.n_neurons)) 162 | new_data.append(weight.flatten()) 163 | # hidden layers 164 | for i in range(config.n_hidden_layers - 1): 165 | weight = torch.zeros((config.n_neurons, config.n_neurons)).to(data) 166 | torch.nn.init.normal_(weight, 0.0, math.sqrt(2) / math.sqrt(config.n_neurons)) 167 | new_data.append(weight.flatten()) 168 | # last layer 169 | weight = torch.zeros((n_output_dims, config.n_neurons)).to(data) 170 | torch.nn.init.normal_(weight, mean=math.sqrt(math.pi) / math.sqrt(config.n_neurons), std=0.0001) 171 | new_data.append(weight.flatten()) 172 | new_data = torch.cat(new_data) 173 | data.copy_(new_data) 174 | 175 | 176 | def get_mlp(n_input_dims, n_output_dims, config): 177 | if config.otype == 'VanillaMLP': 178 | network = VanillaMLP(n_input_dims, n_output_dims, config_to_primitive(config)) 179 | else: 180 | with torch.cuda.device(get_rank()): 181 | network = tcnn.Network(n_input_dims, n_output_dims, config_to_primitive(config)) 182 | if config.get('sphere_init', False): 183 | sphere_init_tcnn_network(n_input_dims, n_output_dims, config, network) 184 | return network 185 | 186 | 187 | class EncodingWithNetwork(nn.Module): 188 | def __init__(self, encoding, network): 189 | super().__init__() 190 | self.encoding, self.network = encoding, network 191 | 192 | def forward(self, x): 193 | return self.network(self.encoding(x)) 194 | 195 | def update_step(self, epoch, global_step): 196 | update_module_step(self.encoding, epoch, global_step) 197 | update_module_step(self.network, epoch, global_step) 198 | 199 | 200 | def get_encoding_with_network(n_input_dims, n_output_dims, encoding_config, network_config): 201 | # input suppose to be range [0, 1] 202 | if encoding_config.otype in ['VanillaFrequency', 'ProgressiveBandHashGrid'] \ 203 | or network_config.otype in ['VanillaMLP']: 204 | encoding = get_encoding(n_input_dims, encoding_config) 205 | network = get_mlp(encoding.n_output_dims, n_output_dims, network_config) 206 | encoding_with_network = EncodingWithNetwork(encoding, network) 207 | else: 208 | with torch.cuda.device(get_rank()): 209 | encoding_with_network = tcnn.NetworkWithInputEncoding( 210 | n_input_dims=n_input_dims, 211 | n_output_dims=n_output_dims, 212 | encoding_config=config_to_primitive(encoding_config), 213 | network_config=config_to_primitive(network_config) 214 | ) 215 | return encoding_with_network 216 | -------------------------------------------------------------------------------- /instant-nsr-pl/utils/mixins.py: -------------------------------------------------------------------------------- 1 | import os 2 | import re 3 | import shutil 4 | import numpy as np 5 | import cv2 6 | import imageio 7 | from matplotlib import cm 8 | from matplotlib.colors import LinearSegmentedColormap 9 | import json 10 | 11 | import torch 12 | 13 | from utils.obj import write_obj 14 | 15 | 16 | class SaverMixin(): 17 | @property 18 | def save_dir(self): 19 | return self.config.save_dir 20 | 21 | def convert_data(self, data): 22 | if isinstance(data, np.ndarray): 23 | return data 24 | elif isinstance(data, torch.Tensor): 25 | return data.cpu().numpy() 26 | elif isinstance(data, list): 27 | return [self.convert_data(d) for d in data] 28 | elif isinstance(data, dict): 29 | return {k: self.convert_data(v) for k, v in data.items()} 30 | else: 31 | raise TypeError('Data must be in type numpy.ndarray, torch.Tensor, list or dict, getting', type(data)) 32 | 33 | def get_save_path(self, filename): 34 | save_path = os.path.join(self.save_dir, filename) 35 | os.makedirs(os.path.dirname(save_path), exist_ok=True) 36 | return save_path 37 | 38 | DEFAULT_RGB_KWARGS = {'data_format': 'CHW', 'data_range': (0, 1)} 39 | DEFAULT_UV_KWARGS = {'data_format': 'CHW', 'data_range': (0, 1), 'cmap': 'checkerboard'} 40 | DEFAULT_GRAYSCALE_KWARGS = {'data_range': None, 'cmap': 'jet'} 41 | 42 | def get_rgb_image_(self, img, data_format, data_range): 43 | img = self.convert_data(img) 44 | assert data_format in ['CHW', 'HWC'] 45 | if data_format == 'CHW': 46 | img = img.transpose(1, 2, 0) 47 | img = img.clip(min=data_range[0], max=data_range[1]) 48 | img = ((img - data_range[0]) / (data_range[1] - data_range[0]) * 255.).astype(np.uint8) 49 | imgs = [img[...,start:start+3] for start in range(0, img.shape[-1], 3)] 50 | imgs = [img_ if img_.shape[-1] == 3 else np.concatenate([img_, np.zeros((img_.shape[0], img_.shape[1], 3 - img_.shape[2]), dtype=img_.dtype)], axis=-1) for img_ in imgs] 51 | img = np.concatenate(imgs, axis=1) 52 | img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) 53 | return img 54 | 55 | def save_rgb_image(self, filename, img, data_format=DEFAULT_RGB_KWARGS['data_format'], data_range=DEFAULT_RGB_KWARGS['data_range']): 56 | img = self.get_rgb_image_(img, data_format, data_range) 57 | cv2.imwrite(self.get_save_path(filename), img) 58 | 59 | def get_uv_image_(self, img, data_format, data_range, cmap): 60 | img = self.convert_data(img) 61 | assert data_format in ['CHW', 'HWC'] 62 | if data_format == 'CHW': 63 | img = img.transpose(1, 2, 0) 64 | img = img.clip(min=data_range[0], max=data_range[1]) 65 | img = (img - data_range[0]) / (data_range[1] - data_range[0]) 66 | assert cmap in ['checkerboard', 'color'] 67 | if cmap == 'checkerboard': 68 | n_grid = 64 69 | mask = (img * n_grid).astype(int) 70 | mask = (mask[...,0] + mask[...,1]) % 2 == 0 71 | img = np.ones((img.shape[0], img.shape[1], 3), dtype=np.uint8) * 255 72 | img[mask] = np.array([255, 0, 255], dtype=np.uint8) 73 | img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) 74 | elif cmap == 'color': 75 | img_ = np.zeros((img.shape[0], img.shape[1], 3), dtype=np.uint8) 76 | img_[..., 0] = (img[..., 0] * 255).astype(np.uint8) 77 | img_[..., 1] = (img[..., 1] * 255).astype(np.uint8) 78 | img_ = cv2.cvtColor(img_, cv2.COLOR_RGB2BGR) 79 | img = img_ 80 | return img 81 | 82 | def save_uv_image(self, filename, img, data_format=DEFAULT_UV_KWARGS['data_format'], data_range=DEFAULT_UV_KWARGS['data_range'], cmap=DEFAULT_UV_KWARGS['cmap']): 83 | img = self.get_uv_image_(img, data_format, data_range, cmap) 84 | cv2.imwrite(self.get_save_path(filename), img) 85 | 86 | def get_grayscale_image_(self, img, data_range, cmap): 87 | img = self.convert_data(img) 88 | img = np.nan_to_num(img) 89 | if data_range is None: 90 | img = (img - img.min()) / (img.max() - img.min()) 91 | else: 92 | img = img.clip(data_range[0], data_range[1]) 93 | img = (img - data_range[0]) / (data_range[1] - data_range[0]) 94 | assert cmap in [None, 'jet', 'magma'] 95 | if cmap == None: 96 | img = (img * 255.).astype(np.uint8) 97 | img = np.repeat(img[...,None], 3, axis=2) 98 | elif cmap == 'jet': 99 | img = (img * 255.).astype(np.uint8) 100 | img = cv2.applyColorMap(img, cv2.COLORMAP_JET) 101 | elif cmap == 'magma': 102 | img = 1. - img 103 | base = cm.get_cmap('magma') 104 | num_bins = 256 105 | colormap = LinearSegmentedColormap.from_list( 106 | f"{base.name}{num_bins}", 107 | base(np.linspace(0, 1, num_bins)), 108 | num_bins 109 | )(np.linspace(0, 1, num_bins))[:,:3] 110 | a = np.floor(img * 255.) 111 | b = (a + 1).clip(max=255.) 112 | f = img * 255. - a 113 | a = a.astype(np.uint16).clip(0, 255) 114 | b = b.astype(np.uint16).clip(0, 255) 115 | img = colormap[a] + (colormap[b] - colormap[a]) * f[...,None] 116 | img = (img * 255.).astype(np.uint8) 117 | return img 118 | 119 | def save_grayscale_image(self, filename, img, data_range=DEFAULT_GRAYSCALE_KWARGS['data_range'], cmap=DEFAULT_GRAYSCALE_KWARGS['cmap']): 120 | img = self.get_grayscale_image_(img, data_range, cmap) 121 | cv2.imwrite(self.get_save_path(filename), img) 122 | 123 | def get_image_grid_(self, imgs): 124 | if isinstance(imgs[0], list): 125 | return np.concatenate([self.get_image_grid_(row) for row in imgs], axis=0) 126 | cols = [] 127 | for col in imgs: 128 | assert col['type'] in ['rgb', 'uv', 'grayscale'] 129 | if col['type'] == 'rgb': 130 | rgb_kwargs = self.DEFAULT_RGB_KWARGS.copy() 131 | rgb_kwargs.update(col['kwargs']) 132 | cols.append(self.get_rgb_image_(col['img'], **rgb_kwargs)) 133 | elif col['type'] == 'uv': 134 | uv_kwargs = self.DEFAULT_UV_KWARGS.copy() 135 | uv_kwargs.update(col['kwargs']) 136 | cols.append(self.get_uv_image_(col['img'], **uv_kwargs)) 137 | elif col['type'] == 'grayscale': 138 | grayscale_kwargs = self.DEFAULT_GRAYSCALE_KWARGS.copy() 139 | grayscale_kwargs.update(col['kwargs']) 140 | cols.append(self.get_grayscale_image_(col['img'], **grayscale_kwargs)) 141 | return np.concatenate(cols, axis=1) 142 | 143 | def save_image_grid(self, filename, imgs): 144 | img = self.get_image_grid_(imgs) 145 | cv2.imwrite(self.get_save_path(filename), img) 146 | 147 | def save_image(self, filename, img): 148 | img = self.convert_data(img) 149 | assert img.dtype == np.uint8 150 | if img.shape[-1] == 3: 151 | img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) 152 | elif img.shape[-1] == 4: 153 | img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGRA) 154 | cv2.imwrite(self.get_save_path(filename), img) 155 | 156 | def save_cubemap(self, filename, img, data_range=(0, 1)): 157 | img = self.convert_data(img) 158 | assert img.ndim == 4 and img.shape[0] == 6 and img.shape[1] == img.shape[2] 159 | 160 | imgs_full = [] 161 | for start in range(0, img.shape[-1], 3): 162 | img_ = img[...,start:start+3] 163 | img_ = np.stack([self.get_rgb_image_(img_[i], 'HWC', data_range) for i in range(img_.shape[0])], axis=0) 164 | size = img_.shape[1] 165 | placeholder = np.zeros((size, size, 3), dtype=np.float32) 166 | img_full = np.concatenate([ 167 | np.concatenate([placeholder, img_[2], placeholder, placeholder], axis=1), 168 | np.concatenate([img_[1], img_[4], img_[0], img_[5]], axis=1), 169 | np.concatenate([placeholder, img_[3], placeholder, placeholder], axis=1) 170 | ], axis=0) 171 | img_full = cv2.cvtColor(img_full, cv2.COLOR_RGB2BGR) 172 | imgs_full.append(img_full) 173 | 174 | imgs_full = np.concatenate(imgs_full, axis=1) 175 | cv2.imwrite(self.get_save_path(filename), imgs_full) 176 | 177 | def save_data(self, filename, data): 178 | data = self.convert_data(data) 179 | if isinstance(data, dict): 180 | if not filename.endswith('.npz'): 181 | filename += '.npz' 182 | np.savez(self.get_save_path(filename), **data) 183 | else: 184 | if not filename.endswith('.npy'): 185 | filename += '.npy' 186 | np.save(self.get_save_path(filename), data) 187 | 188 | def save_state_dict(self, filename, data): 189 | torch.save(data, self.get_save_path(filename)) 190 | 191 | def save_img_sequence(self, filename, img_dir, matcher, save_format='gif', fps=30): 192 | assert save_format in ['gif', 'mp4'] 193 | if not filename.endswith(save_format): 194 | filename += f".{save_format}" 195 | matcher = re.compile(matcher) 196 | img_dir = os.path.join(self.save_dir, img_dir) 197 | imgs = [] 198 | for f in os.listdir(img_dir): 199 | if matcher.search(f): 200 | imgs.append(f) 201 | imgs = sorted(imgs, key=lambda f: int(matcher.search(f).groups()[0])) 202 | imgs = [cv2.imread(os.path.join(img_dir, f)) for f in imgs] 203 | 204 | if save_format == 'gif': 205 | imgs = [cv2.cvtColor(i, cv2.COLOR_BGR2RGB) for i in imgs] 206 | imageio.mimsave(self.get_save_path(filename), imgs, fps=fps, palettesize=256) 207 | elif save_format == 'mp4': 208 | imgs = [cv2.cvtColor(i, cv2.COLOR_BGR2RGB) for i in imgs] 209 | imageio.mimsave(self.get_save_path(filename), imgs, fps=fps) 210 | 211 | def save_mesh(self, filename, v_pos, t_pos_idx, v_tex=None, t_tex_idx=None, v_rgb=None): 212 | v_pos, t_pos_idx = self.convert_data(v_pos), self.convert_data(t_pos_idx) 213 | if v_rgb is not None: 214 | v_rgb = self.convert_data(v_rgb) 215 | 216 | import trimesh 217 | mesh = trimesh.Trimesh( 218 | vertices=v_pos, 219 | faces=t_pos_idx, 220 | vertex_colors=v_rgb 221 | ) 222 | mesh.export(self.get_save_path(filename)) 223 | 224 | def save_file(self, filename, src_path): 225 | shutil.copyfile(src_path, self.get_save_path(filename)) 226 | 227 | def save_json(self, filename, payload): 228 | with open(self.get_save_path(filename), 'w') as f: 229 | f.write(json.dumps(payload)) 230 | -------------------------------------------------------------------------------- /instant-nsr-pl/systems/nerf.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from torch_efficient_distloss import flatten_eff_distloss 5 | 6 | import pytorch_lightning as pl 7 | from pytorch_lightning.utilities.rank_zero import rank_zero_info, rank_zero_debug 8 | from torch.cuda.amp import autocast 9 | from models.ray_utils import get_rays 10 | import systems 11 | from systems.base import BaseSystem 12 | from systems.criterions import PSNR, SSIM 13 | from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity 14 | import os 15 | from systems.loss_utils import ssim as ca_ssim 16 | lpips = LearnedPerceptualImagePatchSimilarity(net_type='vgg').cuda() 17 | 18 | @systems.register('nerf-system') 19 | class NeRFSystem(BaseSystem): 20 | """ 21 | Two ways to print to console: 22 | 1. self.print: correctly handle progress bar 23 | 2. rank_zero_info: use the logging module 24 | """ 25 | def prepare(self): 26 | self.criterions = { 27 | 'psnr': PSNR(), 28 | } 29 | self.train_num_samples = self.config.model.train_num_rays * self.config.model.num_samples_per_ray 30 | self.train_num_rays = self.config.model.train_num_rays 31 | 32 | def forward(self, batch): 33 | return self.model(batch['rays']) 34 | 35 | def preprocess_data(self, batch, stage): 36 | if 'index' in batch: # validation / testing 37 | index = batch['index'] 38 | else: 39 | if self.config.model.batch_image_sampling: 40 | index = torch.randint(0, len(self.dataset.all_images), size=(self.train_num_rays,), device=self.dataset.all_images.device) 41 | else: 42 | index = torch.randint(0, len(self.dataset.all_images), size=(1,), device=self.dataset.all_images.device) 43 | if stage in ['train']: 44 | c2w = self.dataset.all_c2w[index] 45 | x = torch.randint( 46 | 0, self.dataset.w, size=(self.train_num_rays,), device=self.dataset.all_images.device 47 | ) 48 | y = torch.randint( 49 | 0, self.dataset.h, size=(self.train_num_rays,), device=self.dataset.all_images.device 50 | ) 51 | if self.dataset.directions.ndim == 3: # (H, W, 3) 52 | directions = self.dataset.directions[y, x] 53 | elif self.dataset.directions.ndim == 4: # (N, H, W, 3) 54 | directions = self.dataset.directions[index, y, x] 55 | rays_o, rays_d = get_rays(directions, c2w) 56 | rgb = self.dataset.all_images[index, y, x].view(-1, self.dataset.all_images.shape[-1]).to(self.rank) 57 | fg_mask = self.dataset.all_fg_masks[index, y, x].view(-1).to(self.rank) 58 | else: 59 | c2w = self.dataset.all_c2w[index][0] 60 | if self.dataset.directions.ndim == 3: # (H, W, 3) 61 | directions = self.dataset.directions 62 | elif self.dataset.directions.ndim == 4: # (N, H, W, 3) 63 | directions = self.dataset.directions[index][0] 64 | rays_o, rays_d = get_rays(directions, c2w) 65 | rgb = self.dataset.all_images[index].view(-1, self.dataset.all_images.shape[-1]).to(self.rank) 66 | fg_mask = self.dataset.all_fg_masks[index].view(-1).to(self.rank) 67 | 68 | rays = torch.cat([rays_o, F.normalize(rays_d, p=2, dim=-1)], dim=-1) 69 | 70 | if stage in ['train']: 71 | if self.config.model.background_color == 'white': 72 | self.model.background_color = torch.ones((3,), dtype=torch.float32, device=self.rank) 73 | elif self.config.model.background_color == 'black': 74 | self.model.background_color = torch.zeros((3,), dtype=torch.float32, device=self.rank) 75 | elif self.config.model.background_color == 'random': 76 | self.model.background_color = torch.rand((3,), dtype=torch.float32, device=self.rank) 77 | else: 78 | raise NotImplementedError 79 | else: 80 | self.model.background_color = torch.zeros((3,), dtype=torch.float32, device=self.rank) 81 | 82 | if self.dataset.apply_mask: 83 | rgb = rgb * fg_mask[...,None] + self.model.background_color * (1 - fg_mask[...,None]) 84 | 85 | batch.update({ 86 | 'rays': rays, 87 | 'rgb': rgb, 88 | 'fg_mask': fg_mask 89 | }) 90 | 91 | def training_step(self, batch, batch_idx): 92 | out = self(batch) 93 | 94 | loss = 0. 95 | 96 | # update train_num_rays 97 | if self.config.model.dynamic_ray_sampling: 98 | train_num_rays = int(self.train_num_rays * (self.train_num_samples / out['num_samples'].sum().item())) 99 | self.train_num_rays = min(int(self.train_num_rays * 0.9 + train_num_rays * 0.1), self.config.model.max_train_num_rays) 100 | 101 | loss_rgb = F.smooth_l1_loss(out['comp_rgb'][out['rays_valid'][...,0]], batch['rgb'][out['rays_valid'][...,0]]) 102 | self.log('train/loss_rgb', loss_rgb) 103 | loss += loss_rgb * self.C(self.config.system.loss.lambda_rgb) 104 | 105 | # distortion loss proposed in MipNeRF360 106 | # an efficient implementation from https://github.com/sunset1995/torch_efficient_distloss, but still slows down training by ~30% 107 | if self.C(self.config.system.loss.lambda_distortion) > 0: 108 | loss_distortion = flatten_eff_distloss(out['weights'], out['points'], out['intervals'], out['ray_indices']) 109 | self.log('train/loss_distortion', loss_distortion) 110 | loss += loss_distortion * self.C(self.config.system.loss.lambda_distortion) 111 | 112 | losses_model_reg = self.model.regularizations(out) 113 | for name, value in losses_model_reg.items(): 114 | self.log(f'train/loss_{name}', value) 115 | loss_ = value * self.C(self.config.system.loss[f"lambda_{name}"]) 116 | loss += loss_ 117 | 118 | for name, value in self.config.system.loss.items(): 119 | if name.startswith('lambda'): 120 | self.log(f'train_params/{name}', self.C(value)) 121 | 122 | self.log('train/num_rays', float(self.train_num_rays), prog_bar=True) 123 | 124 | return { 125 | 'loss': loss 126 | } 127 | 128 | """ 129 | # aggregate outputs from different devices (DP) 130 | def training_step_end(self, out): 131 | pass 132 | """ 133 | 134 | """ 135 | # aggregate outputs from different iterations 136 | def training_epoch_end(self, out): 137 | pass 138 | """ 139 | 140 | def validation_step(self, batch, batch_idx): 141 | out = self(batch) 142 | psnr = self.criterions['psnr'](out['comp_rgb'].to(batch['rgb']), batch['rgb']) 143 | W, H = self.dataset.img_wh 144 | self.save_image_grid(f"it{self.global_step}-{batch['index'][0].item()}.png", [ 145 | {'type': 'rgb', 'img': batch['rgb'].view(H, W, 3), 'kwargs': {'data_format': 'HWC'}}, 146 | {'type': 'rgb', 'img': out['comp_rgb'].view(H, W, 3), 'kwargs': {'data_format': 'HWC'}}, 147 | {'type': 'grayscale', 'img': out['depth'].view(H, W), 'kwargs': {}}, 148 | {'type': 'grayscale', 'img': out['opacity'].view(H, W), 'kwargs': {'cmap': None, 'data_range': (0, 1)}} 149 | ]) 150 | return { 151 | 'psnr': psnr, 152 | 'index': batch['index'] 153 | } 154 | 155 | 156 | """ 157 | # aggregate outputs from different devices when using DP 158 | def validation_step_end(self, out): 159 | pass 160 | """ 161 | 162 | def validation_epoch_end(self, out): 163 | out = self.all_gather(out) 164 | if self.trainer.is_global_zero: 165 | out_set = {} 166 | for step_out in out: 167 | # DP 168 | if step_out['index'].ndim == 1: 169 | out_set[step_out['index'].item()] = {'psnr': step_out['psnr']} 170 | # DDP 171 | else: 172 | for oi, index in enumerate(step_out['index']): 173 | out_set[index[0].item()] = {'psnr': step_out['psnr'][oi]} 174 | psnr = torch.mean(torch.stack([o['psnr'] for o in out_set.values()])) 175 | self.log('val/psnr', psnr, prog_bar=True, rank_zero_only=True) 176 | 177 | def test_step(self, batch, batch_idx): 178 | out = self(batch) 179 | W, H = self.dataset.img_wh 180 | render_img = out['comp_rgb'].view(H, W, 3).permute(2, 0, 1).unsqueeze(0).to(batch['rgb']) 181 | target_img = batch['rgb'].view(H, W, 3).permute(2, 0, 1).unsqueeze(0) 182 | psnr = self.criterions['psnr'](out['comp_rgb'].to(batch['rgb']), batch['rgb']) 183 | with autocast(enabled=False): 184 | ssim = ca_ssim(render_img, target_img) 185 | lpips_ = lpips(render_img, target_img) 186 | self.save_image_grid(f"it{self.global_step}-test/{batch['index'][0].item()}.png", [ 187 | {'type': 'rgb', 'img': batch['rgb'].view(H, W, 3), 'kwargs': {'data_format': 'HWC'}}, 188 | {'type': 'rgb', 'img': out['comp_rgb'].view(H, W, 3), 'kwargs': {'data_format': 'HWC'}}, 189 | {'type': 'grayscale', 'img': out['depth'].view(H, W), 'kwargs': {}}, 190 | {'type': 'grayscale', 'img': out['opacity'].view(H, W), 'kwargs': {'cmap': None, 'data_range': (0, 1)}} 191 | ]) 192 | return { 193 | 'psnr': psnr, 194 | 'ssim': ssim, 195 | 'lpips': lpips_, 196 | 'index': batch['index'] 197 | } 198 | 199 | def test_epoch_end(self, out): 200 | out = self.all_gather(out) 201 | if self.trainer.is_global_zero: 202 | out_set = {} 203 | for step_out in out: 204 | # DP 205 | if step_out['index'].ndim == 1: 206 | out_set[step_out['index'].item()] = {'psnr': step_out['psnr'], 'ssim': step_out['ssim'], 'lpips': step_out['lpips']} 207 | # DDP 208 | else: 209 | for oi, index in enumerate(step_out['index']): 210 | out_set[index[0].item()] = {'psnr': step_out['psnr'][oi], 'ssim': step_out['ssim'][oi], 'lpips': step_out['lpips'][oi]} 211 | psnr = torch.mean(torch.stack([o['psnr'] for o in out_set.values()])) 212 | minpsnr = torch.min(torch.stack([o['psnr'] for o in out_set.values()])) 213 | maxpsnr = torch.max(torch.stack([o['psnr'] for o in out_set.values()])) 214 | ssim = torch.mean(torch.stack([o['ssim'] for o in out_set.values()])) 215 | lpips = torch.mean(torch.stack([o['lpips'] for o in out_set.values()])) 216 | self.log('test/psnr', psnr, prog_bar=True, rank_zero_only=True) 217 | text = f"{self.config.dataset.object}:instant-nsr-pl:{self.config.dataset.bsdf_name}: PSNR={psnr} [min={minpsnr} max={maxpsnr}]-SSIM={ssim}-lpips={lpips}" 218 | os.makedirs(f'{self.config.render_save_dir}/{self.config.dataset.object}', exist_ok=True) 219 | with open(os.path.join(f'../output', self.config.dataset.object, "instant-nsr-pl-wmask.txt"), "w") as file: 220 | file.write(text + '\n') 221 | self.export() 222 | 223 | def export(self): 224 | mesh = self.model.export(self.config.export) 225 | self.save_mesh( 226 | f"it{self.global_step}-{self.config.model.geometry.isosurface.method}{self.config.model.geometry.isosurface.resolution}.obj", 227 | **mesh 228 | ) 229 | -------------------------------------------------------------------------------- /Openmaterial-main/eval/clean_mesh.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 as cv 3 | import os 4 | from glob import glob 5 | import trimesh 6 | import argparse 7 | import mitsuba as mi 8 | mi.set_variant('cuda_ad_rgb') 9 | from mitsuba import ScalarTransform4f as T 10 | import numpy as np 11 | import json 12 | import math 13 | 14 | """ 15 | This code is adapted from: 16 | https://github.com/xxlong0/SparseNeuS/blob/main/evaluation/clean_mesh.py 17 | 18 | """ 19 | 20 | def gen_w2c(pose): 21 | 22 | pose[:3, :1] = -pose[:3, :1] 23 | pose[:3, 1:2] = -pose[:3, 1:2] # Flip the x+ and y+ to align coordinate system 24 | 25 | R = pose[:3, :3].transpose() 26 | T = -R @ pose[:3, 3:] 27 | return R, T 28 | 29 | def gen_camera_intrinsic(width, height, fov_x, fov_y): 30 | fx = width / 2.0 / math.tan(fov_x / 180 * math.pi / 2.0) 31 | fy = height / 2.0 / math.tan(fov_y / 180 * math.pi / 2.0) 32 | return fx, fy 33 | 34 | def clean_points_by_mask(points, bsdf_name, scene_name, imgs_idx=None, minimal_vis=0, mask_dilated_size=11): 35 | json_path = glob(f'{args.dataset_dir}/{scene_name}/*{bsdf_name}/transforms_train.json')[0] 36 | with open(json_path, 'r', encoding='utf-8') as f: 37 | data = json.load(f) 38 | fov_x = 37.8492 39 | fov_y = 28.8415 40 | width, height = 1600, 1200 41 | # transform to Colmap format 42 | fx, fy = gen_camera_intrinsic(width, height, fov_x, fov_y) 43 | 44 | # use float64 to avoid loss of precision 45 | intrinsic = np.diag([fx, fy, 1.0, 1.0]).astype(np.float64) 46 | # The origin is in the center and not in the upper left corner of the image 47 | intrinsic[0, 2] = width / 2.0 48 | intrinsic[1, 2] = height / 2.0 49 | flip_mat = np.array([ 50 | [-1, 0, 0, 0], 51 | [0, 1, 0, 0], 52 | [0, 0, -1, 0], 53 | [0, 0, 0, 1] 54 | ]) 55 | bottom = np.array([0, 0, 0, 1.]).reshape([1, 4]) 56 | scale_mat = np.diag([1.0, 1.0, 1.0, 1.0]) 57 | 58 | mask_lis = sorted(glob(f'{args.dataset_dir}/{scene_name}/*{bsdf_name}/train/mask/*.png')) 59 | n_images = len(mask_lis) 60 | inside_mask = np.zeros(len(points)) 61 | 62 | if imgs_idx is None: 63 | imgs_idx = [i for i in range(n_images)] 64 | 65 | for i, frame in enumerate(data['frames']): 66 | cam_pose_ = np.matmul(frame['transform_matrix'], flip_mat) 67 | cam_pose = np.array(cam_pose_) 68 | R, T = gen_w2c(cam_pose) 69 | w2c = np.concatenate([np.concatenate([R, T], 1), bottom], 0) 70 | world_mat = intrinsic @ w2c 71 | 72 | P = world_mat 73 | P = P @ scale_mat 74 | P = P[:3, :4] 75 | pts_image = np.matmul(P[None, :3, :3], points[:, :, None]).squeeze() + P[None, :3, 3] 76 | pts_image = pts_image / pts_image[:, 2:] 77 | pts_image = np.round(pts_image).astype(np.int32) + 1 78 | 79 | mask_image = cv.imread(mask_lis[i]) 80 | kernel_size = mask_dilated_size 81 | kernel = cv.getStructuringElement(cv.MORPH_ELLIPSE, (kernel_size, kernel_size)) 82 | mask_image = cv.dilate(mask_image, kernel, iterations=1) 83 | mask_image = (mask_image[:, :, 0] > 128) 84 | 85 | mask_image = np.concatenate([np.ones([1, 1600]), mask_image, np.ones([1, 1600])], axis=0) 86 | mask_image = np.concatenate([np.ones([1202, 1]), mask_image, np.ones([1202, 1])], axis=1) 87 | 88 | in_mask = (pts_image[:, 0] >= 0) * (pts_image[:, 0] <= 1600) * (pts_image[:, 1] >= 0) * ( 89 | pts_image[:, 1] <= 1200) > 0 90 | curr_mask = mask_image[(pts_image[:, 1].clip(0, 1201), pts_image[:, 0].clip(0, 1601))] 91 | 92 | curr_mask = curr_mask.astype(np.float32) * in_mask 93 | 94 | inside_mask += curr_mask 95 | 96 | if i > len(imgs_idx): 97 | break 98 | 99 | return inside_mask > minimal_vis 100 | 101 | 102 | def clean_points_by_visualhull(points, bsdf_name, scene_name, imgs_idx=None, minimal_vis=0, mask_dilated_size=11): 103 | json_path = glob(f'{args.dataset_dir}/{scene_name}/*{bsdf_name}/transforms_train.json')[0] 104 | with open(json_path, 'r', encoding='utf-8') as f: 105 | data = json.load(f) 106 | fov_x = 37.8492 107 | fov_y = 28.8415 108 | width, height = 1600, 1200 109 | 110 | # transform to Colmap format 111 | fx, fy = gen_camera_intrinsic(width, height, fov_x, fov_y) 112 | 113 | # use float64 to avoid loss of precision 114 | intrinsic = np.diag([fx, fy, 1.0, 1.0]).astype(np.float64) 115 | # The origin is in the center and not in the upper left corner of the image 116 | intrinsic[0, 2] = width / 2.0 117 | intrinsic[1, 2] = height / 2.0 118 | flip_mat = np.array([ 119 | [-1, 0, 0, 0], 120 | [0, 1, 0, 0], 121 | [0, 0, -1, 0], 122 | [0, 0, 0, 1] 123 | ]) 124 | bottom = np.array([0, 0, 0, 1.]).reshape([1, 4]) 125 | scale_mat = np.diag([1.0, 1.0, 1.0, 1.0]) 126 | mask_lis = sorted(glob(f'{args.dataset_dir}/{scene_name}/*{bsdf_name}/train/mask/*.png')) 127 | n_images = len(mask_lis) 128 | outside_mask = np.zeros(len(points)) 129 | if imgs_idx is None: 130 | imgs_idx = [i for i in range(n_images)] 131 | for i in imgs_idx: 132 | cam_pose_ = np.matmul(data['frames'][i]['transform_matrix'], flip_mat) 133 | cam_pose = np.array(cam_pose_) 134 | R, T = gen_w2c(cam_pose) 135 | w2c = np.concatenate([np.concatenate([R, T], 1), bottom], 0) 136 | world_mat = intrinsic @ w2c 137 | 138 | P = world_mat 139 | P = P @ scale_mat 140 | P = P[:3, :4] 141 | 142 | pts_image = np.matmul(P[None, :3, :3], points[:, :, None]).squeeze() + P[None, :3, 3] 143 | pts_image = pts_image / pts_image[:, 2:] 144 | pts_image = np.round(pts_image).astype(np.int32) + 1 145 | 146 | mask_image = cv.imread(mask_lis[i]) 147 | kernel_size = mask_dilated_size # default 101 148 | kernel = cv.getStructuringElement(cv.MORPH_ELLIPSE, (kernel_size, kernel_size)) 149 | mask_image = cv.dilate(mask_image, kernel, iterations=1) 150 | mask_image = (mask_image[:, :, 0] < 128) # * outside the mask 151 | 152 | mask_image = np.concatenate([np.ones([1, 1600]), mask_image, np.ones([1, 1600])], axis=0) 153 | mask_image = np.concatenate([np.ones([1202, 1]), mask_image, np.ones([1202, 1])], axis=1) 154 | 155 | border = 50 156 | in_mask = (pts_image[:, 0] >= (0 + border)) * (pts_image[:, 0] <= (1600 - border)) * ( 157 | pts_image[:, 1] >= (0 + border)) * ( 158 | pts_image[:, 1] <= (1200 - border)) > 0 159 | curr_mask = mask_image[(pts_image[:, 1].clip(0, 1201), pts_image[:, 0].clip(0, 1601))] 160 | 161 | curr_mask = curr_mask.astype(np.float32) * in_mask 162 | 163 | outside_mask += curr_mask 164 | 165 | return outside_mask < 5, scale_mat 166 | 167 | def find_closest_point(p1, d1, p2, d2): 168 | # Calculate the direction vectors of the lines 169 | d1_norm = d1 / np.linalg.norm(d1) 170 | d2_norm = d2 / np.linalg.norm(d2) 171 | 172 | # Create the coefficient matrix A and the constant vector b 173 | A = np.vstack((d1_norm, -d2_norm)).T 174 | b = p2 - p1 175 | 176 | # Solve the linear system to find the parameters t1 and t2 177 | t1, t2 = np.linalg.lstsq(A, b, rcond=None)[0] 178 | 179 | # Calculate the closest point on each line 180 | closest_point1 = p1 + d1_norm * t1 181 | closest_point2 = p2 + d2_norm * t2 182 | 183 | # Calculate the average of the two closest points 184 | closest_point = 0.5 * (closest_point1 + closest_point2) 185 | 186 | return closest_point 187 | 188 | def clean_mesh_faces_by_mask(mesh_file, new_mesh_file, bsdf_name, scene_name, imgs_idx, cut_y=-1.0, minimal_vis=0, mask_dilated_size=11): 189 | old_mesh = trimesh.load(mesh_file) 190 | old_vertices = old_mesh.vertices[:] 191 | old_faces = old_mesh.faces[:] 192 | mask = clean_points_by_mask(old_vertices, bsdf_name, scene_name, imgs_idx, minimal_vis, mask_dilated_size) 193 | y_mask = old_vertices[:, 1] >= cut_y 194 | mask = mask & y_mask 195 | indexes = np.ones(len(old_vertices)) * -1 196 | indexes = indexes.astype(np.int64) 197 | indexes[np.where(mask)] = np.arange(len(np.where(mask)[0])) 198 | 199 | faces_mask = mask[old_faces[:, 0]] & mask[old_faces[:, 1]] & mask[old_faces[:, 2]] 200 | new_faces = old_faces[np.where(faces_mask)] 201 | new_faces[:, 0] = indexes[new_faces[:, 0]] 202 | new_faces[:, 1] = indexes[new_faces[:, 1]] 203 | new_faces[:, 2] = indexes[new_faces[:, 2]] 204 | new_vertices = old_vertices[np.where(mask)] 205 | 206 | new_mesh = trimesh.Trimesh(new_vertices, new_faces) 207 | new_mesh.export(new_mesh_file) 208 | 209 | 210 | def clean_mesh_faces_by_visualhull(mesh_file, new_mesh_file, bsdf_name, scene_name, imgs_idx, minimal_vis=0, mask_dilated_size=11): 211 | old_mesh = trimesh.load(mesh_file) 212 | os.remove(mesh_file) 213 | old_vertices = old_mesh.vertices[:] 214 | old_faces = old_mesh.faces[:] 215 | mask, scale = clean_points_by_visualhull(old_vertices, bsdf_name, scene_name, imgs_idx, minimal_vis, mask_dilated_size) 216 | indexes = np.ones(len(old_vertices)) * -1 217 | indexes = indexes.astype(np.int64) 218 | indexes[np.where(mask)] = np.arange(len(np.where(mask)[0])) 219 | 220 | faces_mask = mask[old_faces[:, 0]] & mask[old_faces[:, 1]] & mask[old_faces[:, 2]] 221 | new_faces = old_faces[np.where(faces_mask)] 222 | new_faces[:, 0] = indexes[new_faces[:, 0]] 223 | new_faces[:, 1] = indexes[new_faces[:, 1]] 224 | new_faces[:, 2] = indexes[new_faces[:, 2]] 225 | new_vertices = old_vertices[np.where(mask)] 226 | 227 | new_mesh = trimesh.Trimesh(new_vertices, new_faces) 228 | new_mesh.vertices *= scale[0, 0] 229 | new_mesh.vertices += scale[:3, 3] 230 | # ! if colmap trim=7, comment these 231 | # meshes = new_mesh.split(only_watertight=False) 232 | # new_mesh = meshes[np.argmax([len(mesh.faces) for mesh in meshes])] 233 | 234 | new_mesh.export(new_mesh_file) 235 | 236 | def load_object(scene_dict, file_name): 237 | object_dir = { 238 | 'type': 'ply', 239 | 'id': 'Material_0001', 240 | 'filename': file_name, 241 | 'to_world': T([[1, 0, 0, 0], 242 | [0, 1, 0, 0], 243 | [0, 0, 1, 0], 244 | [0, 0, 0, 1]]), 245 | } 246 | scene_dict['shape_0'] = object_dir 247 | return scene_dict 248 | 249 | def load_integrator(scene_dict): 250 | integrator_dir = { 251 | 'type': 'path', 252 | 'max_depth': 65 253 | } 254 | scene_dict['integrator'] = integrator_dir 255 | return scene_dict 256 | 257 | 258 | if __name__ == "__main__": 259 | parser = argparse.ArgumentParser() 260 | parser.add_argument('--dataset_dir', type=str, default='../datasets/openmaterial') 261 | parser.add_argument('--groundtruth_dir', type=str, default='../datasets/groundtruth') 262 | parser.add_argument('--method', type=str) 263 | parser.add_argument('--directory', type=str) 264 | parser.add_argument('--object_name', type=str) 265 | args = parser.parse_args() 266 | mask_kernel_size = 11 267 | directory=args.directory 268 | dir_list = os.listdir(f'{directory}/meshes') 269 | object_name = args.object_name 270 | 271 | scene_list = os.listdir(f'{directory}/meshes/{object_name}') 272 | print(scene_list) 273 | for scene_name in scene_list: 274 | scene_name = scene_name.split(".")[0] 275 | base_path = f"{directory}/meshes/{object_name}" 276 | print("processing:", scene_name) 277 | 278 | old_mesh_file = glob(os.path.join(base_path, f"{scene_name}*"))[0] 279 | clean_mesh_file = os.path.join(base_path, f"clean_{scene_name}.ply") 280 | os.makedirs("{}/CleanedMesh/{}".format(directory, object_name), exist_ok=True) 281 | visualhull_mesh_file = f'{directory}/CleanedMesh/{object_name}/{scene_name}.ply' 282 | scene_path = glob(os.path.join(f'{args.groundtruth_dir}', object_name, '*.ply'))[0] 283 | scene_path = os.path.abspath(scene_path) 284 | scene_dict= {'type': 'scene'} 285 | scene_dict = load_integrator(scene_dict) 286 | scene_dict = load_object(scene_dict, scene_path) 287 | scene = mi.load_dict(scene_dict) 288 | bbox = scene.bbox() 289 | cut_y = bbox.min.y 290 | 291 | clean_mesh_faces_by_mask(old_mesh_file, clean_mesh_file, scene_name, object_name, None, cut_y=cut_y, minimal_vis=2, 292 | mask_dilated_size=mask_kernel_size) 293 | 294 | clean_mesh_faces_by_visualhull(clean_mesh_file, visualhull_mesh_file, scene_name, object_name, None, minimal_vis=2, 295 | mask_dilated_size=mask_kernel_size + 20) 296 | 297 | print("finish processing ", scene_name) -------------------------------------------------------------------------------- /instant-nsr-pl/models/geometry.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | 6 | from pytorch_lightning.utilities.rank_zero import rank_zero_info 7 | 8 | import models 9 | from models.base import BaseModel 10 | from models.utils import scale_anything, get_activation, cleanup, chunk_batch 11 | from models.network_utils import get_encoding, get_mlp, get_encoding_with_network 12 | from utils.misc import get_rank 13 | from systems.utils import update_module_step 14 | from nerfacc import ContractionType 15 | 16 | 17 | def contract_to_unisphere(x, radius, contraction_type): 18 | if contraction_type == ContractionType.AABB: 19 | x = scale_anything(x, (-radius, radius), (0, 1)) 20 | elif contraction_type == ContractionType.UN_BOUNDED_SPHERE: 21 | x = scale_anything(x, (-radius, radius), (0, 1)) 22 | x = x * 2 - 1 # aabb is at [-1, 1] 23 | mag = x.norm(dim=-1, keepdim=True) 24 | mask = mag.squeeze(-1) > 1 25 | x[mask] = (2 - 1 / mag[mask]) * (x[mask] / mag[mask]) 26 | x = x / 4 + 0.5 # [-inf, inf] is at [0, 1] 27 | else: 28 | raise NotImplementedError 29 | return x 30 | 31 | 32 | class MarchingCubeHelper(nn.Module): 33 | def __init__(self, resolution, use_torch=True): 34 | super().__init__() 35 | self.resolution = resolution 36 | self.use_torch = use_torch 37 | self.points_range = (0, 1) 38 | if self.use_torch: 39 | import torchmcubes 40 | self.mc_func = torchmcubes.marching_cubes 41 | else: 42 | import mcubes 43 | self.mc_func = mcubes.marching_cubes 44 | self.verts = None 45 | 46 | def grid_vertices(self): 47 | if self.verts is None: 48 | x, y, z = torch.linspace(*self.points_range, self.resolution), torch.linspace(*self.points_range, self.resolution), torch.linspace(*self.points_range, self.resolution) 49 | x, y, z = torch.meshgrid(x, y, z, indexing='ij') 50 | verts = torch.cat([x.reshape(-1, 1), y.reshape(-1, 1), z.reshape(-1, 1)], dim=-1).reshape(-1, 3) 51 | self.verts = verts 52 | return self.verts 53 | 54 | def forward(self, level, threshold=0.): 55 | level = level.float().view(self.resolution, self.resolution, self.resolution) 56 | if self.use_torch: 57 | verts, faces = self.mc_func(level.to(get_rank()), threshold) 58 | verts, faces = verts.cpu(), faces.cpu().long() 59 | else: 60 | verts, faces = self.mc_func(-level.numpy(), threshold) # transform to numpy 61 | verts, faces = torch.from_numpy(verts.astype(np.float32)), torch.from_numpy(faces.astype(np.int64)) # transform back to pytorch 62 | verts = verts / (self.resolution - 1.) 63 | return { 64 | 'v_pos': verts, 65 | 't_pos_idx': faces 66 | } 67 | 68 | 69 | class BaseImplicitGeometry(BaseModel): 70 | def __init__(self, config): 71 | super().__init__(config) 72 | if self.config.isosurface is not None: 73 | assert self.config.isosurface.method in ['mc', 'mc-torch'] 74 | if self.config.isosurface.method == 'mc-torch': 75 | raise NotImplementedError("Please do not use mc-torch. It currently has some scaling issues I haven't fixed yet.") 76 | self.helper = MarchingCubeHelper(self.config.isosurface.resolution, use_torch=self.config.isosurface.method=='mc-torch') 77 | self.radius = self.config.radius 78 | self.contraction_type = None # assigned in system 79 | 80 | def forward_level(self, points): 81 | raise NotImplementedError 82 | 83 | def isosurface_(self, vmin, vmax): 84 | def batch_func(x): 85 | x = torch.stack([ 86 | scale_anything(x[...,0], (0, 1), (vmin[0], vmax[0])), 87 | scale_anything(x[...,1], (0, 1), (vmin[1], vmax[1])), 88 | scale_anything(x[...,2], (0, 1), (vmin[2], vmax[2])), 89 | ], dim=-1).to(self.rank) 90 | rv = self.forward_level(x).cpu() 91 | cleanup() 92 | return rv 93 | 94 | level = chunk_batch(batch_func, self.config.isosurface.chunk, True, self.helper.grid_vertices()) 95 | mesh = self.helper(level, threshold=self.config.isosurface.threshold) 96 | mesh['v_pos'] = torch.stack([ 97 | scale_anything(mesh['v_pos'][...,0], (0, 1), (vmin[0], vmax[0])), 98 | scale_anything(mesh['v_pos'][...,1], (0, 1), (vmin[1], vmax[1])), 99 | scale_anything(mesh['v_pos'][...,2], (0, 1), (vmin[2], vmax[2])) 100 | ], dim=-1) 101 | return mesh 102 | 103 | @torch.no_grad() 104 | def isosurface(self): 105 | if self.config.isosurface is None: 106 | raise NotImplementedError 107 | mesh_coarse = self.isosurface_((-self.radius, -self.radius, -self.radius), (self.radius, self.radius, self.radius)) 108 | vmin, vmax = mesh_coarse['v_pos'].amin(dim=0), mesh_coarse['v_pos'].amax(dim=0) 109 | vmin_ = (vmin - (vmax - vmin) * 0.1).clamp(-self.radius, self.radius) 110 | vmax_ = (vmax + (vmax - vmin) * 0.1).clamp(-self.radius, self.radius) 111 | mesh_fine = self.isosurface_(vmin_, vmax_) 112 | return mesh_fine 113 | 114 | 115 | @models.register('volume-density') 116 | class VolumeDensity(BaseImplicitGeometry): 117 | def setup(self): 118 | self.n_input_dims = self.config.get('n_input_dims', 3) 119 | self.n_output_dims = self.config.feature_dim 120 | self.encoding_with_network = get_encoding_with_network(self.n_input_dims, self.n_output_dims, self.config.xyz_encoding_config, self.config.mlp_network_config) 121 | 122 | def forward(self, points): 123 | points = contract_to_unisphere(points, self.radius, self.contraction_type) 124 | out = self.encoding_with_network(points.view(-1, self.n_input_dims)).view(*points.shape[:-1], self.n_output_dims).float() 125 | density, feature = out[...,0], out 126 | if 'density_activation' in self.config: 127 | density = get_activation(self.config.density_activation)(density + float(self.config.density_bias)) 128 | if 'feature_activation' in self.config: 129 | feature = get_activation(self.config.feature_activation)(feature) 130 | return density, feature 131 | 132 | def forward_level(self, points): 133 | points = contract_to_unisphere(points, self.radius, self.contraction_type) 134 | density = self.encoding_with_network(points.reshape(-1, self.n_input_dims)).reshape(*points.shape[:-1], self.n_output_dims)[...,0] 135 | if 'density_activation' in self.config: 136 | density = get_activation(self.config.density_activation)(density + float(self.config.density_bias)) 137 | return -density 138 | 139 | def update_step(self, epoch, global_step): 140 | update_module_step(self.encoding_with_network, epoch, global_step) 141 | 142 | 143 | @models.register('volume-sdf') 144 | class VolumeSDF(BaseImplicitGeometry): 145 | def setup(self): 146 | self.n_output_dims = self.config.feature_dim 147 | encoding = get_encoding(3, self.config.xyz_encoding_config) 148 | network = get_mlp(encoding.n_output_dims, self.n_output_dims, self.config.mlp_network_config) 149 | self.encoding, self.network = encoding, network 150 | self.grad_type = self.config.grad_type 151 | self.finite_difference_eps = self.config.get('finite_difference_eps', 1e-3) 152 | # the actual value used in training 153 | # will update at certain steps if finite_difference_eps="progressive" 154 | self._finite_difference_eps = None 155 | if self.grad_type == 'finite_difference': 156 | rank_zero_info(f"Using finite difference to compute gradients with eps={self.finite_difference_eps}") 157 | 158 | def forward(self, points, with_grad=True, with_feature=True, with_laplace=False): 159 | with torch.inference_mode(torch.is_inference_mode_enabled() and not (with_grad and self.grad_type == 'analytic')): 160 | with torch.set_grad_enabled(self.training or (with_grad and self.grad_type == 'analytic')): 161 | if with_grad and self.grad_type == 'analytic': 162 | if not self.training: 163 | points = points.clone() # points may be in inference mode, get a copy to enable grad 164 | points.requires_grad_(True) 165 | 166 | points_ = points # points in the original scale 167 | points = contract_to_unisphere(points, self.radius, self.contraction_type) # points normalized to (0, 1) 168 | 169 | out = self.network(self.encoding(points.view(-1, 3))).view(*points.shape[:-1], self.n_output_dims).float() 170 | sdf, feature = out[...,0], out 171 | if 'sdf_activation' in self.config: 172 | sdf = get_activation(self.config.sdf_activation)(sdf + float(self.config.sdf_bias)) 173 | if 'feature_activation' in self.config: 174 | feature = get_activation(self.config.feature_activation)(feature) 175 | if with_grad: 176 | if self.grad_type == 'analytic': 177 | grad = torch.autograd.grad( 178 | sdf, points_, grad_outputs=torch.ones_like(sdf), 179 | create_graph=True, retain_graph=True, only_inputs=True 180 | )[0] 181 | elif self.grad_type == 'finite_difference': 182 | eps = self._finite_difference_eps 183 | offsets = torch.as_tensor( 184 | [ 185 | [eps, 0.0, 0.0], 186 | [-eps, 0.0, 0.0], 187 | [0.0, eps, 0.0], 188 | [0.0, -eps, 0.0], 189 | [0.0, 0.0, eps], 190 | [0.0, 0.0, -eps], 191 | ] 192 | ).to(points_) 193 | points_d_ = (points_[...,None,:] + offsets).clamp(-self.radius, self.radius) 194 | points_d = scale_anything(points_d_, (-self.radius, self.radius), (0, 1)) 195 | points_d_sdf = self.network(self.encoding(points_d.view(-1, 3)))[...,0].view(*points.shape[:-1], 6).float() 196 | grad = 0.5 * (points_d_sdf[..., 0::2] - points_d_sdf[..., 1::2]) / eps 197 | 198 | if with_laplace: 199 | laplace = (points_d_sdf[..., 0::2] + points_d_sdf[..., 1::2] - 2 * sdf[..., None]).sum(-1) / (eps ** 2) 200 | 201 | rv = [sdf] 202 | if with_grad: 203 | rv.append(grad) 204 | if with_feature: 205 | rv.append(feature) 206 | if with_laplace: 207 | assert self.config.grad_type == 'finite_difference', "Laplace computation is only supported with grad_type='finite_difference'" 208 | rv.append(laplace) 209 | rv = [v if self.training else v.detach() for v in rv] 210 | return rv[0] if len(rv) == 1 else rv 211 | 212 | def forward_level(self, points): 213 | points = contract_to_unisphere(points, self.radius, self.contraction_type) # points normalized to (0, 1) 214 | sdf = self.network(self.encoding(points.view(-1, 3))).view(*points.shape[:-1], self.n_output_dims)[...,0] 215 | if 'sdf_activation' in self.config: 216 | sdf = get_activation(self.config.sdf_activation)(sdf + float(self.config.sdf_bias)) 217 | return sdf 218 | 219 | def update_step(self, epoch, global_step): 220 | update_module_step(self.encoding, epoch, global_step) 221 | update_module_step(self.network, epoch, global_step) 222 | if self.grad_type == 'finite_difference': 223 | if isinstance(self.finite_difference_eps, float): 224 | self._finite_difference_eps = self.finite_difference_eps 225 | elif self.finite_difference_eps == 'progressive': 226 | hg_conf = self.config.xyz_encoding_config 227 | assert hg_conf.otype == "ProgressiveBandHashGrid", "finite_difference_eps='progressive' only works with ProgressiveBandHashGrid" 228 | current_level = min( 229 | hg_conf.start_level + max(global_step - hg_conf.start_step, 0) // hg_conf.update_steps, 230 | hg_conf.n_levels 231 | ) 232 | grid_res = hg_conf.base_resolution * hg_conf.per_level_scale**(current_level - 1) 233 | grid_size = 2 * self.config.radius / grid_res 234 | if grid_size != self._finite_difference_eps: 235 | rank_zero_info(f"Update finite_difference_eps to {grid_size}") 236 | self._finite_difference_eps = grid_size 237 | else: 238 | raise ValueError(f"Unknown finite_difference_eps={self.finite_difference_eps}") 239 | -------------------------------------------------------------------------------- /instant-nsr-pl/systems/neus.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from torch_efficient_distloss import flatten_eff_distloss 5 | 6 | import pytorch_lightning as pl 7 | from pytorch_lightning.utilities.rank_zero import rank_zero_info, rank_zero_debug 8 | 9 | import models 10 | from models.utils import cleanup 11 | from models.ray_utils import get_rays 12 | import systems 13 | from systems.base import BaseSystem 14 | from systems.criterions import PSNR, binary_cross_entropy 15 | 16 | 17 | @systems.register('neus-system') 18 | class NeuSSystem(BaseSystem): 19 | """ 20 | Two ways to print to console: 21 | 1. self.print: correctly handle progress bar 22 | 2. rank_zero_info: use the logging module 23 | """ 24 | def prepare(self): 25 | self.criterions = { 26 | 'psnr': PSNR() 27 | } 28 | self.train_num_samples = self.config.model.train_num_rays * (self.config.model.num_samples_per_ray + self.config.model.get('num_samples_per_ray_bg', 0)) 29 | self.train_num_rays = self.config.model.train_num_rays 30 | 31 | def forward(self, batch): 32 | return self.model(batch['rays']) 33 | 34 | def preprocess_data(self, batch, stage): 35 | if 'index' in batch: # validation / testing 36 | index = batch['index'] 37 | else: 38 | if self.config.model.batch_image_sampling: 39 | index = torch.randint(0, len(self.dataset.all_images), size=(self.train_num_rays,), device=self.dataset.all_images.device) 40 | else: 41 | index = torch.randint(0, len(self.dataset.all_images), size=(1,), device=self.dataset.all_images.device) 42 | if stage in ['train']: 43 | c2w = self.dataset.all_c2w[index] 44 | x = torch.randint( 45 | 0, self.dataset.w, size=(self.train_num_rays,), device=self.dataset.all_images.device 46 | ) 47 | y = torch.randint( 48 | 0, self.dataset.h, size=(self.train_num_rays,), device=self.dataset.all_images.device 49 | ) 50 | if self.dataset.directions.ndim == 3: # (H, W, 3) 51 | directions = self.dataset.directions[y, x] 52 | elif self.dataset.directions.ndim == 4: # (N, H, W, 3) 53 | directions = self.dataset.directions[index, y, x] 54 | rays_o, rays_d = get_rays(directions, c2w) 55 | rgb = self.dataset.all_images[index, y, x].view(-1, self.dataset.all_images.shape[-1]).to(self.rank) 56 | fg_mask = self.dataset.all_fg_masks[index, y, x].view(-1).to(self.rank) 57 | else: 58 | c2w = self.dataset.all_c2w[index][0] 59 | if self.dataset.directions.ndim == 3: # (H, W, 3) 60 | directions = self.dataset.directions 61 | elif self.dataset.directions.ndim == 4: # (N, H, W, 3) 62 | directions = self.dataset.directions[index][0] 63 | rays_o, rays_d = get_rays(directions, c2w) 64 | rgb = self.dataset.all_images[index].view(-1, self.dataset.all_images.shape[-1]).to(self.rank) 65 | fg_mask = self.dataset.all_fg_masks[index].view(-1).to(self.rank) 66 | 67 | rays = torch.cat([rays_o, F.normalize(rays_d, p=2, dim=-1)], dim=-1) 68 | 69 | if stage in ['train']: 70 | if self.config.model.background_color == 'white': 71 | self.model.background_color = torch.ones((3,), dtype=torch.float32, device=self.rank) 72 | elif self.config.model.background_color == 'random': 73 | self.model.background_color = torch.rand((3,), dtype=torch.float32, device=self.rank) 74 | else: 75 | raise NotImplementedError 76 | else: 77 | self.model.background_color = torch.ones((3,), dtype=torch.float32, device=self.rank) 78 | 79 | if self.dataset.apply_mask: 80 | rgb = rgb * fg_mask[...,None] + self.model.background_color * (1 - fg_mask[...,None]) 81 | 82 | batch.update({ 83 | 'rays': rays, 84 | 'rgb': rgb, 85 | 'fg_mask': fg_mask 86 | }) 87 | 88 | def training_step(self, batch, batch_idx): 89 | out = self(batch) 90 | 91 | loss = 0. 92 | 93 | # update train_num_rays 94 | if self.config.model.dynamic_ray_sampling: 95 | train_num_rays = int(self.train_num_rays * (self.train_num_samples / out['num_samples_full'].sum().item())) 96 | self.train_num_rays = min(int(self.train_num_rays * 0.9 + train_num_rays * 0.1), self.config.model.max_train_num_rays) 97 | 98 | loss_rgb_mse = F.mse_loss(out['comp_rgb_full'][out['rays_valid_full'][...,0]], batch['rgb'][out['rays_valid_full'][...,0]]) 99 | self.log('train/loss_rgb_mse', loss_rgb_mse) 100 | loss += loss_rgb_mse * self.C(self.config.system.loss.lambda_rgb_mse) 101 | 102 | loss_rgb_l1 = F.l1_loss(out['comp_rgb_full'][out['rays_valid_full'][...,0]], batch['rgb'][out['rays_valid_full'][...,0]]) 103 | self.log('train/loss_rgb', loss_rgb_l1) 104 | loss += loss_rgb_l1 * self.C(self.config.system.loss.lambda_rgb_l1) 105 | 106 | loss_eikonal = ((torch.linalg.norm(out['sdf_grad_samples'], ord=2, dim=-1) - 1.)**2).mean() 107 | self.log('train/loss_eikonal', loss_eikonal) 108 | loss += loss_eikonal * self.C(self.config.system.loss.lambda_eikonal) 109 | 110 | opacity = torch.clamp(out['opacity'].squeeze(-1), 1.e-3, 1.-1.e-3) 111 | loss_mask = binary_cross_entropy(opacity, batch['fg_mask'].float()) 112 | self.log('train/loss_mask', loss_mask) 113 | loss += loss_mask * (self.C(self.config.system.loss.lambda_mask) if self.dataset.has_mask else 0.0) 114 | 115 | loss_opaque = binary_cross_entropy(opacity, opacity) 116 | self.log('train/loss_opaque', loss_opaque) 117 | loss += loss_opaque * self.C(self.config.system.loss.lambda_opaque) 118 | 119 | loss_sparsity = torch.exp(-self.config.system.loss.sparsity_scale * out['sdf_samples'].abs()).mean() 120 | self.log('train/loss_sparsity', loss_sparsity) 121 | loss += loss_sparsity * self.C(self.config.system.loss.lambda_sparsity) 122 | 123 | if self.C(self.config.system.loss.lambda_curvature) > 0: 124 | assert 'sdf_laplace_samples' in out, "Need geometry.grad_type='finite_difference' to get SDF Laplace samples" 125 | loss_curvature = out['sdf_laplace_samples'].abs().mean() 126 | self.log('train/loss_curvature', loss_curvature) 127 | loss += loss_curvature * self.C(self.config.system.loss.lambda_curvature) 128 | 129 | # distortion loss proposed in MipNeRF360 130 | # an efficient implementation from https://github.com/sunset1995/torch_efficient_distloss 131 | if self.C(self.config.system.loss.lambda_distortion) > 0: 132 | loss_distortion = flatten_eff_distloss(out['weights'], out['points'], out['intervals'], out['ray_indices']) 133 | self.log('train/loss_distortion', loss_distortion) 134 | loss += loss_distortion * self.C(self.config.system.loss.lambda_distortion) 135 | 136 | if self.config.model.learned_background and self.C(self.config.system.loss.lambda_distortion_bg) > 0: 137 | loss_distortion_bg = flatten_eff_distloss(out['weights_bg'], out['points_bg'], out['intervals_bg'], out['ray_indices_bg']) 138 | self.log('train/loss_distortion_bg', loss_distortion_bg) 139 | loss += loss_distortion_bg * self.C(self.config.system.loss.lambda_distortion_bg) 140 | 141 | losses_model_reg = self.model.regularizations(out) 142 | for name, value in losses_model_reg.items(): 143 | self.log(f'train/loss_{name}', value) 144 | loss_ = value * self.C(self.config.system.loss[f"lambda_{name}"]) 145 | loss += loss_ 146 | 147 | self.log('train/inv_s', out['inv_s'], prog_bar=True) 148 | 149 | for name, value in self.config.system.loss.items(): 150 | if name.startswith('lambda'): 151 | self.log(f'train_params/{name}', self.C(value)) 152 | 153 | self.log('train/num_rays', float(self.train_num_rays), prog_bar=True) 154 | 155 | return { 156 | 'loss': loss 157 | } 158 | 159 | """ 160 | # aggregate outputs from different devices (DP) 161 | def training_step_end(self, out): 162 | pass 163 | """ 164 | 165 | """ 166 | # aggregate outputs from different iterations 167 | def training_epoch_end(self, out): 168 | pass 169 | """ 170 | 171 | def validation_step(self, batch, batch_idx): 172 | out = self(batch) 173 | psnr = self.criterions['psnr'](out['comp_rgb_full'].to(batch['rgb']), batch['rgb']) 174 | W, H = self.dataset.img_wh 175 | self.save_image_grid(f"it{self.global_step}-{batch['index'][0].item()}.png", [ 176 | {'type': 'rgb', 'img': batch['rgb'].view(H, W, 3), 'kwargs': {'data_format': 'HWC'}}, 177 | {'type': 'rgb', 'img': out['comp_rgb_full'].view(H, W, 3), 'kwargs': {'data_format': 'HWC'}} 178 | ] + ([ 179 | {'type': 'rgb', 'img': out['comp_rgb_bg'].view(H, W, 3), 'kwargs': {'data_format': 'HWC'}}, 180 | {'type': 'rgb', 'img': out['comp_rgb'].view(H, W, 3), 'kwargs': {'data_format': 'HWC'}}, 181 | ] if self.config.model.learned_background else []) + [ 182 | {'type': 'grayscale', 'img': out['depth'].view(H, W), 'kwargs': {}}, 183 | {'type': 'rgb', 'img': out['comp_normal'].view(H, W, 3), 'kwargs': {'data_format': 'HWC', 'data_range': (-1, 1)}} 184 | ]) 185 | return { 186 | 'psnr': psnr, 187 | 'index': batch['index'] 188 | } 189 | 190 | 191 | """ 192 | # aggregate outputs from different devices when using DP 193 | def validation_step_end(self, out): 194 | pass 195 | """ 196 | 197 | def validation_epoch_end(self, out): 198 | out = self.all_gather(out) 199 | if self.trainer.is_global_zero: 200 | out_set = {} 201 | for step_out in out: 202 | # DP 203 | if step_out['index'].ndim == 1: 204 | out_set[step_out['index'].item()] = {'psnr': step_out['psnr']} 205 | # DDP 206 | else: 207 | for oi, index in enumerate(step_out['index']): 208 | out_set[index[0].item()] = {'psnr': step_out['psnr'][oi]} 209 | psnr = torch.mean(torch.stack([o['psnr'] for o in out_set.values()])) 210 | self.log('val/psnr', psnr, prog_bar=True, rank_zero_only=True) 211 | 212 | def test_step(self, batch, batch_idx): 213 | # out = self(batch) 214 | # psnr = self.criterions['psnr'](out['comp_rgb_full'].to(batch['rgb']), batch['rgb']) 215 | # W, H = self.dataset.img_wh 216 | # self.save_image_grid(f"it{self.global_step}-test/{batch['index'][0].item()}.png", [ 217 | # {'type': 'rgb', 'img': batch['rgb'].view(H, W, 3), 'kwargs': {'data_format': 'HWC'}}, 218 | # {'type': 'rgb', 'img': out['comp_rgb_full'].view(H, W, 3), 'kwargs': {'data_format': 'HWC'}} 219 | # ] + ([ 220 | # {'type': 'rgb', 'img': out['comp_rgb_bg'].view(H, W, 3), 'kwargs': {'data_format': 'HWC'}}, 221 | # {'type': 'rgb', 'img': out['comp_rgb'].view(H, W, 3), 'kwargs': {'data_format': 'HWC'}}, 222 | # ] if self.config.model.learned_background else []) + [ 223 | # {'type': 'grayscale', 'img': out['depth'].view(H, W), 'kwargs': {}}, 224 | # {'type': 'rgb', 'img': out['comp_normal'].view(H, W, 3), 'kwargs': {'data_format': 'HWC', 'data_range': (-1, 1)}} 225 | # ]) 226 | return { 227 | # 'psnr': psnr, 228 | 'index': batch['index'] 229 | } 230 | 231 | def test_epoch_end(self, out): 232 | """ 233 | Synchronize devices. 234 | Generate image sequence using test outputs. 235 | """ 236 | # out = self.all_gather(out) 237 | if self.trainer.is_global_zero: 238 | # out_set = {} 239 | # for step_out in out: 240 | # # DP 241 | # if step_out['index'].ndim == 1: 242 | # out_set[step_out['index'].item()] = {'psnr': step_out['psnr']} 243 | # # DDP 244 | # else: 245 | # for oi, index in enumerate(step_out['index']): 246 | # out_set[index[0].item()] = {'psnr': step_out['psnr'][oi]} 247 | # psnr = torch.mean(torch.stack([o['psnr'] for o in out_set.values()])) 248 | # self.log('test/psnr', psnr, prog_bar=True, rank_zero_only=True) 249 | 250 | self.export() 251 | 252 | def export(self): 253 | mesh = self.model.export(self.config.export) 254 | self.save_mesh( 255 | f"{self.config.dataset.bsdf_name}.obj", 256 | **mesh 257 | ) 258 | -------------------------------------------------------------------------------- /sum_metrics-ablation.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | import pandas as pd 4 | import re 5 | import argparse 6 | import glob 7 | from tqdm import tqdm 8 | 9 | bsdf_names = [ 10 | 'conductor', 11 | 'dielectric', 12 | 'plastic', 13 | 'roughconductor', 14 | 'roughdielectric', 15 | 'roughplastic', 16 | 'diffuse' 17 | ] 18 | 19 | hdr_names = [ 20 | 'cobblestone_street_night_4k', 21 | 'leadenhall_market_4k', 22 | 'symmetrical_garden_4k' 23 | ] 24 | 25 | object_names = { 26 | 'b14a251fe8ad4a10bbc75f7dd3f6cebb': "vase", 27 | 'fc4f34dae22c4dae95c19b1654c3cb7e': "snail", 28 | '01098ad7973647a9b558f41d2ebc5193': "Boat", 29 | 'df894b66e2e54b558a43497becb94ff0': "Bike", 30 | '5c230ea126b943b8bc1da3f5865d5cd2': "Statue", 31 | } 32 | 33 | parser = argparse.ArgumentParser() 34 | parser.add_argument('--output_dir', type=str, default='output-ablation') 35 | parser.add_argument('--method', type=str,) 36 | parser.add_argument('--eval_mesh', action='store_true') 37 | args = parser.parse_args() 38 | 39 | def sum_instant_nsr_pl(): 40 | count_dir_bsdf, psnr_dir_bsdf, ssim_dir_bsdf, lpips_dir_bsdf = {}, {}, {}, {} 41 | count_dir_hdr, psnr_dir_hdr, ssim_dir_hdr, lpips_dir_hdr = {}, {}, {}, {} 42 | count_dir_obj, psnr_dir_obj, ssim_dir_obj, lpips_dir_obj = {}, {}, {}, {} 43 | folder_path = args.output_dir 44 | for name in bsdf_names: 45 | count_dir_bsdf[name] = 0 46 | psnr_dir_bsdf[name] = 0.0 47 | ssim_dir_bsdf[name] = 0.0 48 | lpips_dir_bsdf[name] = 0.0 49 | 50 | for name in hdr_names: 51 | count_dir_hdr[name] = 0 52 | psnr_dir_hdr[name] = 0.0 53 | ssim_dir_hdr[name] = 0.0 54 | lpips_dir_hdr[name] = 0.0 55 | 56 | for name in object_names.keys(): 57 | count_dir_obj[object_names[name]] = 0 58 | psnr_dir_obj[object_names[name]] = 0.0 59 | ssim_dir_obj[object_names[name]] = 0.0 60 | lpips_dir_obj[object_names[name]] = 0.0 61 | 62 | for file_name in tqdm(object_names.keys(), desc="sum instant-nsr-pl..."): 63 | scene_names_ = os.listdir(os.path.join(folder_path, file_name)) 64 | for scene_name in scene_names_: 65 | if scene_name.endswith("insr"): 66 | txt_path = os.path.join(folder_path, file_name, scene_name, 'instant-nsr-pl-wmask.txt') 67 | 68 | with open(txt_path, 'r') as f: 69 | txt_data = f.readline() 70 | psnr_ = re.search(r'PSNR=(\d+\.\d+)', txt_data).group(1) 71 | ssim_ = re.search(r'SSIM=(\d+\.\d+)', txt_data).group(1) 72 | lpips_ = re.search(r'lpips=(\d+\.\d+)', txt_data).group(1) 73 | bsdf_name = txt_data.split(':')[2] 74 | 75 | for name in bsdf_names: 76 | if bsdf_name.startswith(name): 77 | psnr_dir_bsdf[name] += float(psnr_) 78 | ssim_dir_bsdf[name] += float(ssim_) 79 | lpips_dir_bsdf[name] += float(lpips_) 80 | count_dir_bsdf[name] += 1 81 | 82 | for name in hdr_names: 83 | if scene_name.startswith(name): 84 | psnr_dir_hdr[name] += float(psnr_) 85 | ssim_dir_hdr[name] += float(ssim_) 86 | lpips_dir_hdr[name] += float(lpips_) 87 | count_dir_hdr[name] += 1 88 | 89 | for name in object_names.keys(): 90 | if file_name.startswith(name): 91 | psnr_dir_obj[object_names[name]] += float(psnr_) 92 | ssim_dir_obj[object_names[name]] += float(ssim_) 93 | lpips_dir_obj[object_names[name]] += float(lpips_) 94 | count_dir_obj[object_names[name]] += 1 95 | 96 | for name in bsdf_names: 97 | print(f"[+] {name} result: {count_dir_bsdf[name]}") 98 | if count_dir_bsdf[name] > 0: 99 | psnr_dir_bsdf[name] = psnr_dir_bsdf[name] / count_dir_bsdf[name] 100 | ssim_dir_bsdf[name] = ssim_dir_bsdf[name] / count_dir_bsdf[name] 101 | lpips_dir_bsdf[name] = lpips_dir_bsdf[name] / count_dir_bsdf[name] 102 | for name in hdr_names: 103 | print(f"[+] {name} result: {count_dir_hdr[name]}") 104 | if count_dir_hdr[name] > 0: 105 | psnr_dir_hdr[name] = psnr_dir_hdr[name] / count_dir_hdr[name] 106 | ssim_dir_hdr[name] = ssim_dir_hdr[name] / count_dir_hdr[name] 107 | lpips_dir_hdr[name] = lpips_dir_hdr[name] / count_dir_hdr[name] 108 | for name in object_names.keys(): 109 | print(f"[+] {name} result: {count_dir_obj[object_names[name]]}") 110 | if count_dir_obj[object_names[name]] > 0: 111 | psnr_dir_obj[object_names[name]] = psnr_dir_obj[object_names[name]] / count_dir_obj[object_names[name]] 112 | ssim_dir_obj[object_names[name]] = ssim_dir_obj[object_names[name]] / count_dir_obj[object_names[name]] 113 | lpips_dir_obj[object_names[name]] = lpips_dir_obj[object_names[name]] / count_dir_obj[object_names[name]] 114 | return psnr_dir_bsdf, ssim_dir_bsdf, lpips_dir_bsdf, psnr_dir_hdr, ssim_dir_hdr, lpips_dir_hdr, psnr_dir_obj, ssim_dir_obj, lpips_dir_obj 115 | 116 | 117 | def mesh_cds(method): 118 | cds_dir_bsdf, count_dir_bsdf = {}, {} 119 | cds_dir_hdr, count_dir_hdr = {}, {} 120 | cds_dir_obj, count_dir_obj = {}, {} 121 | folder_path = args.output_dir 122 | cds_dir_bsdf[method] = {} 123 | count_dir_bsdf[method] = {} 124 | cds_dir_hdr[method] = {} 125 | count_dir_hdr[method] = {} 126 | cds_dir_obj[method] = {} 127 | count_dir_obj[method] = {} 128 | 129 | for name in bsdf_names: 130 | count_dir_bsdf[method][name] = 0 131 | cds_dir_bsdf[method][name] = 0.0 132 | 133 | for name in hdr_names: 134 | count_dir_hdr[method][name] = 0 135 | cds_dir_hdr[method][name] = 0.0 136 | 137 | for name in object_names.keys(): 138 | count_dir_obj[method][object_names[name]] = 0 139 | cds_dir_obj[method][object_names[name]] = 0.0 140 | 141 | for file_name in tqdm(object_names.keys(), desc="sum chamfer distance..."): 142 | file_name = file_name.strip() 143 | txt_paths = glob.glob(os.path.join(folder_path, file_name, '*mesh-output.txt')) 144 | for txt_path in txt_paths: 145 | with open(txt_path, 'r') as f: 146 | for line in f: 147 | txt_data = line.strip() 148 | method_name = txt_data.split(':')[1] 149 | bsdf_name = txt_data.split(':')[2] 150 | cds = txt_data.split(':')[-1] 151 | if method_name == "instant-nsr-pl-wmask": 152 | method_name = "insr" 153 | if method_name == method: 154 | for name in bsdf_names: 155 | if bsdf_name.endswith(name): 156 | cds_dir_bsdf[method][name] += float(cds) 157 | count_dir_bsdf[method][name] += 1 158 | 159 | for name in hdr_names: 160 | if bsdf_name.startswith(name): 161 | cds_dir_hdr[method][name] += float(cds) 162 | count_dir_hdr[method][name] += 1 163 | 164 | for name in object_names.keys(): 165 | if file_name.startswith(name): 166 | cds_dir_obj[method][object_names[name]] += float(cds) 167 | count_dir_obj[method][object_names[name]] += 1 168 | 169 | for name in bsdf_names: 170 | if count_dir_bsdf[method][name] > 0: 171 | cds_dir_bsdf[method][name] = cds_dir_bsdf[method][name] / count_dir_bsdf[method][name] 172 | else: 173 | cds_dir_bsdf[method][name] = 0.0 174 | print(f"Warning: No material type: {name}") 175 | for name in hdr_names: 176 | if count_dir_hdr[method][name] > 0: 177 | cds_dir_hdr[method][name] = cds_dir_hdr[method][name] / count_dir_hdr[method][name] 178 | for name in object_names.keys(): 179 | if count_dir_obj[method][object_names[name]] > 0: 180 | cds_dir_obj[method][object_names[name]] = cds_dir_obj[method][object_names[name]] / count_dir_obj[method][object_names[name]] 181 | return cds_dir_bsdf, cds_dir_hdr, cds_dir_obj 182 | 183 | 184 | if __name__ == '__main__': 185 | 186 | result = {} 187 | cds = {} 188 | flag = True 189 | 190 | method_list = [args.method] 191 | if args.eval_mesh: 192 | method_cds = [args.method] 193 | else: 194 | method_cds = [] 195 | 196 | for method in method_list: 197 | if method == "insr": 198 | result['insr'] = sum_instant_nsr_pl() 199 | flag = True 200 | 201 | for method in method_cds: 202 | if method in method_cds: 203 | cds[method] = mesh_cds(method) 204 | 205 | if args.method == None or flag: 206 | print('\n') 207 | print("******************************************") 208 | print(" Novel View Synthesis ") 209 | print("******************************************") 210 | 211 | print("--------------Material Type--------------") 212 | rows = method_list 213 | columns = [name for name in bsdf_names] 214 | PSNR_df = pd.DataFrame(index=rows, columns=columns) 215 | for r in rows: 216 | for c in columns: 217 | PSNR_df.at[r, c] = result[r][0][c] 218 | print("PSNR:") 219 | print(PSNR_df) 220 | SSIM_df = pd.DataFrame(index=rows, columns=columns) 221 | for r in rows: 222 | for c in columns: 223 | SSIM_df.at[r, c] = result[r][1][c] 224 | print("SSIM:") 225 | print(SSIM_df) 226 | LPIPS_df = pd.DataFrame(index=rows, columns=columns) 227 | for r in rows: 228 | for c in columns: 229 | LPIPS_df.at[r, c] = result[r][2][c] 230 | print("LPIPS:") 231 | print(LPIPS_df) 232 | print('\n') 233 | 234 | print("--------------Lighting Type--------------") 235 | rows = method_list 236 | columns = [name for name in hdr_names] 237 | PSNR_df = pd.DataFrame(index=rows, columns=columns) 238 | for r in rows: 239 | for c in columns: 240 | PSNR_df.at[r, c] = result[r][3][c] 241 | print("PSNR:") 242 | print(PSNR_df) 243 | SSIM_df = pd.DataFrame(index=rows, columns=columns) 244 | for r in rows: 245 | for c in columns: 246 | SSIM_df.at[r, c] = result[r][4][c] 247 | print("SSIM:") 248 | print(SSIM_df) 249 | LPIPS_df = pd.DataFrame(index=rows, columns=columns) 250 | for r in rows: 251 | for c in columns: 252 | LPIPS_df.at[r, c] = result[r][5][c] 253 | print("LPIPS:") 254 | print(LPIPS_df) 255 | print('\n') 256 | 257 | print("--------------Object Name--------------") 258 | rows = method_list 259 | columns = [name for name in object_names.values()] 260 | PSNR_df = pd.DataFrame(index=rows, columns=columns) 261 | for r in rows: 262 | for c in columns: 263 | PSNR_df.at[r, c] = result[r][6][c] 264 | print("PSNR:") 265 | print(PSNR_df) 266 | SSIM_df = pd.DataFrame(index=rows, columns=columns) 267 | for r in rows: 268 | for c in columns: 269 | SSIM_df.at[r, c] = result[r][7][c] 270 | print("SSIM:") 271 | print(SSIM_df) 272 | LPIPS_df = pd.DataFrame(index=rows, columns=columns) 273 | for r in rows: 274 | for c in columns: 275 | LPIPS_df.at[r, c] = result[r][8][c] 276 | print("LPIPS:") 277 | print(LPIPS_df) 278 | print('\n') 279 | 280 | if args.eval_mesh or args.method == None: 281 | print("*******************************************") 282 | print(" 3D Reconstruction ") 283 | print("*******************************************") 284 | 285 | print("--------------Material Type--------------") 286 | rows = method_cds 287 | columns = [name for name in bsdf_names] 288 | CDs_df = pd.DataFrame(index=rows, columns=columns) 289 | for r in rows: 290 | for c in columns: 291 | CDs_df.at[r, c] = cds[r][0][r][c] 292 | print("Chamfer Distance:") 293 | print(CDs_df) 294 | print('\n') 295 | 296 | print("--------------Lighting Type--------------") 297 | rows = method_cds 298 | columns = [name for name in hdr_names] 299 | CDs_df = pd.DataFrame(index=rows, columns=columns) 300 | for r in rows: 301 | for c in columns: 302 | CDs_df.at[r, c] = cds[r][1][r][c] 303 | print("Chamfer Distance:") 304 | print(CDs_df) 305 | print('\n') 306 | 307 | print("--------------Object Name--------------") 308 | rows = method_cds 309 | columns = [name for name in object_names.values()] 310 | CDs_df = pd.DataFrame(index=rows, columns=columns) 311 | for r in rows: 312 | for c in columns: 313 | CDs_df.at[r, c] = cds[r][2][r][c] 314 | print("Chamfer Distance:") 315 | print(CDs_df) 316 | print('\n') 317 | -------------------------------------------------------------------------------- /instant-nsr-pl/datasets/colmap_utils.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) 2018, ETH Zurich and UNC Chapel Hill. 2 | # All rights reserved. 3 | # 4 | # Redistribution and use in source and binary forms, with or without 5 | # modification, are permitted provided that the following conditions are met: 6 | # 7 | # * Redistributions of source code must retain the above copyright 8 | # notice, this list of conditions and the following disclaimer. 9 | # 10 | # * Redistributions in binary form must reproduce the above copyright 11 | # notice, this list of conditions and the following disclaimer in the 12 | # documentation and/or other materials provided with the distribution. 13 | # 14 | # * Neither the name of ETH Zurich and UNC Chapel Hill nor the names of 15 | # its contributors may be used to endorse or promote products derived 16 | # from this software without specific prior written permission. 17 | # 18 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 19 | # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 20 | # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE 21 | # ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE 22 | # LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 23 | # CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF 24 | # SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 25 | # INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN 26 | # CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) 27 | # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 28 | # POSSIBILITY OF SUCH DAMAGE. 29 | # 30 | # Author: Johannes L. Schoenberger (jsch at inf.ethz.ch) 31 | 32 | import os 33 | import collections 34 | import numpy as np 35 | import struct 36 | 37 | 38 | CameraModel = collections.namedtuple( 39 | "CameraModel", ["model_id", "model_name", "num_params"]) 40 | Camera = collections.namedtuple( 41 | "Camera", ["id", "model", "width", "height", "params"]) 42 | BaseImage = collections.namedtuple( 43 | "Image", ["id", "qvec", "tvec", "camera_id", "name", "xys", "point3D_ids"]) 44 | Point3D = collections.namedtuple( 45 | "Point3D", ["id", "xyz", "rgb", "error", "image_ids", "point2D_idxs"]) 46 | 47 | class Image(BaseImage): 48 | def qvec2rotmat(self): 49 | return qvec2rotmat(self.qvec) 50 | 51 | 52 | CAMERA_MODELS = { 53 | CameraModel(model_id=0, model_name="SIMPLE_PINHOLE", num_params=3), 54 | CameraModel(model_id=1, model_name="PINHOLE", num_params=4), 55 | CameraModel(model_id=2, model_name="SIMPLE_RADIAL", num_params=4), 56 | CameraModel(model_id=3, model_name="RADIAL", num_params=5), 57 | CameraModel(model_id=4, model_name="OPENCV", num_params=8), 58 | CameraModel(model_id=5, model_name="OPENCV_FISHEYE", num_params=8), 59 | CameraModel(model_id=6, model_name="FULL_OPENCV", num_params=12), 60 | CameraModel(model_id=7, model_name="FOV", num_params=5), 61 | CameraModel(model_id=8, model_name="SIMPLE_RADIAL_FISHEYE", num_params=4), 62 | CameraModel(model_id=9, model_name="RADIAL_FISHEYE", num_params=5), 63 | CameraModel(model_id=10, model_name="THIN_PRISM_FISHEYE", num_params=12) 64 | } 65 | CAMERA_MODEL_IDS = dict([(camera_model.model_id, camera_model) \ 66 | for camera_model in CAMERA_MODELS]) 67 | 68 | 69 | def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"): 70 | """Read and unpack the next bytes from a binary file. 71 | :param fid: 72 | :param num_bytes: Sum of combination of {2, 4, 8}, e.g. 2, 6, 16, 30, etc. 73 | :param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}. 74 | :param endian_character: Any of {@, =, <, >, !} 75 | :return: Tuple of read and unpacked values. 76 | """ 77 | data = fid.read(num_bytes) 78 | return struct.unpack(endian_character + format_char_sequence, data) 79 | 80 | 81 | def read_cameras_text(path): 82 | """ 83 | see: src/base/reconstruction.cc 84 | void Reconstruction::WriteCamerasText(const std::string& path) 85 | void Reconstruction::ReadCamerasText(const std::string& path) 86 | """ 87 | cameras = {} 88 | with open(path, "r") as fid: 89 | while True: 90 | line = fid.readline() 91 | if not line: 92 | break 93 | line = line.strip() 94 | if len(line) > 0 and line[0] != "#": 95 | elems = line.split() 96 | camera_id = int(elems[0]) 97 | model = elems[1] 98 | width = int(elems[2]) 99 | height = int(elems[3]) 100 | params = np.array(tuple(map(float, elems[4:]))) 101 | cameras[camera_id] = Camera(id=camera_id, model=model, 102 | width=width, height=height, 103 | params=params) 104 | return cameras 105 | 106 | 107 | def read_cameras_binary(path_to_model_file): 108 | """ 109 | see: src/base/reconstruction.cc 110 | void Reconstruction::WriteCamerasBinary(const std::string& path) 111 | void Reconstruction::ReadCamerasBinary(const std::string& path) 112 | """ 113 | cameras = {} 114 | with open(path_to_model_file, "rb") as fid: 115 | num_cameras = read_next_bytes(fid, 8, "Q")[0] 116 | for camera_line_index in range(num_cameras): 117 | camera_properties = read_next_bytes( 118 | fid, num_bytes=24, format_char_sequence="iiQQ") 119 | camera_id = camera_properties[0] 120 | model_id = camera_properties[1] 121 | model_name = CAMERA_MODEL_IDS[camera_properties[1]].model_name 122 | width = camera_properties[2] 123 | height = camera_properties[3] 124 | num_params = CAMERA_MODEL_IDS[model_id].num_params 125 | params = read_next_bytes(fid, num_bytes=8*num_params, 126 | format_char_sequence="d"*num_params) 127 | cameras[camera_id] = Camera(id=camera_id, 128 | model=model_name, 129 | width=width, 130 | height=height, 131 | params=np.array(params)) 132 | assert len(cameras) == num_cameras 133 | return cameras 134 | 135 | 136 | def read_images_text(path): 137 | """ 138 | see: src/base/reconstruction.cc 139 | void Reconstruction::ReadImagesText(const std::string& path) 140 | void Reconstruction::WriteImagesText(const std::string& path) 141 | """ 142 | images = {} 143 | with open(path, "r") as fid: 144 | while True: 145 | line = fid.readline() 146 | if not line: 147 | break 148 | line = line.strip() 149 | if len(line) > 0 and line[0] != "#": 150 | elems = line.split() 151 | image_id = int(elems[0]) 152 | qvec = np.array(tuple(map(float, elems[1:5]))) 153 | tvec = np.array(tuple(map(float, elems[5:8]))) 154 | camera_id = int(elems[8]) 155 | image_name = elems[9] 156 | elems = fid.readline().split() 157 | xys = np.column_stack([tuple(map(float, elems[0::3])), 158 | tuple(map(float, elems[1::3]))]) 159 | point3D_ids = np.array(tuple(map(int, elems[2::3]))) 160 | images[image_id] = Image( 161 | id=image_id, qvec=qvec, tvec=tvec, 162 | camera_id=camera_id, name=image_name, 163 | xys=xys, point3D_ids=point3D_ids) 164 | return images 165 | 166 | 167 | def read_images_binary(path_to_model_file): 168 | """ 169 | see: src/base/reconstruction.cc 170 | void Reconstruction::ReadImagesBinary(const std::string& path) 171 | void Reconstruction::WriteImagesBinary(const std::string& path) 172 | """ 173 | images = {} 174 | with open(path_to_model_file, "rb") as fid: 175 | num_reg_images = read_next_bytes(fid, 8, "Q")[0] 176 | for image_index in range(num_reg_images): 177 | binary_image_properties = read_next_bytes( 178 | fid, num_bytes=64, format_char_sequence="idddddddi") 179 | image_id = binary_image_properties[0] 180 | qvec = np.array(binary_image_properties[1:5]) 181 | tvec = np.array(binary_image_properties[5:8]) 182 | camera_id = binary_image_properties[8] 183 | image_name = "" 184 | current_char = read_next_bytes(fid, 1, "c")[0] 185 | while current_char != b"\x00": # look for the ASCII 0 entry 186 | image_name += current_char.decode("utf-8") 187 | current_char = read_next_bytes(fid, 1, "c")[0] 188 | num_points2D = read_next_bytes(fid, num_bytes=8, 189 | format_char_sequence="Q")[0] 190 | x_y_id_s = read_next_bytes(fid, num_bytes=24*num_points2D, 191 | format_char_sequence="ddq"*num_points2D) 192 | xys = np.column_stack([tuple(map(float, x_y_id_s[0::3])), 193 | tuple(map(float, x_y_id_s[1::3]))]) 194 | point3D_ids = np.array(tuple(map(int, x_y_id_s[2::3]))) 195 | images[image_id] = Image( 196 | id=image_id, qvec=qvec, tvec=tvec, 197 | camera_id=camera_id, name=image_name, 198 | xys=xys, point3D_ids=point3D_ids) 199 | return images 200 | 201 | 202 | def read_points3D_text(path): 203 | """ 204 | see: src/base/reconstruction.cc 205 | void Reconstruction::ReadPoints3DText(const std::string& path) 206 | void Reconstruction::WritePoints3DText(const std::string& path) 207 | """ 208 | points3D = {} 209 | with open(path, "r") as fid: 210 | while True: 211 | line = fid.readline() 212 | if not line: 213 | break 214 | line = line.strip() 215 | if len(line) > 0 and line[0] != "#": 216 | elems = line.split() 217 | point3D_id = int(elems[0]) 218 | xyz = np.array(tuple(map(float, elems[1:4]))) 219 | rgb = np.array(tuple(map(int, elems[4:7]))) 220 | error = float(elems[7]) 221 | image_ids = np.array(tuple(map(int, elems[8::2]))) 222 | point2D_idxs = np.array(tuple(map(int, elems[9::2]))) 223 | points3D[point3D_id] = Point3D(id=point3D_id, xyz=xyz, rgb=rgb, 224 | error=error, image_ids=image_ids, 225 | point2D_idxs=point2D_idxs) 226 | return points3D 227 | 228 | 229 | def read_points3d_binary(path_to_model_file): 230 | """ 231 | see: src/base/reconstruction.cc 232 | void Reconstruction::ReadPoints3DBinary(const std::string& path) 233 | void Reconstruction::WritePoints3DBinary(const std::string& path) 234 | """ 235 | points3D = {} 236 | with open(path_to_model_file, "rb") as fid: 237 | num_points = read_next_bytes(fid, 8, "Q")[0] 238 | for point_line_index in range(num_points): 239 | binary_point_line_properties = read_next_bytes( 240 | fid, num_bytes=43, format_char_sequence="QdddBBBd") 241 | point3D_id = binary_point_line_properties[0] 242 | xyz = np.array(binary_point_line_properties[1:4]) 243 | rgb = np.array(binary_point_line_properties[4:7]) 244 | error = np.array(binary_point_line_properties[7]) 245 | track_length = read_next_bytes( 246 | fid, num_bytes=8, format_char_sequence="Q")[0] 247 | track_elems = read_next_bytes( 248 | fid, num_bytes=8*track_length, 249 | format_char_sequence="ii"*track_length) 250 | image_ids = np.array(tuple(map(int, track_elems[0::2]))) 251 | point2D_idxs = np.array(tuple(map(int, track_elems[1::2]))) 252 | points3D[point3D_id] = Point3D( 253 | id=point3D_id, xyz=xyz, rgb=rgb, 254 | error=error, image_ids=image_ids, 255 | point2D_idxs=point2D_idxs) 256 | return points3D 257 | 258 | 259 | def read_model(path, ext): 260 | if ext == ".txt": 261 | cameras = read_cameras_text(os.path.join(path, "cameras" + ext)) 262 | images = read_images_text(os.path.join(path, "images" + ext)) 263 | points3D = read_points3D_text(os.path.join(path, "points3D") + ext) 264 | else: 265 | cameras = read_cameras_binary(os.path.join(path, "cameras" + ext)) 266 | images = read_images_binary(os.path.join(path, "images" + ext)) 267 | points3D = read_points3d_binary(os.path.join(path, "points3D") + ext) 268 | return cameras, images, points3D 269 | 270 | 271 | def qvec2rotmat(qvec): 272 | return np.array([ 273 | [1 - 2 * qvec[2]**2 - 2 * qvec[3]**2, 274 | 2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3], 275 | 2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2]], 276 | [2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3], 277 | 1 - 2 * qvec[1]**2 - 2 * qvec[3]**2, 278 | 2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1]], 279 | [2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2], 280 | 2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1], 281 | 1 - 2 * qvec[1]**2 - 2 * qvec[2]**2]]) 282 | 283 | 284 | def rotmat2qvec(R): 285 | Rxx, Ryx, Rzx, Rxy, Ryy, Rzy, Rxz, Ryz, Rzz = R.flat 286 | K = np.array([ 287 | [Rxx - Ryy - Rzz, 0, 0, 0], 288 | [Ryx + Rxy, Ryy - Rxx - Rzz, 0, 0], 289 | [Rzx + Rxz, Rzy + Ryz, Rzz - Rxx - Ryy, 0], 290 | [Ryz - Rzy, Rzx - Rxz, Rxy - Ryx, Rxx + Ryy + Rzz]]) / 3.0 291 | eigvals, eigvecs = np.linalg.eigh(K) 292 | qvec = eigvecs[[3, 0, 1, 2], np.argmax(eigvals)] 293 | if qvec[0] < 0: 294 | qvec *= -1 295 | return qvec 296 | --------------------------------------------------------------------------------