├── crf ├── __init__.py ├── plot.py ├── emor.py ├── hdr2ldr.py └── model_crf.py ├── utils ├── __init__.py ├── dataset │ ├── scannetpp │ │ ├── __init__.py │ │ ├── process.py │ │ └── render_semantic.py │ ├── __init__.py │ └── neilf │ │ ├── extract.sh │ │ └── extract_geometry.py ├── loss.py ├── extract_emitter_mesh.py ├── ops.py ├── disco_ball.py ├── metric_brdf.py ├── fuse_segmentation.py ├── gen_path.py ├── metric_crf.py ├── export.py ├── video.py └── ray_utils.py ├── assets └── teaser_cvpr25.png ├── .gitignore ├── scripts ├── metric.sh ├── conda_env.sh ├── export.sh ├── extract_emitter.sh ├── albedo.sh ├── scannetpp │ ├── convert_psdf.sh │ ├── room2 │ │ ├── render.sh │ │ └── train.sh │ └── bathroom2 │ │ ├── render.sh │ │ └── train.sh └── fipt │ ├── classroom │ ├── render.sh │ └── train.sh │ ├── bedroom │ ├── render.sh │ └── train.sh │ ├── kitchen │ ├── render.sh │ └── train.sh │ ├── bathroom │ ├── render.sh │ └── train.sh │ ├── conferenceroom │ ├── render.sh │ └── train.sh │ └── livingroom │ ├── render.sh │ └── train.sh ├── const.py ├── configs ├── scannetpp │ ├── bathroom2 │ │ ├── relight_1.yaml │ │ ├── relight_0.yaml │ │ └── insert.yaml │ └── room2 │ │ ├── relight_1.yaml │ │ ├── relight_0.yaml │ │ └── insert.yaml ├── fipt │ ├── conferenceroom │ │ ├── relight_1.yaml │ │ ├── relight_0.yaml │ │ └── insert.yaml │ ├── classroom │ │ ├── relight_0.yaml │ │ ├── relight_1.yaml │ │ └── insert.yaml │ ├── bedroom │ │ ├── relight_1.yaml │ │ ├── relight_0.yaml │ │ └── insert.yaml │ ├── kitchen │ │ ├── relight_0.yaml │ │ ├── relight_1.yaml │ │ └── insert.yaml │ └── bathroom │ │ ├── relight_0.yaml │ │ ├── insert.yaml │ │ └── relight_1.yaml └── config.py ├── LICENSE ├── environment.yml ├── CONTRIBUTING.md ├── model ├── mlps.py ├── slf.py └── fipt_bsdf.py ├── CODE_OF_CONDUCT.md ├── slf_refine.py ├── extract_emitter_ldr.py ├── slf_bake.py ├── refine_shading.py └── bake_shading.py /crf/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /utils/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /utils/dataset/scannetpp/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /assets/teaser_cvpr25.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/iris/HEAD/assets/teaser_cvpr25.png -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | /datas/ 2 | **/__pycache__/ 3 | notebooks/* 4 | **/.ipynb_checkpoints 5 | /logs/ 6 | /outputs/ 7 | /outputs_ldr/ 8 | /checkpoints/ 9 | /pretrained/ 10 | /data/ 11 | lightning_logs/ 12 | *-Copy* 13 | -------------------------------------------------------------------------------- /scripts/metric.sh: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | 8 | python -m utils.metric_crf \ 9 | --crf_gt /hdd/datasets/fipt/indoor_synthetic/livingroom/train/Image/cam/crf.npy \ 10 | --ckpt checkpoints/250127_no_crf_livingroom/last_1.ckpt \ 11 | --dir_save outputs/250127_no_crf_livingroom/ -------------------------------------------------------------------------------- /utils/dataset/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | from .real_ldr import RealDatasetLDR, InvRealDatasetLDR 8 | from .synthetic_ldr import SyntheticDatasetLDR, InvSyntheticDatasetLDR 9 | 10 | __all__ = [RealDatasetLDR, InvRealDatasetLDR, 11 | SyntheticDatasetLDR, InvSyntheticDatasetLDR] 12 | -------------------------------------------------------------------------------- /scripts/conda_env.sh: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | conda env create -f environment.yml 8 | conda activate iris 9 | pip install torch_scatter -f https://data.pyg.org/whl/torch-1.13.0+cu117.html 10 | pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch # tested with tinycudann-1.7 11 | -------------------------------------------------------------------------------- /scripts/export.sh: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | python -m utils.export \ 8 | --mesh /hdd/datasets/scannetpp/data/45b0dac5e3/scans/scene.ply \ 9 | --ckpt /hdd/meta/meta_iris_dev/checkpoints/240506_scannetpp_bathroom2/last_1.ckpt \ 10 | --emitter_path /hdd/meta/meta_iris_dev/outputs/240506_scannetpp_bathroom2/bake \ 11 | --dir_save /hdd/meta/meta_iris_dev/outputs/240506_scannetpp_bathroom2/output/textured_mesh \ -------------------------------------------------------------------------------- /const.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import numpy as np 8 | import torch 9 | 10 | GAMMA = 2.2 11 | SEED=0 12 | 13 | def set_random_seed(): 14 | np.random.seed(SEED) 15 | torch.manual_seed(SEED) 16 | torch.cuda.manual_seed(SEED) 17 | torch.cuda.manual_seed_all(SEED) 18 | 19 | # Ensures deterministic behavior (slightly slower performance) 20 | torch.backends.cudnn.deterministic = True 21 | torch.backends.cudnn.benchmark = False -------------------------------------------------------------------------------- /utils/dataset/neilf/extract.sh: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | # data folder 8 | DATASET_ROOT='/hdd/datasets/scannetpp/' 9 | DATASET='scannetpp' 10 | # scene name 11 | SCENE='45b0dac5e3' 12 | OUTPUT='/hdd/datasets/neilf/scannetpp/bathroom2/' 13 | INPUT=ldr 14 | LDR_IMG_DIR=Image 15 | SPLIT=all 16 | RES_SCALE=0.25 17 | 18 | python -m utils.dataset.neilf.extract_geometry \ 19 | --dataset_root $DATASET_ROOT --scene $SCENE --dataset $DATASET \ 20 | --input $INPUT --split $SPLIT \ 21 | --output $OUTPUT --res_scale $RES_SCALE -------------------------------------------------------------------------------- /scripts/extract_emitter.sh: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | 8 | 9 | python -m utils.extract_emitter_mesh \ 10 | --mesh_scene /hdd/datasets/scannetpp/data/7e09430da7/scans/scene.ply \ 11 | --emitter /hdd/meta/meta_iris_dev/outputs/240506_scannetpp_room2/bake/emitter.pth \ 12 | --mesh_emitter /hdd/meta/meta_iris_dev/outputs/240506_scannetpp_room2/bake/emitter.ply 13 | 14 | # Emitter average radiance: 15 | # bathroom2: [ 9.040693, 9.697464, 10.583247] 16 | # conferenceroom: [2.1285791, 2.1815128, 2.379316 ] 17 | # room2: [8.679976, 8.248041, 2.860989] -------------------------------------------------------------------------------- /scripts/albedo.sh: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | DIR_IMG=/hdd/datasets/scannetpp/data/4a1a3a7dc5/psdf/images/ 8 | DIR_OUTPUT=/hdd/datasets/scannetpp/data/4a1a3a7dc5/psdf/albedo/ 9 | 10 | cd /hdd/meta/IRISFormer_beta 11 | conda activate iris 12 | python train/train_customize.py --task_name DATE-train_mm1_albedo_eval_IIW_RESUME20230517-224958 \ 13 | --if_train False --if_val True --if_vis True --eval_every_iter 4000 \ 14 | --config-file train/configs/train_albedo.yaml --resume 20230517-224958--DATE-train_mm1_albedo \ 15 | --data_root $DIR_IMG \ 16 | --dir_output $DIR_OUTPUT 17 | 18 | cd /hdd/meta/meta_iris_dev 19 | conda activate fipt 20 | -------------------------------------------------------------------------------- /configs/scannetpp/bathroom2/relight_1.yaml: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | type: 'scene' 8 | 9 | PerspectiveCamera: 10 | type: 'perspective' 11 | fov: 45 12 | film: 13 | type: 'hdrfilm' 14 | width: 1920 15 | height: 1080 16 | rfilter: 17 | type: 'box' 18 | 19 | Integrator: 20 | type: 'path' 21 | max_depth: 7 22 | 23 | main_scene: 24 | type: 'ply' 25 | filename: '' 26 | bsdf: 27 | type: 'twosided' 28 | fipt_bsdf: 29 | type: 'fipt' 30 | emitter_path: '' 31 | brdf_path: '' 32 | 33 | disco_ball: 34 | T: 60 35 | position: [0.0 , 0.3 , -0.5] 36 | radius: 0.2 37 | light_intensity: 40 38 | light_num: 40 39 | light_radius_rate: 0.1 40 | spot_intensity: 0.5 41 | spot_cutoff_angle: 20.0 -------------------------------------------------------------------------------- /configs/fipt/conferenceroom/relight_1.yaml: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | type: 'scene' 8 | 9 | PerspectiveCamera: 10 | type: 'perspective' 11 | fov: 45 12 | film: 13 | type: 'hdrfilm' 14 | width: 1920 15 | height: 1080 16 | rfilter: 17 | type: 'box' 18 | 19 | Integrator: 20 | type: 'path' 21 | max_depth: 7 22 | 23 | main_scene: 24 | type: 'ply' 25 | filename: '' 26 | bsdf: 27 | type: 'twosided' 28 | fipt_bsdf: 29 | type: 'fipt' 30 | emitter_path: '' 31 | brdf_path: '' 32 | 33 | disco_ball: 34 | T: 60 35 | position: [-2.90256, -0.923296, 0.603608] 36 | radius: 0.5 37 | light_intensity: 50 38 | light_num: 40 39 | light_radius_rate: 0.1 40 | spot_intensity: 5 41 | spot_cutoff_angle: 20.0 -------------------------------------------------------------------------------- /configs/scannetpp/room2/relight_1.yaml: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | type: 'scene' 8 | 9 | PerspectiveCamera: 10 | type: 'perspective' 11 | fov: 45 12 | film: 13 | type: 'hdrfilm' 14 | width: 1920 15 | height: 1080 16 | rfilter: 17 | type: 'box' 18 | 19 | Integrator: 20 | type: 'path' 21 | max_depth: 7 22 | 23 | main_scene: 24 | type: 'ply' 25 | filename: '' 26 | bsdf: 27 | type: 'twosided' 28 | fipt_bsdf: 29 | type: 'fipt' 30 | emitter_path: '' 31 | brdf_path: '' 32 | 33 | disco_ball: 34 | T: 60 35 | position: [-0.033905, -0.41319, -0.327367] 36 | radius: 0.05 37 | light_intensity: 100 38 | light_num: 40 39 | light_radius_rate: 0.1 40 | spot_intensity: 0.2 41 | spot_cutoff_angle: 20.0 -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | LICENSE 2 | 3 | Creative Commons Attribution-NonCommercial 4.0 International 4 | (CC BY-NC 4.0) 5 | 6 | Copyright (c) [2025] [Meta Platforms, Inc.] 7 | 8 | This work is licensed under the Creative Commons Attribution-NonCommercial 4.0 International License. 9 | To view a copy of this license, visit https://creativecommons.org/licenses/by-nc/4.0/ or send a letter to Creative Commons, PO Box 1866, Mountain View, CA 94042, USA. 10 | 11 | You are free to: 12 | 13 | Share — copy and redistribute the material in any medium or format 14 | 15 | Adapt — remix, transform, and build upon the material 16 | 17 | Under the following terms: 18 | 19 | Attribution — You must give appropriate credit, provide a link to the license, and indicate if changes were made. You may do so in any reasonable manner, but not in any way that suggests the licensor endorses you or your use. 20 | 21 | NonCommercial — You may not use the material for commercial purposes. 22 | 23 | No additional restrictions — You may not apply legal terms or technological measures that legally restrict others from doing anything the license permits. -------------------------------------------------------------------------------- /configs/fipt/classroom/relight_0.yaml: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | type: 'scene' 8 | 9 | PerspectiveCamera: 10 | type: 'perspective' 11 | fov: 45 12 | film: 13 | type: 'hdrfilm' 14 | width: 1920 15 | height: 1080 16 | rfilter: 17 | type: 'box' 18 | 19 | Integrator: 20 | type: 'path' 21 | max_depth: 7 22 | 23 | main_scene: 24 | type: 'ply' 25 | filename: '' 26 | bsdf: 27 | type: 'twosided' 28 | fipt_bsdf: 29 | type: 'fipt' 30 | emitter_path: '' 31 | brdf_path: '' 32 | 33 | light_ball: 34 | type: 'sphere' 35 | to_world: 36 | - type: 'translate' 37 | value: [-3.0, 0.5, 2.5] 38 | - type: 'scale' 39 | value: [0.3, 0.3, 0.3] 40 | bsdf: 41 | type: 'diffuse' 42 | reflectance: 43 | type: 'rgb' 44 | value: [0., 0., 0.] 45 | emitter: 46 | type: 'area' 47 | radiance: 48 | type: 'rgb' 49 | value: [24, 24, 30] -------------------------------------------------------------------------------- /configs/fipt/bedroom/relight_1.yaml: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | type: 'scene' 8 | 9 | PerspectiveCamera: 10 | type: 'perspective' 11 | fov: 45 12 | film: 13 | type: 'hdrfilm' 14 | width: 1920 15 | height: 1080 16 | rfilter: 17 | type: 'box' 18 | 19 | Integrator: 20 | type: 'path' 21 | max_depth: 7 22 | 23 | main_scene: 24 | type: 'ply' 25 | filename: '' 26 | bsdf: 27 | type: 'twosided' 28 | fipt_bsdf: 29 | type: 'fipt' 30 | emitter_path: '' 31 | brdf_path: '' 32 | 33 | light_ball: 34 | type: 'sphere' 35 | to_world: 36 | - type: 'translate' 37 | value: [-1.29994, 1.21888, 0.020946] 38 | - type: 'scale' 39 | value: [0.2, 0.2, 0.2] 40 | bsdf: 41 | type: 'diffuse' 42 | reflectance: 43 | type: 'rgb' 44 | value: [0., 0., 0.] 45 | emitter: 46 | type: 'area' 47 | radiance: 48 | type: 'rgb' 49 | value: [15, 12, 15] -------------------------------------------------------------------------------- /configs/fipt/conferenceroom/relight_0.yaml: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | type: 'scene' 8 | 9 | PerspectiveCamera: 10 | type: 'perspective' 11 | fov: 45 12 | film: 13 | type: 'hdrfilm' 14 | width: 1920 15 | height: 1080 16 | rfilter: 17 | type: 'box' 18 | 19 | Integrator: 20 | type: 'path' 21 | max_depth: 7 22 | 23 | main_scene: 24 | type: 'ply' 25 | filename: '' 26 | bsdf: 27 | type: 'twosided' 28 | fipt_bsdf: 29 | type: 'fipt' 30 | emitter_path: '' 31 | brdf_path: '' 32 | 33 | light_ball: 34 | type: 'sphere' 35 | to_world: 36 | - type: 'translate' 37 | value: [-2.0, -1.2, 0.5] 38 | - type: 'scale' 39 | value: [0.3, 0.3, 0.3] 40 | bsdf: 41 | type: 'diffuse' 42 | reflectance: 43 | type: 'rgb' 44 | value: [0., 0., 0.] 45 | emitter: 46 | type: 'area' 47 | radiance: 48 | type: 'rgb' 49 | value: [15, 15, 12] 50 | -------------------------------------------------------------------------------- /configs/fipt/kitchen/relight_0.yaml: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | type: 'scene' 8 | 9 | PerspectiveCamera: 10 | type: 'perspective' 11 | fov: 45 12 | film: 13 | type: 'hdrfilm' 14 | width: 1920 15 | height: 1080 16 | rfilter: 17 | type: 'box' 18 | 19 | Integrator: 20 | type: 'path' 21 | max_depth: 7 22 | 23 | main_scene: 24 | type: 'ply' 25 | filename: '' 26 | bsdf: 27 | type: 'twosided' 28 | fipt_bsdf: 29 | type: 'fipt' 30 | emitter_path: '' 31 | brdf_path: '' 32 | 33 | light_ball: 34 | type: 'sphere' 35 | to_world: 36 | - type: 'translate' 37 | value: [-1.03821, 1.46279,-1.43346] 38 | - type: 'scale' 39 | value: [0.2, 0.2, 0.2] 40 | bsdf: 41 | type: 'diffuse' 42 | reflectance: 43 | type: 'rgb' 44 | value: [0., 0., 0.] 45 | emitter: 46 | type: 'area' 47 | radiance: 48 | type: 'rgb' 49 | value: [10.2, 9.8, 7.28] -------------------------------------------------------------------------------- /configs/scannetpp/bathroom2/relight_0.yaml: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | type: 'scene' 8 | 9 | PerspectiveCamera: 10 | type: 'perspective' 11 | fov: 45 12 | film: 13 | type: 'hdrfilm' 14 | width: 1920 15 | height: 1080 16 | rfilter: 17 | type: 'box' 18 | 19 | Integrator: 20 | type: 'path' 21 | max_depth: 7 22 | 23 | main_scene: 24 | type: 'ply' 25 | filename: '' 26 | bsdf: 27 | type: 'twosided' 28 | fipt_bsdf: 29 | type: 'fipt' 30 | emitter_path: '' 31 | brdf_path: '' 32 | 33 | light_ball: 34 | type: 'sphere' 35 | to_world: 36 | - type: 'translate' 37 | value: [0.06 , 0.29 , -0.5] 38 | - type: 'scale' 39 | value: [0.1, 0.1, 0.1] 40 | bsdf: 41 | type: 'diffuse' 42 | reflectance: 43 | type: 'rgb' 44 | value: [0., 0., 0.] 45 | emitter: 46 | type: 'area' 47 | radiance: 48 | type: 'rgb' 49 | value: [12., 12., 15.] 50 | -------------------------------------------------------------------------------- /configs/scannetpp/room2/relight_0.yaml: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | type: 'scene' 8 | 9 | PerspectiveCamera: 10 | type: 'perspective' 11 | fov: 45 12 | film: 13 | type: 'hdrfilm' 14 | width: 1920 15 | height: 1080 16 | rfilter: 17 | type: 'box' 18 | 19 | Integrator: 20 | type: 'path' 21 | max_depth: 7 22 | 23 | main_scene: 24 | type: 'ply' 25 | filename: '' 26 | bsdf: 27 | type: 'twosided' 28 | fipt_bsdf: 29 | type: 'fipt' 30 | emitter_path: '' 31 | brdf_path: '' 32 | 33 | light_ball: 34 | type: 'sphere' 35 | to_world: 36 | - type: 'translate' 37 | value: [-0.240946, -0.199994, -0.034389] 38 | - type: 'scale' 39 | value: [0.1, 0.1, 0.1] 40 | bsdf: 41 | type: 'diffuse' 42 | reflectance: 43 | type: 'rgb' 44 | value: [0., 0., 0.] 45 | emitter: 46 | type: 'area' 47 | radiance: 48 | type: 'rgb' 49 | value: [5, 5, 10] 50 | -------------------------------------------------------------------------------- /environment.yml: -------------------------------------------------------------------------------- 1 | name: iris 2 | channels: 3 | - defaults 4 | dependencies: 5 | - python=3.8 6 | - pip 7 | - pip: 8 | # Core Python libraries 9 | - numpy==1.24.4 10 | - torch==1.13.1+cu117 11 | - torchvision==0.14.1+cu117 12 | - --extra-index-url https://download.pytorch.org/whl/cu117 13 | - pytorch-lightning==1.9.0 14 | - mitsuba==3.5.0 15 | - drjit==0.4.4 16 | 17 | # Image processing and visualization 18 | - opencv-python 19 | - imageio 20 | - imageio-ffmpeg 21 | - matplotlib 22 | - pillow 23 | 24 | # Utilities 25 | - tqdm==4.66.1 26 | - argparse==1.4.0 27 | - omegaconf==2.3.0 28 | - scipy==1.10.1 29 | - scikit-image==0.21.0 30 | - trimesh==3.21.4 31 | - pyquaternion==0.9.9 32 | 33 | # Logging and monitoring 34 | - tensorboardx==2.6.2.2 35 | - wandb==0.16.1 36 | 37 | # File and path handling 38 | - pathlib==1.0.1 39 | 40 | # Optional dependencies for CRF and plotting 41 | - lpips==0.1.4 42 | - matplotlib-inline==0.1.6 43 | - git+https://github.com/sbarratt/torch_interpolations.git -------------------------------------------------------------------------------- /utils/dataset/scannetpp/process.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | from argparse import ArgumentParser 8 | import cv2 9 | import os 10 | from tqdm import tqdm 11 | 12 | def main(): 13 | parser = ArgumentParser() 14 | parser.add_argument('--input', type=str, required=True) 15 | parser.add_argument('--output', type=str, required=True) 16 | parser.add_argument('--max_width', type=int, default=1024) 17 | args = parser.parse_args() 18 | 19 | os.makedirs(args.output, exist_ok=True) 20 | img_names = sorted([name for name in os.listdir(args.input)]) 21 | 22 | for name in tqdm(img_names): 23 | img = cv2.imread(os.path.join(args.input, name)) 24 | if img.shape[1] > args.max_width: 25 | w_new = args.max_width 26 | h_new = int(args.max_width / img.shape[1] * img.shape[0]) 27 | img = cv2.resize(img, (w_new, h_new)) 28 | cv2.imwrite(os.path.join(args.output, name), img) 29 | 30 | if __name__ == '__main__': 31 | main() -------------------------------------------------------------------------------- /configs/fipt/bedroom/relight_0.yaml: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | type: 'scene' 8 | 9 | PerspectiveCamera: 10 | type: 'perspective' 11 | fov: 45 12 | film: 13 | type: 'hdrfilm' 14 | width: 1920 15 | height: 1080 16 | rfilter: 17 | type: 'box' 18 | 19 | Integrator: 20 | type: 'path' 21 | max_depth: 7 22 | 23 | main_scene: 24 | type: 'ply' 25 | filename: '' 26 | bsdf: 27 | type: 'twosided' 28 | fipt_bsdf: 29 | type: 'fipt' 30 | emitter_path: '' 31 | brdf_path: '' 32 | 33 | light_plane: 34 | type: 'rectangle' 35 | to_world: 36 | - type: 'translate' 37 | value: [-1.6257, 2.3767, 1.5447] 38 | - type: 'scale' 39 | value: [0.197,1,0.589] 40 | - type: 'rotate' 41 | axis: [1, 0, 0] 42 | angle: 90 43 | bsdf: 44 | type: 'twosided' 45 | bsdf: 46 | type: 'diffuse' 47 | reflectance: 48 | type: 'rgb' 49 | value: [0., 0., 0.] 50 | emitter: 51 | type: 'area' 52 | radiance: 53 | type: 'rgb' 54 | value: [15, 15, 8] -------------------------------------------------------------------------------- /configs/fipt/classroom/relight_1.yaml: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | type: 'scene' 8 | 9 | PerspectiveCamera: 10 | type: 'perspective' 11 | fov: 45 12 | film: 13 | type: 'hdrfilm' 14 | width: 1920 15 | height: 1080 16 | rfilter: 17 | type: 'box' 18 | 19 | Integrator: 20 | type: 'path' 21 | max_depth: 7 22 | 23 | main_scene: 24 | type: 'ply' 25 | filename: '' 26 | bsdf: 27 | type: 'twosided' 28 | fipt_bsdf: 29 | type: 'fipt' 30 | emitter_path: '' 31 | brdf_path: '' 32 | 33 | light_plane: 34 | type: 'rectangle' 35 | to_world: 36 | - type: 'translate' 37 | value: [1.6643, 1.7467, 3.0747] 38 | - type: 'scale' 39 | value: [1.287,1,0.809] 40 | - type: 'rotate' 41 | axis: [1, 0, 0] 42 | angle: 90 43 | bsdf: 44 | type: 'twosided' 45 | bsdf: 46 | type: 'diffuse' 47 | reflectance: 48 | type: 'rgb' 49 | value: [0., 0., 0.] 50 | emitter: 51 | type: 'area' 52 | radiance: 53 | type: 'rgb' 54 | value: [15, 15, 12] -------------------------------------------------------------------------------- /configs/fipt/kitchen/relight_1.yaml: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | type: 'scene' 8 | 9 | PerspectiveCamera: 10 | type: 'perspective' 11 | fov: 45 12 | film: 13 | type: 'hdrfilm' 14 | width: 1920 15 | height: 1080 16 | rfilter: 17 | type: 'box' 18 | 19 | Integrator: 20 | type: 'path' 21 | max_depth: 7 22 | 23 | main_scene: 24 | type: 'ply' 25 | filename: '' 26 | bsdf: 27 | type: 'twosided' 28 | fipt_bsdf: 29 | type: 'fipt' 30 | emitter_path: '' 31 | brdf_path: '' 32 | 33 | light_plane: 34 | type: 'rectangle' 35 | to_world: 36 | - type: 'translate' 37 | value: [0.068418, 3.2243, 0.85067] 38 | - type: 'scale' 39 | value: [0.4, 0.4, 0.4] 40 | - type: 'rotate' 41 | axis: [1, 0, 0] 42 | angle: 90 43 | bsdf: 44 | type: 'twosided' 45 | bsdf: 46 | type: 'diffuse' 47 | reflectance: 48 | type: 'rgb' 49 | value: [0., 0., 0.] 50 | emitter: 51 | type: 'area' 52 | radiance: 53 | type: 'rgb' 54 | value: [20, 20, 20] -------------------------------------------------------------------------------- /configs/fipt/bathroom/relight_0.yaml: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | type: 'scene' 8 | 9 | PerspectiveCamera: 10 | type: 'perspective' 11 | fov: 45 12 | film: 13 | type: 'hdrfilm' 14 | width: 1920 15 | height: 1080 16 | rfilter: 17 | type: 'box' 18 | 19 | Integrator: 20 | type: 'path' 21 | max_depth: 7 22 | 23 | main_scene: 24 | type: 'ply' 25 | filename: '' 26 | bsdf: 27 | type: 'twosided' 28 | fipt_bsdf: 29 | type: 'fipt' 30 | emitter_path: '' 31 | brdf_path: '' 32 | 33 | light_plane: 34 | type: 'rectangle' 35 | to_world: 36 | - type: 'translate' 37 | value: [-1.362, 2.7918, -2.2202] 38 | - type: 'scale' 39 | value: [0.865, 0.865, 0.08] 40 | - type: 'rotate' 41 | axis: [1, 0, 0] 42 | angle: 90 43 | bsdf: 44 | type: 'twosided' 45 | bsdf: 46 | type: 'diffuse' 47 | reflectance: 48 | type: 'rgb' 49 | value: [0., 0., 0.] 50 | emitter: 51 | type: 'area' 52 | radiance: 53 | type: 'rgb' 54 | value: [22, 25, 25] -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # Contributing to iris 2 | We want to make contributing to this project as easy and transparent as 3 | possible. 4 | 5 | ## Pull Requests 6 | We actively welcome your pull requests. 7 | 8 | 1. Fork the repo and create your branch from `main`. 9 | 2. If you've added code that should be tested, add tests. 10 | 3. If you've changed APIs, update the documentation. 11 | 4. Ensure the test suite passes. 12 | 5. Make sure your code lints. 13 | 6. If you haven't already, complete the Contributor License Agreement ("CLA"). 14 | 15 | ## Contributor License Agreement ("CLA") 16 | In order to accept your pull request, we need you to submit a CLA. You only need 17 | to do this once to work on any of Facebook's open source projects. 18 | 19 | Complete your CLA here: 20 | 21 | ## Issues 22 | We use GitHub issues to track public bugs. Please ensure your description is 23 | clear and has sufficient instructions to be able to reproduce the issue. 24 | 25 | Facebook has a [bounty program](https://www.facebook.com/whitehat/) for the safe 26 | disclosure of security bugs. In those cases, please go through the process 27 | outlined on that page and do not file a public issue. 28 | 29 | ## License 30 | By contributing to iris, you agree that your contributions will be licensed 31 | under the LICENSE file in the root directory of this source tree. -------------------------------------------------------------------------------- /utils/loss.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import numpy as np 8 | import math 9 | import torch 10 | import torch.nn as nn 11 | import torch.nn.functional as F 12 | import torch_scatter 13 | 14 | def compute_scale(source, target): 15 | ''' 16 | solve least square error problem: target = source * scale 17 | ''' 18 | source, target = source.view(-1), target.view(-1) #(n, ) 19 | scale = torch.dot(source, target) / torch.dot(source, source) 20 | return scale.item() 21 | 22 | def compute_scale_shift(source, target): 23 | ''' 24 | solve least square error problem: target = source * scale + shift 25 | ''' 26 | source, target = source.view(-1), target.view(-1) #(n, ) 27 | source_one = torch.stack([source, torch.ones_like(source)], dim=-1) #(n, 2) 28 | pseudo_inverse = torch.inverse(source_one.T @ source_one) @ source_one.T #(2, n) 29 | x = pseudo_inverse @ target 30 | scale, shift = x[0].item(), x[1].item() 31 | return scale, shift 32 | 33 | def scale_invariant_mse(source, target): 34 | scale = compute_scale(source, target) 35 | source_transformed = source * scale 36 | loss = F.mse_loss(source_transformed, target) 37 | return loss 38 | 39 | def scale_shift_invariant_mse(source, target): 40 | scale, shift = compute_scale_shift(source, target) 41 | source_transformed = source * scale + shift 42 | loss = F.mse_loss(source_transformed, target) 43 | return loss -------------------------------------------------------------------------------- /configs/fipt/kitchen/insert.yaml: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | type: 'scene' 8 | 9 | PerspectiveCamera: 10 | type: 'perspective' 11 | fov: 45 12 | film: 13 | type: 'hdrfilm' 14 | width: 1920 15 | height: 1080 16 | rfilter: 17 | type: 'box' 18 | 19 | Integrator: 20 | type: 'path' 21 | max_depth: 7 22 | 23 | main_scene: 24 | type: 'ply' 25 | filename: '' 26 | bsdf: 27 | type: 'twosided' 28 | fipt_bsdf: 29 | type: 'fipt' 30 | emitter_path: '' 31 | brdf_path: '' 32 | 33 | light_ball: 34 | type: 'sphere' 35 | to_world: 36 | - type: 'translate' 37 | value: [0.2, 2.25, -1.6] 38 | - type: 'scale' 39 | value: [0.15, 0.15, 0.15] 40 | bsdf: 41 | type: 'diffuse' 42 | reflectance: 43 | type: 'rgb' 44 | value: [0., 0., 0.] 45 | emitter: 46 | type: 'area' 47 | radiance: 48 | type: 'rgb' 49 | value: [25, 25, 25] 50 | 51 | sphere_2: 52 | type: 'sphere' 53 | to_world: 54 | - type: 'translate' 55 | value: [0.125, 1.14, -1.07569] 56 | - type: 'scale' 57 | value: [0.25, 0.25, 0.25] 58 | bsdf: 59 | type: 'conductor' 60 | material: 'none' 61 | 62 | sphere_3: 63 | type: 'sphere' 64 | to_world: 65 | - type: 'translate' 66 | value: [0.125, 1.02819, -2] 67 | - type: 'scale' 68 | value: [0.15, 0.15, 0.15] 69 | bsdf: 70 | type: 'diffuse' 71 | reflectance: 72 | type: 'rgb' 73 | value: [0.2, 0.25, 0.7] -------------------------------------------------------------------------------- /configs/scannetpp/bathroom2/insert.yaml: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | type: 'scene' 8 | 9 | PerspectiveCamera: 10 | type: 'perspective' 11 | fov: 45 12 | film: 13 | type: 'hdrfilm' 14 | width: 1920 15 | height: 1080 16 | rfilter: 17 | type: 'box' 18 | 19 | Integrator: 20 | type: 'path' 21 | max_depth: 7 22 | 23 | main_scene: 24 | type: 'ply' 25 | filename: '' 26 | bsdf: 27 | type: 'twosided' 28 | fipt_bsdf: 29 | type: 'fipt' 30 | emitter_path: '' 31 | brdf_path: '' 32 | 33 | emitter_bathroom2: 34 | type: 'ply' 35 | filename: 'checkpoints/scannetpp_bathroom2/bake/emitter.ply' 36 | to_world: 37 | - type: 'translate' 38 | value: [0.0, 0.0, 0.0] 39 | - type: 'scale' 40 | value: [1.0, 1.0, 1.0] 41 | bsdf: 42 | type: 'diffuse' 43 | reflectance: 44 | type: 'rgb' 45 | value: [0., 0., 0.] 46 | emitter: 47 | type: 'area' 48 | radiance: 49 | type: 'rgb' 50 | value: [ 9.040693, 9.697464, 10.583247] 51 | 52 | spot: 53 | type: 'obj' 54 | filename: 'outputs/00_assets/spot.obj' 55 | to_world: 56 | - type: 'translate' 57 | value: [-0.561803, -0.323403, 0.079234] 58 | - type: 'scale' 59 | value: [0.2, 0.2, 0.2] 60 | - type: 'rotate' 61 | axis: [1, 0, 0] 62 | angle: -90 63 | - type: 'rotate' 64 | axis: [0, 1, 0] 65 | angle: -90 66 | bsdf: 67 | type: 'conductor' 68 | material: 'Au' 69 | 70 | -------------------------------------------------------------------------------- /scripts/scannetpp/convert_psdf.sh: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | DATA_ROOT=/hdd/datasets/scannetpp/data 8 | MAX_W=1024 9 | 10 | scenes=( 11 | '1ada7a0617' 12 | ) 13 | 14 | for scene in "${scenes[@]}"; do 15 | echo "============== $scene ==============" 16 | cp $DATA_ROOT/$scene/dslr/train_test_lists.json $DATA_ROOT/$scene/psdf 17 | 18 | # Extract semantic 19 | python -m utils.dataset.scannetpp.gen_mesh --scene $scene 20 | python -m utils.dataset.scannetpp.render_semantic --scene $scene 21 | 22 | # Estimate albedo 23 | cd /hdd/meta/IRISFormer_beta 24 | conda activate iris 25 | python train/train_customize.py --task_name DATE-train_mm1_albedo_eval_IIW_RESUME20230517-224958 \ 26 | --if_train False --if_val True --if_vis True --eval_every_iter 4000 \ 27 | --config-file train/configs/train_albedo.yaml --resume 20230517-224958--DATE-train_mm1_albedo \ 28 | --data_root $DATA_ROOT/$scene/dslr/undistorted_images \ 29 | --dir_output $DATA_ROOT/$scene/dslr/albedo 30 | 31 | cd /hdd/meta/meta_iris_dev 32 | conda activate fipt 33 | 34 | # resize albedo 35 | python -m utils.dataset.scannetpp.process \ 36 | --input $DATA_ROOT/$scene/dslr/albedo/ \ 37 | --output $DATA_ROOT/$scene/psdf/albedo/ \ 38 | --max_width $MAX_W 39 | 40 | # resize semantic mask 41 | python -m utils.dataset.scannetpp.process \ 42 | --input $DATA_ROOT/$scene/dslr/seg/ \ 43 | --output $DATA_ROOT/$scene/psdf/seg/ \ 44 | --max_width $MAX_W 45 | done 46 | 47 | -------------------------------------------------------------------------------- /scripts/fipt/classroom/render.sh: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | # data folder 8 | DATASET_ROOT='/hdd/datasets/fipt/real/' 9 | DATASET='real' 10 | # scene name 11 | SCENE='classroom' 12 | LDR_IMG_DIR='Image' 13 | EXP='fipt_real_classroom' 14 | VAL_FRAME=1 15 | CRF_BASIS=3 16 | # whether has part segmentation 17 | HAS_PART=0 18 | SPP=256 19 | spp=16 20 | 21 | python render.py --experiment_name $EXP --device 0\ 22 | --ckpt last_1.ckpt \ 23 | --dataset $DATASET $DATASET_ROOT$SCENE \ 24 | --emitter_path checkpoints/$EXP/bake\ 25 | --output_path 'outputs/'$EXP'/output'\ 26 | --split 'val'\ 27 | --ldr_img_dir $LDR_IMG_DIR \ 28 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS 29 | 30 | python render_video.py --experiment_name $EXP --device 0\ 31 | --ckpt last_1.ckpt \ 32 | --dataset $DATASET $DATASET_ROOT$SCENE\ 33 | --emitter_path checkpoints/$EXP/bake\ 34 | --output_path 'outputs/'$EXP'/video'\ 35 | --split 'val'\ 36 | --ldr_img_dir $LDR_IMG_DIR \ 37 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS 38 | 39 | # relighting 40 | python render_relight.py --experiment_name $EXP --device 0\ 41 | --ckpt last_1.ckpt --mode traj\ 42 | --dataset $DATASET $DATASET_ROOT$SCENE\ 43 | --emitter_path checkpoints/$EXP/bake\ 44 | --output_path 'outputs/'$EXP'/relight/video_relight_0'\ 45 | --split 'test'\ 46 | --ldr_img_dir $LDR_IMG_DIR \ 47 | --light_cfg 'configs/fipt/classroom/relight_0.yaml' \ 48 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS 49 | -------------------------------------------------------------------------------- /scripts/fipt/bedroom/render.sh: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | # data folder 8 | DATASET_ROOT='/hdd/datasets/fipt/indoor_synthetic/' 9 | DATASET='synthetic' 10 | # scene name 11 | SCENE='bedroom' 12 | LDR_IMG_DIR='Image' 13 | EXP='fipt_syn_bedroom' 14 | VAL_FRAME=10 15 | CRF_BASIS=3 16 | # whether has part segmentation 17 | HAS_PART=1 18 | SPP=256 19 | spp=16 20 | 21 | python render.py --experiment_name $EXP --device 0\ 22 | --ckpt last_1.ckpt \ 23 | --dataset $DATASET $DATASET_ROOT$SCENE \ 24 | --emitter_path checkpoints/$EXP/bake\ 25 | --output_path 'outputs/'$EXP'/output'\ 26 | --split 'val'\ 27 | --ldr_img_dir $LDR_IMG_DIR \ 28 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS 29 | 30 | python render_video.py --experiment_name $EXP --device 0\ 31 | --ckpt last_1.ckpt \ 32 | --dataset $DATASET $DATASET_ROOT$SCENE\ 33 | --emitter_path checkpoints/$EXP/bake\ 34 | --output_path 'outputs/'$EXP'/video'\ 35 | --split 'val'\ 36 | --ldr_img_dir $LDR_IMG_DIR \ 37 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS 38 | 39 | # relighting 40 | python render_relight.py --experiment_name $EXP --device 0\ 41 | --ckpt last_1.ckpt --mode traj\ 42 | --dataset $DATASET $DATASET_ROOT$SCENE\ 43 | --emitter_path checkpoints/$EXP/bake\ 44 | --output_path 'outputs/'$EXP'/relight/video_relight_0'\ 45 | --split 'test'\ 46 | --ldr_img_dir $LDR_IMG_DIR \ 47 | --light_cfg 'configs/fipt/bedroom/relight_0.yaml' \ 48 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS -------------------------------------------------------------------------------- /scripts/fipt/kitchen/render.sh: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | # data folder 8 | DATASET_ROOT='/hdd/datasets/fipt/indoor_synthetic/' 9 | DATASET='synthetic' 10 | # scene name 11 | SCENE='kitchen' 12 | LDR_IMG_DIR='Image' 13 | EXP='fipt_syn_kitchen' 14 | VAL_FRAME=10 15 | CRF_BASIS=3 16 | # whether has part segmentation 17 | HAS_PART=1 18 | SPP=256 19 | spp=16 20 | 21 | python render.py --experiment_name $EXP --device 0\ 22 | --ckpt last_1.ckpt \ 23 | --dataset $DATASET $DATASET_ROOT$SCENE \ 24 | --emitter_path checkpoints/$EXP/bake\ 25 | --output_path 'outputs/'$EXP'/output'\ 26 | --split 'val'\ 27 | --ldr_img_dir $LDR_IMG_DIR \ 28 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS 29 | 30 | python render_video.py --experiment_name $EXP --device 0\ 31 | --ckpt last_1.ckpt \ 32 | --dataset $DATASET $DATASET_ROOT$SCENE\ 33 | --emitter_path checkpoints/$EXP/bake\ 34 | --output_path 'outputs/'$EXP'/video'\ 35 | --split 'val'\ 36 | --ldr_img_dir $LDR_IMG_DIR \ 37 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS 38 | 39 | # relighting 40 | python render_relight.py --experiment_name $EXP --device 0\ 41 | --ckpt last_1.ckpt --mode traj\ 42 | --dataset $DATASET $DATASET_ROOT$SCENE\ 43 | --emitter_path checkpoints/$EXP/bake\ 44 | --output_path 'outputs/'$EXP'/relight/video_relight_0'\ 45 | --split 'test'\ 46 | --ldr_img_dir $LDR_IMG_DIR \ 47 | --light_cfg 'configs/fipt/kitchen/relight_0.yaml' \ 48 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS -------------------------------------------------------------------------------- /scripts/fipt/bathroom/render.sh: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | # data folder 8 | DATASET_ROOT='/hdd/datasets/fipt/indoor_synthetic/' 9 | DATASET='synthetic' 10 | # scene name 11 | SCENE='bathroom' 12 | LDR_IMG_DIR='Image' 13 | EXP='fipt_syn_bathroom' 14 | VAL_FRAME=10 15 | CRF_BASIS=3 16 | # whether has part segmentation 17 | HAS_PART=1 18 | SPP=256 19 | spp=16 20 | 21 | python render.py --experiment_name $EXP --device 0\ 22 | --ckpt last_1.ckpt \ 23 | --dataset $DATASET $DATASET_ROOT$SCENE \ 24 | --emitter_path checkpoints/$EXP/bake\ 25 | --output_path 'outputs/'$EXP'/output'\ 26 | --split 'val'\ 27 | --ldr_img_dir $LDR_IMG_DIR \ 28 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS 29 | 30 | python render_video.py --experiment_name $EXP --device 0\ 31 | --ckpt last_1.ckpt \ 32 | --dataset $DATASET $DATASET_ROOT$SCENE\ 33 | --emitter_path checkpoints/$EXP/bake\ 34 | --output_path 'outputs/'$EXP'/video'\ 35 | --split 'val'\ 36 | --ldr_img_dir $LDR_IMG_DIR \ 37 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS 38 | 39 | # relighting 40 | python render_relight.py --experiment_name $EXP --device 0\ 41 | --ckpt last_1.ckpt --mode traj\ 42 | --dataset $DATASET $DATASET_ROOT$SCENE\ 43 | --emitter_path checkpoints/$EXP/bake\ 44 | --output_path 'outputs/'$EXP'/relight/video_relight_0'\ 45 | --split 'test'\ 46 | --ldr_img_dir $LDR_IMG_DIR \ 47 | --light_cfg 'configs/fipt/bathroom/relight_0.yaml' \ 48 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS -------------------------------------------------------------------------------- /scripts/fipt/conferenceroom/render.sh: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | # data folder 8 | DATASET_ROOT='/hdd/datasets/fipt/real/' 9 | DATASET='real' 10 | # scene name 11 | SCENE='conferenceroom' 12 | LDR_IMG_DIR='Image' 13 | EXP='fipt_real_conferenceroom' 14 | VAL_FRAME=0 15 | CRF_BASIS=3 16 | # whether has part segmentation 17 | HAS_PART=0 18 | SPP=256 19 | spp=16 20 | 21 | python render.py --experiment_name $EXP --device 0\ 22 | --ckpt last_1.ckpt \ 23 | --dataset $DATASET $DATASET_ROOT$SCENE \ 24 | --emitter_path checkpoints/$EXP/bake\ 25 | --output_path 'outputs/'$EXP'/output'\ 26 | --split 'val'\ 27 | --ldr_img_dir $LDR_IMG_DIR \ 28 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS 29 | 30 | python render_video.py --experiment_name $EXP --device 0\ 31 | --ckpt last_1.ckpt \ 32 | --dataset $DATASET $DATASET_ROOT$SCENE\ 33 | --emitter_path checkpoints/$EXP/bake\ 34 | --output_path 'outputs/'$EXP'/video'\ 35 | --split 'val'\ 36 | --ldr_img_dir $LDR_IMG_DIR \ 37 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS 38 | 39 | # relighting 40 | python render_relight.py --experiment_name $EXP --device 0\ 41 | --ckpt last_1.ckpt --mode traj\ 42 | --dataset $DATASET $DATASET_ROOT$SCENE\ 43 | --emitter_path checkpoints/$EXP/bake\ 44 | --output_path 'outputs/'$EXP'/relight/video_relight_0'\ 45 | --split 'test'\ 46 | --ldr_img_dir $LDR_IMG_DIR \ 47 | --light_cfg 'configs/fipt/conferenceroom/relight_0.yaml' \ 48 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS -------------------------------------------------------------------------------- /configs/fipt/bathroom/insert.yaml: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | type: 'scene' 8 | 9 | PerspectiveCamera: 10 | type: 'perspective' 11 | fov: 45 12 | film: 13 | type: 'hdrfilm' 14 | width: 1920 15 | height: 1080 16 | rfilter: 17 | type: 'box' 18 | 19 | Integrator: 20 | type: 'path' 21 | max_depth: 7 22 | 23 | main_scene: 24 | type: 'ply' 25 | filename: '' 26 | bsdf: 27 | type: 'twosided' 28 | fipt_bsdf: 29 | type: 'fipt' 30 | emitter_path: '' 31 | brdf_path: '' 32 | 33 | light_plane: 34 | type: 'rectangle' 35 | to_world: 36 | - type: 'translate' 37 | value: [-1.362, 2.7918, -2.2202] 38 | - type: 'scale' 39 | value: [0.865, 0.865, 0.08] 40 | - type: 'rotate' 41 | axis: [1, 0, 0] 42 | angle: 90 43 | bsdf: 44 | type: 'twosided' 45 | bsdf: 46 | type: 'diffuse' 47 | reflectance: 48 | type: 'rgb' 49 | value: [0., 0., 0.] 50 | emitter: 51 | type: 'area' 52 | radiance: 53 | type: 'rgb' 54 | value: [22, 25, 25] 55 | 56 | sphere_2: 57 | type: 'sphere' 58 | to_world: 59 | - type: 'translate' 60 | value: [-1.31264, 0.87, -2.1] 61 | - type: 'scale' 62 | value: [0.1, 0.1, 0.1] 63 | bsdf: 64 | type: 'conductor' 65 | material: 'none' 66 | 67 | sphere_3: 68 | type: 'sphere' 69 | to_world: 70 | - type: 'translate' 71 | value: [-0.380948, 0.87, -2.1] 72 | - type: 'scale' 73 | value: [0.1, 0.1, 0.1] 74 | bsdf: 75 | type: 'diffuse' 76 | reflectance: 77 | type: 'rgb' 78 | value: [0.2, 0.25, 0.7] -------------------------------------------------------------------------------- /configs/fipt/bedroom/insert.yaml: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | type: 'scene' 8 | 9 | PerspectiveCamera: 10 | type: 'perspective' 11 | fov: 45 12 | film: 13 | type: 'hdrfilm' 14 | width: 1920 15 | height: 1080 16 | rfilter: 17 | type: 'box' 18 | 19 | Integrator: 20 | type: 'path' 21 | max_depth: 7 22 | 23 | main_scene: 24 | type: 'ply' 25 | filename: '' 26 | bsdf: 27 | type: 'twosided' 28 | fipt_bsdf: 29 | type: 'fipt' 30 | emitter_path: '' 31 | brdf_path: '' 32 | 33 | light_plane: 34 | type: 'rectangle' 35 | to_world: 36 | - type: 'translate' 37 | value: [-1.6257, 2.3767, 1.5447] 38 | - type: 'scale' 39 | value: [0.197,1,0.589] 40 | - type: 'rotate' 41 | axis: [1, 0, 0] 42 | angle: 90 43 | bsdf: 44 | type: 'twosided' 45 | bsdf: 46 | type: 'diffuse' 47 | reflectance: 48 | type: 'rgb' 49 | value: [0., 0., 0.] 50 | emitter: 51 | type: 'area' 52 | radiance: 53 | type: 'rgb' 54 | value: [20, 23, 21] 55 | 56 | sphere_2: 57 | type: 'sphere' 58 | to_world: 59 | - type: 'translate' 60 | value: [-0.374835, 0.665015, 0.420669] 61 | - type: 'scale' 62 | value: [0.2, 0.2, 0.2] 63 | bsdf: 64 | type: 'conductor' 65 | material: 'none' 66 | 67 | sphere_3: 68 | type: 'sphere' 69 | to_world: 70 | - type: 'translate' 71 | value: [0.337497, 0.665015, 0] 72 | - type: 'scale' 73 | value: [0.2, 0.2, 0.2] 74 | bsdf: 75 | type: 'diffuse' 76 | reflectance: 77 | type: 'rgb' 78 | value: [0.2, 0.25, 0.7] -------------------------------------------------------------------------------- /configs/fipt/classroom/insert.yaml: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | type: 'scene' 8 | 9 | PerspectiveCamera: 10 | type: 'perspective' 11 | fov: 45 12 | film: 13 | type: 'hdrfilm' 14 | width: 1920 15 | height: 1080 16 | rfilter: 17 | type: 'box' 18 | 19 | Integrator: 20 | type: 'path' 21 | max_depth: 7 22 | 23 | main_scene: 24 | type: 'ply' 25 | filename: '' 26 | bsdf: 27 | type: 'twosided' 28 | fipt_bsdf: 29 | type: 'fipt' 30 | emitter_path: '' 31 | brdf_path: '' 32 | 33 | light_plane: 34 | type: 'rectangle' 35 | to_world: 36 | - type: 'translate' 37 | value: [1.6643, 1.7467, 3.0747] 38 | - type: 'scale' 39 | value: [1.287,1,0.809] 40 | - type: 'rotate' 41 | axis: [1, 0, 0] 42 | angle: 90 43 | bsdf: 44 | type: 'twosided' 45 | bsdf: 46 | type: 'diffuse' 47 | reflectance: 48 | type: 'rgb' 49 | value: [0., 0., 0.] 50 | emitter: 51 | type: 'area' 52 | radiance: 53 | type: 'rgb' 54 | value: [15, 15, 18] 55 | 56 | sphere_2: 57 | type: 'sphere' 58 | to_world: 59 | - type: 'translate' 60 | value: [0.31224, -0.768295, 2.71416] 61 | - type: 'scale' 62 | value: [0.3, 0.3, 0.3] 63 | bsdf: 64 | type: 'conductor' 65 | material: 'none' 66 | 67 | sphere_3: 68 | type: 'sphere' 69 | to_world: 70 | - type: 'translate' 71 | value: [-1.1346, -0.768295, 2.71416] 72 | - type: 'scale' 73 | value: [0.3, 0.3, 0.3] 74 | bsdf: 75 | type: 'diffuse' 76 | reflectance: 77 | type: 'rgb' 78 | value: [0.2, 0.25, 0.7] -------------------------------------------------------------------------------- /scripts/scannetpp/room2/render.sh: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | # data folder 8 | DATASET_ROOT='/hdd/datasets/scannetpp/' 9 | DATASET='scannetpp' 10 | # scene name 11 | SCENE='7e09430da7' 12 | LDR_IMG_DIR='Image' 13 | EXP='scannetpp_room2' 14 | VAL_FRAME=3 15 | CRF_BASIS=3 16 | RES_SCALE=0.5 17 | # whether has part segmentation 18 | HAS_PART=0 19 | SPP=256 20 | spp=16 21 | 22 | # render 23 | python render.py --experiment_name $EXP --device 0\ 24 | --ckpt last_1.ckpt \ 25 | --dataset $DATASET $DATASET_ROOT --scene $SCENE\ 26 | --res_scale $RES_SCALE\ 27 | --emitter_path checkpoints/$EXP/bake\ 28 | --output_path 'outputs/'$EXP'/output'\ 29 | --split 'test'\ 30 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS 31 | 32 | python render_video.py --experiment_name $EXP --device 0\ 33 | --ckpt last_1.ckpt \ 34 | --dataset $DATASET $DATASET_ROOT --scene $SCENE \ 35 | --res_scale $RES_SCALE\ 36 | --emitter_path checkpoints/$EXP/bake\ 37 | --output_path 'outputs/'$EXP'/video'\ 38 | --split 'test'\ 39 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS 40 | 41 | # relighting 42 | python render_relight.py --experiment_name $EXP --device 0\ 43 | --ckpt last_1.ckpt --mode traj\ 44 | --dataset $DATASET $DATASET_ROOT --scene $SCENE \ 45 | --res_scale $RES_SCALE \ 46 | --emitter_path checkpoints/$EXP/bake\ 47 | --output_path 'outputs/'$EXP'/relight/video_relight_0'\ 48 | --split 'test'\ 49 | --light_cfg 'configs/scannetpp/room2/relight_0.yaml' \ 50 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS -------------------------------------------------------------------------------- /scripts/scannetpp/bathroom2/render.sh: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | # data folder 8 | DATASET_ROOT='/hdd/datasets/scannetpp/' 9 | DATASET='scannetpp' 10 | # scene name 11 | SCENE='45b0dac5e3' 12 | LDR_IMG_DIR='Image' 13 | EXP='scannetpp_bathroom2' 14 | VAL_FRAME=0 15 | CRF_BASIS=3 16 | RES_SCALE=0.5 17 | # whether has part segmentation 18 | HAS_PART=0 19 | SPP=256 20 | spp=16 21 | 22 | # render 23 | python render.py --experiment_name $EXP --device 0\ 24 | --ckpt last_1.ckpt \ 25 | --dataset $DATASET $DATASET_ROOT --scene $SCENE\ 26 | --res_scale $RES_SCALE\ 27 | --emitter_path checkpoints/$EXP/bake\ 28 | --output_path 'outputs/'$EXP'/output'\ 29 | --split 'test'\ 30 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS 31 | 32 | python render_video.py --experiment_name $EXP --device 0\ 33 | --ckpt last_1.ckpt \ 34 | --dataset $DATASET $DATASET_ROOT --scene $SCENE \ 35 | --res_scale $RES_SCALE\ 36 | --emitter_path checkpoints/$EXP/bake\ 37 | --output_path 'outputs/'$EXP'/video'\ 38 | --split 'test'\ 39 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS 40 | 41 | # relighting 42 | python render_relight.py --experiment_name $EXP --device 0\ 43 | --ckpt last_1.ckpt --mode traj\ 44 | --dataset $DATASET $DATASET_ROOT --scene $SCENE \ 45 | --res_scale $RES_SCALE \ 46 | --emitter_path checkpoints/$EXP/bake\ 47 | --output_path 'outputs/'$EXP'/relight/video_relight_0'\ 48 | --split 'test'\ 49 | --light_cfg 'configs/scannetpp/bathroom2/relight_0.yaml' \ 50 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS -------------------------------------------------------------------------------- /scripts/fipt/livingroom/render.sh: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | # data folder 8 | DATASET_ROOT='/hdd/datasets/fipt/indoor_synthetic/' 9 | DATASET='synthetic' 10 | # scene name 11 | SCENE='livingroom' 12 | LDR_IMG_DIR='Image' 13 | EXP='fipt_syn_livingroom' 14 | VAL_FRAME=10 # bathroom: 10, bedroom: 10, livingroom:10, classroom:64, conferenceroom: 4 15 | CRF_BASIS=3 16 | # whether has part segmentation 17 | HAS_PART=1 18 | SPP=256 19 | spp=16 20 | 21 | python render.py --experiment_name $EXP --device 0\ 22 | --ckpt last_1.ckpt \ 23 | --dataset $DATASET $DATASET_ROOT$SCENE \ 24 | --emitter_path checkpoints/$EXP/bake\ 25 | --output_path 'outputs/'$EXP'/output'\ 26 | --split 'val'\ 27 | --ldr_img_dir $LDR_IMG_DIR \ 28 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS 29 | 30 | python render_video.py --experiment_name $EXP --device 0\ 31 | --ckpt last_1.ckpt \ 32 | --dataset $DATASET $DATASET_ROOT$SCENE\ 33 | --emitter_path checkpoints/$EXP/bake\ 34 | --output_path 'outputs/'$EXP'/video'\ 35 | --split 'val'\ 36 | --ldr_img_dir $LDR_IMG_DIR \ 37 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS 38 | 39 | # relighting 40 | python render_relight.py --experiment_name $EXP --device 0\ 41 | --ckpt last_1.ckpt --mode traj\ 42 | --dataset $DATASET $DATASET_ROOT$SCENE\ 43 | --emitter_path checkpoints/$EXP/bake\ 44 | --output_path 'outputs/'$EXP'/relight/video_relight_0'\ 45 | --split 'test'\ 46 | --ldr_img_dir $LDR_IMG_DIR \ 47 | --light_cfg 'configs/fipt/livingroom/relight_0.yaml' \ 48 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS -------------------------------------------------------------------------------- /utils/extract_emitter_mesh.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import os 8 | import numpy as np 9 | import trimesh 10 | import torch 11 | from argparse import ArgumentParser 12 | 13 | def main(): 14 | parser = ArgumentParser() 15 | parser.add_argument('--mesh_scene', type=str, required=True) 16 | parser.add_argument('--emitter', type=str, required=True) 17 | parser.add_argument('--mesh_emitter', type=str, required=True) 18 | args = parser.parse_args() 19 | 20 | mesh = trimesh.load(args.mesh_scene) 21 | vertices = np.array(mesh.vertices).astype(np.float32) 22 | faces = np.array(mesh.faces).astype(np.int32) 23 | 24 | emitter = torch.load(args.emitter) 25 | is_emitter = emitter['is_emitter'].numpy() 26 | emitter_vertices = emitter['emitter_vertices'].numpy() 27 | emitter_area = emitter['emitter_area'].numpy() 28 | emitter_radiance = emitter['emitter_radiance'].numpy() 29 | 30 | # calculate average radiance weighted by area 31 | emitter_radiance = emitter_radiance[:np.sum(is_emitter)] 32 | emission_avg = np.sum(emitter_radiance * emitter_area[..., None], axis=0) / np.sum(emitter_area) 33 | print('Average radiance:', emission_avg) 34 | 35 | # Export emitter mesh 36 | ef = faces[is_emitter] 37 | ev = vertices[ef] 38 | n_ef = ef.shape[0] 39 | ef = ef.flatten() 40 | u, i = np.unique(ef, return_inverse=True) 41 | # print(np.abs(u[i] - ef).max()) 42 | 43 | v_new = vertices[u] 44 | f_new = i.reshape(n_ef, 3).astype(np.int32) 45 | emitter_mesh = trimesh.Trimesh(vertices=v_new, faces=f_new) 46 | emitter_mesh.export(args.mesh_emitter) 47 | print('Exported emitter mesh to', args.mesh_emitter) 48 | 49 | 50 | if __name__ == '__main__': 51 | main() -------------------------------------------------------------------------------- /crf/plot.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import os 8 | import torch 9 | import numpy as np 10 | from matplotlib import pyplot as plt 11 | 12 | def plot_crfs(crfs_pred, crfs_gt, path): 13 | crfs_pred = crfs_pred.detach().cpu().numpy() 14 | fig, (ax0, ax1, ax2) = plt.subplots(1, 3, figsize=(12, 4)) 15 | x = np.linspace(0, 1, crfs_gt.shape[1]) 16 | ax0.title.set_text('CRF (R)') 17 | ax0.set_ylabel('Pixel intensity') 18 | ax0.set_xlabel('Irradiance') 19 | ax0.plot(x, crfs_gt[0], c='b') 20 | ax0.plot(x, crfs_pred[0], c='r') 21 | ax1.title.set_text('CRF (G)') 22 | ax1.set_xlabel('Irradiance') 23 | ax1.plot(x, crfs_gt[1], c='b') 24 | ax1.plot(x, crfs_pred[1], c='r') 25 | ax2.title.set_text('CRF (B)') 26 | ax2.set_xlabel('Irradiance') 27 | ax2.plot(x, crfs_gt[2], c='b', label='GT') 28 | ax2.plot(x, crfs_pred[2], c='r', label='pred') 29 | ax2.legend() 30 | plt.savefig(path) 31 | plt.close() 32 | 33 | def plot_weights(weights_pred, weights_gt, path): 34 | weights_pred = weights_pred.detach().cpu().numpy() 35 | fig, (ax0, ax1, ax2) = plt.subplots(1, 3, figsize=(12, 4)) 36 | x = np.linspace(0, 1, weights_gt.shape[1]) 37 | interval = 0.05 38 | ax0.title.set_text('Basis coefficients (R)') 39 | ax0.set_ylabel('value') 40 | ax0.bar(x-interval, weights_gt[0], interval*2) 41 | ax0.bar(x+interval, weights_pred[0], interval*2) 42 | ax1.title.set_text('Basis coefficients (G)') 43 | ax1.bar(x-interval, weights_gt[1], interval*2) 44 | ax1.bar(x+interval, weights_pred[1], interval*2) 45 | ax2.title.set_text('Basis coefficients (B)') 46 | ax2.bar(x-interval, weights_gt[2], interval*2, label='GT') 47 | ax2.bar(x+interval, weights_pred[2], interval*2, label='pred') 48 | ax2.legend() 49 | plt.savefig(path) 50 | plt.close() 51 | -------------------------------------------------------------------------------- /configs/scannetpp/room2/insert.yaml: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | type: 'scene' 8 | 9 | PerspectiveCamera: 10 | type: 'perspective' 11 | fov: 45 12 | film: 13 | type: 'hdrfilm' 14 | width: 1920 15 | height: 1080 16 | rfilter: 17 | type: 'box' 18 | 19 | Integrator: 20 | type: 'path' 21 | max_depth: 7 22 | 23 | main_scene: 24 | type: 'ply' 25 | filename: '' 26 | bsdf: 27 | type: 'twosided' 28 | fipt_bsdf: 29 | type: 'fipt' 30 | emitter_path: '' 31 | brdf_path: '' 32 | 33 | emitter_room2: 34 | type: 'ply' 35 | filename: 'checkpoints/scannetpp_room2/bake/emitter.ply' 36 | to_world: 37 | - type: 'translate' 38 | value: [0.0, 0.0, 0.0] 39 | - type: 'scale' 40 | value: [1.0, 1.0, 1.0] 41 | bsdf: 42 | type: 'diffuse' 43 | reflectance: 44 | type: 'rgb' 45 | value: [0., 0., 0.] 46 | emitter: 47 | type: 'area' 48 | radiance: 49 | type: 'rgb' 50 | value: [8.679976, 8.248041, 2.860989] 51 | 52 | andersen: 53 | type: 'obj' 54 | filename: 'outputs/00_assets/hans-christian-andersen/source/HansChristianAndersen/HansChristianAndersen-80k.obj' 55 | to_world: 56 | - type: 'translate' 57 | value: [-0.192385 , -0.926694, 0.643368] 58 | - type: 'scale' 59 | value: [0.015, 0.015, 0.015] 60 | - type: 'rotate' 61 | axis: [1, 0, 0] 62 | angle: 180 63 | bsdf: 64 | type: 'roughconductor' 65 | distribution: 'ggx' 66 | alpha_u: 0.05 67 | alpha_v: 0.3 68 | eta: 69 | type: 'rgb' 70 | value: [0.47, 0.35, 0.29] 71 | k: 72 | type: 'rgb' 73 | value: [0.332, 0.239, 0.235] 74 | 75 | bunny: 76 | type: 'obj' 77 | filename: 'outputs/00_assets/stanford-bunny.obj' 78 | to_world: 79 | - type: 'translate' 80 | value: [-0.515105, -0.626611, 0.3113] 81 | - type: 'scale' 82 | value: [2, 2, 2] 83 | - type: 'rotate' 84 | axis: [1, 0, 0] 85 | angle: -90 86 | - type: 'rotate' 87 | axis: [0, 1, 0] 88 | angle: 90 89 | bsdf: 90 | type: 'conductor' 91 | material: 'Cu' -------------------------------------------------------------------------------- /configs/fipt/conferenceroom/insert.yaml: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | type: 'scene' 8 | 9 | PerspectiveCamera: 10 | type: 'perspective' 11 | fov: 45 12 | film: 13 | type: 'hdrfilm' 14 | width: 1920 15 | height: 1080 16 | rfilter: 17 | type: 'box' 18 | 19 | Integrator: 20 | type: 'path' 21 | max_depth: 7 22 | 23 | main_scene: 24 | type: 'ply' 25 | filename: '' 26 | bsdf: 27 | type: 'twosided' 28 | fipt_bsdf: 29 | type: 'fipt' 30 | emitter_path: '' 31 | brdf_path: '' 32 | 33 | emitter_conferenceroom: 34 | type: 'ply' 35 | filename: 'checkpoints/fipt_real_conferenceroom/bake/emitter.ply' 36 | to_world: 37 | - type: 'translate' 38 | value: [0.0, 0.0, 0.0] 39 | - type: 'scale' 40 | value: [1.0, 1.0, 1.0] 41 | bsdf: 42 | type: 'diffuse' 43 | reflectance: 44 | type: 'rgb' 45 | value: [0., 0., 0.] 46 | emitter: 47 | type: 'area' 48 | radiance: 49 | type: 'rgb' 50 | value: [5, 5, 5.4] 51 | 52 | teapot: 53 | type: 'obj' 54 | filename: 'outputs/00_assets/teapot.obj' 55 | to_world: 56 | - type: 'translate' 57 | value: [-1.37901, 0.303306, -0.869783] 58 | - type: 'scale' 59 | value: [0.15, 0.15, 0.15] 60 | - type: 'rotate' 61 | axis: [1, 0, 0] 62 | angle: 180 63 | - type: 'rotate' 64 | axis: [0, 1, 0] 65 | angle: -45 66 | bsdf: 67 | type: 'conductor' 68 | material: 'Ag' 69 | 70 | lucy: 71 | type: 'ply' 72 | filename: 'outputs/00_assets/lucy.ply' 73 | to_world: 74 | - type: 'translate' 75 | value: [-1.00135, -0.148423, -1.13769] 76 | - type: 'scale' 77 | value: [0.001, 0.001, 0.001] 78 | - type: 'rotate' 79 | axis: [1, 0, 0] 80 | angle: 90 81 | - type: 'rotate' 82 | axis: [0, 1, 0] 83 | angle: -12 84 | - type: 'rotate' 85 | axis: [0, 0, 1] 86 | angle: 90 87 | bsdf: 88 | type: 'roughconductor' 89 | distribution: 'ggx' 90 | alpha_u: 0.05 91 | alpha_v: 0.3 92 | eta: 93 | type: 'rgb' 94 | value: [0.47, 0.35, 0.29] 95 | k: 96 | type: 'rgb' 97 | value: [0.332, 0.239, 0.235] -------------------------------------------------------------------------------- /configs/fipt/bathroom/relight_1.yaml: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | type: 'scene' 8 | 9 | PerspectiveCamera: 10 | type: 'perspective' 11 | fov: 45 12 | film: 13 | type: 'hdrfilm' 14 | width: 1920 15 | height: 1080 16 | rfilter: 17 | type: 'box' 18 | 19 | Integrator: 20 | type: 'path' 21 | max_depth: 7 22 | 23 | main_scene: 24 | type: 'ply' 25 | filename: '' 26 | bsdf: 27 | type: 'twosided' 28 | fipt_bsdf: 29 | type: 'fipt' 30 | emitter_path: '' 31 | brdf_path: '' 32 | 33 | light_ball_0: 34 | type: 'sphere' 35 | to_world: 36 | - type: 'translate' 37 | value: [-1.72946, 1.95877, -2.19141] 38 | - type: 'scale' 39 | value: [0.05, 0.05, 0.05] 40 | bsdf: 41 | type: 'diffuse' 42 | reflectance: 43 | type: 'rgb' 44 | value: [0., 0., 0.] 45 | emitter: 46 | type: 'area' 47 | radiance: 48 | type: 'rgb' 49 | value: [18, 16, 15] 50 | 51 | light_ball_1: 52 | type: 'sphere' 53 | to_world: 54 | - type: 'translate' 55 | value: [-1.48904, 1.95877, -2.19659] 56 | - type: 'scale' 57 | value: [0.05, 0.05, 0.05] 58 | bsdf: 59 | type: 'diffuse' 60 | reflectance: 61 | type: 'rgb' 62 | value: [0., 0., 0.] 63 | emitter: 64 | type: 'area' 65 | radiance: 66 | type: 'rgb' 67 | value: [15, 15, 12] 68 | 69 | light_ball_2: 70 | type: 'sphere' 71 | to_world: 72 | - type: 'translate' 73 | value: [-1.23297, 1.96617, -2.20512] 74 | - type: 'scale' 75 | value: [0.05, 0.05, 0.05] 76 | bsdf: 77 | type: 'diffuse' 78 | reflectance: 79 | type: 'rgb' 80 | value: [0., 0., 0.] 81 | emitter: 82 | type: 'area' 83 | radiance: 84 | type: 'rgb' 85 | value: [18, 16, 15] 86 | 87 | light_ball_3: 88 | type: 'sphere' 89 | to_world: 90 | - type: 'translate' 91 | value: [-0.982932, 1.96617, -2.21057] 92 | - type: 'scale' 93 | value: [0.05, 0.05, 0.05] 94 | bsdf: 95 | type: 'diffuse' 96 | reflectance: 97 | type: 'rgb' 98 | value: [0., 0., 0.] 99 | emitter: 100 | type: 'area' 101 | radiance: 102 | type: 'rgb' 103 | value: [15, 15, 12] -------------------------------------------------------------------------------- /crf/emor.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import numpy as np 8 | import os 9 | import pathlib 10 | from matplotlib import pyplot as plt 11 | from const import set_random_seed 12 | set_random_seed() 13 | 14 | repo_root = pathlib.Path().resolve() 15 | emor_path = os.path.join(repo_root, 'crf', 'emor.txt') 16 | invemor_path = os.path.join(repo_root, 'crf', 'invemor.txt') 17 | curves_path = os.path.join(repo_root, 'crf', 'dorfCurves.txt') 18 | 19 | def parse_emor_file(inv=True): 20 | file_path = invemor_path if inv else emor_path 21 | with open(file_path, 'r') as file: 22 | lines = file.readlines() 23 | lines = [line.strip() for line in lines] 24 | 25 | stride = 1 + 256 26 | names = [] 27 | vectors = [] 28 | for i in range(len(lines)//stride): 29 | name_line = lines[i*stride] 30 | name = name_line.split('=')[0].strip() 31 | names.append(name) 32 | lines_numbers = lines[i*stride+1: (i+1)*stride] 33 | numbers = [line.split() for line in lines_numbers] 34 | vector = np.float32(numbers).reshape(-1) 35 | vectors.append(vector) 36 | names = np.array(names) 37 | vectors = np.stack(vectors) 38 | return names, vectors 39 | 40 | def parse_dorf_curves(): 41 | with open(curves_path, 'r') as file: 42 | lines = file.readlines() 43 | lines = [line.strip() for line in lines] 44 | stride = 6 45 | names = [] 46 | vectors = [] 47 | for i in range(len(lines)//stride): 48 | line_sample = lines[i*stride: (i+1)*stride] 49 | n_i = '{}-{}-{}'.format(line_sample[0], line_sample[1], line_sample[2][0]) 50 | n_b = '{}-{}-{}'.format(line_sample[0], line_sample[1], line_sample[4][0]) 51 | names += [n_b] 52 | v_i = np.float32(line_sample[3].split()) 53 | v_b = np.float32(line_sample[5].split()) 54 | vectors += [v_b] 55 | names = np.array(names) 56 | vectors = np.stack(vectors) 57 | return names, vectors #(201, 1024) 58 | 59 | def get_dorf_mean_basis(top=25): 60 | names, curves = parse_dorf_curves() 61 | mean = np.mean(curves, 0) 62 | curves = curves - mean[None] 63 | u, s, vh = np.linalg.svd(curves) 64 | scaled_basis = s[:top, None] * vh[:top] 65 | # scaled_basis = vh[:top] 66 | return mean, scaled_basis 67 | 68 | def mono_increase_constraint(crf): 69 | diff = crf[1:] - crf[:-1] 70 | gap = -1 * np.min([0.0, diff.min()]) 71 | diff += gap 72 | diff /= diff.sum() 73 | crf = np.cumsum(diff) 74 | crf = np.concatenate([np.zeros((1)), crf]) 75 | return crf 76 | 77 | -------------------------------------------------------------------------------- /model/mlps.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import torch 8 | import torch.nn as nn 9 | import math 10 | from const import set_random_seed 11 | set_random_seed() 12 | 13 | """ mlps """ 14 | 15 | class PositionalEncoding(nn.Module): 16 | def __init__(self, L): 17 | """ L: number of frequency bands """ 18 | super(PositionalEncoding, self).__init__() 19 | self.L= L 20 | 21 | def forward(self, inputs): 22 | L = self.L 23 | encoded = [inputs] 24 | for l in range(L): 25 | encoded.append(torch.sin((2 ** l * math.pi) * inputs)) 26 | encoded.append(torch.cos((2 ** l * math.pi) * inputs)) 27 | return torch.cat(encoded, -1) 28 | 29 | 30 | def mlp(dim_in, dims, dim_out): 31 | """ create an MLP in format: dim_in->dims[0]->...->dims[-1]->dim_out""" 32 | lists = [] 33 | dims = [dim_in] + dims 34 | 35 | for i in range(len(dims)-1): 36 | lists.append(nn.Linear(dims[i],dims[i+1])) 37 | lists.append(nn.ReLU(inplace=True)) 38 | lists.append(nn.Linear(dims[-1], dim_out)) 39 | 40 | return nn.Sequential(*lists) 41 | 42 | 43 | class ImplicitMLP(nn.Module): 44 | """ NeRF style coordinate based MLP """ 45 | def __init__(self, D, C, S, dim_out, dim_enc,dim_in=3): 46 | """ 47 | Args: 48 | D: depth of the MLP 49 | C: intermediate feature dim 50 | S: skip layers 51 | dim_out: out feature dimension 52 | dim_enc: frequency band of positional encdoing for the vertex location 53 | """ 54 | super(ImplicitMLP, self).__init__() 55 | self.input_ch = dim_enc * 2 * dim_in + dim_in 56 | 57 | self.D = D 58 | self.C = C 59 | self.skips = S 60 | self.dim_out = dim_out 61 | 62 | self.point_encode = PositionalEncoding(dim_enc) 63 | 64 | self.pts_linears = nn.ModuleList( 65 | [nn.Linear(self.input_ch, self.C)] + 66 | [ 67 | nn.Linear(self.C, self.C) 68 | if i not in self.skips 69 | else nn.Linear(self.C + self.input_ch, self.C) 70 | for i in range(1, self.D) 71 | ] 72 | ) 73 | self.feature_linear = nn.Linear(self.C, self.dim_out) 74 | 75 | def forward(self, x): 76 | """ 77 | Args: 78 | x: Bx3 3D vertex locations 79 | Return: 80 | Bx3 corresponding feature vectors 81 | """ 82 | points = self.point_encode(x) 83 | p = points 84 | for i, l in enumerate(self.pts_linears): 85 | if i in self.skips: 86 | p = torch.cat([points, p], 1) 87 | p = l(p) 88 | p = torch.relu(p) 89 | feature = self.feature_linear(p) 90 | return feature -------------------------------------------------------------------------------- /model/slf.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import torch 8 | import torch.nn as nn 9 | import torch.nn.functional as NF 10 | from const import set_random_seed 11 | set_random_seed() 12 | 13 | """ surface light field model """ 14 | 15 | 16 | class VoxelSLF(nn.Module): 17 | """ voxel grid based surface light field """ 18 | def __init__(self,mask,voxel_min,voxel_max): 19 | """ 20 | mask: NxNxN voxel occupancy mask 21 | voxel_min,voxel_max: voxel bounding box 22 | """ 23 | super(VoxelSLF,self).__init__() 24 | H = mask.shape[0] 25 | self.H = H 26 | self.voxel_min = voxel_min 27 | self.voxel_max = voxel_max 28 | 29 | # find coordinates for occupied voxels 30 | kk,jj,ii = torch.where(mask) 31 | inds = -torch.ones(H,H,H,dtype=torch.long) 32 | inds[kk,jj,ii] = torch.arange(len(ii)) 33 | 34 | 35 | self.register_buffer('inds',inds) 36 | self.register_buffer( 37 | 'radiance', torch.zeros(len(ii),3)*1e-1) 38 | self.register_buffer( 39 | 'count', torch.zeros(len(ii),dtype=torch.long)) # number of entries, used for mean pooling 40 | 41 | def spatial_idx(self,x): 42 | """ get voxel entry index for input location 43 | Args: 44 | x: Bx3 3D position 45 | Return: 46 | B indices 47 | """ 48 | # map to voxel grid coordinates 49 | x_ = (x-self.voxel_min)/(self.voxel_max-self.voxel_min) 50 | x_ = (x_*self.H).long().clamp(0,self.H-1) 51 | 52 | # find entry indices 53 | idx = self.inds[x_[...,2],x_[...,1],x_[...,0]] 54 | return idx 55 | 56 | def scatter_add(self,x,radiance): 57 | """ scatter add radiance to voxel grid 58 | """ 59 | idx = self.spatial_idx(x) 60 | self.radiance.scatter_add_(0,idx[...,None].expand_as(radiance),radiance) 61 | self.count.scatter_add_(0,idx,torch.ones_like(idx)) 62 | 63 | def forward(self,x): 64 | """ query surface light field """ 65 | idx = self.spatial_idx(x) 66 | radiance = self.radiance[idx] 67 | radiance[idx==-1] = 0 # if hit empty space, return zero radiance 68 | return { 69 | 'rgb': radiance 70 | } 71 | 72 | 73 | class TextureSLF(nn.Module): 74 | """ textured mesh based surface light field (unused) """ 75 | def __init__(self,res,texture=None,Co=3): 76 | super(ExplicitSLF,self).__init__() 77 | self.res = res 78 | self.Co = Co 79 | if texture is None: 80 | texture = torch.randn(Co,res,res)*0.1 + 0.5 81 | self.register_parameter('feature',nn.Parameter(texture)) 82 | 83 | def texture(self,uv): 84 | """ uv: Bx2""" 85 | uv = uv*2-1 86 | B,_ = uv.shape 87 | feat = NF.grid_sample(self.feature[None],uv.reshape(1,B,1,2), 88 | mode='bilinear',align_corners=True).reshape(self.Co,B).T 89 | return feat 90 | 91 | def forward(self,uv,wo): 92 | feat = self.texture(uv) 93 | rgb = feat[...,:3] 94 | return { 95 | 'rgb': rgb 96 | } -------------------------------------------------------------------------------- /utils/ops.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import torch 8 | import torch.nn.functional as NF 9 | import math 10 | 11 | 12 | def get_normal_space(normal): 13 | """ get matrix transform shading space to normal spanned space 14 | Args: 15 | normal: Bx3 16 | Return: 17 | Bx3x3 transformation matrix 18 | """ 19 | v1 = torch.zeros_like(normal) 20 | tangent = v1.clone() 21 | v1[...,0] = 1.0 22 | tangent[...,1] = 1.0 23 | 24 | mask = (v1*normal).sum(-1).abs() <= 1e-1 25 | tangent[mask] = NF.normalize(torch.cross(v1[mask],normal[mask],dim=-1),dim=-1) 26 | mask = ~mask 27 | tangent[mask] = NF.normalize(torch.cross(tangent[mask],normal[mask],dim=-1),dim=-1) 28 | 29 | bitangent = torch.cross(normal,tangent,dim=-1) 30 | return torch.stack([tangent,bitangent,normal],dim=-1) 31 | 32 | def angle2xyz(theta,phi): 33 | """ spherical coordinates to euclidean 34 | Args: 35 | theta,phi: B 36 | Return: 37 | Bx3 euclidean coordinates 38 | """ 39 | sin_theta = torch.sin(theta) 40 | x = sin_theta*torch.cos(phi) 41 | y = sin_theta*torch.sin(phi) 42 | z = torch.cos(theta) 43 | ret = torch.stack([x,y,z],dim=-1) 44 | return NF.normalize(ret,dim=-1) 45 | 46 | def G1_GGX_Schlick(NoV, eta): 47 | """ G term of schlick GGX 48 | eta: roughness 49 | """ 50 | r = eta 51 | k = (r+1) 52 | k = k*k/8 53 | denom = NoV*(1-k)+k 54 | return 1 /denom 55 | 56 | def G_Smith(NoV,NoL,eta): 57 | """ Smith shadow masking divided by (NoV*NoL) 58 | eta: roughness 59 | """ 60 | g1_l = G1_GGX_Schlick(NoL,eta) 61 | g1_v = G1_GGX_Schlick(NoV,eta) 62 | return g1_l*g1_v 63 | 64 | def fresnelSchlick(VoH,F0): 65 | """ schlick fresnel """ 66 | x = (1-VoH).pow(5) 67 | return F0 + (1-F0)*x 68 | 69 | def fresnelSchlick_sep(VoH): 70 | """ two terms of schlick fresnel """ 71 | x = (1-VoH).pow(5) 72 | return (1-x),x 73 | 74 | def D_GGX(cos_h,eta): 75 | """GGX normal distribution 76 | eta: roughness 77 | """ 78 | alpha = eta*eta 79 | alpha2 = alpha*alpha 80 | denom = (cos_h*cos_h*(alpha2-1.0)+1.0) 81 | denom = math.pi * denom*denom 82 | return alpha2/denom 83 | 84 | 85 | def double_sided(V,N): 86 | """ double sided normal 87 | Args: 88 | V: Bx3 viewing direction 89 | N: Bx3 normal direction 90 | Return: 91 | Bx3 flipped normal towards camera direction 92 | """ 93 | NoV = (N*V).sum(-1) 94 | flipped = NoV<0 95 | N[flipped] = -N[flipped] 96 | return N 97 | 98 | 99 | def lerp_specular(specular,roughness): 100 | """ interpolate specular shadings by roughness 101 | Args: 102 | specular: Bx6x3 specular shadings 103 | roughness: Bx1 roughness in [0.02,1.0] 104 | Return: 105 | Bx3 interpolated specular shading 106 | """ 107 | # remap roughness from to [0,1] 108 | r_min,r_max = 0.02,1.0 109 | r_num = specular.shape[-2] 110 | r = (roughness-r_min)/(r_max-r_min)*(r_num-1) 111 | 112 | 113 | r1 = r.ceil().long() 114 | r0 = r.floor().long() 115 | r_ = (r-r0) 116 | s0 = torch.gather(specular,1,r0[...,None].expand(r0.shape[0],1,3))[:,0] 117 | s1 = torch.gather(specular,1,r1[...,None].expand(r1.shape[0],1,3))[:,0] 118 | s = s0*(1-r_) + s1*r_ 119 | return s -------------------------------------------------------------------------------- /utils/dataset/scannetpp/render_semantic.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import os 8 | import numpy as np 9 | import cv2 10 | import torch 11 | import mitsuba 12 | mitsuba.set_variant('cuda_ad_rgb') 13 | import trimesh 14 | from .dataset_dslr import Scannetpp 15 | from utils.path_tracing import ray_intersect 16 | from tqdm import tqdm 17 | import argparse 18 | from const import set_random_seed 19 | set_random_seed() 20 | 21 | def main(): 22 | parser = argparse.ArgumentParser() 23 | parser.add_argument('--data_root', default='/hdd/datasets/scannetpp/') 24 | parser.add_argument('--scene') 25 | args = parser.parse_args() 26 | dataset = Scannetpp(args.data_root, args.scene, split='all', pixel=False) 27 | 28 | mesh_path = os.path.join(args.data_root, 'data', args.scene, 'scans', 'mesh.ply') 29 | mesh_type = 'ply' 30 | scene = mitsuba.load_dict({ 31 | 'type': 'scene', 32 | 'shape_id':{ 33 | 'type': mesh_type, 34 | 'filename': mesh_path, 35 | } 36 | }) 37 | 38 | sem_mesh_path = os.path.join(args.data_root, 'data', args.scene, 'scans', 'mesh_aligned_0.05_semantic.ply') 39 | sem_mesh = trimesh.load(sem_mesh_path) 40 | vertices = np.array(sem_mesh.vertices) # (v, 3) 41 | faces = np.array(sem_mesh.faces) # (f, 3) 42 | 43 | vtx_colors = (sem_mesh.visual.vertex_colors).astype(np.int64) 44 | vtx_values = vtx_colors[:, 0] + vtx_colors[:, 1]*(2**8) + vtx_colors[:, 2]*(2**16) 45 | seg_idxs, inv_idxs = np.unique(vtx_values, return_inverse=True) 46 | # face_colors = (sem_mesh.visual.face_colors).astype(np.int64) 47 | # face_values = face_colors[:, 0] + face_colors[:, 1]*(2**8) + face_colors[:, 2]*(2**16) 48 | # seg_idxs, inv_idxs = np.unique(face_values, return_inverse=True) 49 | seg_idxs_new = np.arange(len(seg_idxs)) 50 | vtx_labels = seg_idxs_new[inv_idxs] 51 | 52 | 53 | img_name_list = dataset.img_name_list 54 | dir_sem_out = os.path.join(args.data_root, 'data', args.scene, 'dslr', 'seg') 55 | os.makedirs(dir_sem_out, exist_ok=True) 56 | 57 | h, w = dataset.img_hw 58 | device = torch.device(0) 59 | colors = np.random.rand(len(seg_idxs), 3) 60 | for i in tqdm(range(len(dataset))): 61 | name = img_name_list[i] 62 | batch = dataset[i] 63 | rays = batch['rays'] 64 | rays_x,rays_d = rays[...,:3].to(device),rays[...,3:6].to(device) 65 | _, _, _, face_idxs, valid = ray_intersect(scene,rays_x,rays_d) 66 | 67 | face_idxs = face_idxs.cpu().numpy() 68 | valid = valid.cpu().numpy() 69 | valid_face_idxs = face_idxs[valid] # (n, ) 70 | valid_faces = faces[valid_face_idxs] # (n, 3) 71 | 72 | valid_labels = vtx_labels[valid_faces] # (n, 3) 73 | valid_labels = np.max(valid_labels, axis=1) #choose max id as label 74 | labels = np.zeros((h*w)) 75 | labels[~valid] = -1 76 | labels[valid] = valid_labels 77 | labels = labels.reshape(h, w).astype(np.uint8) 78 | save_path = os.path.join(dir_sem_out, name) 79 | cv2.imwrite(save_path, labels) 80 | 81 | labels_vis = np.zeros((h*w, 3)) 82 | labels_vis[valid] = colors[valid_labels] 83 | labels_vis = (labels_vis*255).reshape(h, w, -1).astype(np.uint8) 84 | save_path = os.path.join(dir_sem_out, name.split('.')[0] + '_vis.jpg') 85 | cv2.imwrite(save_path, labels_vis) 86 | 87 | if __name__ == '__main__': 88 | main() -------------------------------------------------------------------------------- /CODE_OF_CONDUCT.md: -------------------------------------------------------------------------------- 1 | # Code of Conduct 2 | 3 | ## Our Pledge 4 | 5 | In the interest of fostering an open and welcoming environment, we as 6 | contributors and maintainers pledge to make participation in our project and 7 | our community a harassment-free experience for everyone, regardless of age, body 8 | size, disability, ethnicity, sex characteristics, gender identity and expression, 9 | level of experience, education, socio-economic status, nationality, personal 10 | appearance, race, religion, or sexual identity and orientation. 11 | 12 | ## Our Standards 13 | 14 | Examples of behavior that contributes to creating a positive environment 15 | include: 16 | 17 | * Using welcoming and inclusive language 18 | * Being respectful of differing viewpoints and experiences 19 | * Gracefully accepting constructive criticism 20 | * Focusing on what is best for the community 21 | * Showing empathy towards other community members 22 | 23 | Examples of unacceptable behavior by participants include: 24 | 25 | * The use of sexualized language or imagery and unwelcome sexual attention or 26 | advances 27 | * Trolling, insulting/derogatory comments, and personal or political attacks 28 | * Public or private harassment 29 | * Publishing others' private information, such as a physical or electronic 30 | address, without explicit permission 31 | * Other conduct which could reasonably be considered inappropriate in a 32 | professional setting 33 | 34 | ## Our Responsibilities 35 | 36 | Project maintainers are responsible for clarifying the standards of acceptable 37 | behavior and are expected to take appropriate and fair corrective action in 38 | response to any instances of unacceptable behavior. 39 | 40 | Project maintainers have the right and responsibility to remove, edit, or 41 | reject comments, commits, code, wiki edits, issues, and other contributions 42 | that are not aligned to this Code of Conduct, or to ban temporarily or 43 | permanently any contributor for other behaviors that they deem inappropriate, 44 | threatening, offensive, or harmful. 45 | 46 | ## Scope 47 | 48 | This Code of Conduct applies within all project spaces, and it also applies when 49 | an individual is representing the project or its community in public spaces. 50 | Examples of representing a project or community include using an official 51 | project e-mail address, posting via an official social media account, or acting 52 | as an appointed representative at an online or offline event. Representation of 53 | a project may be further defined and clarified by project maintainers. 54 | 55 | This Code of Conduct also applies outside the project spaces when there is a 56 | reasonable belief that an individual's behavior may have a negative impact on 57 | the project or its community. 58 | 59 | ## Enforcement 60 | 61 | Instances of abusive, harassing, or otherwise unacceptable behavior may be 62 | reported by contacting the project team at . All 63 | complaints will be reviewed and investigated and will result in a response that 64 | is deemed necessary and appropriate to the circumstances. The project team is 65 | obligated to maintain confidentiality with regard to the reporter of an incident. 66 | Further details of specific enforcement policies may be posted separately. 67 | 68 | Project maintainers who do not follow or enforce the Code of Conduct in good 69 | faith may face temporary or permanent repercussions as determined by other 70 | members of the project's leadership. 71 | 72 | ## Attribution 73 | 74 | This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4, 75 | available at https://www.contributor-covenant.org/version/1/4/code-of-conduct.html 76 | 77 | [homepage]: https://www.contributor-covenant.org 78 | 79 | For answers to common questions about this code of conduct, see 80 | https://www.contributor-covenant.org/faq 81 | -------------------------------------------------------------------------------- /configs/config.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | default_options = { 8 | # dataset config 9 | 'batch_size':{ 10 | 'type': int, 11 | 'default': 1024*8 12 | }, 13 | 'dataset': { 14 | 'type': str, 15 | 'nargs': 2, 16 | 'default': ['synthetic','../data/indoor_synthetic/kitchen'] 17 | }, 18 | 'scene': { 19 | 'type': str, 20 | 'default': '' 21 | }, 22 | 'voxel_path': { 23 | 'type': str, 24 | 'default': 'outputs/kitchen/vslf.npz' 25 | }, 26 | 'num_workers': { 27 | 'type': int, 28 | 'default': 12 29 | }, 30 | 'dir_val': { # small batch of samples per point 31 | 'type': str, 32 | 'default': 'val' 33 | }, 34 | 'val_step': { 35 | 'type': int, 36 | 'default': 250 37 | }, 38 | 39 | # whether has part segmentation 40 | 'has_part': { 41 | 'type': int, 42 | 'default': 1 43 | }, 44 | 45 | # image resolution scaling factor, 1.0 = original size 46 | 'res_scale': { 47 | 'type': float, 48 | 'default': 1.0 49 | }, 50 | 51 | 52 | # optimizer config 53 | 'optimizer': { 54 | 'type': str, 55 | 'choices': ['SGD', 'Ranger', 'Adam'], 56 | 'default': 'Adam' 57 | }, 58 | 'learning_rate': { 59 | 'type': float, 60 | 'default': 1e-3 61 | }, 62 | 'weight_decay': { 63 | 'type': float, 64 | 'default': 0 65 | }, 66 | 67 | 'scheduler_rate':{ 68 | 'type': float, 69 | 'default': 0.5 70 | }, 71 | 'milestones':{ 72 | 'type': int, 73 | 'nargs': '*', 74 | 'default': [1000] 75 | }, 76 | 77 | 78 | # reuglarization config 79 | 'le': { 80 | 'type': float, 81 | 'default': 1.0 82 | }, 83 | 'ld': { 84 | 'type': float, 85 | 'default': 5e-4 86 | }, 87 | 'lp': { 88 | 'type': float, 89 | 'default': 5e-3 90 | }, 91 | 'ls': { 92 | 'type': float, 93 | 'default': 1e-3 94 | }, 95 | 'la': { 96 | 'type': float, 97 | 'default': 0.0 98 | }, 99 | 'sigma_albedo': { 100 | 'type': float, 101 | 'default': 0.05/3.0 102 | }, 103 | 'sigma_pos': { 104 | 'type': float, 105 | 'default': 0.3/3.0 106 | }, 107 | 108 | # model params 109 | 'ckpt_path': { 110 | 'type': str, 111 | 'default': None 112 | }, 113 | 'emitter_path': { 114 | 'type': str, 115 | 'default': None 116 | }, 117 | 'freeze_emitter': { 118 | 'type': int, 119 | 'default': 0 120 | }, 121 | 'freeze_crf': { 122 | 'type': int, 123 | 'default': 0 124 | }, 125 | 'indir_depth': { 126 | 'type': int, 127 | 'default': 5 128 | }, 129 | 'SPP': { # Total samples per point 130 | 'type': int, 131 | 'default': 512 132 | }, 133 | 'spp': { # small batch of samples per point 134 | 'type': int, 135 | 'default': 8 136 | }, 137 | 138 | # LDR images with varying exposure 139 | 'ldr_img_dir': { 140 | 'type': str, 141 | 'default': None 142 | }, 143 | 'crf_basis': { 144 | 'type': int, 145 | 'default': 3 146 | }, 147 | 'load_crf': { 148 | 'type': int, 149 | 'default': 0 150 | }, 151 | 'l_crf_increasing': { 152 | 'type': float, 153 | 'default': 0.1 154 | }, 155 | 'l_crf_weight': { 156 | 'type': float, 157 | 'default': 0.001 158 | }, 159 | } 160 | -------------------------------------------------------------------------------- /utils/disco_ball.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import numpy as np 8 | import mitsuba 9 | 10 | def sample_points_on_sphere(num_points, phase=0): 11 | """Uniformly sample points on a unit sphere using the Fibonacci lattice method.""" 12 | points = [] 13 | phi = (1 + np.sqrt(5)) / 2 # Golden ratio 14 | 15 | for i in range(num_points): 16 | theta = 2 * np.pi * i / phi # Longitude angle 17 | z = 1 - (2 * i + 1) / num_points # Latitude adjustment for uniform spacing 18 | radius = np.sqrt(1 - z * z) # Radius of the circle at latitude z 19 | 20 | x = radius * np.cos(theta + phase) 21 | y = radius * np.sin(theta + phase) 22 | points.append([x, y, z]) 23 | 24 | return np.array(points) 25 | 26 | def make_disco_ball( 27 | scene_dict: dict, 28 | position, 29 | radius, 30 | light_intensity, 31 | light_num=20, 32 | light_radius_rate=0.1, 33 | spot_intensity=10, 34 | spot_cutoff_angle=20.0, 35 | phase=0): 36 | 37 | colors = np.array([ 38 | [1.0, 0.0, 0.0], 39 | [0.0, 1.0, 0.0], 40 | [0.0, 0.0, 1.0], 41 | [1.0, 1.0, 0.0], 42 | [1.0, 0.0, 1.0], 43 | [0.0, 1.0, 1.0], 44 | ]) 45 | position = np.array(position) 46 | points_unit = sample_points_on_sphere(light_num, phase) 47 | light_radius = radius * light_radius_rate 48 | distance_from_center = radius - light_radius * 0.6 49 | points_sphere = points_unit * distance_from_center + position 50 | 51 | # make center ball 52 | scene_dict.update({ 53 | 'disco_ball': { 54 | 'type': 'sphere', 55 | 'to_world': mitsuba.ScalarTransform4f\ 56 | .translate(position.tolist())\ 57 | .scale([radius, radius, radius]), 58 | 'bsdf': { 59 | 'type': 'diffuse', 60 | 'reflectance': { 61 | 'type': 'rgb', 62 | 'value': [0.2, 0.2, 0.2] 63 | } 64 | }, 65 | } 66 | }) 67 | 68 | # make lights 69 | for i in range(light_num): 70 | light_pos = points_sphere[i] 71 | light_value = colors[i % colors.shape[0]] * light_intensity 72 | 73 | light_config = { 74 | 'type': 'sphere', 75 | 'to_world': mitsuba.ScalarTransform4f\ 76 | .translate(light_pos.tolist())\ 77 | .scale([light_radius, light_radius, light_radius]), 78 | 'emitter': { 79 | 'type': 'area', 80 | 'radiance': { 81 | 'type': 'rgb', 82 | 'value': light_value.tolist() 83 | } 84 | } 85 | } 86 | 87 | spot_o = points_unit[i] * (radius + light_radius) + position 88 | spot_t = spot_o + points_unit[i] 89 | spot_value = colors[i % colors.shape[0]] * spot_intensity 90 | spot_config = { 91 | 'type': 'spot', 92 | 'to_world': mitsuba.ScalarTransform4f.look_at( 93 | origin=spot_o.tolist(), 94 | target=spot_t.tolist(), 95 | up=[0, 0, 1] 96 | ), # Light direction 97 | 'cutoff_angle': spot_cutoff_angle, # Defines the angle of the light cone # Inner angle with full intensity 98 | 'intensity': { 99 | 'type': 'rgb', 100 | 'value': spot_value.tolist() 101 | } 102 | } 103 | 104 | 105 | scene_dict.update({ 106 | 'light_{}'.format(i): light_config, 107 | 'spot_{}'.format(i): spot_config 108 | }) -------------------------------------------------------------------------------- /utils/metric_brdf.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import torch 8 | import torch.nn.functional as NF 9 | 10 | import math 11 | import os 12 | os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1" 13 | import cv2 14 | 15 | import numpy as np 16 | import pandas as pd 17 | 18 | import matplotlib.pyplot as plt 19 | from tqdm import tqdm 20 | 21 | METHOD = 'fipt_syn_kitchen' 22 | GT_PATH = '/hdd/datasets/fipt/indoor_synthetic/kitchen/train' 23 | METHOD_PATH = os.path.join('outputs', METHOD, 'output', 'train') 24 | 25 | image_num = len([f for f in os.listdir(os.path.join(GT_PATH,'Image')) if f[0] != '.' and f[-4:]=='.exr']) 26 | mse_roughness = [] 27 | mse_albedo = [] 28 | mse_diff = [] 29 | iou_emission = [] 30 | mse_emission = [] 31 | for i in tqdm(range(image_num)): 32 | emission_gt = cv2.imread(os.path.join(GT_PATH,'Emit','{:03d}_0001.exr'.format(i)),-1)[...,[2,1,0]] 33 | emission_gt = torch.from_numpy(emission_gt).float() 34 | emission_mask = emission_gt.sum(-1) > 0 35 | albedo_gt = cv2.imread(os.path.join(GT_PATH,'albedo','{:03d}.exr'.format(i)),-1)[...,[2,1,0]] 36 | albedo_gt = torch.from_numpy(albedo_gt).float().clamp(0,1).mul(255).long().float()/255 37 | albedo_gt[emission_mask] = 0 38 | 39 | kd_gt = cv2.imread(os.path.join(GT_PATH,'DiffCol','{:03d}_0001.exr'.format(i)),-1)[...,[2,1,0]] 40 | kd_gt = torch.from_numpy(kd_gt).float().clamp(0,1).mul(255).long().float()/255 41 | kd_gt[emission_mask] = 0 42 | 43 | roughness_gt = cv2.imread(os.path.join(GT_PATH,'Roughness','{:03d}_0001.exr'.format(i)),-1)[...,0] 44 | roughness_gt = (torch.from_numpy(roughness_gt).float().mul(255).long().float()/255).clamp(0.2,1) 45 | roughness_gt[emission_mask] = 0 46 | 47 | diff_mask =roughness_gt==1 48 | kd_gt[~diff_mask] = 0 49 | 50 | 51 | emission = cv2.imread(os.path.join(METHOD_PATH, 'emission', '{:05d}_emission.exr'.format(i)),-1)[...,[2,1,0]] 52 | emission = torch.from_numpy(emission).float() 53 | 54 | albedo = cv2.imread(os.path.join(METHOD_PATH, 'a_prime','{:05d}_a_prime.png'.format(i)),-1)[...,[2,1,0]] 55 | albedo = torch.from_numpy(albedo).float()/255 56 | albedo[emission_mask] = 0 57 | 58 | kd = cv2.imread(os.path.join(METHOD_PATH, 'diffuse', '{:05d}_kd.png'.format(i)),-1)[...,[2,1,0]] 59 | kd = torch.from_numpy(kd).float()/255 60 | kd[emission_mask] = 0 61 | kd[~diff_mask] = 0 62 | 63 | roughness = cv2.imread(os.path.join(METHOD_PATH, 'roughness','{:05d}_roughness.png'.format(i)),-1) 64 | roughness = roughness[:, :, 0] 65 | roughness = (torch.from_numpy(roughness).float()/255).clamp(0.2,1) 66 | 67 | roughness[emission_mask] = 0 68 | 69 | 70 | emission_mask_est = emission.sum(-1)>0 71 | if emission_mask.any(): 72 | iou = (emission_mask&emission_mask_est).sum()*1.0/(emission_mask|emission_mask_est).sum() 73 | iou_emission.append(iou) 74 | intersect = emission_mask&emission_mask_est 75 | mse_emission.append(NF.mse_loss(torch.log(emission+1),torch.log(emission_gt+1))) 76 | #mse_emission.append(NF.mse_loss(emission.pow(1/2.2),emission_gt.pow(1/2.2))) 77 | 78 | mse_roughness.append(NF.mse_loss(roughness,roughness_gt)) 79 | mse_albedo.append(NF.mse_loss(albedo,albedo_gt)) 80 | mse_diff.append(NF.mse_loss(kd,kd_gt)) 81 | mse_roughness = torch.tensor(mse_roughness) 82 | mse_diff = torch.tensor(mse_diff) 83 | mse_albedo = torch.tensor(mse_albedo) 84 | mse_emission = torch.tensor(mse_emission) 85 | iou_emission = torch.tensor(iou_emission) 86 | 87 | print(METHOD) 88 | print('kd: ',-10*torch.log10(mse_diff.clamp_min(1e-10)).mean().item()) 89 | print('albedo: ',-10*torch.log10(mse_albedo.clamp_min(1e-10)).mean().item()) 90 | print('roughness: ',(-10*torch.log10(mse_roughness.clamp_min(1e-10))).mean().item()) 91 | print('emit_iou: ',iou_emission.mean().item()) 92 | print('emit_log_mse:',mse_emission.mean().item()) -------------------------------------------------------------------------------- /utils/fuse_segmentation.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | 8 | import torch 9 | 10 | import mitsuba 11 | mitsuba.set_variant('cuda_ad_rgb') 12 | 13 | import os 14 | os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1" 15 | import cv2 16 | import trimesh 17 | import numpy as np 18 | 19 | import sys 20 | sys.path.append('..') 21 | from utils.dataset import RealDataset, SyntheticDataset 22 | from utils.path_tracing import ray_intersect 23 | 24 | from argparse import ArgumentParser 25 | from tqdm import tqdm 26 | 27 | 28 | if __name__ == '__main__': 29 | parser = ArgumentParser() 30 | parser.add_argument('--scene',type=str,required=True,help='scene path') 31 | parser.add_argument('--dataset',type=str,required=True,help='scene type') 32 | 33 | 34 | args = parser.parse_args() 35 | device = torch.device(0) 36 | 37 | if args.dataset == 'real': 38 | dataset_fn = RealDataset 39 | elif args.dataset == 'synthetic': 40 | dataset_fn = SyntheticDataset 41 | 42 | # load dataset 43 | dataset = dataset_fn(args.scene,split='train',pixel=False,ray_diff=False) 44 | img_hw = dataset.img_hw 45 | 46 | # load scene geometry 47 | scene = mitsuba.load_dict({ 48 | 'type': 'scene', 49 | 'shape_id':{ 50 | 'type': 'obj', 51 | 'filename': os.path.join(args.scene, 'scene.obj') 52 | } 53 | }) 54 | mesh = trimesh.load_mesh(os.path.join(args.scene,'scene.obj')) 55 | vertices = torch.from_numpy(np.array(mesh.vertices)).float() 56 | faces = torch.from_numpy(np.array(mesh.faces)) 57 | 58 | # segmentation root folder 59 | if args.dataset == 'synthetic': 60 | args.scene = os.path.join(args.scene,'train') 61 | 62 | # find max segmentation indices 63 | seg_num = 0 64 | for i in tqdm(range(len(dataset))): 65 | segmentation = cv2.imread(os.path.join(args.scene,'segmentation/{:03d}.exr'.format(i)),-1) 66 | seg_num = max(int(segmentation.max()),seg_num) 67 | 68 | # build histogram of segmentation labels for each triangle 69 | labels = torch.zeros(len(faces)*seg_num,device=device,dtype=torch.long) 70 | for i in tqdm(range(len(dataset))): 71 | segmentation = cv2.imread(os.path.join(args.scene,'segmentation/{:03d}.exr'.format(i)),-1) 72 | segmentation = torch.from_numpy(segmentation[...,0]).long().reshape(-1) 73 | batch = dataset[i] 74 | rays = batch['rays'].to(device) 75 | xs,ds = rays[...,:3],rays[...,3:6] 76 | 77 | 78 | positions,normals,_,triangle_idx,valid = ray_intersect(scene,xs,ds) 79 | # flattened indices for histogram entries 80 | inds = triangle_idx[valid]*int(seg_num) + segmentation.to(device)[valid] 81 | labels.scatter_add_(0,inds,torch.ones_like(inds)) 82 | 83 | # assign triangle label by maximum occurance 84 | new_label_num,new_label = labels.reshape(-1,seg_num)[:,1:].max(-1) 85 | new_label += 1 86 | # if the label occured 0 times, assign 0 label 87 | new_label[new_label_num==0] = 0 88 | 89 | 90 | # write fused segmentation 91 | OUTPUT_PATH = os.path.join(args.scene,'segmentation-new') 92 | os.makedirs(OUTPUT_PATH,exist_ok=True) 93 | im_id =0 94 | for batch in tqdm(dataset): 95 | rays = batch['rays'].to(device) 96 | xs = rays[...,:3] 97 | ds = rays[...,3:6] 98 | _,_,_,triangle_idx,valid = ray_intersect(scene,xs,ds) 99 | segmentation = new_label[triangle_idx] 100 | segmentation[~valid] = 0 101 | segmentation = segmentation.cpu().reshape(*img_hw,1).expand(*img_hw,3) 102 | cv2.imwrite(os.path.join(OUTPUT_PATH,'{:03d}.exr'.format(im_id)),segmentation.float().numpy()) 103 | im_id += 1 104 | 105 | # move file folder 106 | os.system('mv {}/segmentation {}/segmentation-old'.format(args.scene,args.scene)) 107 | os.system('mv {}/segmentation-new {}/segmentation'.format(args.scene,args.scene)) -------------------------------------------------------------------------------- /crf/hdr2ldr.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import numpy as np 8 | import os 9 | os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1" 10 | import cv2 11 | from PIL import Image 12 | from matplotlib import pyplot as plt 13 | from argparse import ArgumentParser 14 | from tqdm import tqdm 15 | from scipy import interpolate 16 | from .emor import parse_dorf_curves, parse_emor_file 17 | 18 | def open_exr(file): 19 | img = cv2.imread(file,cv2.IMREAD_UNCHANGED)[...,[2,1,0]] 20 | # img = img.astype(np.float32) 21 | return img 22 | 23 | def apply_crf(image, curves): 24 | h, w, _ = image.shape 25 | image_ldr = [] 26 | for i in range(3): 27 | ch = image[:, :, i].reshape(-1) 28 | crf = curves[i] 29 | x = np.linspace(0, 1, len(crf)) 30 | crf_inter = interpolate.interp1d(x, crf) 31 | ch_ldr = crf_inter(ch).reshape(h, w) 32 | image_ldr.append(ch_ldr) 33 | image_ldr = np.stack(image_ldr, axis=-1) 34 | return image_ldr 35 | 36 | def main(): 37 | parser = ArgumentParser() 38 | parser.add_argument('--dir_src', help='dir path to HDR images') 39 | parser.add_argument('--dir_tgt', help='dir path to output LDR images') 40 | parser.add_argument('--curve_idx', nargs='+', type=int, default=0) 41 | args = parser.parse_args() 42 | os.makedirs(args.dir_tgt, exist_ok=True) 43 | 44 | level = np.array([-2, -1, 0, 1, 2]) 45 | exposure = (np.ones(len(level))*2)**level 46 | 47 | names_all, curves_all = parse_dorf_curves() 48 | idx = args.curve_idx 49 | names = names_all[[idx[0], idx[1], idx[2]]] 50 | curves = curves_all[[idx[0], idx[1], idx[2]]] 51 | dir_cam = os.path.join(args.dir_tgt, 'cam') 52 | os.makedirs(dir_cam, exist_ok=True) 53 | for i in range(len(names)): 54 | n = names[i] 55 | crf = curves[i] 56 | x = np.linspace(0, 1, len(crf)) 57 | plt.title(n) 58 | plt.plot(x, crf) 59 | channel = 'R' if i==0 else 'G' if i==1 else 'B' 60 | out_path = os.path.join(dir_cam, 'crf_{}.png'.format(channel)) 61 | plt.savefig(out_path) 62 | plt.close() 63 | np.save(os.path.join(dir_cam, 'crf.npy'), curves) 64 | 65 | save_sorted_exposure(args.dir_src, args.dir_tgt, exposure, curves) 66 | 67 | def save_all_exposure(dir_src, dir_tgt, exposure, curves): 68 | hdr_paths = sorted([os.path.join(dir_src, name) for name in os.listdir(dir_src) if name.endswith('.exr')]) 69 | for path in tqdm(hdr_paths): 70 | hdr = open_exr(path) 71 | prefix = path.split('/')[-1].split('.')[0] 72 | for i, exp in enumerate(exposure): 73 | irr = hdr * exp 74 | irr = np.clip(irr, 0, 1) 75 | ldr = apply_crf(irr, curves) 76 | out_path = os.path.join(dir_tgt, '{}_{}.png'.format(prefix, i)) 77 | img = Image.fromarray((ldr*255).astype(np.uint8)) 78 | img.save(out_path) 79 | 80 | def save_sorted_exposure(dir_src, dir_tgt, exposure, curves): 81 | hdr_paths = sorted([os.path.join(dir_src, name) for name in os.listdir(dir_src) if name.endswith('.exr')]) 82 | img_means = [] 83 | for path in hdr_paths: 84 | img = open_exr(path) 85 | mean = np.mean(img) 86 | img_means.append(mean) 87 | img_means = np.array(img_means) 88 | 89 | argsort = np.argsort(img_means) 90 | exposure = np.sort(exposure)[::-1] #big to small 91 | img_exp = np.zeros_like(img_means) 92 | step = len(img_exp) // len(exposure) 93 | for i, exp in enumerate(exposure): 94 | idx = argsort[i*step:(i+1)*step] 95 | img_exp[idx] = exp 96 | idx = argsort[step*len(exposure):] 97 | img_exp[idx] = exposure[-1] 98 | np.save(os.path.join(dir_tgt, 'cam', 'exposure.npy'), img_exp) 99 | 100 | for i in tqdm(range(len(hdr_paths))): 101 | path = hdr_paths[i] 102 | hdr = open_exr(path) 103 | exp = img_exp[i] 104 | irr = np.clip(hdr * exp, 0, 1) 105 | ldr = apply_crf(irr, curves) 106 | prefix = path.split('/')[-1].split('.')[0] 107 | ldr_path = os.path.join(dir_tgt, '{}.png'.format(prefix)) 108 | img = Image.fromarray((ldr*255).astype(np.uint8)) 109 | img.save(ldr_path) 110 | 111 | if __name__ == '__main__': 112 | main() 113 | -------------------------------------------------------------------------------- /crf/model_crf.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import numpy as np 8 | import torch 9 | import torch.nn as nn 10 | import torch.nn.functional as F 11 | import torch_interpolations 12 | import os 13 | import pathlib 14 | from tqdm import tqdm 15 | import scipy 16 | from scipy import interpolate 17 | from matplotlib import pyplot as plt 18 | from .emor import parse_dorf_curves, parse_emor_file 19 | from const import set_random_seed 20 | set_random_seed() 21 | 22 | def mono_increase_constraint(crf): 23 | diff = crf[1:] - crf[:-1] 24 | diff_min = diff.min() 25 | gap = -diff_min if diff_min < 0 else 0 26 | diff += gap 27 | diff /= diff.sum() 28 | crf = torch.cumsum(diff, dim=0) 29 | crf = torch.cat([torch.zeros((1), device=crf.device), crf]) 30 | return crf 31 | 32 | class EmorCRF(nn.Module): 33 | def __init__(self, dim=11): 34 | super().__init__() 35 | self.dim = dim 36 | names, vectors = parse_emor_file(inv=False) 37 | self.register_buffer('f0', torch.FloatTensor(vectors[1])[None]) 38 | self.register_buffer('basis', torch.FloatTensor(vectors[2:2+dim])) 39 | self.weight = nn.Parameter(torch.zeros(3, dim)) 40 | 41 | def get_crf(self): 42 | crf = self.f0 + self.weight @ self.basis 43 | return crf 44 | 45 | def get_inv_crf(self): 46 | crf = self.get_crf() 47 | inv_crf = [] 48 | for i in range(3): 49 | crf_ch = mono_increase_constraint(crf[i]) 50 | x = torch.linspace(0, 1, len(crf_ch)).to(self.weight.device) 51 | interp_func = torch_interpolations.RegularGridInterpolator([crf_ch], x) 52 | inv_crf_ch = interp_func([x]) 53 | inv_crf.append(inv_crf_ch) 54 | inv_crf = torch.stack(inv_crf, dim=0) 55 | return inv_crf 56 | 57 | def initialize_weight(self, crf): 58 | weight = self.cal_weight_fitting_crf(crf) #(3, dim) 59 | self.weight = nn.Parameter(torch.FloatTensor(weight).to(self.weight.device)) 60 | 61 | def cal_weight_fitting_crf(self, crf): 62 | f0 = self.f0.detach().cpu().numpy() 63 | basis = self.basis.detach().cpu().numpy().T 64 | pseudo_inverse = np.linalg.inv(basis.T @ basis) @ basis.T 65 | weight = pseudo_inverse @ (crf - f0).T 66 | return weight.T 67 | 68 | def forward(self, hdr, exposure): 69 | ''' 70 | Input: 71 | hdr: (n, 3) 72 | Return: 73 | ldr: (n, 3) 74 | ''' 75 | hdr = torch.clip(hdr*exposure, 0, 1) 76 | crf = self.get_crf() 77 | x = torch.linspace(0, 1, crf.size(1)).to(self.weight.device) 78 | ldr = [] 79 | for i in range(3): 80 | hdr_ch = hdr[:, i] 81 | crf_ch = crf[i] 82 | interp_func = torch_interpolations.RegularGridInterpolator([x], crf_ch) 83 | ldr_ch = interp_func([hdr_ch]) 84 | ldr.append(ldr_ch) 85 | ldr = torch.stack(ldr, dim=-1) 86 | return ldr 87 | 88 | def inverse(self, ldr, exposure): 89 | ''' 90 | Input: 91 | ldr: (n, 3) 92 | Return: 93 | hdr: (n, 3) 94 | ''' 95 | ldr = torch.clip(ldr, 0, 1) 96 | inv_crf = self.get_inv_crf() 97 | x = torch.linspace(0, 1, inv_crf.size(1)).to(self.weight.device) 98 | hdr = [] 99 | for i in range(3): 100 | ldr_ch = ldr[:, i] 101 | inv_crf_ch = inv_crf[i] 102 | interp_func = torch_interpolations.RegularGridInterpolator([x], inv_crf_ch) 103 | hdr_ch = interp_func([ldr_ch]) 104 | hdr.append(hdr_ch) 105 | hdr = torch.stack(hdr, dim=-1) / exposure 106 | return hdr 107 | 108 | def reg_weight(self): 109 | loss = torch.mean(self.weight ** 2) 110 | return loss 111 | 112 | def reg_monotonically_increasing(self): 113 | crf = self.get_crf() #(3, 1024) 114 | diff = crf[:, 1:] - crf[:, :-1] # should be all positive 115 | loss = torch.sum(F.relu(-diff)) 116 | return loss 117 | 118 | def reg_smoothness(self): 119 | crf = self.get_crf() 120 | smoothness = crf[:, :-2] + crf[:, 2:] - 2 * crf[:, 1:-1] 121 | loss = torch.mean(smoothness ** 2) 122 | return loss -------------------------------------------------------------------------------- /slf_refine.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import torch 8 | import mitsuba 9 | mitsuba.set_variant('cuda_ad_rgb') 10 | import os 11 | os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1" 12 | import cv2 13 | 14 | from utils.dataset import SyntheticDatasetLDR,RealDatasetLDR 15 | from utils.dataset.scannetpp.dataset import Scannetpp 16 | from utils.path_tracing import ray_intersect 17 | from model.slf import VoxelSLF 18 | from model.brdf import BaseBRDF 19 | from crf.model_crf import EmorCRF 20 | from pathlib import Path 21 | from tqdm import tqdm 22 | from argparse import ArgumentParser 23 | from const import set_random_seed 24 | set_random_seed() 25 | 26 | if __name__ == '__main__': 27 | parser = ArgumentParser() 28 | parser.add_argument('--dataset_root', type=str, help='dataset root') 29 | parser.add_argument('--scene', type=str, required=True, help='dataset folder') 30 | parser.add_argument('--output', type=str, required=True, help='output path') 31 | parser.add_argument('--load', type=str, default='vslf.npz') 32 | parser.add_argument('--save', type=str, default='vslf.npz') 33 | parser.add_argument('--dataset', type=str,required=True, help='dataset type') 34 | parser.add_argument('--voxel_num', type=int,default=256, help='resolution for voxel radiance cache') 35 | parser.add_argument('--ldr_img_dir', type=str, default=None) 36 | parser.add_argument('--ckpt', type=str, help='checkpoint to load CRF model') 37 | parser.add_argument('--crf_basis', type=int, help='number of CRF basis') 38 | parser.add_argument('--res_scale', type=float, default=1.0) 39 | args = parser.parse_args() 40 | 41 | device = torch.device(0) # use gpu device 0 42 | 43 | DATASET_PATH = args.scene 44 | OUTPUT_PATH = args.output 45 | os.makedirs(OUTPUT_PATH,exist_ok=True) 46 | # load mesh 47 | if args.dataset in ['synthetic', 'real']: 48 | mesh_path = os.path.join(DATASET_PATH,'scene.obj') 49 | mesh_type = 'obj' 50 | elif args.dataset == 'scannetpp': 51 | mesh_path = os.path.join(args.dataset_root, 'data', args.scene, 'scans', 'scene.ply') 52 | mesh_type = 'ply' 53 | assert Path(mesh_path).exists(), 'mesh not found: '+mesh_path 54 | 55 | scene = mitsuba.load_dict({ 56 | 'type': 'scene', 57 | 'shape_id':{ 58 | 'type': mesh_type, 59 | 'filename': mesh_path, 60 | } 61 | }) 62 | 63 | # load dataset 64 | if args.dataset == 'synthetic': 65 | dataset = SyntheticDatasetLDR(DATASET_PATH, img_dir=args.ldr_img_dir, split='train', pixel=False) 66 | elif args.dataset == 'real': 67 | dataset = RealDatasetLDR(DATASET_PATH, img_dir=args.ldr_img_dir, split='train', pixel=False) 68 | elif args.dataset == 'scannetpp': 69 | dataset = Scannetpp(args.dataset_root, args.scene, split='train', pixel=False, res_scale=args.res_scale) 70 | img_hw = dataset.img_hw 71 | 72 | model_crf = EmorCRF(args.crf_basis) 73 | if args.ckpt: 74 | state_dict = torch.load(args.ckpt, map_location='cpu')['state_dict'] 75 | weight = {} 76 | for k,v in state_dict.items(): 77 | if 'model_crf.' in k: 78 | weight[k.replace('model_crf.','')]=v 79 | model_crf.load_state_dict(weight) 80 | model_crf = model_crf.to(device) 81 | 82 | vslf_load = os.path.join(args.output, args.load) 83 | vslf_save = os.path.join(args.output, args.save) 84 | 85 | with torch.no_grad(): 86 | device = model_crf.weight.device 87 | state_dict = torch.load(vslf_load,map_location='cpu') 88 | vslf = VoxelSLF(state_dict['mask'], state_dict['voxel_min'],state_dict['voxel_max']) 89 | for idx in tqdm(range(len(dataset)), postfix='Update VSLF'): 90 | batch = dataset[idx] 91 | rays = batch['rays'].to(device) 92 | radiance = batch['rgbs'].to(device) 93 | xs = rays[...,:3] 94 | ds = rays[...,3:6] 95 | 96 | exposure = batch['exposure'] 97 | radiance = model_crf.inverse(radiance, exposure) 98 | 99 | positions,_,_,_,valid = ray_intersect(scene,xs,ds) 100 | if not valid.any(): 101 | continue 102 | 103 | vslf.scatter_add(positions[valid].cpu(),radiance.to(device)[valid].cpu()) 104 | 105 | # average pooling the radiance 106 | vslf.radiance = vslf.radiance/vslf.count[...,None].float().clamp_min(1) 107 | state_dict['weight'] = vslf.state_dict() 108 | torch.save(state_dict, vslf_save) -------------------------------------------------------------------------------- /scripts/fipt/classroom/train.sh: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | # data folder 8 | DATASET_ROOT='/hdd/datasets/fipt/real/' 9 | DATASET='real' 10 | # scene name 11 | SCENE='classroom' 12 | LDR_IMG_DIR='Image' 13 | EXP='fipt_real_classroom' 14 | VAL_FRAME=1 15 | CRF_BASIS=3 16 | # whether has part segmentation 17 | HAS_PART=0 18 | 19 | # bake surface light field (SLF) 20 | python slf_bake.py --scene $DATASET_ROOT$SCENE\ 21 | --output checkpoints/$EXP/bake \ 22 | --dataset $DATASET --ldr_img_dir $LDR_IMG_DIR 23 | 24 | # extract emitter mask 25 | python extract_emitter_ldr.py --scene $DATASET_ROOT$SCENE\ 26 | --output checkpoints/$EXP/bake --dataset $DATASET\ 27 | --ldr_img_dir $LDR_IMG_DIR --threshold 0.99 28 | 29 | python initialize.py --experiment_name $EXP --max_epochs 3 \ 30 | --dataset $DATASET $DATASET_ROOT$SCENE \ 31 | --voxel_path checkpoints/$EXP/bake/vslf.npz \ 32 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 33 | --has_part $HAS_PART --ldr_img_dir $LDR_IMG_DIR --val_frame $VAL_FRAME\ 34 | --SPP 128 --spp 32 --crf_basis $CRF_BASIS 35 | 36 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/init.ckpt 37 | 38 | # extract emitters 39 | python extract_emitter_ldr.py --mode update\ 40 | --scene $DATASET_ROOT$SCENE\ 41 | --output checkpoints/$EXP/bake\ 42 | --ckpt checkpoints/$EXP/init.ckpt\ 43 | --dataset $DATASET --ldr_img_dir $LDR_IMG_DIR 44 | 45 | python bake_shading.py \ 46 | --scene $DATASET_ROOT$SCENE --dataset $DATASET \ 47 | --ldr_img_dir $LDR_IMG_DIR \ 48 | --slf_path checkpoints/$EXP/bake/vslf.npz \ 49 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 50 | --output outputs/$EXP/shading 51 | 52 | # optimize BRDF, CRF 53 | python train_brdf_crf.py --experiment_name $EXP \ 54 | --max_epochs 2 --dir_val val_0 \ 55 | --ckpt_path checkpoints/$EXP/init.ckpt \ 56 | --voxel_path checkpoints/$EXP/bake/vslf.npz \ 57 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 58 | --cache_dir outputs/$EXP/shading \ 59 | --dataset $DATASET $DATASET_ROOT$SCENE \ 60 | --has_part $HAS_PART --ldr_img_dir $LDR_IMG_DIR --val_frame $VAL_FRAME\ 61 | --SPP 128 --spp 32 --lp 0.005 --la 0.01 --l_crf_weight 0.01 --crf_basis $CRF_BASIS 62 | 63 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/last_0.ckpt 64 | 65 | # refine SLF 66 | python slf_refine.py --scene $DATASET_ROOT$SCENE \ 67 | --output checkpoints/$EXP/bake --load vslf.npz --save vslf_0.npz \ 68 | --dataset $DATASET --ldr_img_dir $LDR_IMG_DIR \ 69 | --ckpt checkpoints/$EXP/last_0.ckpt --crf_basis $CRF_BASIS 70 | 71 | # refine emitter 72 | python train_emitter.py --experiment_name $EXP \ 73 | --max_epochs 1 --dir_val val_0_emitter \ 74 | --ckpt_path checkpoints/$EXP/last_0.ckpt \ 75 | --voxel_path checkpoints/$EXP/bake/vslf_0.npz \ 76 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 77 | --dataset $DATASET $DATASET_ROOT$SCENE \ 78 | --has_part $HAS_PART --ldr_img_dir $LDR_IMG_DIR --val_frame $VAL_FRAME\ 79 | --SPP 128 --spp 32 --crf_basis $CRF_BASIS 80 | 81 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/last_0.ckpt 82 | 83 | # extract emitter 84 | python extract_emitter_ldr.py --mode update\ 85 | --scene $DATASET_ROOT$SCENE\ 86 | --output checkpoints/$EXP/bake\ 87 | --ckpt checkpoints/$EXP/last_0.ckpt\ 88 | --dataset $DATASET --ldr_img_dir $LDR_IMG_DIR 89 | 90 | # # refine shading 91 | python refine_shading.py \ 92 | --scene $DATASET_ROOT$SCENE --dataset $DATASET \ 93 | --ldr_img_dir $LDR_IMG_DIR \ 94 | --slf_path checkpoints/$EXP/bake/vslf_0.npz \ 95 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 96 | --ckpt checkpoints/$EXP/last_0.ckpt \ 97 | --output outputs/$EXP/shading 98 | 99 | # optimize BRDF, CRF 100 | python train_brdf_crf.py --experiment_name $EXP \ 101 | --max_epochs 2 --dir_val val_1 \ 102 | --ckpt_path checkpoints/$EXP/init.ckpt \ 103 | --voxel_path checkpoints/$EXP/bake/vslf_0.npz \ 104 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 105 | --cache_dir outputs/$EXP/shading \ 106 | --dataset $DATASET $DATASET_ROOT$SCENE \ 107 | --has_part $HAS_PART --ldr_img_dir $LDR_IMG_DIR --val_frame $VAL_FRAME\ 108 | --SPP 128 --spp 32 --lp 0.005 --la 0.01 --l_crf_weight 0.01 --crf_basis $CRF_BASIS 109 | 110 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/last_1.ckpt -------------------------------------------------------------------------------- /model/fipt_bsdf.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import torch 8 | import torch.nn.functional as NF 9 | 10 | import mitsuba 11 | mitsuba.set_variant('cuda_ad_rgb') 12 | 13 | import os 14 | import sys 15 | sys.path.append('..') 16 | from model.brdf import NGPBRDF 17 | from model.emitter import SLFEmitter 18 | 19 | class FIPTBSDF(mitsuba.BSDF): 20 | def __init__(self, props): 21 | mitsuba.BSDF.__init__(self, props) 22 | # default device for mitsuba 23 | device = torch.device(0) 24 | 25 | # load BRDF and emission mask 26 | mask = torch.load(os.path.join(props['emitter_path'],'vslf.npz'),map_location='cpu') 27 | state_dict = torch.load(props['brdf_path'],map_location='cpu')['state_dict'] 28 | weight = {} 29 | for k,v in state_dict.items(): 30 | if 'material.' in k: 31 | weight[k.replace('material.','')]=v 32 | material_net = NGPBRDF(mask['voxel_min'],mask['voxel_max']) 33 | material_net.load_state_dict(weight) 34 | material_net.to(device) 35 | for p in material_net.parameters(): 36 | p.requires_grad=False 37 | self.material_net = material_net 38 | self.is_emitter = torch.load(os.path.join(props['emitter_path'],'emitter.pth'))['is_emitter'].to(device) 39 | 40 | # specify flags 41 | reflection_flags = mitsuba.BSDFFlags.SpatiallyVarying|mitsuba.BSDFFlags.DiffuseReflection|mitsuba.BSDFFlags.FrontSide | mitsuba.BSDFFlags.BackSide 42 | self.m_components = [reflection_flags] 43 | self.m_flags = reflection_flags 44 | 45 | def sample(self, ctx, si, sample1, sample2, active): 46 | wi = si.to_world(si.wi).torch() 47 | normal = si.n.torch() 48 | position = si.p.torch() 49 | triangle_idx = mitsuba.Int(si.prim_index).torch().long() 50 | 51 | mat = self.material_net(position) 52 | wo,pdf,brdf_weight = self.material_net.sample_brdf( 53 | sample1.torch().reshape(-1),sample2.torch(), 54 | wi,normal,mat 55 | ) 56 | brdf_weight[self.is_emitter[triangle_idx]] = 1.0 # increase from 0 to 1 to fill the emitter region color 57 | 58 | pdf_mi = mitsuba.Float(pdf.squeeze(-1)) 59 | wo_mi = mitsuba.Vector3f(wo[...,0],wo[...,1],wo[...,2]) 60 | wo_mi = si.to_local(wo_mi) 61 | value_mi = mitsuba.Vector3f(brdf_weight[...,0],brdf_weight[...,1],brdf_weight[...,2]) 62 | 63 | bs = mitsuba.BSDFSample3f() 64 | bs.pdf = pdf_mi 65 | bs.sampled_component = mitsuba.UInt32(0) 66 | bs.sampled_type = mitsuba.UInt32(+self.m_flags) 67 | bs.wo = wo_mi 68 | bs.eta = 1.0 69 | 70 | return (bs,value_mi) 71 | 72 | def eval(self, ctx, si, wo, active): 73 | wo = si.to_world(wo).torch() 74 | wi = si.to_world(si.wi).torch() 75 | triangle_idx = mitsuba.Int(si.prim_index).torch().long() 76 | 77 | normal = si.n.torch() 78 | position = si.p.torch() 79 | 80 | mat = self.material_net(position) 81 | 82 | brdf,_ = self.material_net.eval_brdf(wo,wi,normal,mat) 83 | brdf[self.is_emitter[triangle_idx]]=0 84 | brdf = mitsuba.Vector3f(brdf[...,0],brdf[...,1],brdf[...,2]) 85 | 86 | return brdf 87 | 88 | 89 | def pdf(self, ctx, si, wo,active): 90 | wo = si.to_world(wo).torch() 91 | wi = si.to_world(si.wi).torch() 92 | 93 | normal = si.n.torch() 94 | position = si.p.torch() 95 | 96 | mat = self.material_net(position) 97 | _,pdf = self.material_net.eval_brdf(wo,wi,normal,mat) 98 | pdf = mitsuba.Float(pdf.squeeze(-1)) 99 | return pdf 100 | 101 | 102 | def eval_pdf(self, ctx, si, wo, active=True): 103 | wo = si.to_world(wo).torch() 104 | wi = si.to_world(si.wi).torch() 105 | triangle_idx = mitsuba.Int(si.prim_index).torch().long() 106 | 107 | normal = si.n.torch() 108 | position = si.p.torch() 109 | 110 | mat = self.material_net(position) 111 | 112 | brdf,pdf = self.material_net.eval_brdf(wo,wi,normal,mat) 113 | brdf[self.is_emitter[triangle_idx]] = 0 114 | brdf = mitsuba.Vector3f(brdf[...,0],brdf[...,1],brdf[...,2]) 115 | pdf = mitsuba.Float(pdf.squeeze(-1)) 116 | 117 | return brdf,pdf 118 | def to_string(self,): 119 | return 'FIPTBSDF' 120 | 121 | 122 | mitsuba.register_bsdf("fipt", lambda props: FIPTBSDF(props)) -------------------------------------------------------------------------------- /scripts/fipt/bathroom/train.sh: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | # data folder 8 | DATASET_ROOT='/hdd/datasets/fipt/indoor_synthetic/' 9 | DATASET='synthetic' 10 | # scene name 11 | SCENE='bathroom' 12 | LDR_IMG_DIR='Image' 13 | EXP='fipt_syn_bathroom' 14 | VAL_FRAME=10 15 | CRF_BASIS=3 16 | # whether has part segmentation 17 | HAS_PART=1 18 | SPP=128 19 | spp=32 20 | 21 | # bake surface light field (SLF) 22 | python slf_bake.py --scene $DATASET_ROOT$SCENE\ 23 | --output checkpoints/$EXP/bake \ 24 | --dataset $DATASET --ldr_img_dir $LDR_IMG_DIR 25 | 26 | # extract emitter mask 27 | python extract_emitter_ldr.py --scene $DATASET_ROOT$SCENE\ 28 | --output checkpoints/$EXP/bake --dataset $DATASET\ 29 | --ldr_img_dir $LDR_IMG_DIR --threshold 0.99 30 | 31 | python initialize.py --experiment_name $EXP --max_epochs 6 \ 32 | --dataset $DATASET $DATASET_ROOT$SCENE \ 33 | --voxel_path checkpoints/$EXP/bake/vslf.npz \ 34 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 35 | --has_part $HAS_PART --ldr_img_dir $LDR_IMG_DIR --val_frame $VAL_FRAME\ 36 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS 37 | 38 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/init.ckpt 39 | 40 | # extract emitters 41 | python extract_emitter_ldr.py --mode update\ 42 | --scene $DATASET_ROOT$SCENE\ 43 | --output checkpoints/$EXP/bake\ 44 | --ckpt checkpoints/$EXP/init.ckpt\ 45 | --dataset $DATASET --ldr_img_dir $LDR_IMG_DIR 46 | 47 | python bake_shading.py \ 48 | --scene $DATASET_ROOT$SCENE --dataset $DATASET \ 49 | --ldr_img_dir $LDR_IMG_DIR \ 50 | --slf_path checkpoints/$EXP/bake/vslf.npz \ 51 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 52 | --output outputs/$EXP/shading 53 | 54 | # optimize BRDF, CRF 55 | python train_brdf_crf.py --experiment_name $EXP \ 56 | --max_epochs 4 --dir_val val_0 \ 57 | --ckpt_path checkpoints/$EXP/init.ckpt \ 58 | --voxel_path checkpoints/$EXP/bake/vslf.npz \ 59 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 60 | --cache_dir outputs/$EXP/shading \ 61 | --dataset $DATASET $DATASET_ROOT$SCENE \ 62 | --has_part $HAS_PART --ldr_img_dir $LDR_IMG_DIR --val_frame $VAL_FRAME\ 63 | --SPP $SPP --spp $spp --lp 0.005 --la 0.01 --l_crf_weight 0.001 --crf_basis $CRF_BASIS 64 | 65 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/last_0.ckpt 66 | 67 | # refine SLF 68 | python slf_refine.py --scene $DATASET_ROOT$SCENE \ 69 | --output checkpoints/$EXP/bake --load vslf.npz --save vslf_0.npz \ 70 | --dataset $DATASET --ldr_img_dir $LDR_IMG_DIR \ 71 | --ckpt checkpoints/$EXP/last_0.ckpt --crf_basis $CRF_BASIS 72 | 73 | # refine emitter 74 | python train_emitter.py --experiment_name $EXP \ 75 | --max_epochs 1 --dir_val val_0_emitter \ 76 | --ckpt_path checkpoints/$EXP/last_0.ckpt \ 77 | --voxel_path checkpoints/$EXP/bake/vslf_0.npz \ 78 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 79 | --dataset $DATASET $DATASET_ROOT$SCENE \ 80 | --has_part $HAS_PART --ldr_img_dir $LDR_IMG_DIR --val_frame $VAL_FRAME\ 81 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS 82 | 83 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/last_0.ckpt 84 | 85 | # extract emitter 86 | python extract_emitter_ldr.py --mode update\ 87 | --scene $DATASET_ROOT$SCENE\ 88 | --output checkpoints/$EXP/bake\ 89 | --ckpt checkpoints/$EXP/last_0.ckpt\ 90 | --dataset $DATASET --ldr_img_dir $LDR_IMG_DIR 91 | 92 | # # refine shading 93 | python refine_shading.py \ 94 | --scene $DATASET_ROOT$SCENE --dataset $DATASET \ 95 | --ldr_img_dir $LDR_IMG_DIR \ 96 | --slf_path checkpoints/$EXP/bake/vslf_0.npz \ 97 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 98 | --ckpt checkpoints/$EXP/last_0.ckpt \ 99 | --output outputs/$EXP/shading 100 | 101 | # optimize BRDF, CRF 102 | python train_brdf_crf.py --experiment_name $EXP \ 103 | --max_epochs 4 --dir_val val_1 \ 104 | --ckpt_path checkpoints/$EXP/init.ckpt \ 105 | --voxel_path checkpoints/$EXP/bake/vslf_0.npz \ 106 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 107 | --cache_dir outputs/$EXP/shading \ 108 | --dataset $DATASET $DATASET_ROOT$SCENE \ 109 | --has_part $HAS_PART --ldr_img_dir $LDR_IMG_DIR --val_frame $VAL_FRAME\ 110 | --SPP $SPP --spp $spp --lp 0.005 --la 0.01 --l_crf_weight 0.001 --crf_basis $CRF_BASIS 111 | 112 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/last_1.ckpt -------------------------------------------------------------------------------- /scripts/fipt/bedroom/train.sh: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | # data folder 8 | DATASET_ROOT='/hdd/datasets/fipt/indoor_synthetic/' 9 | DATASET='synthetic' 10 | # scene name 11 | SCENE='bedroom' 12 | LDR_IMG_DIR='Image' 13 | EXP='fipt_syn_bedroom' 14 | VAL_FRAME=10 15 | CRF_BASIS=3 16 | # whether has part segmentation 17 | HAS_PART=1 18 | SPP=128 19 | spp=32 20 | 21 | # bake surface light field (SLF) 22 | python slf_bake.py --scene $DATASET_ROOT$SCENE\ 23 | --output checkpoints/$EXP/bake \ 24 | --dataset $DATASET --ldr_img_dir $LDR_IMG_DIR 25 | 26 | # extract emitter mask 27 | python extract_emitter_ldr.py --scene $DATASET_ROOT$SCENE\ 28 | --output checkpoints/$EXP/bake --dataset $DATASET\ 29 | --ldr_img_dir $LDR_IMG_DIR --threshold 0.99 30 | 31 | python initialize.py --experiment_name $EXP --max_epochs 3 \ 32 | --dataset $DATASET $DATASET_ROOT$SCENE \ 33 | --voxel_path checkpoints/$EXP/bake/vslf.npz \ 34 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 35 | --has_part $HAS_PART --ldr_img_dir $LDR_IMG_DIR --val_frame $VAL_FRAME\ 36 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS 37 | 38 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/init.ckpt 39 | 40 | # extract emitters 41 | python extract_emitter_ldr.py --mode update\ 42 | --scene $DATASET_ROOT$SCENE\ 43 | --output checkpoints/$EXP/bake\ 44 | --ckpt checkpoints/$EXP/init.ckpt\ 45 | --dataset $DATASET --ldr_img_dir $LDR_IMG_DIR 46 | 47 | python bake_shading.py \ 48 | --scene $DATASET_ROOT$SCENE --dataset $DATASET \ 49 | --ldr_img_dir $LDR_IMG_DIR \ 50 | --slf_path checkpoints/$EXP/bake/vslf.npz \ 51 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 52 | --output outputs/$EXP/shading 53 | 54 | # optimize BRDF, CRF 55 | python train_brdf_crf.py --experiment_name $EXP \ 56 | --max_epochs 2 --dir_val val_0 \ 57 | --ckpt_path checkpoints/$EXP/init.ckpt \ 58 | --voxel_path checkpoints/$EXP/bake/vslf.npz \ 59 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 60 | --cache_dir outputs/$EXP/shading \ 61 | --dataset $DATASET $DATASET_ROOT$SCENE \ 62 | --has_part $HAS_PART --ldr_img_dir $LDR_IMG_DIR --val_frame $VAL_FRAME\ 63 | --SPP $SPP --spp $spp --lp 0.005 --la 0.01 --l_crf_weight 0.001 --crf_basis $CRF_BASIS 64 | 65 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/last_0.ckpt 66 | 67 | # refine SLF 68 | python slf_refine.py --scene $DATASET_ROOT$SCENE \ 69 | --output checkpoints/$EXP/bake --load vslf.npz --save vslf_0.npz \ 70 | --dataset $DATASET --ldr_img_dir $LDR_IMG_DIR \ 71 | --ckpt checkpoints/$EXP/last_0.ckpt --crf_basis $CRF_BASIS 72 | 73 | # refine emitter 74 | python train_emitter.py --experiment_name $EXP \ 75 | --max_epochs 1 --dir_val val_0_emitter \ 76 | --ckpt_path checkpoints/$EXP/last_0.ckpt \ 77 | --voxel_path checkpoints/$EXP/bake/vslf_0.npz \ 78 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 79 | --dataset $DATASET $DATASET_ROOT$SCENE \ 80 | --has_part $HAS_PART --ldr_img_dir $LDR_IMG_DIR --val_frame $VAL_FRAME\ 81 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS 82 | 83 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/last_0.ckpt 84 | 85 | # extract emitter 86 | python extract_emitter_ldr.py --mode update\ 87 | --scene $DATASET_ROOT$SCENE\ 88 | --output checkpoints/$EXP/bake\ 89 | --ckpt checkpoints/$EXP/last_0.ckpt\ 90 | --dataset $DATASET --ldr_img_dir $LDR_IMG_DIR 91 | 92 | # # refine shading 93 | python refine_shading.py \ 94 | --scene $DATASET_ROOT$SCENE --dataset $DATASET \ 95 | --ldr_img_dir $LDR_IMG_DIR \ 96 | --slf_path checkpoints/$EXP/bake/vslf_0.npz \ 97 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 98 | --ckpt checkpoints/$EXP/last_0.ckpt \ 99 | --output outputs/$EXP/shading 100 | 101 | # optimize BRDF, CRF 102 | python train_brdf_crf.py --experiment_name $EXP \ 103 | --max_epochs 2 --dir_val val_1 \ 104 | --ckpt_path checkpoints/$EXP/init.ckpt \ 105 | --voxel_path checkpoints/$EXP/bake/vslf_0.npz \ 106 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 107 | --cache_dir outputs/$EXP/shading \ 108 | --dataset $DATASET $DATASET_ROOT$SCENE \ 109 | --has_part $HAS_PART --ldr_img_dir $LDR_IMG_DIR --val_frame $VAL_FRAME\ 110 | --SPP $SPP --spp $spp --lp 0.005 --la 0.01 --l_crf_weight 0.001 --crf_basis $CRF_BASIS 111 | 112 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/last_1.ckpt -------------------------------------------------------------------------------- /scripts/fipt/conferenceroom/train.sh: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | # data folder 8 | DATASET_ROOT='/hdd/datasets/fipt/real/' 9 | DATASET='real' 10 | # scene name 11 | SCENE='conferenceroom' 12 | LDR_IMG_DIR='Image' 13 | EXP='fipt_real_conferenceroom' 14 | VAL_FRAME=0 15 | CRF_BASIS=3 16 | # whether has part segmentation 17 | HAS_PART=0 18 | SPP=128 19 | spp=32 20 | 21 | # bake surface light field (SLF) 22 | python slf_bake.py --scene $DATASET_ROOT$SCENE\ 23 | --output checkpoints/$EXP/bake \ 24 | --dataset $DATASET --ldr_img_dir $LDR_IMG_DIR 25 | 26 | # extract emitter mask 27 | python extract_emitter_ldr.py --scene $DATASET_ROOT$SCENE\ 28 | --output checkpoints/$EXP/bake --dataset $DATASET\ 29 | --ldr_img_dir $LDR_IMG_DIR --threshold 0.99 30 | 31 | python initialize.py --experiment_name $EXP --max_epochs 3 \ 32 | --dataset $DATASET $DATASET_ROOT$SCENE \ 33 | --voxel_path checkpoints/$EXP/bake/vslf.npz \ 34 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 35 | --has_part $HAS_PART --ldr_img_dir $LDR_IMG_DIR --val_frame $VAL_FRAME\ 36 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS 37 | 38 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/init.ckpt 39 | 40 | # extract emitters 41 | python extract_emitter_ldr.py --mode update\ 42 | --scene $DATASET_ROOT$SCENE\ 43 | --output checkpoints/$EXP/bake\ 44 | --ckpt checkpoints/$EXP/init.ckpt\ 45 | --dataset $DATASET --ldr_img_dir $LDR_IMG_DIR 46 | 47 | python bake_shading.py \ 48 | --scene $DATASET_ROOT$SCENE --dataset $DATASET \ 49 | --ldr_img_dir $LDR_IMG_DIR \ 50 | --slf_path checkpoints/$EXP/bake/vslf.npz \ 51 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 52 | --output outputs/$EXP/shading 53 | 54 | # optimize BRDF, CRF 55 | python train_brdf_crf.py --experiment_name $EXP \ 56 | --max_epochs 2 --dir_val val_0 \ 57 | --ckpt_path checkpoints/$EXP/init.ckpt \ 58 | --voxel_path checkpoints/$EXP/bake/vslf.npz \ 59 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 60 | --cache_dir outputs/$EXP/shading \ 61 | --dataset $DATASET $DATASET_ROOT$SCENE \ 62 | --has_part $HAS_PART --ldr_img_dir $LDR_IMG_DIR --val_frame $VAL_FRAME\ 63 | --SPP $SPP --spp $spp --lp 0.005 --la 0.01 --l_crf_weight 0.001 --crf_basis $CRF_BASIS 64 | 65 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/last_0.ckpt 66 | 67 | # refine SLF 68 | python slf_refine.py --scene $DATASET_ROOT$SCENE \ 69 | --output checkpoints/$EXP/bake --load vslf.npz --save vslf_0.npz \ 70 | --dataset $DATASET --ldr_img_dir $LDR_IMG_DIR \ 71 | --ckpt checkpoints/$EXP/last_0.ckpt --crf_basis $CRF_BASIS 72 | 73 | # refine emitter 74 | python train_emitter.py --experiment_name $EXP \ 75 | --max_epochs 1 --dir_val val_0_emitter \ 76 | --ckpt_path checkpoints/$EXP/last_0.ckpt \ 77 | --voxel_path checkpoints/$EXP/bake/vslf_0.npz \ 78 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 79 | --dataset $DATASET $DATASET_ROOT$SCENE \ 80 | --has_part $HAS_PART --ldr_img_dir $LDR_IMG_DIR --val_frame $VAL_FRAME\ 81 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS 82 | 83 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/last_0.ckpt 84 | 85 | # extract emitter 86 | python extract_emitter_ldr.py --mode update\ 87 | --scene $DATASET_ROOT$SCENE\ 88 | --output checkpoints/$EXP/bake\ 89 | --ckpt checkpoints/$EXP/last_0.ckpt\ 90 | --dataset $DATASET --ldr_img_dir $LDR_IMG_DIR 91 | 92 | # # refine shading 93 | python refine_shading.py \ 94 | --scene $DATASET_ROOT$SCENE --dataset $DATASET \ 95 | --ldr_img_dir $LDR_IMG_DIR \ 96 | --slf_path checkpoints/$EXP/bake/vslf_0.npz \ 97 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 98 | --ckpt checkpoints/$EXP/last_0.ckpt \ 99 | --output outputs/$EXP/shading 100 | 101 | # optimize BRDF, CRF 102 | python train_brdf_crf.py --experiment_name $EXP \ 103 | --max_epochs 2 --dir_val val_1 \ 104 | --ckpt_path checkpoints/$EXP/init.ckpt \ 105 | --voxel_path checkpoints/$EXP/bake/vslf_0.npz \ 106 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 107 | --cache_dir outputs/$EXP/shading \ 108 | --dataset $DATASET $DATASET_ROOT$SCENE \ 109 | --has_part $HAS_PART --ldr_img_dir $LDR_IMG_DIR --val_frame $VAL_FRAME\ 110 | --SPP $SPP --spp $spp --lp 0.005 --la 0.01 --l_crf_weight 0.001 --crf_basis $CRF_BASIS 111 | 112 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/last_1.ckpt -------------------------------------------------------------------------------- /scripts/fipt/kitchen/train.sh: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | # data folder 8 | DATASET_ROOT='/hdd/datasets/fipt/indoor_synthetic/' 9 | DATASET='synthetic' 10 | # scene name 11 | SCENE='kitchen' 12 | LDR_IMG_DIR='Image' 13 | EXP='fipt_syn_kitchen' 14 | VAL_FRAME=10 15 | CRF_BASIS=3 16 | # whether has part segmentation 17 | HAS_PART=1 18 | SPP=128 19 | spp=32 20 | 21 | # bake surface light field (SLF) 22 | python slf_bake.py --scene $DATASET_ROOT$SCENE\ 23 | --output checkpoints/$EXP/bake \ 24 | --dataset $DATASET --ldr_img_dir $LDR_IMG_DIR 25 | 26 | # extract emitter mask 27 | python extract_emitter_ldr.py --scene $DATASET_ROOT$SCENE\ 28 | --output checkpoints/$EXP/bake --dataset $DATASET\ 29 | --ldr_img_dir $LDR_IMG_DIR --threshold 0.99 30 | 31 | python initialize.py --experiment_name $EXP --max_epochs 3 \ 32 | --dataset $DATASET $DATASET_ROOT$SCENE \ 33 | --voxel_path checkpoints/$EXP/bake/vslf.npz \ 34 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 35 | --has_part $HAS_PART --ldr_img_dir $LDR_IMG_DIR --val_frame $VAL_FRAME\ 36 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS 37 | 38 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/init.ckpt 39 | 40 | # extract emitters 41 | python extract_emitter_ldr.py --mode update\ 42 | --scene $DATASET_ROOT$SCENE\ 43 | --output checkpoints/$EXP/bake\ 44 | --ckpt checkpoints/$EXP/init.ckpt\ 45 | --dataset $DATASET --ldr_img_dir $LDR_IMG_DIR 46 | 47 | python bake_shading.py \ 48 | --scene $DATASET_ROOT$SCENE --dataset $DATASET \ 49 | --ldr_img_dir $LDR_IMG_DIR \ 50 | --slf_path checkpoints/$EXP/bake/vslf.npz \ 51 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 52 | --output outputs/$EXP/shading 53 | 54 | # optimize BRDF, CRF 55 | python train_brdf_crf.py --experiment_name $EXP \ 56 | --max_epochs 2 --dir_val val_0 \ 57 | --ckpt_path checkpoints/$EXP/init.ckpt \ 58 | --voxel_path checkpoints/$EXP/bake/vslf.npz \ 59 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 60 | --cache_dir outputs/$EXP/shading \ 61 | --dataset $DATASET $DATASET_ROOT$SCENE \ 62 | --has_part $HAS_PART --ldr_img_dir $LDR_IMG_DIR --val_frame $VAL_FRAME\ 63 | --SPP $SPP --spp $spp --lp 0.005 --la 0.01 --l_crf_weight 0.001 --crf_basis $CRF_BASIS 64 | 65 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/last_0.ckpt 66 | 67 | # refine SLF 68 | python slf_refine.py --scene $DATASET_ROOT$SCENE \ 69 | --output checkpoints/$EXP/bake --load vslf.npz --save vslf_0.npz \ 70 | --dataset $DATASET --ldr_img_dir $LDR_IMG_DIR \ 71 | --ckpt checkpoints/$EXP/last_0.ckpt --crf_basis $CRF_BASIS 72 | 73 | # refine emitter 74 | python train_emitter.py --experiment_name $EXP \ 75 | --max_epochs 1 --dir_val val_0_emitter \ 76 | --ckpt_path checkpoints/$EXP/last_0.ckpt \ 77 | --voxel_path checkpoints/$EXP/bake/vslf_0.npz \ 78 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 79 | --dataset $DATASET $DATASET_ROOT$SCENE \ 80 | --has_part $HAS_PART --ldr_img_dir $LDR_IMG_DIR --val_frame $VAL_FRAME\ 81 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS 82 | 83 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/last_0.ckpt 84 | 85 | # extract emitter 86 | python extract_emitter_ldr.py --mode update\ 87 | --scene $DATASET_ROOT$SCENE\ 88 | --output checkpoints/$EXP/bake\ 89 | --ckpt checkpoints/$EXP/last_0.ckpt\ 90 | --dataset $DATASET --ldr_img_dir $LDR_IMG_DIR 91 | 92 | # # refine shading 93 | python refine_shading.py \ 94 | --scene $DATASET_ROOT$SCENE --dataset $DATASET \ 95 | --ldr_img_dir $LDR_IMG_DIR \ 96 | --slf_path checkpoints/$EXP/bake/vslf_0.npz \ 97 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 98 | --ckpt checkpoints/$EXP/last_0.ckpt \ 99 | --output outputs/$EXP/shading 100 | 101 | # optimize BRDF, CRF 102 | python train_brdf_crf.py --experiment_name $EXP \ 103 | --max_epochs 2 --dir_val val_1 \ 104 | --ckpt_path checkpoints/$EXP/init.ckpt \ 105 | --voxel_path checkpoints/$EXP/bake/vslf_0.npz \ 106 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 107 | --cache_dir outputs/$EXP/shading \ 108 | --dataset $DATASET $DATASET_ROOT$SCENE \ 109 | --has_part $HAS_PART --ldr_img_dir $LDR_IMG_DIR --val_frame $VAL_FRAME\ 110 | --SPP $SPP --spp $spp --lp 0.005 --la 0.01 --l_crf_weight 0.001 --crf_basis $CRF_BASIS 111 | 112 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/last_1.ckpt -------------------------------------------------------------------------------- /scripts/fipt/livingroom/train.sh: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | # data folder 8 | DATASET_ROOT='/hdd/datasets/fipt/indoor_synthetic/' 9 | DATASET='synthetic' 10 | # scene name 11 | SCENE='livingroom' 12 | LDR_IMG_DIR='Image' 13 | EXP='fipt_syn_livingroom' 14 | VAL_FRAME=10 # bathroom: 10, bedroom: 10, livingroom:10, classroom:64, conferenceroom: 4 15 | CRF_BASIS=3 16 | # whether has part segmentation 17 | HAS_PART=1 18 | SPP=128 19 | spp=32 20 | 21 | # bake surface light field (SLF) 22 | python slf_bake.py --scene $DATASET_ROOT$SCENE\ 23 | --output checkpoints/$EXP/bake \ 24 | --dataset $DATASET --ldr_img_dir $LDR_IMG_DIR 25 | 26 | # extract emitter mask 27 | python extract_emitter_ldr.py --scene $DATASET_ROOT$SCENE\ 28 | --output checkpoints/$EXP/bake --dataset $DATASET\ 29 | --ldr_img_dir $LDR_IMG_DIR --threshold 0.99 30 | 31 | python initialize.py --experiment_name $EXP --max_epochs 3 \ 32 | --dataset $DATASET $DATASET_ROOT$SCENE \ 33 | --voxel_path checkpoints/$EXP/bake/vslf.npz \ 34 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 35 | --has_part $HAS_PART --ldr_img_dir $LDR_IMG_DIR --val_frame $VAL_FRAME\ 36 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS 37 | 38 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/init.ckpt 39 | 40 | # extract emitters 41 | python extract_emitter_ldr.py --mode update\ 42 | --scene $DATASET_ROOT$SCENE\ 43 | --output checkpoints/$EXP/bake\ 44 | --ckpt checkpoints/$EXP/init.ckpt\ 45 | --dataset $DATASET --ldr_img_dir $LDR_IMG_DIR 46 | 47 | python bake_shading.py \ 48 | --scene $DATASET_ROOT$SCENE --dataset $DATASET \ 49 | --ldr_img_dir $LDR_IMG_DIR \ 50 | --slf_path checkpoints/$EXP/bake/vslf.npz \ 51 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 52 | --output outputs/$EXP/shading 53 | 54 | # optimize BRDF, CRF 55 | python train_brdf_crf.py --experiment_name $EXP \ 56 | --max_epochs 2 --dir_val val_0 \ 57 | --ckpt_path checkpoints/$EXP/init.ckpt \ 58 | --voxel_path checkpoints/$EXP/bake/vslf.npz \ 59 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 60 | --cache_dir outputs/$EXP/shading \ 61 | --dataset $DATASET $DATASET_ROOT$SCENE \ 62 | --has_part $HAS_PART --ldr_img_dir $LDR_IMG_DIR --val_frame $VAL_FRAME\ 63 | --SPP $SPP --spp $spp --lp 0.005 --la 0.01 --l_crf_weight 0.001 --crf_basis $CRF_BASIS 64 | 65 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/last_0.ckpt 66 | 67 | # refine SLF 68 | python slf_refine.py --scene $DATASET_ROOT$SCENE \ 69 | --output checkpoints/$EXP/bake --load vslf.npz --save vslf_0.npz \ 70 | --dataset $DATASET --ldr_img_dir $LDR_IMG_DIR \ 71 | --ckpt checkpoints/$EXP/last_0.ckpt --crf_basis $CRF_BASIS 72 | 73 | # refine emitter 74 | python train_emitter.py --experiment_name $EXP \ 75 | --max_epochs 1 --dir_val val_0_emitter \ 76 | --ckpt_path checkpoints/$EXP/last_0.ckpt \ 77 | --voxel_path checkpoints/$EXP/bake/vslf_0.npz \ 78 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 79 | --dataset $DATASET $DATASET_ROOT$SCENE \ 80 | --has_part $HAS_PART --ldr_img_dir $LDR_IMG_DIR --val_frame $VAL_FRAME\ 81 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS 82 | 83 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/last_0.ckpt 84 | 85 | # extract emitter 86 | python extract_emitter_ldr.py --mode update\ 87 | --scene $DATASET_ROOT$SCENE\ 88 | --output checkpoints/$EXP/bake\ 89 | --ckpt checkpoints/$EXP/last_0.ckpt\ 90 | --dataset $DATASET --ldr_img_dir $LDR_IMG_DIR 91 | 92 | # # refine shading 93 | python refine_shading.py \ 94 | --scene $DATASET_ROOT$SCENE --dataset $DATASET \ 95 | --ldr_img_dir $LDR_IMG_DIR \ 96 | --slf_path checkpoints/$EXP/bake/vslf_0.npz \ 97 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 98 | --ckpt checkpoints/$EXP/last_0.ckpt \ 99 | --output outputs/$EXP/shading 100 | 101 | # optimize BRDF, CRF 102 | python train_brdf_crf.py --experiment_name $EXP \ 103 | --max_epochs 2 --dir_val val_1 \ 104 | --ckpt_path checkpoints/$EXP/init.ckpt \ 105 | --voxel_path checkpoints/$EXP/bake/vslf_0.npz \ 106 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 107 | --cache_dir outputs/$EXP/shading \ 108 | --dataset $DATASET $DATASET_ROOT$SCENE \ 109 | --has_part $HAS_PART --ldr_img_dir $LDR_IMG_DIR --val_frame $VAL_FRAME\ 110 | --SPP $SPP --spp $spp --lp 0.005 --la 0.01 --l_crf_weight 0.001 --crf_basis $CRF_BASIS 111 | 112 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/last_1.ckpt -------------------------------------------------------------------------------- /utils/gen_path.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import os 8 | import numpy as np 9 | from .ray_utils import generate_interpolated_path 10 | from .dataset.synthetic_ldr import SyntheticDatasetLDR 11 | from .dataset.real_ldr import RealDatasetLDR 12 | from .dataset.scannetpp.dataset import Scannetpp 13 | 14 | def generate_path_kitchen(): 15 | root_dir = '/hdd/datasets/fipt/indoor_synthetic/kitchen' 16 | dataset = SyntheticDatasetLDR(root_dir, split='val', pixel=True) 17 | poses = np.array(dataset.poses) 18 | # print('poses:', poses.shape) #(n, 3, 4) 19 | 20 | key_idx = [11, 12] 21 | pose_keyframe = poses[key_idx] 22 | render_traj = generate_interpolated_path(pose_keyframe, n_interp=300) 23 | output_path = os.path.join(root_dir, 'render_traj.npy') 24 | np.save(output_path, render_traj) 25 | print('Save', output_path) 26 | 27 | def generate_path_bathroom(): 28 | root_dir = '/hdd/datasets/fipt/indoor_synthetic/bathroom' 29 | dataset = SyntheticDatasetLDR(root_dir, split='val', pixel=True) 30 | poses = np.array(dataset.poses) 31 | 32 | key_idx = [10, 12] 33 | pose_keyframe = poses[key_idx] 34 | render_traj = generate_interpolated_path(pose_keyframe, n_interp=300) 35 | output_path = os.path.join(root_dir, 'render_traj.npy') 36 | np.save(output_path, render_traj) 37 | print('Save', output_path) 38 | 39 | def generate_path_bedroom(): 40 | root_dir = '/hdd/datasets/fipt/indoor_synthetic/bedroom' 41 | dataset = SyntheticDatasetLDR(root_dir, split='val', pixel=True) 42 | poses = np.array(dataset.poses) 43 | 44 | key_idx = [12, 10] 45 | pose_keyframe = poses[key_idx] 46 | render_traj = generate_interpolated_path(pose_keyframe, n_interp=300) 47 | output_path = os.path.join(root_dir, 'render_traj.npy') 48 | np.save(output_path, render_traj) 49 | print('Save', output_path) 50 | 51 | def generate_path_livingroom(): 52 | root_dir = '/hdd/datasets/fipt/indoor_synthetic/livingroom' 53 | dataset = SyntheticDatasetLDR(root_dir, split='val', pixel=True) 54 | poses = np.array(dataset.poses) 55 | 56 | key_idx = [10, 13] 57 | pose_keyframe = poses[key_idx] 58 | render_traj = generate_interpolated_path(pose_keyframe, n_interp=300) 59 | output_path = os.path.join(root_dir, 'render_traj.npy') 60 | np.save(output_path, render_traj) 61 | print('Save', output_path) 62 | 63 | def generate_path_classroom(): 64 | root_dir = '/hdd/datasets/fipt/real/classroom' 65 | dataset = RealDatasetLDR(root_dir, split='train', pixel=True) 66 | poses = np.array(dataset.C2Ws) 67 | 68 | key_idx = [76, 37] 69 | pose_keyframe = poses[key_idx] 70 | render_traj = generate_interpolated_path(pose_keyframe, n_interp=300) 71 | output_path = os.path.join(root_dir, 'render_traj.npy') 72 | np.save(output_path, render_traj) 73 | print('Save', output_path) 74 | 75 | def generate_path_conferenceroom(): 76 | root_dir = '/hdd/datasets/fipt/real/conferenceroom' 77 | dataset = RealDatasetLDR(root_dir, split='train', pixel=True) 78 | poses = np.array(dataset.C2Ws) 79 | 80 | key_idx = [30, 42, 108] 81 | pose_keyframe = poses[key_idx] 82 | render_traj = generate_interpolated_path(pose_keyframe, n_interp=150) 83 | output_path = os.path.join(root_dir, 'render_traj.npy') 84 | np.save(output_path, render_traj) 85 | print('Save', output_path) 86 | 87 | def generate_path_scannetpp_bathroom2(): 88 | root_dir = '/hdd/datasets/scannetpp' 89 | scene_id = '45b0dac5e3' 90 | split = 'all' 91 | n_frames = 150 92 | key_idx = [6, 227, 147] 93 | dataset = Scannetpp(root_dir, scene_id, split=split) 94 | poses = np.array(dataset.C2Ws) 95 | 96 | pose_keyframe = poses[key_idx] 97 | render_traj = generate_interpolated_path(pose_keyframe, n_interp=n_frames) 98 | output_path = os.path.join(root_dir, 'data', scene_id, 'psdf', 'render_traj.npy') 99 | np.save(output_path, render_traj) 100 | print('Save', output_path) 101 | 102 | def generate_path_scannetpp_room2(): 103 | root_dir = '/hdd/datasets/scannetpp' 104 | scene_id = '7e09430da7' 105 | split = 'train' 106 | n_frames = 225 107 | key_idx = [200, 121, 96] 108 | dataset = Scannetpp(root_dir, scene_id, split=split) 109 | poses = np.array(dataset.C2Ws) 110 | 111 | pose_keyframe = poses[key_idx] 112 | render_traj = generate_interpolated_path(pose_keyframe, n_interp=n_frames)[-300:] 113 | output_path = os.path.join(root_dir, 'data', scene_id, 'psdf', 'render_traj.npy') 114 | np.save(output_path, render_traj) 115 | print('Save', output_path) 116 | 117 | if __name__ == '__main__': 118 | generate_path_scannetpp_room2() -------------------------------------------------------------------------------- /scripts/scannetpp/room2/train.sh: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | # data folder 8 | DATASET_ROOT='/hdd/datasets/scannetpp/' 9 | DATASET='scannetpp' 10 | # scene name 11 | SCENE='7e09430da7' 12 | LDR_IMG_DIR='Image' 13 | EXP='scannetpp_room2' 14 | VAL_FRAME=3 15 | CRF_BASIS=3 16 | RES_SCALE=0.5 17 | # whether has part segmentation 18 | HAS_PART=0 19 | SPP=128 20 | spp=32 21 | 22 | # bake surface light field (SLF) 23 | python slf_bake.py --dataset_root $DATASET_ROOT --scene $SCENE\ 24 | --output checkpoints/$EXP/bake --res_scale $RES_SCALE\ 25 | --dataset $DATASET 26 | 27 | # extract emitter mask 28 | python extract_emitter_ldr.py \ 29 | --dataset_root $DATASET_ROOT --scene $SCENE\ 30 | --output checkpoints/$EXP/bake --dataset $DATASET --res_scale $RES_SCALE\ 31 | --threshold 0.99 32 | 33 | python initialize.py --experiment_name $EXP --max_epochs 5 \ 34 | --dataset $DATASET $DATASET_ROOT --scene $SCENE \ 35 | --voxel_path checkpoints/$EXP/bake/vslf.npz \ 36 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 37 | --has_part $HAS_PART --val_frame $VAL_FRAME\ 38 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS --res_scale $RES_SCALE 39 | 40 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/init.ckpt 41 | 42 | # extract emitters 43 | python extract_emitter_ldr.py --mode update\ 44 | --dataset_root $DATASET_ROOT --scene $SCENE\ 45 | --output checkpoints/$EXP/bake --res_scale $RES_SCALE\ 46 | --ckpt checkpoints/$EXP/init.ckpt\ 47 | --dataset $DATASET 48 | 49 | python bake_shading.py \ 50 | --dataset_root $DATASET_ROOT --scene $SCENE \ 51 | --dataset $DATASET --res_scale $RES_SCALE\ 52 | --slf_path checkpoints/$EXP/bake/vslf.npz \ 53 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 54 | --output outputs/$EXP/shading 55 | 56 | # optimize BRDF, CRF 57 | python train_brdf_crf.py --experiment_name $EXP \ 58 | --dataset $DATASET $DATASET_ROOT --scene $SCENE\ 59 | --has_part $HAS_PART --val_frame $VAL_FRAME --res_scale $RES_SCALE\ 60 | --max_epochs 2 --dir_val val_0 \ 61 | --ckpt_path checkpoints/$EXP/init.ckpt \ 62 | --voxel_path checkpoints/$EXP/bake/vslf.npz \ 63 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 64 | --cache_dir outputs/$EXP/shading \ 65 | --SPP $SPP --spp $spp --lp 0.005 --la 0.01 --l_crf_weight 0.001 --crf_basis $CRF_BASIS 66 | 67 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/last_0.ckpt 68 | 69 | # refine SLF 70 | python slf_refine.py --dataset_root $DATASET_ROOT --scene $SCENE \ 71 | --output checkpoints/$EXP/bake --load vslf.npz --save vslf_0.npz \ 72 | --dataset $DATASET --res_scale $RES_SCALE\ 73 | --ckpt checkpoints/$EXP/last_0.ckpt --crf_basis $CRF_BASIS 74 | 75 | # refine emitter 76 | python train_emitter.py --experiment_name $EXP \ 77 | --dataset $DATASET $DATASET_ROOT --scene $SCENE\ 78 | --has_part $HAS_PART --val_frame $VAL_FRAME --res_scale $RES_SCALE\ 79 | --max_epochs 1 --dir_val val_0_emitter \ 80 | --ckpt_path checkpoints/$EXP/last_0.ckpt \ 81 | --voxel_path checkpoints/$EXP/bake/vslf_0.npz \ 82 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 83 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS 84 | 85 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/last_0.ckpt 86 | 87 | # extract emitter 88 | python extract_emitter_ldr.py --mode update\ 89 | --dataset_root $DATASET_ROOT --scene $SCENE\ 90 | --output checkpoints/$EXP/bake --res_scale $RES_SCALE\ 91 | --ckpt checkpoints/$EXP/last_0.ckpt\ 92 | --dataset $DATASET --ldr_img_dir $LDR_IMG_DIR 93 | 94 | # refine shading 95 | python refine_shading.py \ 96 | --dataset_root $DATASET_ROOT --scene $SCENE \ 97 | --dataset $DATASET --res_scale $RES_SCALE\ 98 | --slf_path checkpoints/$EXP/bake/vslf_0.npz \ 99 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 100 | --ckpt checkpoints/$EXP/last_0.ckpt \ 101 | --output outputs/$EXP/shading 102 | 103 | # optimize BRDF, CRF 104 | python train_brdf_crf.py --experiment_name $EXP \ 105 | --dataset $DATASET $DATASET_ROOT --scene $SCENE\ 106 | --has_part $HAS_PART --val_frame $VAL_FRAME --res_scale $RES_SCALE\ 107 | --max_epochs 2 --dir_val val_1 \ 108 | --ckpt_path checkpoints/$EXP/init.ckpt \ 109 | --voxel_path checkpoints/$EXP/bake/vslf_0.npz \ 110 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 111 | --cache_dir outputs/$EXP/shading \ 112 | --SPP $SPP --spp $spp --lp 0.005 --la 0.01 --l_crf_weight 0.001 --crf_basis $CRF_BASIS 113 | 114 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/last_1.ckpt -------------------------------------------------------------------------------- /scripts/scannetpp/bathroom2/train.sh: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | # data folder 8 | DATASET_ROOT='/hdd/datasets/scannetpp/' 9 | DATASET='scannetpp' 10 | # scene name 11 | SCENE='45b0dac5e3' 12 | LDR_IMG_DIR='Image' 13 | EXP='scannetpp_bathroom2' 14 | VAL_FRAME=0 15 | CRF_BASIS=3 16 | RES_SCALE=0.5 17 | # whether has part segmentation 18 | HAS_PART=0 19 | SPP=128 20 | spp=32 21 | 22 | # bake surface light field (SLF) 23 | python slf_bake.py --dataset_root $DATASET_ROOT --scene $SCENE\ 24 | --output checkpoints/$EXP/bake --res_scale $RES_SCALE\ 25 | --dataset $DATASET 26 | 27 | # extract emitter mask 28 | python extract_emitter_ldr.py \ 29 | --dataset_root $DATASET_ROOT --scene $SCENE\ 30 | --output checkpoints/$EXP/bake --dataset $DATASET --res_scale $RES_SCALE\ 31 | --threshold 0.99 32 | 33 | python initialize.py --experiment_name $EXP --max_epochs 5 \ 34 | --dataset $DATASET $DATASET_ROOT --scene $SCENE \ 35 | --voxel_path checkpoints/$EXP/bake/vslf.npz \ 36 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 37 | --has_part $HAS_PART --val_frame $VAL_FRAME\ 38 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS --res_scale $RES_SCALE 39 | 40 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/init.ckpt 41 | 42 | # extract emitters 43 | python extract_emitter_ldr.py --mode update\ 44 | --dataset_root $DATASET_ROOT --scene $SCENE\ 45 | --output checkpoints/$EXP/bake --res_scale $RES_SCALE\ 46 | --ckpt checkpoints/$EXP/init.ckpt\ 47 | --dataset $DATASET 48 | 49 | python bake_shading.py \ 50 | --dataset_root $DATASET_ROOT --scene $SCENE \ 51 | --dataset $DATASET --res_scale $RES_SCALE\ 52 | --slf_path checkpoints/$EXP/bake/vslf.npz \ 53 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 54 | --output outputs/$EXP/shading 55 | 56 | # optimize BRDF, CRF 57 | python train_brdf_crf.py --experiment_name $EXP \ 58 | --dataset $DATASET $DATASET_ROOT --scene $SCENE\ 59 | --has_part $HAS_PART --val_frame $VAL_FRAME --res_scale $RES_SCALE\ 60 | --max_epochs 2 --dir_val val_0 \ 61 | --ckpt_path checkpoints/$EXP/init.ckpt \ 62 | --voxel_path checkpoints/$EXP/bake/vslf.npz \ 63 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 64 | --cache_dir outputs/$EXP/shading \ 65 | --SPP $SPP --spp $spp --lp 0.005 --la 0.01 --l_crf_weight 0.001 --crf_basis $CRF_BASIS 66 | 67 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/last_0.ckpt 68 | 69 | # refine SLF 70 | python slf_refine.py --dataset_root $DATASET_ROOT --scene $SCENE \ 71 | --output checkpoints/$EXP/bake --load vslf.npz --save vslf_0.npz \ 72 | --dataset $DATASET --res_scale $RES_SCALE\ 73 | --ckpt checkpoints/$EXP/last_0.ckpt --crf_basis $CRF_BASIS 74 | 75 | # refine emitter 76 | python train_emitter.py --experiment_name $EXP \ 77 | --dataset $DATASET $DATASET_ROOT --scene $SCENE\ 78 | --has_part $HAS_PART --val_frame $VAL_FRAME --res_scale $RES_SCALE\ 79 | --max_epochs 1 --dir_val val_0_emitter \ 80 | --ckpt_path checkpoints/$EXP/last_0.ckpt \ 81 | --voxel_path checkpoints/$EXP/bake/vslf_0.npz \ 82 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 83 | --SPP $SPP --spp $spp --crf_basis $CRF_BASIS 84 | 85 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/last_0.ckpt 86 | 87 | # extract emitter 88 | python extract_emitter_ldr.py --mode update\ 89 | --dataset_root $DATASET_ROOT --scene $SCENE\ 90 | --output checkpoints/$EXP/bake --res_scale $RES_SCALE\ 91 | --ckpt checkpoints/$EXP/last_0.ckpt\ 92 | --dataset $DATASET --ldr_img_dir $LDR_IMG_DIR 93 | 94 | # refine shading 95 | python refine_shading.py \ 96 | --dataset_root $DATASET_ROOT --scene $SCENE \ 97 | --dataset $DATASET --res_scale $RES_SCALE\ 98 | --slf_path checkpoints/$EXP/bake/vslf_0.npz \ 99 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 100 | --ckpt checkpoints/$EXP/last_0.ckpt \ 101 | --output outputs/$EXP/shading 102 | 103 | # optimize BRDF, CRF 104 | python train_brdf_crf.py --experiment_name $EXP \ 105 | --dataset $DATASET $DATASET_ROOT --scene $SCENE\ 106 | --has_part $HAS_PART --val_frame $VAL_FRAME --res_scale $RES_SCALE\ 107 | --max_epochs 2 --dir_val val_1 \ 108 | --ckpt_path checkpoints/$EXP/init.ckpt \ 109 | --voxel_path checkpoints/$EXP/bake/vslf_0.npz \ 110 | --emitter_path checkpoints/$EXP/bake/emitter.pth \ 111 | --cache_dir outputs/$EXP/shading \ 112 | --SPP $SPP --spp $spp --lp 0.005 --la 0.01 --l_crf_weight 0.001 --crf_basis $CRF_BASIS 113 | 114 | mv checkpoints/$EXP/last.ckpt checkpoints/$EXP/last_1.ckpt -------------------------------------------------------------------------------- /utils/dataset/neilf/extract_geometry.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import torch 8 | 9 | import mitsuba 10 | mitsuba.set_variant('cuda_ad_rgb') 11 | import numpy as np 12 | import torch 13 | import torch.nn.functional as F 14 | import math 15 | 16 | import os 17 | os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1" 18 | import cv2 19 | 20 | from utils.dataset import SyntheticDataset,RealDataset,SyntheticDatasetLDR,RealDatasetLDR 21 | from utils.dataset.scannetpp.dataset import Scannetpp 22 | from utils.path_tracing import ray_intersect 23 | from model.brdf import BaseBRDF 24 | from model.emitter import SLFEmitter 25 | from pathlib import Path 26 | from tqdm import tqdm 27 | from argparse import ArgumentParser 28 | import time 29 | 30 | if __name__ == '__main__': 31 | parser = ArgumentParser() 32 | parser.add_argument('--dataset_root', type=str, help='dataset root') 33 | parser.add_argument('--scene', type=str, required=True, help='dataset folder') 34 | parser.add_argument('--output', type=str, required=True, help='output path') 35 | parser.add_argument('--dataset', type=str,required=True, help='dataset type') 36 | parser.add_argument('--input', type=str, default='ldr', choices=['hdr', 'ldr']) 37 | parser.add_argument('--split', type=str, default='train') 38 | parser.add_argument('--ldr_img_dir', type=str, default=None) 39 | parser.add_argument('--res_scale', type=float, default=1.0) 40 | args = parser.parse_args() 41 | 42 | device = torch.device(0) # use gpu device 0 43 | 44 | DATASET_PATH = args.scene 45 | OUTPUT_PATH = args.output 46 | os.makedirs(OUTPUT_PATH,exist_ok=True) 47 | dir_position = os.path.join(OUTPUT_PATH, 'position_maps') 48 | dir_normal = os.path.join(OUTPUT_PATH, 'normal_maps') 49 | dir_image = os.path.join(OUTPUT_PATH, 'images') 50 | dir_depth = os.path.join(OUTPUT_PATH, 'depth_maps') 51 | os.makedirs(dir_position, exist_ok=True) 52 | os.makedirs(dir_normal, exist_ok=True) 53 | os.makedirs(dir_image, exist_ok=True) 54 | os.makedirs(dir_depth, exist_ok=True) 55 | 56 | # load mesh 57 | if args.dataset in ['synthetic', 'real']: 58 | mesh_path = os.path.join(DATASET_PATH,'scene.obj') 59 | mesh_type = 'obj' 60 | elif args.dataset == 'scannetpp': 61 | mesh_path = os.path.join(args.dataset_root, 'data', args.scene, 'scans', 'scene.ply') 62 | mesh_type = 'ply' 63 | assert Path(mesh_path).exists(), 'mesh not found: '+mesh_path 64 | 65 | scene = mitsuba.load_dict({ 66 | 'type': 'scene', 67 | 'shape_id':{ 68 | 'type': mesh_type, 69 | 'filename': mesh_path, 70 | } 71 | }) 72 | 73 | # load dataset 74 | if args.input == 'hdr': 75 | if args.dataset == 'synthetic': 76 | dataset = SyntheticDataset(DATASET_PATH, split=args.split, pixel=False) 77 | elif args.dataset == 'real': 78 | dataset = RealDataset(DATASET_PATH, split=args.split, pixel=False) 79 | else: 80 | if args.dataset == 'synthetic': 81 | dataset = SyntheticDatasetLDR(DATASET_PATH, img_dir=args.ldr_img_dir, split=args.split, pixel=False) 82 | elif args.dataset == 'real': 83 | dataset = RealDatasetLDR(DATASET_PATH, img_dir=args.ldr_img_dir, split=args.split, pixel=False) 84 | elif args.dataset == 'scannetpp': 85 | dataset = Scannetpp(args.dataset_root, args.scene, split=args.split, pixel=False, res_scale=args.res_scale) 86 | img_hw = dataset.img_hw 87 | img_hw = dataset.img_hw 88 | 89 | 90 | denoiser = mitsuba.OptixDenoiser(img_hw[::-1]) 91 | 92 | im_id = 0 93 | for i in tqdm(range(len(dataset))): 94 | name = dataset.img_name_list[i].split('.')[0] 95 | batch = dataset[i] 96 | rays = batch['rays'] 97 | xs = rays[...,:3] 98 | ds = rays[...,3:6] 99 | rgb = batch['rgbs'] 100 | rgb = rgb.reshape(*img_hw, 3).cpu().numpy() 101 | rgb = (rgb*255).astype(np.uint8) 102 | rgb_path = os.path.join(dir_image, '{}.png'.format(name)) 103 | cv2.imwrite(rgb_path, cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)) 104 | 105 | positions,normals,_,_,valid = ray_intersect(scene,xs.to(device),ds.to(device)) 106 | positions = positions.reshape(*img_hw, 3).cpu().numpy() 107 | normals = normals.reshape(*img_hw, 3).cpu().numpy() 108 | ds = F.normalize(ds, dim=-1).reshape(*img_hw, 3).cpu().numpy() 109 | xs = xs.reshape(*img_hw, 3).cpu().numpy() 110 | depths = np.sum((positions - xs) * ds, axis=-1) 111 | 112 | position_path = os.path.join(dir_position, '{}.exr'.format(name)) 113 | cv2.imwrite(position_path, positions) 114 | normal_path = os.path.join(dir_normal, '{}.exr'.format(name)) 115 | cv2.imwrite(normal_path, normals) 116 | depth_path = os.path.join(dir_depth, '{}.npy'.format(name)) 117 | np.save(depth_path, depths) -------------------------------------------------------------------------------- /extract_emitter_ldr.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import torch 8 | import torch.nn.functional as NF 9 | import mitsuba 10 | mitsuba.set_variant('cuda_ad_rgb') 11 | 12 | from utils.dataset import InvRealDatasetLDR,InvSyntheticDatasetLDR 13 | from utils.dataset.scannetpp.dataset import InvScannetpp 14 | from utils.path_tracing import ray_intersect 15 | import os 16 | import numpy as np 17 | import trimesh 18 | from tqdm import tqdm 19 | from argparse import ArgumentParser 20 | import torch_scatter 21 | from pathlib import Path 22 | from const import set_random_seed 23 | set_random_seed() 24 | 25 | def main(): 26 | parser = ArgumentParser() 27 | parser.add_argument('--dataset_root', type=str, help='dataset root') 28 | parser.add_argument('--scene', type=str, required=True, help= 'dataset folder') 29 | parser.add_argument('--output', type=str, required=True, help='output path') 30 | parser.add_argument('--dataset', type=str,required=True, help='dataset type') 31 | parser.add_argument('--mode', type=str, default='export', choices=['export', 'update', 'test']) 32 | parser.add_argument('--ckpt', type=str,help='checkpoint path') 33 | parser.add_argument('--spp', type=int,default=100, help='number of samples for each triangle emitter') 34 | parser.add_argument('--threshold',type=float,default=0.99,help='threshold for emitter') 35 | parser.add_argument('--ldr_img_dir', type=str, default=None) 36 | parser.add_argument('--res_scale', type=float, default=1.0) 37 | args = parser.parse_args() 38 | 39 | device = torch.device(0) 40 | 41 | SCENE = args.scene 42 | OUTPUT =args.output 43 | os.makedirs(OUTPUT, exist_ok=True) 44 | 45 | # load geometry 46 | if args.dataset in ['synthetic', 'real']: 47 | mesh_path = os.path.join(args.scene,'scene.obj') 48 | mesh_type = 'obj' 49 | elif args.dataset == 'scannetpp': 50 | mesh_path = os.path.join(args.dataset_root, 'data', args.scene, 'scans', 'scene.ply') 51 | mesh_type = 'ply' 52 | assert Path(mesh_path).exists(), 'mesh not found: '+mesh_path 53 | 54 | scene = mitsuba.load_dict({ 55 | 'type': 'scene', 56 | 'shape_id':{ 57 | 'type': mesh_type, 58 | 'filename': mesh_path, 59 | } 60 | }) 61 | 62 | if args.dataset == 'synthetic': 63 | dataset = InvSyntheticDatasetLDR(SCENE, img_dir=args.ldr_img_dir, split='train', pixel=False) 64 | elif args.dataset == 'real': 65 | dataset = InvRealDatasetLDR(SCENE, img_dir=args.ldr_img_dir, split='train', pixel=False) 66 | elif args.dataset == 'scannetpp': 67 | dataset = InvScannetpp(args.dataset_root, args.scene, split='train', pixel=False, res_scale=args.res_scale) 68 | img_hw = dataset.img_hw 69 | 70 | 71 | # get mesh vertices and triangles 72 | if args.mode == 'export': 73 | mesh = trimesh.load_mesh(mesh_path) 74 | vertices = torch.from_numpy(np.array(mesh.vertices)).float() #(v, 3) 75 | faces = torch.from_numpy(np.array(mesh.faces)) #(f, 3) 76 | 77 | n_face = len(faces) 78 | triangle_radiance = torch.zeros(n_face, 3) 79 | triangle_count = torch.zeros(n_face) 80 | for batch in tqdm(dataset): 81 | rays = batch['rays'] 82 | rays_x,rays_d = rays[...,:3].to(device),rays[...,3:6].to(device) 83 | positions,normals,uvs,triangle_idxs,valid = ray_intersect(scene,rays_x,rays_d) 84 | 85 | triangle_idxs = triangle_idxs[valid].cpu() 86 | radiance = batch['rgbs'][valid.cpu()] 87 | # segmentation = batch['segmentation'][valid.cpu()] 88 | # seg_idxs, inv_idxs = segmentation.unique(return_inverse=True) 89 | 90 | triangle_radiance = torch_scatter.scatter( 91 | radiance, triangle_idxs, 0, triangle_radiance, reduce='sum' 92 | ) 93 | triangle_count = torch_scatter.scatter( 94 | torch.ones(len(triangle_idxs)), triangle_idxs, 0, triangle_count, reduce='sum' 95 | ) 96 | 97 | triangle_radiance_mean = triangle_radiance / triangle_count.unsqueeze(-1).clamp_min(1) #(f, 3) 98 | triangle_radiance_mean = torch.max(triangle_radiance_mean, dim=-1)[0] # max value among 3 channels 99 | 100 | is_emitter = triangle_radiance_mean > args.threshold 101 | emitter_vertices = vertices[faces[is_emitter]] 102 | emitter_area = torch.cross(emitter_vertices[:,1]-emitter_vertices[:,0], 103 | emitter_vertices[:,2]-emitter_vertices[:,0],-1) 104 | emitter_normal = NF.normalize(emitter_area,dim=-1) 105 | emitter_area = emitter_area.norm(dim=-1)/2.0 106 | emitter_radiance = torch.zeros(n_face, 3) 107 | 108 | # print('is_emitter ratio:', torch.sum(is_emitter) / is_emitter.numel()) 109 | torch.save({ 110 | 'is_emitter': is_emitter, 111 | 'emitter_vertices': emitter_vertices, 112 | 'emitter_area': emitter_area, 113 | 'emitter_normal': emitter_normal, 114 | 'emitter_radiance': emitter_radiance 115 | }, os.path.join(OUTPUT,'emitter.pth')) 116 | 117 | if args.mode == 'update': 118 | ckpt_state = torch.load(args.ckpt, map_location='cpu')['state_dict'] 119 | emitter_path = os.path.join(OUTPUT, 'emitter.pth') 120 | emitter_state = torch.load(emitter_path, map_location='cpu') 121 | emitter_state['emitter_radiance'] = ckpt_state['emitter.radiance'] 122 | torch.save(emitter_state, emitter_path) 123 | 124 | if __name__ == '__main__': 125 | main() -------------------------------------------------------------------------------- /slf_bake.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import torch 8 | import mitsuba 9 | mitsuba.set_variant('cuda_ad_rgb') 10 | import os 11 | os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1" 12 | 13 | from utils.dataset import SyntheticDatasetLDR,RealDatasetLDR 14 | from utils.dataset.scannetpp.dataset import Scannetpp 15 | from utils.path_tracing import ray_intersect 16 | from model.slf import VoxelSLF 17 | from crf.model_crf import EmorCRF 18 | from pathlib import Path 19 | from tqdm import tqdm 20 | from argparse import ArgumentParser 21 | import time 22 | from const import set_random_seed 23 | set_random_seed() 24 | 25 | if __name__ == '__main__': 26 | parser = ArgumentParser() 27 | parser.add_argument('--dataset_root', type=str, help='dataset root') 28 | parser.add_argument('--scene', type=str, required=True, help='dataset folder') 29 | parser.add_argument('--output', type=str, required=True, help='output path') 30 | parser.add_argument('--dataset', type=str,required=True, help='dataset type') 31 | parser.add_argument('--voxel_num', type=int,default=256, help='resolution for voxel radiance cache') 32 | parser.add_argument('--ldr_img_dir', type=str, default=None) 33 | parser.add_argument('--res_scale', type=float, default=1.0) 34 | args = parser.parse_args() 35 | 36 | device = torch.device(0) # use gpu device 0 37 | 38 | DATASET_PATH = args.scene 39 | OUTPUT_PATH = args.output 40 | os.makedirs(OUTPUT_PATH,exist_ok=True) 41 | # load mesh 42 | if args.dataset in ['synthetic', 'real']: 43 | mesh_path = os.path.join(DATASET_PATH,'scene.obj') 44 | mesh_type = 'obj' 45 | elif args.dataset == 'scannetpp': 46 | mesh_path = os.path.join(args.dataset_root, 'data', args.scene, 'scans', 'scene.ply') 47 | mesh_type = 'ply' 48 | assert Path(mesh_path).exists(), 'mesh not found: '+mesh_path 49 | 50 | scene = mitsuba.load_dict({ 51 | 'type': 'scene', 52 | 'shape_id':{ 53 | 'type': mesh_type, 54 | 'filename': mesh_path, 55 | } 56 | }) 57 | 58 | # load dataset 59 | if args.dataset == 'synthetic': 60 | dataset = SyntheticDatasetLDR(DATASET_PATH, img_dir=args.ldr_img_dir, split='train', pixel=False) 61 | elif args.dataset == 'real': 62 | dataset = RealDatasetLDR(DATASET_PATH, img_dir=args.ldr_img_dir, split='train', pixel=False) 63 | elif args.dataset == 'scannetpp': 64 | dataset = Scannetpp(args.dataset_root, args.scene, split='train', pixel=False, res_scale=args.res_scale) 65 | img_hw = dataset.img_hw 66 | 67 | model_crf = EmorCRF() 68 | 69 | start_time = time.time() 70 | # extract scene bounding box 71 | print('find scene bound') 72 | voxel_min = 1000. 73 | voxel_max = 0.0 74 | for idx in tqdm(range(len(dataset))): 75 | batch = dataset[idx] 76 | rays = batch['rays'] 77 | xs = rays[...,:3] 78 | ds = rays[...,3:6] 79 | 80 | positions,_,_,_,valid = ray_intersect(scene,xs.to(device),ds.to(device)) 81 | if not valid.any(): 82 | continue 83 | position = positions[valid] 84 | voxel_min = min(voxel_min,position.min()) 85 | voxel_max = max(voxel_max,position.max()) 86 | 87 | if args.dataset in ['synthetic', 'real']: 88 | voxel_min = 1.1*voxel_min 89 | voxel_max = 1.1*voxel_max 90 | else: 91 | voxel_c = (voxel_min + voxel_max) 92 | voxel_min = voxel_c + (voxel_min - voxel_c) * 1.1 93 | voxel_max = voxel_c + (voxel_max - voxel_c) * 1.1 94 | 95 | # find voxels that are not occupied 96 | print('find visible voxels') 97 | res_spatial = args.voxel_num 98 | SpatialHist = torch.zeros(res_spatial**3,device=device) 99 | for idx in tqdm(range(len(dataset))): 100 | batch = dataset[idx] 101 | rays = batch['rays'] 102 | xs = rays[...,:3] 103 | ds = rays[...,3:6] 104 | 105 | positions,_,_,_,valid = ray_intersect(scene,xs.to(device),ds.to(device)) 106 | if not valid.any(): 107 | continue 108 | 109 | position = (positions[valid]-voxel_min)/(voxel_max-voxel_min) 110 | position = (position*res_spatial).long().clamp(0,res_spatial-1) 111 | inds = position[...,0] + position[...,1]*res_spatial\ 112 | + position[...,2]*res_spatial*res_spatial 113 | SpatialHist.scatter_add_(0,inds,torch.ones_like(inds).float()) 114 | SpatialHist = SpatialHist.reshape(res_spatial,res_spatial,res_spatial) 115 | 116 | mask = (SpatialHist>0) 117 | 118 | # create voxle surface light field 119 | print('bake voxel surface light field') 120 | vslf = VoxelSLF(mask.cpu(),voxel_min.item(),voxel_max.item()) 121 | for idx in tqdm(range(len(dataset))): 122 | batch = dataset[idx] 123 | rays = batch['rays'] 124 | radiance = batch['rgbs'] 125 | xs = rays[...,:3] 126 | ds = rays[...,3:6] 127 | 128 | exposure = batch['exposure'] 129 | radiance = model_crf.inverse(radiance, exposure) 130 | 131 | positions,_,_,_,valid = ray_intersect(scene,xs.to(device),ds.to(device)) 132 | if not valid.any(): 133 | continue 134 | 135 | vslf.scatter_add(positions[valid].cpu(),radiance.to(device)[valid].cpu()) 136 | 137 | # average pooling the radiance 138 | vslf.radiance = vslf.radiance/vslf.count[...,None].float().clamp_min(1) 139 | 140 | torch.save({ 141 | 'mask': (SpatialHist>0), 142 | 'voxel_min': voxel_min.item(), 143 | 'voxel_max': voxel_max.item(), 144 | 'weight':vslf.state_dict() 145 | },os.path.join(OUTPUT_PATH,'vslf.npz')) 146 | -------------------------------------------------------------------------------- /utils/metric_crf.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import torch 8 | import torch.nn.functional as NF 9 | 10 | import math 11 | import os 12 | os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1" 13 | import cv2 14 | 15 | import numpy as np 16 | import pandas as pd 17 | 18 | import matplotlib.pyplot as plt 19 | from tqdm import tqdm 20 | from crf.model_crf import EmorCRF 21 | import plotly.graph_objects as go 22 | import argparse 23 | 24 | parser = argparse.ArgumentParser() 25 | parser.add_argument('--crf_gt', help='path to LDR_IMG_DIR/cam/crf.npy') 26 | parser.add_argument('--ckpt', help='path to checkpoint last_1.ckpt') 27 | parser.add_argument('--dir_save', help='dir for saving image') 28 | args = parser.parse_args() 29 | 30 | # SCENE = 'kitchen' 31 | # METHOD = '1104_kitchen_single_base' 32 | # LDR_IMG_DIR='Image' 33 | # GT_PATH = '/hdd/datasets/fipt/indoor_synthetic/' 34 | # path_gt_crf = os.path.join(GT_PATH, SCENE, 'train', LDR_IMG_DIR, 'cam', 'crf.npy') 35 | crf_gt = np.load(args.crf_gt) 36 | last_ckpt = os.path.join(args.ckpt) 37 | 38 | n_basis = 3 39 | model_crf = EmorCRF(dim=n_basis) 40 | state_dict = torch.load(last_ckpt, map_location='cpu')['state_dict'] 41 | weight = {} 42 | for k,v in state_dict.items(): 43 | if 'model_crf.' in k: 44 | weight[k.replace('model_crf.','')]=v 45 | model_crf.load_state_dict(weight) 46 | crf_pred = model_crf.get_crf().detach().numpy() 47 | 48 | l2 = np.sqrt(np.sum((crf_gt - crf_pred)**2)) 49 | l2 = np.linalg.norm(crf_gt - crf_pred) 50 | print('L2: {:.5f}'.format(l2)) 51 | 52 | def plot_plotly(): 53 | x = np.linspace(0, 1, 1024) 54 | y1 = crf_gt[0] 55 | y2 = crf_pred[0] 56 | trace1 = go.Scatter(x=x, y=y1, mode='lines', name='GT', line=dict(color='blue', width=2, dash='dash')) 57 | trace2 = go.Scatter(x=x, y=y2, mode='lines', name='Pred.', line=dict(color='red', width=2)) 58 | 59 | # Create layout 60 | layout = go.Layout(title='', xaxis=dict(title='Irradiance', range=[0, 1]), yaxis=dict(title='LDR', range=[0, 1])) 61 | 62 | # Create a figure 63 | fig = go.Figure(data=[trace1, trace2], layout=layout) 64 | 65 | # Show the plot 66 | fig.show() 67 | 68 | def plot_test_dark(): 69 | x = np.linspace(0, 1, 1024) 70 | line = crf_pred[0] 71 | 72 | fig, ax = plt.subplots() 73 | ax.set_facecolor('black') 74 | ax.plot(x, line, color='yellow') 75 | # ax.xaxis.label.set_color('white') 76 | # ax.yaxis.label.set_color('white') 77 | ax.set_xlabel('Irradiance', size=18) 78 | ax.set_ylabel('Pixel Value', size=18) 79 | # ax.tick_params(axis='x', labelsize=12) 80 | # ax.tick_params(axis='y', labelsize=12) 81 | ax.tick_params(which='both', bottom=False, left=False) 82 | ax.set_xlim(0, 1) 83 | ax.set_ylim(0, 1) 84 | ax.grid(color='white', linestyle='-', linewidth=0.5) 85 | ax.set_aspect('equal') 86 | plt.show() 87 | 88 | def plot_test(): 89 | x = np.linspace(0, 1, 1024) 90 | line = crf_pred[0] 91 | 92 | fig, ax = plt.subplots(figsize=(8, 8)) 93 | ax.set_facecolor((0.9, 0.9, 0.9)) 94 | ax.plot(x, line, color=(0, 0, 1), linestyle='--', linewidth=2.5) 95 | # ax.xaxis.label.set_color('white') 96 | # ax.yaxis.label.set_color('white') 97 | ax.set_xlabel('Irradiance', size=24) 98 | ax.set_ylabel('Pixel Value', size=24) 99 | # ax.tick_params(axis='x', labelsize=12) 100 | # ax.tick_params(axis='y', labelsize=12) 101 | ax.tick_params(which='both', bottom=False, left=False, labelsize=15) 102 | ax.set_xlim(0, 1) 103 | ax.set_ylim(0, 1) 104 | ax.grid(color=(0, 0, 0), linestyle='--', linewidth=0.5) 105 | ax.set_aspect('equal') 106 | plt.show() 107 | 108 | def plot_single(): 109 | x = np.linspace(0, 1, 1024) 110 | 111 | fig, ax = plt.subplots(figsize=(10, 10)) 112 | # ax.set_facecolor((0.95, 0.95, 0.95)) 113 | ax.plot(x, crf_pred[0], color=(0, 0.1882, 0.6), linewidth=8.0, label='Pred.') 114 | ax.plot(x, crf_gt[0], color=(0.97, 0.5, 0), linestyle='--', linewidth=8.0, label='GT') 115 | ax.legend(fontsize=32) 116 | # ax.xaxis.label.set_color('white') 117 | # ax.yaxis.label.set_color('white') 118 | ax.set_xlabel('Irradiance', size=48) 119 | ax.set_ylabel('Pixel Value', size=48) 120 | # ax.tick_params(axis='x', labelsize=12) 121 | # ax.tick_params(axis='y', labelsize=12) 122 | ax.tick_params(which='both', bottom=False, left=False, labelsize=20) 123 | ax.set_xlim(0, 1) 124 | ax.set_ylim(0, 1) 125 | ax.grid(color=(0, 0, 0), linestyle='--', linewidth=1.0) 126 | ax.set_aspect('equal') 127 | path = os.path.join(args.dir_save, 'crf.png') 128 | plt.savefig(path) 129 | print('Saved to', path) 130 | 131 | def plot_all(): 132 | x = np.linspace(0, 1, 1024) 133 | 134 | fig, ax = plt.subplots(figsize=(10, 10)) 135 | # ax.set_facecolor((0.95, 0.95, 0.95)) 136 | ax.plot(x, crf_pred[0], color=(1, 0, 0), linewidth=4.0, label='R Pred.') 137 | ax.plot(x, crf_gt[0], color=(1, 0, 0), linewidth=4.0, label='R GT', linestyle='--') 138 | ax.plot(x, crf_pred[1], color=(0, 1, 0), linewidth=4.0, label='G Pred.') 139 | ax.plot(x, crf_gt[1], color=(0, 1, 0), linewidth=4.0, label='G GT', linestyle='--') 140 | ax.plot(x, crf_pred[2], color=(0, 0, 1), linewidth=4.0, label='B Pred.') 141 | ax.plot(x, crf_gt[2], color=(0, 0, 1), linewidth=4.0, label='B GT', linestyle='--') 142 | ax.legend(fontsize=32) 143 | # ax.xaxis.label.set_color('white') 144 | # ax.yaxis.label.set_color('white') 145 | ax.set_xlabel('Irradiance', size=48) 146 | ax.set_ylabel('Pixel Value', size=48) 147 | # ax.tick_params(axis='x', labelsize=12) 148 | # ax.tick_params(axis='y', labelsize=12) 149 | ax.tick_params(which='both', bottom=False, left=False, labelsize=20) 150 | ax.set_xlim(0, 1) 151 | ax.set_ylim(0, 1) 152 | ax.grid(color=(0, 0, 0), linestyle='--', linewidth=1.0) 153 | ax.set_aspect('equal') 154 | path = os.path.join(args.dir_save, 'crf_all.png') 155 | plt.savefig(path) 156 | print('Saved to', path) 157 | 158 | # plot_single() 159 | # plot_all() -------------------------------------------------------------------------------- /utils/export.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import mitsuba 8 | mitsuba.set_variant('cuda_ad_rgb') 9 | 10 | import os 11 | os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1" 12 | import cv2 13 | import imageio 14 | import numpy as np 15 | import torch 16 | import open3d as o3d 17 | import trimesh 18 | import xatlas 19 | import nvdiffrast.torch as dr 20 | from PIL import Image 21 | from argparse import ArgumentParser 22 | from model.brdf import NGPBRDF 23 | 24 | def main(): 25 | parser = ArgumentParser() 26 | parser.add_argument('--mesh') 27 | parser.add_argument('--ckpt') 28 | parser.add_argument('--emitter_path') 29 | parser.add_argument('--dir_save') 30 | parser.add_argument('--device', default='cuda') 31 | parser.add_argument('--tex_res', type=int, default=2048) 32 | parser.add_argument('--chunk_size', type=int, default=160000) 33 | args = parser.parse_args() 34 | 35 | dir_save = args.dir_save 36 | os.makedirs(dir_save, exist_ok=True) 37 | 38 | # Load material network 39 | device = torch.device(args.device) 40 | mask = torch.load(os.path.join(args.emitter_path,'vslf.npz'),map_location='cpu') 41 | material_net = NGPBRDF(mask['voxel_min'],mask['voxel_max']) 42 | state_dict = torch.load(args.ckpt, map_location='cpu')['state_dict'] 43 | weight = {} 44 | for k,v in state_dict.items(): 45 | if 'material.' in k: 46 | weight[k.replace('material.','')]=v 47 | material_net.load_state_dict(weight) 48 | material_net.to(device) 49 | print(f'[INFO] loaded material network from {args.ckpt}') 50 | 51 | mesh = trimesh.load(args.mesh) 52 | v_np = np.array(mesh.vertices).astype(np.float32) 53 | f_np = np.array(mesh.faces).astype(np.int32) 54 | 55 | # unwrap uvs 56 | print(f'[INFO] running xatlas to unwrap UVs for mesh: v={v_np.shape} f={f_np.shape}') 57 | atlas = xatlas.Atlas() 58 | atlas.add_mesh(v_np, f_np) 59 | 60 | # xatlas to unwarp 61 | path_ft_np = os.path.join(dir_save, 'ft.npy') 62 | path_vt_np = os.path.join(dir_save, 'vt.npy') 63 | if os.path.exists(path_ft_np) and os.path.exists(path_vt_np): 64 | print(f'[INFO] found existing UVs, loading from {path_ft_np} and {path_vt_np}') 65 | ft_np = np.load(path_ft_np) 66 | vt_np = np.load(path_vt_np) 67 | else: 68 | chart_options = xatlas.ChartOptions() 69 | chart_options.max_iterations = 0 # disable merge_chart for faster unwrap... 70 | pack_options = xatlas.PackOptions() 71 | atlas.generate(chart_options=chart_options, pack_options=pack_options) 72 | _, ft_np, vt_np = atlas[0] # [N], [M, 3], [N, 2] 73 | np.save(path_ft_np, ft_np) 74 | np.save(path_vt_np, vt_np) 75 | print(f'[INFO] finished: xatlas unwraps UVs for mesh: v={v_np.shape} f={f_np.shape} vt={vt_np.shape} ft={ft_np.shape}') 76 | 77 | vt = torch.from_numpy(vt_np.astype(np.float32)).float().to(device) 78 | ft = torch.from_numpy(ft_np.astype(np.int64)).int().to(device) 79 | # padding 80 | uv = vt * 2.0 - 1.0 # uvs to range [-1, 1] 81 | uv = torch.cat((uv, torch.zeros_like(uv[..., :1]), torch.ones_like(uv[..., :1])), dim=-1) # [N, 4] 82 | 83 | glctx = dr.RasterizeGLContext(output_db=False) 84 | tex_res = args.tex_res 85 | h, w = tex_res, tex_res 86 | # rasterize 2d texture vertices to texture image 87 | rast, _ = dr.rasterize(glctx, uv.unsqueeze(0), ft, (h, w)) # [1, h, w, 4]rast 88 | # interpolate to get the corresponding 3D location of each pixel 89 | v = torch.from_numpy(v_np).to(device) 90 | f = torch.from_numpy(f_np).to(device) 91 | xyzs, _ = dr.interpolate(v.unsqueeze(0), rast, f) # [1, h, w, 3] 92 | mask, _ = dr.interpolate(torch.ones_like(v[:, :1]).unsqueeze(0), rast, f) # [1, h, w, 1] 93 | vt = vt.cpu() 94 | ft = ft.cpu() 95 | uv = uv.cpu() 96 | # masked query 97 | xyzs = xyzs.view(-1, 3).cpu() 98 | mask = (mask > 0).view(-1).cpu() 99 | 100 | # masked query 101 | xyzs = xyzs.view(-1, 3).cpu() 102 | mask = (mask > 0).view(-1).cpu() 103 | 104 | mat_dim = 5 105 | feats = torch.zeros(h * w, mat_dim, dtype=torch.float32).cpu() 106 | if mask.any(): 107 | with torch.no_grad(): 108 | xyzs = xyzs[mask] # [M, 3] 109 | chunk_size = args.chunk_size 110 | # batched inference to avoid OOM 111 | all_feats = torch.zeros((xyzs.shape[0], mat_dim)).cpu() 112 | head = 0 113 | while head < xyzs.shape[0]: 114 | tail = min(head + chunk_size, xyzs.shape[0]) 115 | with torch.cuda.amp.autocast(enabled=False): 116 | slice_xyzs = xyzs[head:tail].clone().detach().cuda() 117 | pred = material_net(slice_xyzs) 118 | slice_mats = torch.cat([pred['albedo'], pred['roughness'], pred['metallic']], dim=-1) 119 | all_feats[head:tail] = slice_mats.cpu().float() 120 | slice_xyzs = slice_xyzs.cpu() 121 | del slice_xyzs 122 | head += chunk_size 123 | feats[mask] = all_feats 124 | 125 | feats = feats.view(h, w, -1) 126 | mask = mask.view(h, w) 127 | # quantize [0.0, 1.0] to [0, 255] 128 | feats = feats.cpu().numpy() 129 | feats = (feats * 255).astype(np.uint8) 130 | mask = mask.cpu().numpy() 131 | 132 | albedo = np.zeros((h, w, 3), dtype=np.uint8) 133 | roughness_metallic = np.zeros((h, w, 3), dtype=np.uint8) 134 | albedo[:, :, :3] = feats[:, :, :3] 135 | roughness_metallic[:, :, :2] = feats[:, :, 3:] 136 | 137 | img_albedo = Image.fromarray(albedo) 138 | path_albedo = os.path.join(dir_save, 'albedo.png') 139 | img_albedo.save(path_albedo) 140 | img_roughness_metallic = Image.fromarray(roughness_metallic) 141 | path_roughness_metallic = os.path.join(dir_save, 'rm.png') 142 | img_roughness_metallic.save(path_roughness_metallic) 143 | print(f'[INFO] saved albedo to {path_albedo}, saved roughness and metallic to {path_roughness_metallic}') 144 | 145 | if __name__ == '__main__': 146 | main() -------------------------------------------------------------------------------- /refine_shading.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import torch 8 | import mitsuba 9 | mitsuba.set_variant('cuda_ad_rgb') 10 | import math 11 | 12 | from utils.dataset import RealDatasetLDR,SyntheticDatasetLDR 13 | from utils.dataset.scannetpp.dataset import Scannetpp 14 | from utils.ops import * 15 | from utils.path_tracing import * 16 | from model.emitter import SLFEmitter 17 | from model.brdf import NGPBRDF 18 | import os 19 | os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1" 20 | import cv2 21 | from pathlib import Path 22 | 23 | from tqdm import tqdm 24 | from argparse import ArgumentParser 25 | import time 26 | from const import set_random_seed 27 | set_random_seed() 28 | 29 | if __name__ == '__main__': 30 | parser = ArgumentParser() 31 | parser.add_argument('--dataset_root', type=str, help='dataset root') 32 | parser.add_argument('--scene', type=str, required=True, help='dataset folder') 33 | parser.add_argument('--slf_path', type=str, required=True) 34 | parser.add_argument('--emitter_path', type=str, required=True) 35 | parser.add_argument('--output', type=str, required=True, help='last shading folder') 36 | parser.add_argument('--ckpt', type=str, required=True, help='checkpoint path') 37 | parser.add_argument('--dataset',type=str,required=True, help='dataset type') 38 | parser.add_argument('--ldr_img_dir', type=str, default=None) 39 | parser.add_argument('--res_scale', type=float, default=1.0) 40 | args = parser.parse_args() 41 | device = torch.device(0) 42 | 43 | DATASET_PATH = args.scene 44 | OUTPUT_PATH = args.output 45 | os.makedirs(OUTPUT_PATH,exist_ok=True) 46 | 47 | # load geometry 48 | if args.dataset in ['synthetic', 'real']: 49 | mesh_path = os.path.join(DATASET_PATH,'scene.obj') 50 | mesh_type = 'obj' 51 | elif args.dataset == 'scannetpp': 52 | mesh_path = os.path.join(args.dataset_root, 'data', args.scene, 'scans', 'scene.ply') 53 | mesh_type = 'ply' 54 | assert Path(mesh_path).exists(), 'mesh not found: '+mesh_path 55 | 56 | scene = mitsuba.load_dict({ 57 | 'type': 'scene', 58 | 'shape_id':{ 59 | 'type': mesh_type, 60 | 'filename': mesh_path, 61 | } 62 | }) 63 | 64 | # load dataset 65 | if args.dataset == 'synthetic': 66 | dataset = SyntheticDatasetLDR(DATASET_PATH, img_dir=args.ldr_img_dir, split='train', pixel=False) 67 | elif args.dataset == 'real': 68 | dataset = RealDatasetLDR(DATASET_PATH, img_dir=args.ldr_img_dir, split='train', pixel=False) 69 | elif args.dataset == 'scannetpp': 70 | dataset = Scannetpp(args.dataset_root, args.scene, split='train', pixel=False, res_scale=args.res_scale) 71 | img_hw = dataset.img_hw 72 | 73 | # load emitter 74 | emitter = SLFEmitter(args.emitter_path, args.slf_path) 75 | for p in emitter.parameters(): 76 | p.requires_grad=False 77 | emitter.to(device) 78 | for p in emitter.parameters(): 79 | p.requires_grad=False 80 | 81 | # load brdf 82 | mask = torch.load(args.slf_path, map_location='cpu') 83 | material_net = NGPBRDF(mask['voxel_min'],mask['voxel_max']) 84 | state_dict = torch.load(args.ckpt,map_location='cpu')['state_dict'] 85 | weight = {} 86 | for k,v in state_dict.items(): 87 | if 'material.' in k: 88 | weight[k.replace('material.','')] = v 89 | material_net.load_state_dict(weight) 90 | material_net.to(device) 91 | for p in material_net.parameters(): 92 | p.requires_grad=False 93 | 94 | # set up denoiser 95 | denoiser = mitsuba.OptixDenoiser(img_hw[::-1]) 96 | 97 | start_time = time.time() 98 | 99 | # refine diffuse shading 100 | print('refine diffuse') 101 | output_path = os.path.join(OUTPUT_PATH,'diffuse') 102 | os.makedirs(output_path,exist_ok=True) 103 | spp = 128 104 | indir_depth = 5 105 | 106 | # batched process 107 | batch_size = 10240*128//spp 108 | im_id = 0 109 | for batch in tqdm(dataset): 110 | rays = batch['rays'] 111 | rays_x,rays_d = rays[...,:3].to(device),rays[...,3:6].to(device) 112 | positions,normals,uvs,triangle_idxs,valid = ray_intersect(scene,rays_x,rays_d) 113 | wi = rays_d 114 | B = len(positions) 115 | L = torch.zeros(B,3,device=device) 116 | for b in range(math.ceil(B*1.0/batch_size)): 117 | b0 = b*batch_size 118 | b1 = min(b0+batch_size,B) 119 | L[b0:b1] = path_tracing_det_diff(scene,emitter,material_net, 120 | positions[b0:b1],wi[b0:b1],normals[b0:b1], 121 | uvs[b0:b1],triangle_idxs[b0:b1], 122 | spp,indir_depth) 123 | assert L.isnan().any() == False 124 | L = denoiser(mitsuba.TensorXf(L.reshape(*img_hw,3))).numpy() 125 | cv2.imwrite(os.path.join(output_path,'{:03d}.exr'.format(im_id)),L[:,:,[2,1,0]]) 126 | im_id += 1 127 | 128 | 129 | print('[refine_shading - diffuse] time (s): ', time.time()-start_time) 130 | start_time = time.time() 131 | 132 | # refine spacular shadings 133 | print('refine specular') 134 | output_path = os.path.join(OUTPUT_PATH,'specular') 135 | os.makedirs(output_path,exist_ok=True) 136 | spp = 64 137 | 138 | batch_size = 10240*128//spp 139 | im_id = 0 140 | 141 | # 6 roughness level 142 | roughness_level = torch.linspace(0.02,1.0,6) 143 | for batch in tqdm(dataset): 144 | rays = batch['rays'] 145 | rays_x,rays_d = rays[...,:3].to(device),rays[...,3:6].to(device) 146 | positions,normals,uvs,triangle_idxs,valid = ray_intersect(scene,rays_x,rays_d) 147 | wi = rays_d 148 | B = len(positions) 149 | L0 = torch.zeros(B,3,device=device) 150 | L1 = L0.clone() 151 | 152 | for r_idx,roughness in enumerate(roughness_level): 153 | # BxSx3 154 | 155 | B = len(positions) 156 | L0 = torch.zeros(B,3,device=device) 157 | L1 = L0.clone() 158 | 159 | for b in range(math.ceil(B*1.0/batch_size)): 160 | b0 = b*batch_size 161 | b1 = min(b0+batch_size,B) 162 | L0_,L1_ = path_tracing_det_spec(scene,emitter,material_net, 163 | roughness, 164 | positions[b0:b1],wi[b0:b1],normals[b0:b1], 165 | uvs[b0:b1],triangle_idxs[b0:b1], 166 | spp,indir_depth) 167 | L0[b0:b1] = L0_ 168 | L1[b0:b1] = L1_ 169 | assert L0.isnan().any() == False 170 | assert L1.isnan().any() == False 171 | L0 = denoiser(mitsuba.TensorXf(L0.reshape(*img_hw,3))).numpy() 172 | L1 = denoiser(mitsuba.TensorXf(L1.reshape(*img_hw,3))).numpy() 173 | cv2.imwrite(os.path.join(output_path,'{:03d}_0_{}.exr'.format(im_id,r_idx)),L0[:,:,[2,1,0]]) 174 | cv2.imwrite(os.path.join(output_path,'{:03d}_1_{}.exr'.format(im_id,r_idx)),L1[:,:,[2,1,0]]) 175 | im_id += 1 176 | 177 | print('[refine_shading - specular] time (s): ', time.time()-start_time) 178 | -------------------------------------------------------------------------------- /utils/video.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import os 8 | import cv2 9 | try: 10 | # for backward compatibility 11 | import imageio.v2 as imageio 12 | except ModuleNotFoundError: 13 | import imageio 14 | from PIL import Image 15 | from tqdm import tqdm 16 | import numpy as np 17 | import argparse 18 | 19 | def extract_frames(): 20 | parser = argparse.ArgumentParser() 21 | parser.add_argument('-video', help='path to video') 22 | parser.add_argument('-outdir', help='output folder') 23 | args = parser.parse_args() 24 | 25 | os.makedirs(args.outdir, exist_ok=True) 26 | vidcap = cv2.VideoCapture(args.video) 27 | i = 0 28 | while True: 29 | success, image = vidcap.read() 30 | if not success: 31 | break 32 | path = os.path.join(args.outdir, '{:0>5d}.png'.format(i)) 33 | cv2.imwrite(path, image) 34 | i += 1 35 | 36 | def read_video_frames(video_path): 37 | vidcap = cv2.VideoCapture(video_path) 38 | frames = [] 39 | while True: 40 | success, image = vidcap.read() 41 | if not success: 42 | break 43 | else: 44 | image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) 45 | frames.append(image) 46 | if len(frames) == 0: 47 | print('ERROR: {} does not exist'.format(video_path)) 48 | return frames 49 | 50 | def is_image_name(name): 51 | valid_names = ['.jpg', '.png', '.JPG', '.PNG'] 52 | for v in valid_names: 53 | if name.endswith(v): 54 | return True 55 | return False 56 | 57 | def generate_video(): 58 | parser = argparse.ArgumentParser() 59 | parser.add_argument('-dir') 60 | parser.add_argument('-out') 61 | parser.add_argument('-fps', type=int, default=30) 62 | args = parser.parse_args() 63 | 64 | imgs = sorted([os.path.join(args.dir, img) for img in os.listdir(args.dir) if is_image_name(img)]) 65 | imgs = [np.array(Image.open(img)) for img in imgs] 66 | h, w = imgs[0].shape[:2] 67 | imgs = [img[:h-h%2,:w-w%2] for img in imgs] 68 | imgs += imgs[::-1] 69 | imageio.mimsave(args.out, imgs, fps=args.fps, macro_block_size=1) 70 | 71 | # imgs = [Image.open(img) for img in imgs] 72 | # imgs += imgs[::-1] 73 | # imgs[0].save(args.out, format='GIF', append_images=imgs, 74 | # save_all=True, duration=10, loop=0) 75 | 76 | def add_text(): 77 | parser = argparse.ArgumentParser() 78 | parser.add_argument('--video_in') 79 | parser.add_argument('--video_out') 80 | parser.add_argument('--text') 81 | parser.add_argument('--font_size', type=float, default=2.0) 82 | parser.add_argument('--font_thickness', type=int, default=4) 83 | parser.add_argument('--right', dest='right', action='store_true') 84 | parser.add_argument('--bottom', dest='bottom', action='store_true') 85 | parser.add_argument('--fps', type=int, default=30) 86 | args = parser.parse_args() 87 | 88 | background_color = (0, 0, 0) # black color in BGR format 89 | text = args.text 90 | font = cv2.FONT_HERSHEY_COMPLEX 91 | font_size = args.font_size 92 | font_color = (255, 255, 255) # White color in BGR format 93 | thickness = args.font_thickness 94 | 95 | frames_in = read_video_frames(args.video_in) 96 | h, w, _ = frames_in[0].shape 97 | 98 | text_size = cv2.getTextSize(text, font, font_size, thickness)[0] 99 | border = 10 # to image 100 | buffer = 30 # inside box 101 | x, y = border, border # Position of the text 102 | x2, y2 = x + text_size[0] + buffer, y + text_size[1] + buffer 103 | if args.right: 104 | x2 = w - border 105 | x = x2 - text_size[0] - buffer 106 | if args.bottom: 107 | y2 = h - border 108 | y = y2 - text_size[1] - buffer 109 | 110 | text_shift_x = 15/2 * args.font_size 111 | text_shift_y = 55/2 * args.font_size 112 | frames_out = [] 113 | for frame in frames_in: 114 | cv2.rectangle(frame, (x, y), (x2, y2), background_color, -1) 115 | cv2.putText(frame, text, (int(x + text_shift_x), int(y + text_shift_y)), font, font_size, font_color, thickness) 116 | frames_out.append(frame) 117 | 118 | imageio.mimsave(args.video_out, 119 | frames_out, 120 | fps=args.fps, macro_block_size=1) 121 | 122 | def switch_video(): 123 | parser = argparse.ArgumentParser() 124 | parser.add_argument('--video_in', nargs='+') 125 | parser.add_argument('--video_out') 126 | parser.add_argument('--mid', type=int) 127 | parser.add_argument('--slope', type=float, default=1.0) 128 | parser.add_argument('--window', type=int, default=30) 129 | parser.add_argument('--linewidth', type=int, default=0) 130 | parser.add_argument('--flip', dest='flip', action='store_true') 131 | parser.add_argument('-fps', type=int, default=30) 132 | args = parser.parse_args() 133 | 134 | vdo_path_0, vdo_path_1 = args.video_in 135 | frames_0 = read_video_frames(vdo_path_0) 136 | frames_1 = read_video_frames(vdo_path_1) 137 | h, w, _ = frames_0[0].shape 138 | 139 | # to be removed 140 | frames_0 = frames_0[:300] 141 | frames_1 = frames_1[:300] 142 | 143 | v_start = 0 144 | v_end = (w-1) + (h-1) * args.slope 145 | v_slope = (v_end - v_start) / args.window 146 | if args.flip: 147 | v_slope *= -1 148 | v_const = (v_end+v_start)/2 - args.mid * v_slope 149 | 150 | grid_y, grid_x = np.meshgrid(np.arange(w), np.arange(h)) 151 | grid_value = grid_y + grid_x * args.slope 152 | 153 | frames_out = [] 154 | vdo_len = len(frames_0) 155 | for i in tqdm(range(vdo_len)): 156 | v_threshold = i * v_slope + v_const 157 | mask = grid_value > v_threshold 158 | 159 | f_0 = frames_0[i] 160 | f_1 = frames_1[i] 161 | f_out = np.zeros_like(f_0) 162 | f_out[mask] = f_0[mask] 163 | f_out[~mask] = f_1[~mask] 164 | frames_out.append(f_out) 165 | 166 | imageio.mimsave(args.video_out, 167 | frames_out, 168 | fps=args.fps, macro_block_size=1) 169 | 170 | def merge_video(): 171 | parser = argparse.ArgumentParser() 172 | parser.add_argument('-first') 173 | parser.add_argument('-second') 174 | parser.add_argument('-out') 175 | parser.add_argument('-axis', type=int, default=0) 176 | parser.add_argument('-fps', type=int, default=30) 177 | args = parser.parse_args() 178 | 179 | frames_first = read_video_frames(args.first) 180 | frames_second = read_video_frames(args.second) 181 | frames_out = [] 182 | 183 | frame_len = min(len(frames_first), len(frames_second)) 184 | for i in range(frame_len): 185 | left = frames_first[i] 186 | right = frames_second[i] 187 | out = np.concatenate([left, right], axis=args.axis) 188 | frames_out.append(out) 189 | 190 | imageio.mimsave(args.out, 191 | frames_out, 192 | fps=args.fps, macro_block_size=1) 193 | print('Export video:', args.out) 194 | 195 | def loop(): 196 | parser = argparse.ArgumentParser() 197 | parser.add_argument('--video_in') 198 | parser.add_argument('--video_out') 199 | parser.add_argument('-fps', type=int, default=30) 200 | args = parser.parse_args() 201 | 202 | frames = read_video_frames(args.video_in) 203 | frames += frames[::-1] 204 | imageio.mimsave(args.video_out, 205 | frames, 206 | fps=args.fps, macro_block_size=1) 207 | 208 | 209 | if __name__ == '__main__': 210 | # extract_frames() 211 | generate_video() 212 | # merge_video() 213 | # add_text() 214 | # switch_video() 215 | # loop() 216 | -------------------------------------------------------------------------------- /utils/ray_utils.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import torch 8 | import numpy as np 9 | from kornia import create_meshgrid 10 | from einops import rearrange 11 | import scipy 12 | 13 | # @torch.cuda.amp.autocast(dtype=torch.float32) 14 | def axisangle_to_R(v): 15 | """ 16 | Convert an axis-angle vector to rotation matrix 17 | from https://github.com/ActiveVisionLab/nerfmm/blob/main/utils/lie_group_helper.py#L47 18 | 19 | Inputs: 20 | v: (B, 3) 21 | 22 | Outputs: 23 | R: (B, 3, 3) 24 | """ 25 | v_ndim = v.ndim 26 | if v_ndim==1: 27 | v = rearrange(v, 'c -> 1 c') 28 | zero = torch.zeros_like(v[:, :1]) # (B, 1) 29 | skew_v0 = torch.cat([zero, -v[:, 2:3], v[:, 1:2]], 1) # (B, 3) 30 | skew_v1 = torch.cat([v[:, 2:3], zero, -v[:, 0:1]], 1) 31 | skew_v2 = torch.cat([-v[:, 1:2], v[:, 0:1], zero], 1) 32 | skew_v = torch.stack([skew_v0, skew_v1, skew_v2], dim=1) # (B, 3, 3) 33 | 34 | norm_v = rearrange(torch.norm(v, dim=1)+1e-7, 'b -> b 1 1') 35 | eye = torch.eye(3, device=v.device) 36 | R = eye + (torch.sin(norm_v)/norm_v)*skew_v + \ 37 | ((1-torch.cos(norm_v))/norm_v**2)*(skew_v@skew_v) 38 | if v_ndim==1: 39 | R = rearrange(R, '1 c d -> c d') 40 | return R 41 | 42 | 43 | def normalize(v): 44 | """Normalize a vector.""" 45 | return v/np.linalg.norm(v) 46 | 47 | 48 | def average_poses(poses, pts3d=None): 49 | """ 50 | Calculate the average pose, which is then used to center all poses 51 | using @center_poses. Its computation is as follows: 52 | 1. Compute the center: the average of 3d point cloud (if None, center of cameras). 53 | 2. Compute the z axis: the normalized average z axis. 54 | 3. Compute axis y': the average y axis. 55 | 4. Compute x' = y' cross product z, then normalize it as the x axis. 56 | 5. Compute the y axis: z cross product x. 57 | 58 | Note that at step 3, we cannot directly use y' as y axis since it's 59 | not necessarily orthogonal to z axis. We need to pass from x to y. 60 | Inputs: 61 | poses: (N_images, 3, 4) 62 | pts3d: (N, 3) 63 | 64 | Outputs: 65 | pose_avg: (3, 4) the average pose 66 | """ 67 | # 1. Compute the center 68 | if pts3d is not None: 69 | center = pts3d.mean(0) 70 | else: 71 | center = poses[..., 3].mean(0) 72 | 73 | # 2. Compute the z axis 74 | z = normalize(poses[..., 2].mean(0)) # (3) 75 | 76 | # 3. Compute axis y' (no need to normalize as it's not the final output) 77 | y_ = poses[..., 1].mean(0) # (3) 78 | 79 | # 4. Compute the x axis 80 | x = normalize(np.cross(y_, z)) # (3) 81 | 82 | # 5. Compute the y axis (as z and x are normalized, y is already of norm 1) 83 | y = np.cross(z, x) # (3) 84 | 85 | pose_avg = np.stack([x, y, z, center], 1) # (3, 4) 86 | 87 | return pose_avg 88 | 89 | 90 | def center_poses(poses, pts3d=None): 91 | """ 92 | See https://github.com/bmild/nerf/issues/34 93 | Inputs: 94 | poses: (N_images, 3, 4) 95 | pts3d: (N, 3) reconstructed point cloud 96 | 97 | Outputs: 98 | poses_centered: (N_images, 3, 4) the centered poses 99 | pts3d_centered: (N, 3) centered point cloud 100 | """ 101 | 102 | pose_avg = average_poses(poses, pts3d) # (3, 4) 103 | pose_avg_homo = np.eye(4) 104 | pose_avg_homo[:3] = pose_avg # convert to homogeneous coordinate for faster computation 105 | # by simply adding 0, 0, 0, 1 as the last row 106 | pose_avg_inv = np.linalg.inv(pose_avg_homo) 107 | last_row = np.tile(np.array([0, 0, 0, 1]), (len(poses), 1, 1)) # (N_images, 1, 4) 108 | poses_homo = \ 109 | np.concatenate([poses, last_row], 1) # (N_images, 4, 4) homogeneous coordinate 110 | 111 | poses_centered = pose_avg_inv @ poses_homo # (N_images, 4, 4) 112 | poses_centered = poses_centered[:, :3] # (N_images, 3, 4) 113 | 114 | if pts3d is not None: 115 | pts3d_centered = pts3d @ pose_avg_inv[:, :3].T + pose_avg_inv[:, 3:].T 116 | return poses_centered, pts3d_centered 117 | 118 | return poses_centered 119 | 120 | def create_spheric_poses(radius, mean_h, n_poses=120): 121 | """ 122 | Create circular poses around z axis. 123 | Inputs: 124 | radius: the (negative) height and the radius of the circle. 125 | mean_h: mean camera height 126 | Outputs: 127 | spheric_poses: (n_poses, 3, 4) the poses in the circular path 128 | """ 129 | def spheric_pose(theta, phi, radius): 130 | trans_t = lambda t : np.array([ 131 | [1,0,0,0], 132 | [0,1,0,2*mean_h], 133 | [0,0,1,-t] 134 | ]) 135 | 136 | rot_phi = lambda phi : np.array([ 137 | [1,0,0], 138 | [0,np.cos(phi),-np.sin(phi)], 139 | [0,np.sin(phi), np.cos(phi)] 140 | ]) 141 | 142 | rot_theta = lambda th : np.array([ 143 | [np.cos(th),0,-np.sin(th)], 144 | [0,1,0], 145 | [np.sin(th),0, np.cos(th)] 146 | ]) 147 | 148 | c2w = rot_theta(theta) @ rot_phi(phi) @ trans_t(radius) 149 | c2w = np.array([[-1,0,0],[0,0,1],[0,1,0]]) @ c2w 150 | return c2w 151 | 152 | spheric_poses = [] 153 | for th in np.linspace(0, 2*np.pi, n_poses+1)[:-1]: 154 | spheric_poses += [spheric_pose(th, -np.pi/12, radius)] 155 | return np.stack(spheric_poses, 0) 156 | 157 | def viewmatrix(lookdir: np.ndarray, up: np.ndarray, 158 | position: np.ndarray) -> np.ndarray: 159 | """Construct lookat view matrix.""" 160 | vec2 = normalize(lookdir) 161 | vec0 = normalize(np.cross(up, vec2)) 162 | vec1 = normalize(np.cross(vec2, vec0)) 163 | m = np.stack([vec0, vec1, vec2, position], axis=1) 164 | return m 165 | 166 | def generate_interpolated_path(poses: np.ndarray, 167 | n_interp: int, 168 | spline_degree: int = 5, 169 | smoothness: float = .03, 170 | rot_weight: float = .1): 171 | """Creates a smooth spline path between input keyframe camera poses. 172 | Adapted from https://github.com/google-research/multinerf/blob/main/internal/camera_utils.py 173 | Spline is calculated with poses in format (position, lookat-point, up-point). 174 | 175 | Args: 176 | poses: (n, 3, 4) array of input pose keyframes. OPENCV, cam2world 177 | n_interp: returned path will have n_interp * (n - 1) total poses. 178 | spline_degree: polynomial degree of B-spline. 179 | smoothness: parameter for spline smoothing, 0 forces exact interpolation. 180 | rot_weight: relative weighting of rotation/translation in spline solve. 181 | 182 | Returns: 183 | Array of new camera poses with shape (n_interp * (n - 1), 3, 4). 184 | """ 185 | 186 | def poses_to_points(poses, dist): 187 | """Converts from pose matrices to (position, lookat, up) format.""" 188 | pos = poses[:, :3, -1] 189 | lookat = poses[:, :3, -1] - dist * poses[:, :3, 2] 190 | up = poses[:, :3, -1] + dist * poses[:, :3, 1] 191 | return np.stack([pos, lookat, up], 1) 192 | 193 | def points_to_poses(points): 194 | """Converts from (position, lookat, up) format to pose matrices.""" 195 | return np.array([viewmatrix(p - l, u - p, p) for p, l, u in points]) 196 | 197 | def interp(points, n, k, s): 198 | """Runs multidimensional B-spline interpolation on the input points.""" 199 | sh = points.shape 200 | pts = np.reshape(points, (sh[0], -1)) 201 | k = min(k, sh[0] - 1) 202 | tck, _ = scipy.interpolate.splprep(pts.T, k=k, s=s) 203 | u = np.linspace(0, 1, n, endpoint=False) 204 | new_points = np.array(scipy.interpolate.splev(u, tck)) 205 | new_points = np.reshape(new_points.T, (n, sh[1], sh[2])) 206 | return new_points 207 | 208 | points = poses_to_points(poses, dist=rot_weight) 209 | new_points = interp(points, 210 | n_interp * (points.shape[0] - 1), 211 | k=spline_degree, 212 | s=smoothness) 213 | return points_to_poses(new_points) -------------------------------------------------------------------------------- /bake_shading.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates. 2 | # All rights reserved. 3 | 4 | # This source code is licensed under the license found in the 5 | # LICENSE file in the root directory of this source tree. 6 | 7 | import torch 8 | import mitsuba 9 | mitsuba.set_variant('cuda_ad_rgb') 10 | import math 11 | import os 12 | os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1" 13 | import cv2 14 | from utils.dataset import SyntheticDatasetLDR,RealDatasetLDR 15 | from utils.dataset.scannetpp.dataset import Scannetpp 16 | from utils.path_tracing import ray_intersect 17 | from model.slf import VoxelSLF 18 | from model.brdf import BaseBRDF 19 | from model.emitter import SLFEmitter 20 | from crf.model_crf import EmorCRF 21 | from pathlib import Path 22 | from tqdm import tqdm 23 | from argparse import ArgumentParser 24 | import time 25 | from const import set_random_seed 26 | set_random_seed() 27 | 28 | 29 | if __name__ == '__main__': 30 | parser = ArgumentParser() 31 | parser.add_argument('--dataset_root', type=str, help='dataset root') 32 | parser.add_argument('--scene', type=str, required=True, help='dataset folder') 33 | parser.add_argument('--slf_path', type=str, required=True) 34 | parser.add_argument('--emitter_path', type=str, required=True) 35 | parser.add_argument('--output', type=str, required=True, help='output path') 36 | parser.add_argument('--dataset', type=str,required=True, help='dataset type') 37 | parser.add_argument('--ldr_img_dir', type=str, default=None) 38 | parser.add_argument('--res_scale', type=float, default=1.0) 39 | args = parser.parse_args() 40 | 41 | device = torch.device(0) # use gpu device 0 42 | 43 | DATASET_PATH = args.scene 44 | OUTPUT_PATH = args.output 45 | os.makedirs(OUTPUT_PATH,exist_ok=True) 46 | # load mesh 47 | if args.dataset in ['synthetic', 'real']: 48 | mesh_path = os.path.join(DATASET_PATH,'scene.obj') 49 | mesh_type = 'obj' 50 | elif args.dataset == 'scannetpp': 51 | mesh_path = os.path.join(args.dataset_root, 'data', args.scene, 'scans', 'scene.ply') 52 | mesh_type = 'ply' 53 | assert Path(mesh_path).exists(), 'mesh not found: '+mesh_path 54 | 55 | scene = mitsuba.load_dict({ 56 | 'type': 'scene', 57 | 'shape_id':{ 58 | 'type': mesh_type, 59 | 'filename': mesh_path, 60 | } 61 | }) 62 | 63 | # load dataset 64 | if args.dataset == 'synthetic': 65 | dataset = SyntheticDatasetLDR(DATASET_PATH, img_dir=args.ldr_img_dir, split='train', pixel=False) 66 | elif args.dataset == 'real': 67 | dataset = RealDatasetLDR(DATASET_PATH, img_dir=args.ldr_img_dir, split='train', pixel=False) 68 | elif args.dataset == 'scannetpp': 69 | dataset = Scannetpp(args.dataset_root, args.scene, split='train', pixel=False, res_scale=args.res_scale) 70 | img_hw = dataset.img_hw 71 | 72 | 73 | # create voxle surface light field 74 | emitter = SLFEmitter(args.emitter_path, args.slf_path) 75 | for p in emitter.parameters(): 76 | p.requires_grad=False 77 | emitter.to(device) 78 | 79 | material_net = BaseBRDF() 80 | 81 | denoiser = mitsuba.OptixDenoiser(img_hw[::-1]) 82 | 83 | start_time = time.time() 84 | 85 | # bake diffuse shading 86 | print('bake diffuse') 87 | output_path = os.path.join(OUTPUT_PATH,'diffuse') 88 | os.makedirs(output_path,exist_ok=True) 89 | 90 | spp = 256 91 | 92 | im_id = 0 93 | for batch in tqdm(dataset): 94 | rays = batch['rays'] 95 | xs = rays[...,:3] 96 | ds = rays[...,3:6] 97 | 98 | positions,normals,_,_,valid = ray_intersect(scene,xs.to(device),ds.to(device)) 99 | position = positions[valid] 100 | normal = normals[valid] 101 | ds = ds.to(device)[valid] 102 | 103 | B = ds.shape[0] 104 | Ld_ = torch.zeros(B,3,device=device) 105 | batch_size = 10240*64//spp 106 | 107 | # batched diffuse shading calculation 108 | for b in range(math.ceil(B*1.0/batch_size)): 109 | b0 = b*batch_size 110 | b1 = min(b0+batch_size,B) 111 | 112 | # importance sampling wi 113 | wi,_,_, = material_net.sample_diffuse(torch.rand((b1-b0)*spp,2,device=device), 114 | normal[b0:b1].repeat_interleave(spp,0)) 115 | 116 | p_next,_,_,tri_next,valid_next = ray_intersect(scene, 117 | position[b0:b1].repeat_interleave(spp,0)+mitsuba.math.RayEpsilon*wi,# prevent self intersection 118 | wi.reshape(-1,3)) 119 | 120 | # query surface light field 121 | roughness_one = torch.ones_like(tri_next)[:, None] 122 | Le, _, _ = emitter.eval_emitter(p_next, wi, tri_next, roughness_one, trace_roughness=0.0) 123 | Ld_[b0:b1] = Le.reshape(b1-b0,spp,3).mean(1) 124 | 125 | # denoiser renderings 126 | Ld = torch.zeros_like(xs) 127 | Ld[valid.cpu()] = Ld_.cpu() 128 | Ld = Ld.reshape(*img_hw,3).numpy() 129 | Ld = denoiser(Ld).numpy() 130 | 131 | cv2.imwrite(os.path.join(output_path,'{:03d}.exr'.format(im_id)),Ld[:,:,[2,1,0]]) 132 | im_id += 1 133 | 134 | print('[bake_shading - diffuse] time (s): ', time.time()-start_time) 135 | start_time = time.time() 136 | 137 | # bake specular shadings 138 | print('bake specular') 139 | output_path = os.path.join(OUTPUT_PATH,'specular') 140 | os.makedirs(output_path,exist_ok=True) 141 | 142 | 143 | spps = [64,128,128,128,128,128] # use different sampling rate 144 | im_id = 0 145 | 146 | # 6 roughness level 147 | roughness_level = torch.linspace(0.02,1.0,6) 148 | 149 | for batch in tqdm(dataset): 150 | rays = batch['rays'] 151 | xs = rays[...,:3] 152 | ds = rays[...,3:6] 153 | 154 | positions,normals,_,_,valid = ray_intersect(scene,xs.to(device),ds.to(device)) 155 | position = positions[valid] 156 | normal = normals[valid] 157 | wo = -ds.to(device)[valid] 158 | 159 | B = position.shape[0] 160 | # caculate for each roughness value 161 | for r_idx,roughness in enumerate(roughness_level): 162 | spp = spps[r_idx] 163 | Ls0_ = torch.zeros(B,3,device=device) 164 | Ls1_ = torch.zeros(B,3,device=device) 165 | 166 | # batched specular shading calculation 167 | batch_size = 10240*64//spp 168 | for b in range(math.ceil(B*1.0/batch_size)): 169 | b0 = b*batch_size 170 | b1 = min(b0+batch_size,B) 171 | 172 | # importance sampling wi 173 | wi,_,g0,g1 = material_net.sample_specular( 174 | torch.rand((b1-b0)*spp,2,device=device), 175 | wo[b0:b1].repeat_interleave(spp,0), 176 | normal[b0:b1].repeat_interleave(spp,0),roughness 177 | ) 178 | 179 | p_next,_,_, tri_next, valid_next = ray_intersect(scene, 180 | position[b0:b1].repeat_interleave(spp,0)+mitsuba.math.RayEpsilon*wi,# prevent self intersection 181 | wi.reshape(-1,3)) 182 | 183 | # query surface light field 184 | roughness_one = torch.ones_like(tri_next)[:, None] 185 | Le, _, _ = emitter.eval_emitter(p_next, wi, tri_next, roughness_one, trace_roughness=0.0) 186 | 187 | Ls0_[b0:b1] = (Le*g0).reshape(b1-b0,spp,3).mean(1) 188 | Ls1_[b0:b1] = (Le*g1).reshape(b1-b0,spp,3).mean(1) 189 | 190 | Ls0 = torch.zeros_like(xs) 191 | Ls1 = torch.zeros_like(xs) 192 | Ls0[valid.cpu()] = Ls0_.cpu() 193 | Ls1[valid.cpu()] = Ls1_.cpu() 194 | 195 | Ls0 = Ls0.reshape(*img_hw,3).numpy() 196 | Ls1 = Ls1.reshape(*img_hw,3).numpy() 197 | 198 | if r_idx > 0: # no need for denoise of low roughness 199 | Ls0 = denoiser(Ls0).numpy() 200 | Ls1 = denoiser(Ls1).numpy() 201 | 202 | cv2.imwrite(os.path.join(output_path,'{:03d}_0_{}.exr'.format(im_id,r_idx)),Ls0[:,:,[2,1,0]]) 203 | cv2.imwrite(os.path.join(output_path,'{:03d}_1_{}.exr'.format(im_id,r_idx)),Ls1[:,:,[2,1,0]]) 204 | im_id += 1 205 | 206 | print('[bake_shading - specular] time (s): ', time.time()-start_time) --------------------------------------------------------------------------------