## NeRF
* [Data - Synthetic-NeRF](https://drive.google.com/drive/folders/128yBriW1IG_3NJ5Rp7APSTZsJqdJdfc1)
* [Data-Tanks&Temples](https://dl.fbaipublicfiles.com/nsvf/dataset/TanksAndTemple.zip)
The training script can be found at `train_per_scene.py`:
```python
python train_per_scene.py configs/nerf.yaml defaults.expname=lego dataset.datadir=./data/nerf_synthetic/lego
```
## Generalization NeRF
* [Data - Google Scanned Objects](https://drive.google.com/file/d/1w1Cs0yztH6kE3JIz7mdggvPGCwIKkVi2/view)
```python
python train_across_scene.py configs/nerf_set.yaml
```
## More examples
Command explanation with a nerf example:
* `model.basis_dims=[4, 4, 4, 2, 2, 2]` adjusts the number of levels and channels at each level, with a total of 6 levels and 18 channels.
* `model.basis_resos=[32, 51, 70, 89, 108, 128]` represents the resolution of the feature embeddings.
* `model.freq_bands=[2.0, 3.2, 4.4, 5.6, 6.8, 8.0]` indicates the frequency parameters applied at each level of the coordinate transformation function.
* `model.coeff_type` represents the coefficient field representations and can be one of the following: [none, x, grid, mlp, vec, cp, vm].
* `model.basis_type` represents the basis field representation and can be one of the following: [none, x, grid, mlp, vec, cp, vm, hash].
* `model.basis_mapping` represents the coordinate transformation and can be one of the following: [x, triangle, sawtooth, trigonometric]. Please note that if you want to use orthogonal projection, choose the cp or vm basis type, as they automatically utilize the orthogonal projection functions.
* `model.total_params` controls the total model size. It is important to note that the model's size capability is determined by model.basis_resos and model.basis_dims. The total_params parameter mainly affects the capability of the coefficients.
* `exportation.render_only` you can rendering item after training by setting this label to 1. Please also specify the `defaults.ckpt` label.
* `exportation....` you can specify whether to render the items of `[render_test, render_train, render_path, export_mesh]` after training by enable the corressponding label to 1.
Some pre-defined configurations (such as occNet, DVGO, nerf, iNGP, EG3D) can be found in `README_FactorField.py`.
## COPY RIGHT
* [Summer Day](https://www.rijksmuseum.nl/en/collection/SK-A-3005) - Credit goes to Johan Hendrik Weissenbruch and rijksmuseum.
* [Mars](https://solarsystem.nasa.gov/resources/933/true-colors-of-pluto/) - Credit goes to NASA.
* [Albert](https://cdn.loc.gov/service/pnp/cph/3b40000/3b46000/3b46000/3b46036v.jpg) - Credit goes to Orren Jack Turner.
* [Girl With a Pearl Earring](http://profoundism.com/free_licenses.html) - Renovation copyright Koorosh Orooj (CC BY-SA 4.0).
## Citation
If you find our code or paper helpful, please consider citing both of these papers:
```
@article{Chen2023factor,
title={Factor Fields: A Unified Framework for Neural Fields and Beyond},
author={Chen, Anpei and Xu, Zexiang and Wei, Xinyue and Tang, Siyu and Su, Hao and Geiger, Andreas},
journal={arXiv preprint arXiv:2302.01226},
year={2023}
}
@article{Chen2023SIGGRAPH,
title={{Dictionary Fields: Learning a Neural Basis Decomposition}},
author={Anpei, Chen and Zexiang, Xu and Xinyue, Wei and Siyu, Tang and Hao, Su and Andreas, Geiger},
booktitle={International Conference on Computer Graphics and Interactive Techniques (SIGGRAPH)},
year={2023}}
```
================================================
FILE: README_FactorField.md
================================================
## nerf reconstruction with Dictionary field
```python
for scene in ['ship', 'mic', 'chair', 'lego', 'drums', 'ficus', 'hotdog', 'materials']:
cmd = f'python train_basis.py configs/nerf.yaml defaults.expname={scene} ' \
f'dataset.datadir=./data/nerf_synthetic/{scene} '
```
## different model design choices
```python
choice_dict = {
'-grid': '', \
'-DVGO-like': 'model.basis_type=none model.coeff_reso=80', \
'-noC': 'model.coeff_type=none', \
'-SL':'model.basis_dims=[18] model.basis_resos=[70] model.freq_bands=[8.]', \
'-CP': f'model.coeff_type=vec model.basis_type=cp model.freq_bands=[1.,1.,1.,1.,1.,1.] model.basis_resos=[512,512,512,512,512,512] model.basis_dims=[32,32,32,32,32,32]', \
'-iNGP-like': 'model.basis_type=hash model.coeff_type=none', \
'-hash': f'model.basis_type=hash model.coef_init=1.0 ', \
'-sinc': f'model.basis_mapping=sinc', \
'-tria': f'model.basis_mapping=triangle', \
'-vm': f'model.coeff_type=vm model.basis_type=vm', \
'-mlpB': 'model.basis_type=mlp', \
'-mlpC': 'model.coeff_type=mlp', \
'-occNet': f'model.basis_type=x model.coeff_type=none model.basis_mapping=x model.num_layers=8 model.hidden_dim=256 ', \
'-nerf': f'model.basis_type=x model.coeff_type=none model.basis_mapping=trigonometric ' \
f'model.num_layers=8 model.hidden_dim=256 ' \
f'model.freq_bands=[1.,2.,4.,8.,16.,32.,64,128,256.,512.] model.basis_dims=[1,1,1,1,1,1,1,1,1,1] model.basis_resos=[1024,512,256,128,64,32,16,8,4,2]', \
'-hash-sl': f'model.basis_type=hash model.coef_init=1.0 model.basis_dims=[16] model.freq_bands=[8.] model.basis_resos=[64] ', \
'-vm-sl': f'model.coeff_type=vm model.basis_type=vm model.coef_init=1.0 model.basis_dims=[18] model.freq_bands=[1.] model.basis_resos=[64] model.total_params=1308416 ', \
'-DCT':'model.basis_type=fix-grid', \
}
for name in choice_dict.keys():
for scene in [ 'ship', 'mic', 'chair', 'lego', 'drums', 'ficus', 'hotdog', 'materials']:
cmd = f"python train_per_scene.py configs/nerf.yaml defaults.expname={scene}{name} dataset.datadir=./data/nerf_synthetic/{scene} {config}"
```
## generalized nerf
Your can choice of the the following design choice for testing.
```python
choice_dict = {
'-grid': '', \
'-DVGO-like': 'model.basis_type=none model.coeff_reso=48',
'-SL':'model.basis_dims=[72] model.basis_resos=[48] model.freq_bands=[6.]', \
'-CP': f'model.coeff_type=vec model.basis_type=cp model.freq_bands=[1.,1.,1.,1.,1.,1.] model.basis_resos=[512,512,512,512,512,512] model.basis_dims=[32,32,32,32,32,32]', \
'-hash': f'model.basis_type=hash model.coef_init=1.0 ', \
'-sinc': f'model.basis_mapping=sinc', \
'-tria': f'model.basis_mapping=triangle', \
'-vm': f'model.coeff_type=vm model.basis_type=vm', \
'-mlpB': 'model.basis_type=mlp', \
'-mlpC': 'model.coeff_type=mlp', \
'-hash-sl': f'model.basis_type=hash model.coef_init=1.0 model.basis_dims=[16] model.freq_bands=[8.] model.basis_resos=[64] ', \
'-vm-sl': f'model.coeff_type=vm model.basis_type=vm model.coef_init=1.0 model.basis_dims=[18] model.freq_bands=[1.] model.basis_resos=[64] model.total_params=1308416 ', \
'-DCT':'model.basis_type=fix-grid', \
}
for name in choice_dict.keys(): #
cmd = f'python train_across_scene.py configs/nerf_set.yaml defaults.expname=google-obj{name} {config} ' \
f'training.volume_resoFinal=128 dataset.datadir=./data/google_scanned_objects/'
```
You can also fine tune of the trained model for a new scene:
```python
for views in [5]:#3,
for name in choice_dict.keys(): #
for scene in [183]:#183,199,298,467,957,244,963,527,
cmd = f'python train_across_scene_ft.py configs/nerf_ft.yaml defaults.expname=google_objs_{name}_{scene}_{views}_views ' \
f'{config} training.n_iters=10000 ' \
f'dataset.train_views={views} ' \
f'dataset.train_scene_list=[{scene}] ' \
f'dataset.test_scene_list=[{scene}] ' \
f'dataset.datadir=./data/google_scanned_objects/ ' \
f'defaults.ckpt=./logs/google-obj{name}//google-obj{name}.th'
```
# render path after optimization
```python
for views in [5]:
for name in choice_dict.keys(): #
config = commands[name].replace(",", "','")
for scene in [183]:#183,199,298,467,957,244,963,527,681,948
cmd = f'python train_across_scene.py configs/nerf_ft.yaml defaults.expname=google_objs_{name}_{scene}_{views}_views ' \
f'{config} training.n_iters=10000 ' \
f'dataset.train_views={views} exporation.render_only=True exporation.render_path=True exporation.render_test=False ' \
f'dataset.train_scene_list=[{scene}] ' \
f'dataset.test_scene_list=[{scene}] ' \
f'dataset.datadir=./data/google_scanned_objects/ ' \
f'defaults.ckpt=./logs/google_objs_{name}_{scene}_{views}_views//google_objs_{name}_{scene}_{views}_views.th'
```
================================================
FILE: configs/360_v2.yaml
================================================
defaults:
expname: basis_room_real_mask
logdir: ./logs
ckpt: null # help='specific weights npy file to reload for coarse network'
model:
basis_dims: [5,5,5,2,2,2]
basis_resos: [ 64, 83, 102, 121, 140, 160]
coeff_reso: 16
coef_init: 0.01
phases: [0.0]
coef_mode: bilinear
basis_mode: bilinear
freq_bands: [ 1.0000, 1.7689, 2.3526, 3.1290, 4.1616, 6.]
kernel_mapping_type: 'sawtooth'
in_dim: 3
out_dim: 32
num_layers: 2
hidden_dim: 128
dataset:
# loader options
dataset_name: llff # choices=['blender', 'llff', 'nsvf', 'dtu','tankstemple', 'own_data']
datadir: /home/anpei/code/NeuBasis/data/360_v2/room/
ndc_ray: 0
is_unbound: True
with_depth: 0
downsample_train: 4.0
downsample_test: 4.0
N_vis: 5
vis_every: 5000
training:
n_iters: 30000
batch_size: 4096
volume_resoInit: 128 # 128**3：
volume_resoFinal: 320 # 300**3
upsamp_list: [2000,3000,4000,5500]
update_AlphaMask_list: [2500]
shrinking_list: [-1]
L1_weight_inital: 0.0
L1_weight_rest: 0.0
TV_weight_density: 0.0
TV_weight_app: 0.00
exportation:
render_only: 0
render_test: 1
render_train: 0
render_path: 0
export_mesh: 0
export_mesh_only: 0
renderer:
shadingMode: MLP_Fea
num_layers: 3
hidden_dim: 128
fea2denseAct: 'relu'
density_shift: -10
distance_scale: 25.0
view_pe: 6
fea_pe: 2
lindisp: 0
perturb: 1 # help='set to 0. for no jitter, 1. for jitter'
step_ratio: 0.5
max_samples: 1600
alphaMask_thres: 0.04
rayMarch_weight_thres: 1e-3
================================================
FILE: configs/defaults.yaml
================================================
defaults:
expname: basis_lego
logedir: ./logs
mode: 'reconstruction'
progress_refresh_rate: 10
add_timestamp: 0
model:
basis_dims: [4,4,4,2,2,2]
basis_resos: [32,51,70,89,108,128]
coeff_reso: 32
total_params: 10744166
T_basis: 0
T_coeff: 0
coef_init: 1.0
coef_mode: bilinear
basis_mode: bilinear
freq_bands: [1.0000, 1.3300, 1.7689, 2.3526, 3.1290, 4.1616]
basis_mapping: 'sawtooth'
with_dropout: False
in_dim: 3
out_dim: 32
num_layers: 2
hidden_dim: 128
dataset:
# loader options
dataset_name: blender # choices=['blender', 'llff', 'nsvf', 'dtu','tankstemple', 'own_data']
datadir: ./data/nerf_synthetic/lego
with_depth: 0
downsample_train: 1.0
downsample_test: 1.0
is_unbound: False
training:
# training options
batch_size: 4096
n_iters: 30000
# learning rate
lr_small: 0.001
lr_large: 0.02
lr_decay_iters: -1
lr_decay_target_ratio: 0.1 # help = 'number of iterations the lr will decay to the target ratio; -1 will set it to n_iters'
lr_upsample_reset: 1 # help='reset lr to inital after upsampling'
# loss
L1_weight_inital: 0.0 # help='loss weight'
L1_weight_rest: 0
Ortho_weight: 0.0
TV_weight_density: 0.0
TV_weight_app: 0.0
# optimiziable
coeff: True
basis: True
linear_mat: True
renderModule: True
================================================
FILE: configs/image.yaml
================================================
defaults:
expname: basis_image
logdir: ./logs
mode: 'image'
ckpt: null # help='specific weights npy file to reload for coarse network'
model:
basis_dims: [32,32,32,16,16,16]
basis_resos: [32,51,70,89,108,128]
freq_bands: [2. , 3.2, 4.4, 5.6, 6.8, 8.]
total_params: 1426063 # albert
# total_params: 61445328 # pluto
# total_params: 71848800 #Girl_with_a_Pearl_Earring
# total_params: 37138096 # Weissenbruch_Jan_Hendrik_The_Shipping_Canal_at_Rijswijk.jpeg_base
coeff_type: 'grid'
basis_type: 'grid'
coef_init: 0.001
coef_mode: nearest
basis_mode: nearest
basis_mapping: 'sawtooth'
in_dim: 2
out_dim: 3
num_layers: 2
hidden_dim: 64
with_dropout: False
dataset:
# loader options
dataset_name: image
datadir: "../data/image/albert.exr"
# datadir: "../data/image//pluto.jpeg"
# datadir: "../data/image//Girl_with_a_Pearl_Earring.jpeg"
# datadir: "../data/image//Weissenbruch_Jan_Hendrik_The_Shipping_Canal_at_Rijswijk.jpeg"
training:
n_iters: 10000
batch_size: 102400
# learning rate
lr_small: 0.002
lr_large: 0.002
================================================
FILE: configs/image_intro.yaml
================================================
defaults:
expname: basis_image
logdir: ./logs
mode: 'demo'
ckpt: null # help='specific weights npy file to reload for coarse network'
model:
in_dim: 2
out_dim: 1
basis_dims: [32,32,32,16,16,16]
basis_resos: [32,51,70,89,108,128]
freq_bands: [2. , 3.2, 4.4, 5.6, 6.8, 8.]
# occNet
coeff_type: 'none'
basis_type: 'x'
basis_mapping: 'x'
num_layers: 8
hidden_dim: 256
# coef_init: 0.001
# coef_mode: nearest
# basis_mode: nearest
# basis_mapping: 'sawtooth'
with_dropout: False
dataset:
# loader options
dataset_name: image
datadir: ../data/image/cat_occupancy.png
training:
n_iters: 10000
batch_size: 102400
# learning rate
lr_small: 0.0002
lr_large: 0.0002
================================================
FILE: configs/image_set.yaml
================================================
defaults:
expname: basis_image
logdir: ./logs
mode: 'images'
ckpt: null # help='specific weights npy file to reload for coarse network'
model:
basis_dims: [32,32,32,16,16,16]
basis_resos: [32,51,70,89,108,128]
freq_bands: [2. , 3.2, 4.4, 5.6, 6.8, 8.]
coeff_reso: 32
total_params: 1024000
coef_init: 0.001
coef_mode: bilinear
basis_mode: bilinear
coeff_type: 'grid'
basis_type: 'grid'
in_dim: 3
out_dim: 3
num_layers: 2
hidden_dim: 64
with_dropout: True
dataset:
# loader options
dataset_name: images
datadir: data/ffhq/ffhq_512.npy
training:
n_iters: 300000
batch_size: 40960
# learning rate
lr_small: 0.002
lr_large: 0.002
================================================
FILE: configs/nerf.yaml
================================================
defaults:
expname: basis_lego
logdir: ./logs
mode: 'reconstruction'
ckpt: null # help='specific weights npy file to reload for coarse network'
model:
coeff_reso: 32
basis_dims: [4,4,4,2,2,2]
basis_resos: [32,51,70,89,108,128]
freq_bands: [2. , 3.2, 4.4, 5.6, 6.8, 8.]
coef_init: 1.0
phases: [0.0]
total_params: 5308416
coef_mode: bilinear
basis_mode: bilinear
coeff_type: 'grid'
basis_type: 'grid'
basis_mapping: 'sawtooth'
in_dim: 3
out_dim: 32
num_layers: 2
hidden_dim: 64
dataset:
# loader options
dataset_name: blender # choices=['blender', 'llff', 'nsvf', 'dtu','tankstemple', 'own_data']
datadir: ./data/nerf_synthetic/lego
ndc_ray: 0
with_depth: 0
downsample_train: 1.0
downsample_test: 1.0
N_vis: 5
vis_every: 100000
scene_reso: 768
training:
n_iters: 30000
batch_size: 4096
volume_resoInit: 128 # 128**3：
volume_resoFinal: 300 # 300**3
upsamp_list: [2000,3000,4000,5500,7000]
update_AlphaMask_list: [2500,4000]
shrinking_list: [500]
L1_weight_inital: 0.0
L1_weight_rest: 0.0
TV_weight_density: 0.000
TV_weight_app: 0.00
exportation:
render_only: 0
render_test: 1
render_train: 0
render_path: 0
export_mesh: 0
export_mesh_only: 0
renderer:
shadingMode: MLP_Fea
num_layers: 3
hidden_dim: 128
fea2denseAct: 'softplus'
density_shift: -10
distance_scale: 25.0
view_pe: 6
fea_pe: 2
lindisp: 0
perturb: 1 # help='set to 0. for no jitter, 1. for jitter'
step_ratio: 0.5
max_samples: 1200
alphaMask_thres: 0.02
rayMarch_weight_thres: 1e-3
================================================
FILE: configs/nerf_ft.yaml
================================================
defaults:
expname: basis
logdir: ./logs
mode: 'reconstructions'
ckpt: null # help='specific weights npy file to reload for coarse network'
model:
coeff_reso: 16
basis_dims: [16,16,16,8,8,8]
basis_resos: [32,51,70,89,108,128]
freq_bands: [2. , 3.2, 4.4, 5.6, 6.8, 8.]
with_dropout: True
coef_init: 1.0
phases: [0.0]
total_params: 5308416
coef_mode: bilinear
basis_mode: bilinear
coeff_type: 'grid'
basis_type: 'grid'
basis_mapping: 'sawtooth'
in_dim: 3
out_dim: 32
num_layers: 2
hidden_dim: 64
dataset:
# loader options
dataset_name: google_objs # choices=['blender', 'llff', 'nsvf', 'dtu','tankstemple', 'own_data']
datadir: /vlg-nfs/anpei/dataset/google_scanned_objects
ndc_ray: 0
train_scene_list: [100]
test_scene_list: [100]
train_views: 5
with_depth: 0
downsample_train: 1.0
downsample_test: 1.0
N_vis: 5
vis_every: 100000
scene_reso: 768
training:
n_iters: 5000
batch_size: 4096
volume_resoInit: 128 # 128**3：
volume_resoFinal: 300 # 300**3
upsamp_list: [2000,3000,4000]
update_AlphaMask_list: [1500]
shrinking_list: [-1]
L1_weight_inital: 0.0
L1_weight_rest: 0.0
TV_weight_density: 0.000
TV_weight_app: 0.00
# optimiziable
coeff: True
basis: False
linear_mat: False
renderModule: False
exportation:
render_only: 0
render_test: 1
render_train: 0
render_path: 0
export_mesh: 0
export_mesh_only: 0
renderer:
shadingMode: MLP_Fea
num_layers: 3
hidden_dim: 128
fea2denseAct: 'softplus'
density_shift: -10
distance_scale: 25.0
view_pe: 6
fea_pe: 2
lindisp: 0
perturb: 1 # help='set to 0. for no jitter, 1. for jitter'
step_ratio: 0.5
max_samples: 1200
alphaMask_thres: 0.02
rayMarch_weight_thres: 1e-3
================================================
FILE: configs/nerf_set.yaml
================================================
defaults:
expname: basis_no_relu_lego
logdir: ./logs
mode: 'reconstructions'
ckpt: null # help='specific weights npy file to reload for coarse network'
model:
coeff_reso: 16
basis_dims: [16,16,16,8,8,8]
# basis_resos: [32,51,70,89,108,128]
basis_resos: [32,51,70,89,108,128]
# freq_bands: [1.0000, 1.7689, 2.3526, 3.1290, 4.1616, 6.]
freq_bands: [2. , 3.2, 4.4, 5.6, 6.8, 8.]
with_dropout: True
coef_init: 1.0
phases: [0.0]
total_params: 5308416
coef_mode: bilinear
basis_mode: bilinear
coeff_type: 'grid'
basis_type: 'grid'
basis_mapping: 'sawtooth'
in_dim: 3
out_dim: 32
num_layers: 2
hidden_dim: 64
dataset:
# loader options
dataset_name: google_objs # choices=['blender', 'llff', 'nsvf', 'dtu','tankstemple', 'own_data']
datadir: /vlg-nfs/anpei/dataset/google_scanned_objects
ndc_ray: 0
train_scene_list: [0,100]
test_scene_list: [0]
train_views: 100
with_depth: 0
downsample_train: 1.0
downsample_test: 1.0
N_vis: 5
vis_every: 100000
scene_reso: 768
training:
n_iters: 50000
batch_size: 4096
volume_resoInit: 128 # 128**3：
volume_resoFinal: 256 # 300**3
upsamp_list: [2000,3000,4000,5500,7000]
update_AlphaMask_list: [-1]
shrinking_list: [-1]
L1_weight_inital: 0.0
L1_weight_rest: 0.0
TV_weight_density: 0.000
TV_weight_app: 0.00
exportation:
render_only: 0
render_test: 0
render_train: 0
render_path: 0
export_mesh: 0
export_mesh_only: 0
renderer:
shadingMode: MLP_Fea
num_layers: 3
hidden_dim: 128
fea2denseAct: 'softplus'
density_shift: -10
distance_scale: 25.0
view_pe: 6
fea_pe: 2
lindisp: 0
perturb: 1 # help='set to 0. for no jitter, 1. for jitter'
step_ratio: 0.5
max_samples: 1200
alphaMask_thres: 0.005
rayMarch_weight_thres: 1e-3
================================================
FILE: configs/sdf.yaml
================================================
defaults:
expname: basis_sdf
logdir: ./logs
mode: 'sdf'
ckpt: null # help='specific weights npy file to reload for coarse network'
model:
basis_dims: [4,4,4,2,2,2]
basis_resos: [32,51,70,89,108,128]
freq_bands: [2. , 3.2, 4.4, 5.6, 6.8, 8.]
total_params: 5313942
coeff_reso: 32
coef_init: 0.05
coef_mode: bilinear
basis_mode: bilinear
coeff_type: 'grid'
basis_type: 'grid'
kernel_mapping_type: 'sawtooth'
in_dim: 3
out_dim: 1
num_layers: 1
hidden_dim: 64
dataset:
# loader options
dataset_name: sdf
datadir: "../data/mesh/statuette_close.npy"
scene_reso: 384
training:
n_iters: 10000
batch_size: 40960
# learning rate
lr_small: 0.002
lr_large: 0.02
================================================
FILE: configs/tnt.yaml
================================================
defaults:
expname: basis_truck
logdir: ./logs
ckpt: null # help='specific weights npy file to reload for coarse network'
model:
coeff_reso: 32
# basis_dims: [8,4,2]
# basis_resos: [32,64,128]
# freq_bands: [3.0,4.7,6.8]
## 32.88
# basis_dims: [3, 3, 3, 3, 3, 3, 3]
# basis_resos: [64, 64, 64, 64, 64, 64, 64]
## freq_bands: [1.52727273, 2.58181818, 3.63636364, 4.69090909, 5.21818182, 6.27272727, 6.8]
# freq_bands: [2., 3., 4.,5.,6.,7.,8.]
# basis_dims: [5,5,5,2,2,2]
# basis_resos: [32,51,70,89,108,128]
# coeff_reso: 26
# freq_bands: [1.0000, 1.7689, 2.3526, 3.1290, 4.1616, 6.]
basis_dims: [4,4,4,2,2,2]
basis_resos: [32,51,70,89,108,128]
freq_bands: [2. , 3.2, 4.4, 5.6, 6.8, 8.]
# freq_bands: [2. , 2.8, 3.6, 4.4, 5.2, 6.]
coef_init: 1.0
phases: [0.0]
total_params: 5744166
coef_mode: bilinear
basis_mode: bilinear
kernel_mapping_type: 'sawtooth'
in_dim: 3
out_dim: 32
num_layers: 2
hidden_dim: 64
dataset:
# loader options
dataset_name: tankstemple # choices=['blender', 'llff', 'nsvf', 'dtu','tankstemple', 'own_data']
datadir: ./data/TanksAndTemple/Truck
ndc_ray: 0
with_depth: 0
downsample_train: 1.0
downsample_test: 1.0
N_vis: 5
vis_every: 100000
scene_reso: 768
training:
n_iters: 30000
batch_size: 4096
volume_resoInit: 128 # 128**3：
volume_resoFinal: 320 # 300**3
# upsamp_list: [2000,4000,7000]
# update_AlphaMask_list: [2000,3000]
upsamp_list: [2000,3000,4000,5500,7000]
update_AlphaMask_list: [2500,4000]
shrinking_list: [500]
L1_weight_inital: 0.0
L1_weight_rest: 0.0
TV_weight_density: 0.0
TV_weight_app: 0.00
exportation:
render_only: 0
render_test: 1
render_train: 0
render_path: 0
export_mesh: 0
export_mesh_only: 0
renderer:
shadingMode: MLP_Fea
num_layers: 3
hidden_dim: 128
fea2denseAct: 'softplus'
density_shift: -10
distance_scale: 25.0
view_pe: 2
fea_pe: 2
lindisp: 0
perturb: 1 # help='set to 0. for no jitter, 1. for jitter'
step_ratio: 0.5
max_samples: 1200
alphaMask_thres: 0.005
rayMarch_weight_thres: 1e-3
================================================
FILE: dataLoader/__init__.py
================================================
from .llff import LLFFDataset
from .blender import BlenderDataset
from .nsvf import NSVF
from .tankstemple import TanksTempleDataset
from .your_own_data import YourOwnDataset
from .image import ImageDataset
from .image_set import ImageSetDataset
from .colmap import ColmapDataset
from .sdf import SDFDataset
from .blender_set import BlenderDatasetSet
from .google_objs import GoogleObjsDataset
from .dtu_objs import DTUDataset
dataset_dict = {'blender': BlenderDataset,
'blender_set': BlenderDatasetSet,
'llff':LLFFDataset,
'tankstemple':TanksTempleDataset,
'nsvf':NSVF,
'own_data':YourOwnDataset,
'image':ImageDataset,
'images':ImageSetDataset,
'sdf':SDFDataset,
'colmap':ColmapDataset,
'google_objs':GoogleObjsDataset,
'dtu':DTUDataset}
================================================
FILE: dataLoader/blender.py
================================================
import torch, cv2
from torch.utils.data import Dataset
import json
from tqdm import tqdm
import os
from PIL import Image
from torchvision import transforms as T
from .ray_utils import *
class BlenderDataset(Dataset):
def __init__(self, cfg, split='train', batch_size=4096, is_stack=None):
# self.N_vis = N_vis
self.split = split
self.batch_size = batch_size
self.root_dir = cfg.datadir
self.is_stack = is_stack if is_stack is not None else 'train'!=split
self.downsample = cfg.get(f'downsample_{self.split}')
self.img_wh = (int(800 / self.downsample), int(800 / self.downsample))
self.define_transforms()
self.rot = torch.tensor([[0.65561799, -0.65561799, 0.37460659],
[0.73729737, 0.44876192, -0.50498052],
[0.16296514, 0.60727077, 0.77760181]])
self.scene_bbox = (np.array([[-1.0, -1.0, -1.0], [1.0, 1.0, 1.0]])).tolist()
# self.scene_bbox = [[-0.8,-0.8,-0.22],[0.8,0.8,0.2]]
self.blender2opencv = np.array([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
self.read_meta()
self.define_proj_mat()
self.white_bg = True
self.near_far = [2.0, 6.0]
# self.center = torch.mean(self.scene_bbox, axis=0).float().view(1, 1, 3)
# self.radius = (self.scene_bbox[1] - self.center).float().view(1, 1, 3)
def read_depth(self, filename):
depth = np.array(read_pfm(filename)[0], dtype=np.float32) # (800, 800)
return depth
def read_meta(self):
with open(os.path.join(self.root_dir, f"transforms_{self.split}.json"), 'r') as f:
self.meta = json.load(f)
w, h = self.img_wh
self.focal = 0.5 * 800 / np.tan(0.5 * self.meta['camera_angle_x']) # original focal length
self.focal *= self.img_wh[0] / 800 # modify focal length to match size self.img_wh
# ray directions for all pixels, same for all images (same H, W, focal)
self.directions = get_ray_directions(h, w, [self.focal, self.focal]) # (h, w, 3)
self.directions = self.directions / torch.norm(self.directions, dim=-1, keepdim=True)
self.intrinsics = torch.tensor([[self.focal, 0, w / 2], [0, self.focal, h / 2], [0, 0, 1]]).float()
self.image_paths = []
self.poses = []
self.all_rays = []
self.all_rgbs = []
self.all_masks = []
self.all_depth = []
self.downsample = 1.0
img_eval_interval = 1 #if self.N_vis < 0 else len(self.meta['frames']) // self.N_vis
idxs = list(range(0, len(self.meta['frames']), img_eval_interval))
# idxs = idxs[:10] if self.split=='train' else idxs
for i in tqdm(idxs, desc=f'Loading data {self.split} ({len(idxs)})'): # img_list:#
frame = self.meta['frames'][i]
pose = np.array(frame['transform_matrix']) @ self.blender2opencv
c2w = torch.FloatTensor(pose)
c2w[:3,-1] /= 1.5
self.poses += [c2w]
image_path = os.path.join(self.root_dir, f"{frame['file_path']}.png")
self.image_paths += [image_path]
img = Image.open(image_path)
if self.downsample != 1.0:
img = img.resize(self.img_wh, Image.LANCZOS)
img = self.transform(img) # (4, h, w)
img = img.view(4, -1).permute(1, 0) # (h*w, 4) RGBA
img = img[:, :3] * img[:, -1:] + (1 - img[:, -1:]) # blend A to RGB
self.all_rgbs += [img]
rays_o, rays_d = get_rays(self.directions, c2w) # both (h*w, 3)
# rays_o, rays_d = rays_o@self.rot, rays_d@self.rot
self.all_rays += [torch.cat([rays_o, rays_d], 1)] # (h*w, 6)
self.poses = torch.stack(self.poses)
if not self.is_stack:
self.all_rays = torch.cat(self.all_rays, 0) # (len(self.meta['frames])*h*w, 3)
self.all_rgbs = torch.cat(self.all_rgbs, 0) # (len(self.meta['frames])*h*w, 3)
# self.all_depth = torch.cat(self.all_depth, 0) # (len(self.meta['frames])*h*w, 3)
else:
self.all_rays = torch.stack(self.all_rays, 0) # (len(self.meta['frames]),h*w, 3)
self.all_rgbs = torch.stack(self.all_rgbs, 0).reshape(-1, *self.img_wh[::-1],3) # (len(self.meta['frames]),h,w,3)
# self.all_masks = torch.stack(self.all_masks, 0).reshape(-1,*self.img_wh[::-1]) # (len(self.meta['frames]),h,w,3)
self.sampler = SimpleSampler(np.prod(self.all_rgbs.shape[:-1]),self.batch_size)
def define_transforms(self):
self.transform = T.ToTensor()
def define_proj_mat(self):
self.proj_mat = self.intrinsics.unsqueeze(0) @ torch.inverse(self.poses)[:, :3]
# def world2ndc(self,points,lindisp=None):
# device = points.device
# return (points - self.center.to(device)) / self.radius.to(device)
def __len__(self):
return len(self.all_rgbs) if self.split=='test' else 300000
def __getitem__(self, idx):
idx_rand = self.sampler.nextids() #torch.randint(0,len(self.all_rays),(self.batch_size,))
sample = {'rays': self.all_rays[idx_rand], 'rgbs': self.all_rgbs[idx_rand]}
return sample
================================================
FILE: dataLoader/blender_set.py
================================================
import torch, cv2
from torch.utils.data import Dataset
import json
from tqdm import tqdm
import os
from PIL import Image
from torchvision import transforms as T
from .ray_utils import *
class BlenderDatasetSet(Dataset):
def __init__(self, cfg, split='train'):
# self.N_vis = N_vis
self.root_dir = cfg.datadir
self.split = split
self.is_stack = False if 'train'==split else True
self.downsample = cfg.get(f'downsample_{self.split}')
self.img_wh = (int(800 / self.downsample), int(800 / self.downsample))
self.define_transforms()
self.rot = torch.tensor([[0.65561799, -0.65561799, 0.37460659],
[0.73729737, 0.44876192, -0.50498052],
[0.16296514, 0.60727077, 0.77760181]])
self.scene_bbox = (np.array([[-1.0, -1.0, -1.0, 0], [1.0, 1.0, 1.0, 2]])).tolist()
# self.scene_bbox = [[-0.8,-0.8,-0.22],[0.8,0.8,0.2]]
self.blender2opencv = np.array([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
self.read_meta()
self.define_proj_mat()
self.white_bg = True
self.near_far = [2.0, 6.0]
# self.center = torch.mean(self.scene_bbox, axis=0).float().view(1, 1, 3)
# self.radius = (self.scene_bbox[1] - self.center).float().view(1, 1, 3)
def read_depth(self, filename):
depth = np.array(read_pfm(filename)[0], dtype=np.float32) # (800, 800)
return depth
def read_meta(self):
with open(os.path.join(self.root_dir, f"transforms_{self.split}.json"), 'r') as f:
self.meta = json.load(f)
w, h = self.img_wh
self.focal = 0.5 * 800 / np.tan(0.5 * self.meta['camera_angle_x']) # original focal length
self.focal *= self.img_wh[0] / 800 # modify focal length to match size self.img_wh
# ray directions for all pixels, same for all images (same H, W, focal)
self.directions = get_ray_directions(h, w, [self.focal, self.focal]) # (h, w, 3)
self.directions = self.directions / torch.norm(self.directions, dim=-1, keepdim=True)
self.intrinsics = torch.tensor([[self.focal, 0, w / 2], [0, self.focal, h / 2], [0, 0, 1]]).float()
self.image_paths = []
self.poses = []
self.all_rays = []
self.all_rgbs = []
self.all_masks = []
self.all_depth = []
self.downsample = 1.0
img_eval_interval = 1 #if self.N_vis < 0 else len(self.meta['frames']) // self.N_vis
idxs = list(range(0, len(self.meta['frames']), img_eval_interval))
for i in tqdm(idxs, desc=f'Loading data {self.split} ({len(idxs)})'): # img_list:#
frame = self.meta['frames'][i]
pose = np.array(frame['transform_matrix']) @ self.blender2opencv
c2w = torch.FloatTensor(pose)
c2w[:3,-1] /= 1.5
self.poses += [c2w]
image_path = os.path.join(self.root_dir, f"{frame['file_path']}.png")
self.image_paths += [image_path]
img = Image.open(image_path)
if self.downsample != 1.0:
img = img.resize(self.img_wh, Image.LANCZOS)
img = self.transform(img) # (4, h, w)
img = img.view(4, -1).permute(1, 0) # (h*w, 4) RGBA
img = img[:, :3] * img[:, -1:] + (1 - img[:, -1:]) # blend A to RGB
self.all_rgbs += [img]
rays_o, rays_d = get_rays(self.directions, c2w) # both (h*w, 3)
# rays_o, rays_d = rays_o@self.rot, rays_d@self.rot
scene_id = torch.ones_like(rays_o[...,:1])*0
self.all_rays += [torch.cat([rays_o, rays_d, scene_id], 1)] # (h*w, 6)
self.poses = torch.stack(self.poses)
if not self.is_stack:
self.all_rays = torch.cat(self.all_rays, 0) # (len(self.meta['frames])*h*w, 3)
self.all_rgbs = torch.cat(self.all_rgbs, 0) # (len(self.meta['frames])*h*w, 3)
# self.all_depth = torch.cat(self.all_depth, 0) # (len(self.meta['frames])*h*w, 3)
else:
self.all_rays = torch.stack(self.all_rays, 0) # (len(self.meta['frames]),h*w, 3)
self.all_rgbs = torch.stack(self.all_rgbs, 0).reshape(-1, *self.img_wh[::-1],3) # (len(self.meta['frames]),h,w,3)
# self.all_masks = torch.stack(self.all_masks, 0).reshape(-1,*self.img_wh[::-1]) # (len(self.meta['frames]),h,w,3)
def define_transforms(self):
self.transform = T.ToTensor()
def define_proj_mat(self):
self.proj_mat = self.intrinsics.unsqueeze(0) @ torch.inverse(self.poses)[:, :3]
# def world2ndc(self,points,lindisp=None):
# device = points.device
# return (points - self.center.to(device)) / self.radius.to(device)
def __len__(self):
return len(self.all_rgbs)
def __getitem__(self, idx):
rays = torch.cat((self.all_rays[idx],torch.tensor([0+0.5])),dim=-1)
sample = {'rays': rays, 'rgbs': self.all_rgbs[idx]}
return sample
================================================
FILE: dataLoader/colmap.py
================================================
import torch, cv2
from torch.utils.data import Dataset
from tqdm import tqdm
import os
from PIL import Image
from torchvision import transforms as T
from .ray_utils import *
class ColmapDataset(Dataset):
def __init__(self, cfg, split='train'):
self.cfg = cfg
self.root_dir = cfg.datadir
self.split = split
self.is_stack = False if 'train'==split else True
self.downsample = cfg.get(f'downsample_{self.split}')
self.define_transforms()
self.img_eval_interval = 8
self.blender2opencv = np.array([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])#np.eye(4)#
self.read_meta()
self.white_bg = cfg.get('white_bg')
# self.near_far = [0.1,2.0]
def read_meta(self):
if os.path.exists(f'{self.root_dir}/transforms.json'):
self.meta = load_json(f'{self.root_dir}/transforms.json')
i_test = np.arange(0, len(self.meta['frames']), self.img_eval_interval) # [np.argmin(dists)]
idxs = i_test if self.split != 'train' else list(set(np.arange(len(self.meta['frames']))) - set(i_test))
else:
self.meta = load_json(f'{self.root_dir}/transforms_{self.split}.json')
idxs = np.arange(0, len(self.meta['frames'])) # [np.argmin(dists)]
inv_split = 'train' if self.split!='train' else 'test'
self.meta['frames'] += load_json(f'{self.root_dir}/transforms_{inv_split}.json')['frames']
print(len(self.meta['frames']),len(idxs))
self.scale = self.meta.get('scale', 0.5)
self.offset = torch.FloatTensor(self.meta.get('offset', [0.0,0.0,0.0]))
# self.scene_bbox = (torch.tensor([[-6.,-7.,-10.0],[6.,7.,10.]])/5).tolist()
# self.scene_bbox = [[-1., -1., -1.0], [1., 1., 1.]]
# center, radius = torch.tensor([-0.082157, 2.415426,-3.703080]), torch.tensor([7.36916, 11.34958, 20.1616])/2
# self.scene_bbox = torch.stack([center-radius, center+radius]).tolist()
h, w = int(self.meta.get('h')), int(self.meta.get('w'))
cx, cy = self.meta.get('cx'), self.meta.get('cy')
self.focal = [self.meta.get('fl_x'), self.meta.get('fl_y')]
# ray directions for all pixels, same for all images (same H, W, focal)
self.directions = get_ray_directions(h, w, self.focal, center=[cx, cy]) # (h, w, 3)
self.directions = self.directions / torch.norm(self.directions, dim=-1, keepdim=True)
# self.intrinsics = torch.FloatTensor([[self.focal[0], 0, cx], [0, self.focal[1], cy], [0, 0, 1]])
# self.intrinsics[:2] /= self.downsample
poses = pose_from_json(self.meta, self.blender2opencv)
poses, self.scene_bbox = orientation(poses, f'{self.root_dir}/colmap_text/points3D.txt')
self.image_paths = []
self.poses = []
self.all_rays = []
self.all_rgbs = []
self.all_masks = []
self.all_depth = []
self.img_wh = [w,h]
for i in tqdm(idxs, desc=f'Loading data {self.split} ({len(idxs)})'): # img_list:#
frame = self.meta['frames'][i]
c2w = torch.FloatTensor(poses[i])
# c2w[:3,3] = (c2w[:3,3]*self.scale + self.offset)*2-1
self.poses += [c2w]
image_path = os.path.join(self.root_dir, frame['file_path'])
self.image_paths += [image_path]
img = Image.open(image_path)
if self.downsample != 1.0:
img = img.resize(self.img_wh, Image.LANCZOS)
img = self.transform(img)
if img.shape[0]==4:
img = img[:3] * img[-1:] + (1 - img[-1:]) # blend A to RGB
img = img.view(3, -1).permute(1, 0)
self.all_rgbs += [img]
rays_o, rays_d = get_rays(self.directions, c2w) # both (h*w, 3)
self.all_rays += [torch.cat([rays_o, rays_d], 1)] # (h*w, 6)
self.poses = torch.stack(self.poses)
if not self.is_stack:
self.all_rays = torch.cat(self.all_rays, 0) # (len(self.meta['frames])*h*w, 3)
self.all_rgbs = torch.cat(self.all_rgbs, 0) # (len(self.meta['frames])*h*w, 3)
else:
self.all_rays = torch.stack(self.all_rays, 0) # (len(self.meta['frames]),h*w, 3)
self.all_rgbs = torch.stack(self.all_rgbs, 0).reshape(-1, *self.img_wh[::-1],3) # (len(self.meta['frames]),h,w,3)
def define_transforms(self):
self.transform = T.ToTensor()
def __len__(self):
return len(self.all_rgbs)
def __getitem__(self, idx):
if self.split == 'train': # use data in the buffers
sample = {'rays': self.all_rays[idx],
'rgbs': self.all_rgbs[idx]}
else: # create data for each image separately
img = self.all_rgbs[idx]
rays = self.all_rays[idx]
sample = {'rays': rays,
'rgbs': img}
return sample
================================================
FILE: dataLoader/colmap2nerf.py
================================================
#!/usr/bin/env python3
# Copyright (c) 2020-2022, NVIDIA CORPORATION. All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto. Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
import argparse
import os
from pathlib import Path, PurePosixPath
import numpy as np
import json
import sys
import math
import cv2
import os
import shutil
def parse_args():
parser = argparse.ArgumentParser(description="convert a text colmap export to nerf format transforms.json; optionally convert video to images, and optionally run colmap in the first place")
parser.add_argument("--video_in", default="", help="run ffmpeg first to convert a provided video file into a set of images. uses the video_fps parameter also")
parser.add_argument("--video_fps", default=2)
parser.add_argument("--time_slice", default="", help="time (in seconds) in the format t1,t2 within which the images should be generated from the video. eg: \"--time_slice '10,300'\" will generate images only from 10th second to 300th second of the video")
parser.add_argument("--run_colmap", action="store_true", help="run colmap first on the image folder")
parser.add_argument("--colmap_matcher", default="sequential", choices=["exhaustive","sequential","spatial","transitive","vocab_tree"], help="select which matcher colmap should use. sequential for videos, exhaustive for adhoc images")
parser.add_argument("--colmap_db", default="colmap.db", help="colmap database filename")
parser.add_argument("--colmap_camera_model", default="OPENCV", choices=["SIMPLE_PINHOLE", "PINHOLE", "SIMPLE_RADIAL", "RADIAL","OPENCV"], help="camera model")
parser.add_argument("--colmap_camera_params", default="", help="intrinsic parameters, depending on the chosen model. Format: fx,fy,cx,cy,dist")
parser.add_argument("--images", default="images", help="input path to the images")
parser.add_argument("--text", default="colmap_text", help="input path to the colmap text files (set automatically if run_colmap is used)")
parser.add_argument("--aabb_scale", default=16, choices=["1","2","4","8","16"], help="large scene scale factor. 1=scene fits in unit cube; power of 2 up to 16")
parser.add_argument("--skip_early", default=0, help="skip this many images from the start")
parser.add_argument("--keep_colmap_coords", action="store_true", help="keep transforms.json in COLMAP's original frame of reference (this will avoid reorienting and repositioning the scene for preview and rendering)")
parser.add_argument("--out", default="transforms.json", help="output path")
parser.add_argument("--vocab_path", default="", help="vocabulary tree path")
args = parser.parse_args()
return args
def do_system(arg):
print(f"==== running: {arg}")
err = os.system(arg)
if err:
print("FATAL: command failed")
sys.exit(err)
def run_ffmpeg(args):
if not os.path.isabs(args.images):
args.images = os.path.join(os.path.dirname(args.video_in), args.images)
images = args.images
video = args.video_in
fps = float(args.video_fps) or 1.0
print(f"running ffmpeg with input video file={video}, output image folder={images}, fps={fps}.")
if (input(f"warning! folder '{images}' will be deleted/replaced. continue? (Y/n)").lower().strip()+"y")[:1] != "y":
sys.exit(1)
try:
shutil.rmtree(images)
except:
pass
do_system(f"mkdir {images}")
time_slice_value = ""
time_slice = args.time_slice
if time_slice:
start, end = time_slice.split(",")
time_slice_value = f",select='between(t\,{start}\,{end})'"
do_system(f"ffmpeg -i {video} -qscale:v 1 -qmin 1 -vf \"fps={fps}{time_slice_value}\" {images}/%04d.jpg")
def run_colmap(args):
db = args.colmap_db
images = "\"" + args.images + "\""
db_noext=str(Path(db).with_suffix(""))
if args.text=="text":
args.text=db_noext+"_text"
text=args.text
sparse=db_noext+"_sparse"
print(f"running colmap with:\n\tdb={db}\n\timages={images}\n\tsparse={sparse}\n\ttext={text}")
if (input(f"warning! folders '{sparse}' and '{text}' will be deleted/replaced. continue? (Y/n)").lower().strip()+"y")[:1] != "y":
sys.exit(1)
if os.path.exists(db):
os.remove(db)
do_system(f"colmap feature_extractor --ImageReader.camera_model {args.colmap_camera_model} --ImageReader.camera_params \"{args.colmap_camera_params}\" --SiftExtraction.estimate_affine_shape=true --SiftExtraction.domain_size_pooling=true --ImageReader.single_camera 1 --database_path {db} --image_path {images}")
match_cmd = f"colmap {args.colmap_matcher}_matcher --SiftMatching.guided_matching=true --database_path {db}"
if args.vocab_path:
match_cmd += f" --VocabTreeMatching.vocab_tree_path {args.vocab_path}"
do_system(match_cmd)
try:
shutil.rmtree(sparse)
except:
pass
do_system(f"mkdir {sparse}")
do_system(f"colmap mapper --database_path {db} --image_path {images} --output_path {sparse}")
do_system(f"colmap bundle_adjuster --input_path {sparse}/0 --output_path {sparse}/0 --BundleAdjustment.refine_principal_point 1")
try:
shutil.rmtree(text)
except:
pass
do_system(f"mkdir {text}")
do_system(f"colmap model_converter --input_path {sparse}/0 --output_path {text} --output_type TXT")
def variance_of_laplacian(image):
return cv2.Laplacian(image, cv2.CV_64F).var()
def sharpness(imagePath):
image = cv2.imread(imagePath)
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
fm = variance_of_laplacian(gray)
return fm
def qvec2rotmat(qvec):
return np.array([
[
1 - 2 * qvec[2]**2 - 2 * qvec[3]**2,
2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3],
2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2]
], [
2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3],
1 - 2 * qvec[1]**2 - 2 * qvec[3]**2,
2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1]
], [
2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2],
2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1],
1 - 2 * qvec[1]**2 - 2 * qvec[2]**2
]
])
def rotmat(a, b):
a, b = a / np.linalg.norm(a), b / np.linalg.norm(b)
v = np.cross(a, b)
c = np.dot(a, b)
s = np.linalg.norm(v)
kmat = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])
return np.eye(3) + kmat + kmat.dot(kmat) * ((1 - c) / (s ** 2 + 1e-10))
def closest_point_2_lines(oa, da, ob, db): # returns point closest to both rays of form o+t*d, and a weight factor that goes to 0 if the lines are parallel
da = da / np.linalg.norm(da)
db = db / np.linalg.norm(db)
c = np.cross(da, db)
denom = np.linalg.norm(c)**2
t = ob - oa
ta = np.linalg.det([t, db, c]) / (denom + 1e-10)
tb = np.linalg.det([t, da, c]) / (denom + 1e-10)
if ta > 0:
ta = 0
if tb > 0:
tb = 0
return (oa+ta*da+ob+tb*db) * 0.5, denom
############ orientation ##############
def normalize(x):
return x / np.linalg.norm(x)
def rotation_matrix_from_vectors(vec1, vec2):
""" Find the rotation matrix that aligns vec1 to vec2
:param vec1: A 3d "source" vector
:param vec2: A 3d "destination" vector
:return mat: A transform matrix (3x3) which when applied to vec1, aligns it with vec2.
"""
a, b = (vec1 / np.linalg.norm(vec1)).reshape(3), (vec2 / np.linalg.norm(vec2)).reshape(3)
v = np.cross(a, b)
if any(v): # if not all zeros then
c = np.dot(a, b)
s = np.linalg.norm(v)
kmat = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])
return np.eye(3) + kmat + kmat.dot(kmat) * ((1 - c) / (s ** 2))
else:
return np.eye(3) # cross of all zeros only occurs on identical directions
def rotation_up(poses):
up = normalize(np.linalg.lstsq(poses[:, :3, 1],np.ones((poses.shape[0],1)),rcond=None)[0])
rot = rotation_matrix_from_vectors(up,np.array([0.,1.,0.]))
return rot
def search_orientation(points):
from scipy.spatial.transform import Rotation as R
bbox_sizes,rot_mats,bboxs = [],[],[]
for y_angle in np.linspace(-45,45,15):
rotvec = np.array([0,y_angle,0])/180*np.pi
rot = R.from_rotvec(rotvec).as_matrix()
point_orientation = rot@points
bbox = np.max(point_orientation,axis=1) - np.min(point_orientation,axis=1)
bbox_sizes.append(np.prod(bbox))
rot_mats.append(rot)
bboxs.append(bbox)
rot = rot_mats[np.argmin(bbox_sizes)]
bbox = bboxs[np.argmin(bbox_sizes)]
return rot,bbox
def load_point_txt(path):
points = []
with open(path, "r") as f:
for line in f:
if line[0] == "#":
continue
els = line.split(" ")
points.append([float(els[1]),float(els[2]),float(els[3])])
return np.stack(points)
# def load_c2ws(frames):
# c2ws =
# for f in frames:
# f["transform_matrix"] -= totp
#
# def oritation(transform_matrix, point_txt):
if __name__ == "__main__":
args = parse_args()
if args.video_in != "":
run_ffmpeg(args)
if args.run_colmap:
run_colmap(args)
AABB_SCALE = int(args.aabb_scale)
SKIP_EARLY = int(args.skip_early)
IMAGE_FOLDER = args.images
TEXT_FOLDER = args.text
OUT_PATH = args.out
print(f"outputting to {OUT_PATH}...")
with open(os.path.join(TEXT_FOLDER,"cameras.txt"), "r") as f:
angle_x = math.pi / 2
for line in f:
# 1 SIMPLE_RADIAL 2048 1536 1580.46 1024 768 0.0045691
# 1 OPENCV 3840 2160 3178.27 3182.09 1920 1080 0.159668 -0.231286 -0.00123982 0.00272224
# 1 RADIAL 1920 1080 1665.1 960 540 0.0672856 -0.0761443
if line[0] == "#":
continue
els = line.split(" ")
w = float(els[2])
h = float(els[3])
fl_x = float(els[4])
fl_y = float(els[4])
k1 = 0
k2 = 0
p1 = 0
p2 = 0
cx = w / 2
cy = h / 2
if els[1] == "SIMPLE_PINHOLE":
cx = float(els[5])
cy = float(els[6])
elif els[1] == "PINHOLE":
fl_y = float(els[5])
cx = float(els[6])
cy = float(els[7])
elif els[1] == "SIMPLE_RADIAL":
cx = float(els[5])
cy = float(els[6])
k1 = float(els[7])
elif els[1] == "RADIAL":
cx = float(els[5])
cy = float(els[6])
k1 = float(els[7])
k2 = float(els[8])
elif els[1] == "OPENCV":
fl_y = float(els[5])
cx = float(els[6])
cy = float(els[7])
k1 = float(els[8])
k2 = float(els[9])
p1 = float(els[10])
p2 = float(els[11])
else:
print("unknown camera model ", els[1])
# fl = 0.5 * w / tan(0.5 * angle_x);
angle_x = math.atan(w / (fl_x * 2)) * 2
angle_y = math.atan(h / (fl_y * 2)) * 2
fovx = angle_x * 180 / math.pi
fovy = angle_y * 180 / math.pi
print(f"camera:\n\tres={w,h}\n\tcenter={cx,cy}\n\tfocal={fl_x,fl_y}\n\tfov={fovx,fovy}\n\tk={k1,k2} p={p1,p2} ")
with open(os.path.join(TEXT_FOLDER,"images.txt"), "r") as f:
i = 0
bottom = np.array([0.0, 0.0, 0.0, 1.0]).reshape([1, 4])
out = {
"camera_angle_x": angle_x,
"camera_angle_y": angle_y,
"fl_x": fl_x,
"fl_y": fl_y,
"k1": k1,
"k2": k2,
"p1": p1,
"p2": p2,
"cx": cx,
"cy": cy,
"w": w,
"h": h,
"aabb_scale": AABB_SCALE,
"frames": [],
}
up = np.zeros(3)
for line in f:
line = line.strip()
if line[0] == "#":
continue
i = i + 1
if i < SKIP_EARLY*2:
continue
if i % 2 == 1:
elems=line.split(" ") # 1-4 is quat, 5-7 is trans, 9ff is filename (9, if filename contains no spaces)
#name = str(PurePosixPath(Path(IMAGE_FOLDER, elems[9])))
# why is this requireing a relitive path while using ^
image_rel = os.path.relpath(IMAGE_FOLDER)
name = str(f"./{image_rel}/{'_'.join(elems[9:])}")
b=sharpness(name)
print(name, "sharpness=",b)
image_id = int(elems[0])
qvec = np.array(tuple(map(float, elems[1:5])))
tvec = np.array(tuple(map(float, elems[5:8])))
R = qvec2rotmat(-qvec)
t = tvec.reshape([3,1])
m = np.concatenate([np.concatenate([R, t], 1), bottom], 0)
c2w = np.linalg.inv(m)
# c2w[0:3,2] *= -1 # flip the y and z axis
# c2w[0:3,1] *= -1
# c2w = c2w[[1,0,2,3],:] # swap y and z
# c2w[2,:] *= -1 # flip whole world upside down
# up += c2w[0:3,1]
frame={"file_path":name,"sharpness":b,"transform_matrix": c2w}
out["frames"].append(frame)
# up = up / np.linalg.norm(up)
# print("up vector was", up)
# R = rotmat(up,[0,0,1]) # rotate up vector to [0,0,1]
# R = np.pad(R,[0,1])
# R[-1, -1] = 1
# for f in out["frames"]:
# f["transform_matrix"] = np.matmul(R, f["transform_matrix"]) # rotate up to be the z axis
nframes = len(out["frames"])
# find a central point they are all looking at
print("computing center of attention...")
totw = 0.0
totp = np.array([0.0, 0.0, 0.0])
for f in out["frames"]:
mf = f["transform_matrix"][0:3,:]
for g in out["frames"]:
mg = g["transform_matrix"][0:3,:]
p, w = closest_point_2_lines(mf[:,3], mf[:,2], mg[:,3], mg[:,2])
if w > 0.01:
totp += p*w
totw += w
totp /= totw
print(totp) # the cameras are looking at totp
for f in out["frames"]:
f["transform_matrix"][0:3,3] -= totp
avglen = 0.
for f in out["frames"]:
avglen += np.linalg.norm(f["transform_matrix"][0:3,3])
avglen /= nframes
print("avg camera distance from origin", avglen)
for f in out["frames"]:
f["transform_matrix"][0:3,3] *= 4.0 / avglen # scale to "nerf sized"
for f in out["frames"]:
f["transform_matrix"] = f["transform_matrix"].tolist()
print(nframes,"frames")
print(f"writing {OUT_PATH}")
with open(OUT_PATH, "w") as outfile:
json.dump(out, outfile, indent=2)
================================================
FILE: dataLoader/dtu_objs.py
================================================
import torch
import cv2 as cv
import numpy as np
import os
from glob import glob
from .ray_utils import *
from torch.utils.data import Dataset
# This function is borrowed from IDR: https://github.com/lioryariv/idr
def load_K_Rt_from_P(filename, P=None):
if P is None:
lines = open(filename).read().splitlines()
if len(lines) == 4:
lines = lines[1:]
lines = [[x[0], x[1], x[2], x[3]] for x in (x.split(" ") for x in lines)]
P = np.asarray(lines).astype(np.float32).squeeze()
out = cv.decomposeProjectionMatrix(P)
K = out[0]
R = out[1]
t = out[2]
K = K / K[2, 2]
intrinsics = np.eye(4)
intrinsics[:3, :3] = K
pose = np.eye(4, dtype=np.float32)
pose[:3, :3] = R.transpose()
pose[:3, 3] = (t[:3] / t[3])[:, 0]
return intrinsics, pose
def fps_downsample(points, n_points_to_sample):
selected_points = np.zeros((n_points_to_sample, 3))
selected_idxs = []
dist = np.ones(points.shape[0]) * 100
for i in range(n_points_to_sample):
idx = np.argmax(dist)
selected_points[i] = points[idx]
selected_idxs.append(idx)
dist_ = ((points - selected_points[i]) ** 2).sum(-1)
dist = np.minimum(dist, dist_)
return selected_idxs
class DTUDataset(Dataset):
def __init__(self, cfg, split='train', batch_size=4096, is_stack=None):
"""
img_wh should be set to a tuple ex: (1152, 864) to enable test mode!
"""
# self.N_vis = N_vis
self.split = split
self.batch_size = batch_size
self.root_dir = cfg.datadir
self.is_stack = is_stack if is_stack is not None else 'train'!=split
self.downsample = cfg.get(f'downsample_{self.split}')
self.img_wh = (int(400 / self.downsample), int(300 / self.downsample))
train_scene_idxs = sorted(cfg.train_scene_list)
test_scene_idxs = cfg.test_scene_list
if len(train_scene_idxs)==2:
train_scene_idxs = list(range(train_scene_idxs[0],train_scene_idxs[1]))
self.scene_idxs = train_scene_idxs if self.split=='train' else test_scene_idxs
print(self.scene_idxs)
self.train_views = cfg.train_views
self.scene_num = len(self.scene_idxs)
self.test_index = test_scene_idxs
# if 'test' == self.split:
# self.test_index = train_scene_idxs.index(test_scene_idxs[0])
self.scene_path_list = [os.path.join(self.root_dir, f"scan{i}") for i in self.scene_idxs]
# self.scene_path_list = sorted(glob(os.path.join(self.root_dir, "scan*")))
self.read_meta()
self.white_bg = False
def read_meta(self):
self.aabbs = []
self.all_rgb_files,self.all_pose_files,self.all_intrinsics_files = {},{},{}
for i, scene_idx in enumerate(self.scene_idxs):
scene_path = self.scene_path_list[i]
camera_dict = np.load(os.path.join(scene_path, 'cameras.npz'))
self.all_rgb_files[scene_idx] = [
os.path.join(scene_path, "image", f)
for f in sorted(os.listdir(os.path.join(scene_path, "image")))
]
# world_mat is a projection matrix from world to image
n_images = len(self.all_rgb_files[scene_idx])
world_mats_np = [camera_dict['world_mat_%d' % idx].astype(np.float32) for idx in range(n_images)]
scale_mats_np = [camera_dict['scale_mat_%d' % idx].astype(np.float32) for idx in range(n_images)]
object_scale_mat = camera_dict['scale_mat_0']
self.aabbs.append(self.get_bbox(scale_mats_np,object_scale_mat))
# W,H = self.img_wh
intrinsics_scene, poses_scene = [], []
for img_idx, (scale_mat, world_mat) in enumerate(zip(scale_mats_np, world_mats_np)):
P = world_mat @ scale_mat
P = P[:3, :4]
intrinsic, c2w = load_K_Rt_from_P(None, P)
c2w = torch.from_numpy(c2w).float()
intrinsic = torch.from_numpy(intrinsic).float()
intrinsic[:2] /= self.downsample
poses_scene.append(c2w)
intrinsics_scene.append(intrinsic)
self.all_pose_files[scene_idx] = np.stack(poses_scene)
self.all_intrinsics_files[scene_idx] = np.stack(intrinsics_scene)
self.aabbs[0][0].append(0)
self.aabbs[0][1].append(self.scene_num)
self.scene_bbox = self.aabbs[0]
print(self.scene_bbox)
if self.split=='test' or self.scene_num==1:
self.load_data(self.scene_idxs[0],range(49))
def load_data(self, scene_idx, img_idx=None):
self.all_rays = []
n_views = len(self.all_pose_files[scene_idx])
cam_xyzs = self.all_pose_files[scene_idx][:,:3, -1]
idxs = fps_downsample(cam_xyzs, min(self.train_views, n_views)) if img_idx is None else img_idx
# if "test" == self.split:
# idxs = [item for item in list(range(n_views)) if item not in idxs]
# if len(idxs)==0:
# idxs = list(range(n_views))
images_np = np.stack([cv.resize(cv.imread(self.all_rgb_files[scene_idx][idx]), self.img_wh) for idx in idxs]) / 255.0
self.all_rgbs = torch.from_numpy(images_np.astype(np.float32)[..., [2, 1, 0]]) # [n_images, H, W, 3]
for c2w,intrinsic in zip(self.all_pose_files[scene_idx][idxs],self.all_intrinsics_files[scene_idx][idxs]):
rays_o, rays_d = self.gen_rays_at(torch.from_numpy(intrinsic).float(), torch.from_numpy(c2w).float())
self.all_rays += [torch.cat([rays_o, rays_d], 1)] # (h*w, 6)
if not self.is_stack:
self.all_rays = torch.cat(self.all_rays, 0) # (len(self.meta['frames])*h*w, 3)
self.all_rgbs = self.all_rgbs.reshape(-1, 3)
else:
self.all_rays = torch.stack(self.all_rays, 0) # (len(self.meta['frames]),h*w, 3)
self.all_rgbs = self.all_rgbs.reshape(-1, *self.img_wh[::-1], 3) # (len(self.meta['frames]),h,w,3)
# self.sampler = SimpleSampler(np.prod(self.all_rgbs.shape[:-1]), self.batch_size)
# def read_meta(self):
#
# images_lis = sorted(glob(os.path.join(self.root_dir, 'image/*.png')))
# images_np = np.stack([cv.resize(cv.imread(im_name), self.img_wh) for im_name in images_lis]) / 255.0
# # masks_lis = sorted(glob(os.path.join(self.root_dir, 'mask/*.png')))
# # masks_np = np.stack([cv.resize(cv.imread(im_name),self.img_wh) for im_name in masks_lis])>128
#
# self.all_rgbs = torch.from_numpy(images_np.astype(np.float32)[..., [2, 1, 0]]) # [n_images, H, W, 3]
# # self.all_masks = torch.from_numpy(masks_np>0) # [n_images, H, W, 3]
# self.img_wh = [self.all_rgbs.shape[2], self.all_rgbs.shape[1]]
#
# # world_mat is a projection matrix from world to image
# n_images = len(images_lis)
# world_mats_np = [self.camera_dict['world_mat_%d' % idx].astype(np.float32) for idx in range(n_images)]
# self.scale_mats_np = [self.camera_dict['scale_mat_%d' % idx].astype(np.float32) for idx in range(n_images)]
#
# # W,H = self.img_wh
# self.all_rays = []
# self.intrinsics, self.poses = [], []
# for img_idx, (scale_mat, world_mat) in enumerate(zip(self.scale_mats_np, world_mats_np)):
# P = world_mat @ scale_mat
# P = P[:3, :4]
# intrinsic, c2w = load_K_Rt_from_P(None, P)
#
# c2w = torch.from_numpy(c2w).float()
# intrinsic = torch.from_numpy(intrinsic).float()
# intrinsic[:2] /= self.downsample
#
# self.poses.append(c2w)
# self.intrinsics.append(intrinsic)
#
# rays_o, rays_d = self.gen_rays_at(intrinsic, c2w)
# self.all_rays += [torch.cat([rays_o, rays_d], 1)] # (h*w, 6)
#
# self.intrinsics, self.poses = torch.stack(self.intrinsics), torch.stack(self.poses)
#
# # self.all_rgbs[~self.all_masks] = 1.0
# if not self.is_stack:
# self.all_rays = torch.cat(self.all_rays, 0) # (len(self.meta['frames])*h*w, 3)
# self.all_rgbs = self.all_rgbs.reshape(-1, 3)
# else:
# self.all_rays = torch.stack(self.all_rays, 0) # (len(self.meta['frames]),h*w, 3)
# self.all_rgbs = self.all_rgbs.reshape(-1, *self.img_wh[::-1], 3) # (len(self.meta['frames]),h,w,3)
#
# self.sampler = SimpleSampler(np.prod(self.all_rgbs.shape[:-1]), self.batch_size)
def get_bbox(self, scale_mats_np, object_scale_mat):
object_bbox_min = np.array([-1.0, -1.0, -1.0, 1.0])
object_bbox_max = np.array([ 1.0, 1.0, 1.0, 1.0])
# Object scale mat: region of interest to **extract mesh**
# object_scale_mat = np.load(os.path.join(scene_path, 'cameras.npz'))
object_bbox_min = np.linalg.inv(scale_mats_np[0]) @ object_scale_mat @ object_bbox_min[:, None]
object_bbox_max = np.linalg.inv(scale_mats_np[0]) @ object_scale_mat @ object_bbox_max[:, None]
return [object_bbox_min[:3, 0].tolist(),object_bbox_max[:3, 0].tolist()]
# self.near_far = [2.125, 4.525]
def gen_rays_at(self, intrinsic, c2w, resolution_level=1):
"""
Generate rays at world space from one camera.
"""
l = resolution_level
W,H = self.img_wh
tx = torch.linspace(0, W - 1, W // l)+0.5
ty = torch.linspace(0, H - 1, H // l)+0.5
pixels_x, pixels_y = torch.meshgrid(tx, ty)
p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1) # W, H, 3
intrinsic_inv = torch.inverse(intrinsic)
p = torch.matmul(intrinsic_inv[None, None, :3, :3], p[:, :, :, None]).squeeze() # W, H, 3
rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True) # W, H, 3
rays_v = torch.matmul(c2w[None, None, :3, :3], rays_v[:, :, :, None]).squeeze() # W, H, 3
rays_o = c2w[None, None, :3, 3].expand(rays_v.shape) # W, H, 3
return rays_o.transpose(0, 1).reshape(-1,3), rays_v.transpose(0, 1).reshape(-1,3)
def __len__(self):
return 1000000 #len(self.all_rays)
def __getitem__(self, idx):
idx = torch.randint(self.scene_num,(1,)).item()
if self.scene_num >= 1:
scene_name = self.scene_idxs[idx]
img_idx = np.random.choice(len(self.all_rgb_files[scene_name]), size=6)
self.load_data(scene_name, img_idx)
idxs = np.random.choice(self.all_rays.shape[0], size=self.batch_size)
return {'rays': self.all_rays[idxs], 'rgbs': self.all_rgbs[idxs], 'idx': idx}
================================================
FILE: dataLoader/dtu_objs2.py
================================================
import torch
import cv2 as cv
import numpy as np
import os
from glob import glob
from .ray_utils import *
from torch.utils.data import Dataset
# This function is borrowed from IDR: https://github.com/lioryariv/idr
def load_K_Rt_from_P(filename, P=None):
if P is None:
lines = open(filename).read().splitlines()
if len(lines) == 4:
lines = lines[1:]
lines = [[x[0], x[1], x[2], x[3]] for x in (x.split(" ") for x in lines)]
P = np.asarray(lines).astype(np.float32).squeeze()
out = cv.decomposeProjectionMatrix(P)
K = out[0]
R = out[1]
t = out[2]
K = K / K[2, 2]
intrinsics = np.eye(4)
intrinsics[:3, :3] = K
pose = np.eye(4, dtype=np.float32)
pose[:3, :3] = R.transpose()
pose[:3, 3] = (t[:3] / t[3])[:, 0]
return intrinsics, pose
class DTUDataset(Dataset):
def __init__(self, cfg, split='train', batch_size=4096, is_stack=None):
"""
img_wh should be set to a tuple ex: (1152, 864) to enable test mode!
"""
# self.N_vis = N_vis
self.split = split
self.batch_size = batch_size
self.root_dir = cfg.datadir
self.is_stack = is_stack if is_stack is not None else 'train'!=split
self.downsample = cfg.get(f'downsample_{self.split}')
self.img_wh = (int(400 / self.downsample), int(300 / self.downsample))
self.white_bg = False
self.camera_dict = np.load(os.path.join(self.root_dir, 'cameras.npz'))
self.read_meta()
self.get_bbox()
# def define_transforms(self):
# self.transform = T.ToTensor()
def get_bbox(self):
object_bbox_min = np.array([-1.0, -1.0, -1.0, 1.0])
object_bbox_max = np.array([ 1.0, 1.0, 1.0, 1.0])
# Object scale mat: region of interest to **extract mesh**
object_scale_mat = np.load(os.path.join(self.root_dir, 'cameras.npz'))['scale_mat_0']
object_bbox_min = np.linalg.inv(self.scale_mats_np[0]) @ object_scale_mat @ object_bbox_min[:, None]
object_bbox_max = np.linalg.inv(self.scale_mats_np[0]) @ object_scale_mat @ object_bbox_max[:, None]
self.scene_bbox = [object_bbox_min[:3, 0].tolist(),object_bbox_max[:3, 0].tolist()]
self.scene_bbox[0].append(0)
self.scene_bbox[1].append(1)
def gen_rays_at(self, intrinsic, c2w, resolution_level=1):
"""
Generate rays at world space from one camera.
"""
l = resolution_level
W,H = self.img_wh
tx = torch.linspace(0, W - 1, W // l)+0.5
ty = torch.linspace(0, H - 1, H // l)+0.5
pixels_x, pixels_y = torch.meshgrid(tx, ty)
p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1) # W, H, 3
intrinsic_inv = torch.inverse(intrinsic)
p = torch.matmul(intrinsic_inv[None, None, :3, :3], p[:, :, :, None]).squeeze() # W, H, 3
rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True) # W, H, 3
rays_v = torch.matmul(c2w[None, None, :3, :3], rays_v[:, :, :, None]).squeeze() # W, H, 3
rays_o = c2w[None, None, :3, 3].expand(rays_v.shape) # W, H, 3
return rays_o.transpose(0, 1).reshape(-1,3), rays_v.transpose(0, 1).reshape(-1,3)
def read_meta(self):
images_lis = sorted(glob(os.path.join(self.root_dir, 'image/*.png')))
images_np = np.stack([cv.resize(cv.imread(im_name),self.img_wh) for im_name in images_lis]) / 255.0
# masks_lis = sorted(glob(os.path.join(self.root_dir, 'mask/*.png')))
# masks_np = np.stack([cv.resize(cv.imread(im_name),self.img_wh) for im_name in masks_lis])>128
self.all_rgbs = torch.from_numpy(images_np.astype(np.float32)[...,[2,1,0]]) # [n_images, H, W, 3]
# self.all_masks = torch.from_numpy(masks_np>0) # [n_images, H, W, 3]
self.img_wh = [self.all_rgbs.shape[2],self.all_rgbs.shape[1]]
# world_mat is a projection matrix from world to image
n_images = len(images_lis)
world_mats_np = [self.camera_dict['world_mat_%d' % idx].astype(np.float32) for idx in range(n_images)]
self.scale_mats_np = [self.camera_dict['scale_mat_%d' % idx].astype(np.float32) for idx in range(n_images)]
# W,H = self.img_wh
self.all_rays = []
self.intrinsics, self.poses = [],[]
for img_idx, (scale_mat, world_mat) in enumerate(zip(self.scale_mats_np, world_mats_np)):
P = world_mat @ scale_mat
P = P[:3, :4]
intrinsic, c2w = load_K_Rt_from_P(None, P)
c2w = torch.from_numpy(c2w).float()
intrinsic = torch.from_numpy(intrinsic).float()
intrinsic[:2] /= self.downsample
self.poses.append(c2w)
self.intrinsics.append(intrinsic)
rays_o, rays_d = self.gen_rays_at(intrinsic,c2w)
self.all_rays += [torch.cat([rays_o, rays_d], 1)] # (h*w, 6)
self.intrinsics, self.poses = torch.stack(self.intrinsics), torch.stack(self.poses)
# self.all_rgbs[~self.all_masks] = 1.0
if not self.is_stack:
self.all_rays = torch.cat(self.all_rays, 0) # (len(self.meta['frames])*h*w, 3)
self.all_rgbs = self.all_rgbs.reshape(-1,3)
else:
self.all_rays = torch.stack(self.all_rays, 0) # (len(self.meta['frames]),h*w, 3)
self.all_rgbs = self.all_rgbs.reshape(-1, *self.img_wh[::-1],3) # (len(self.meta['frames]),h,w,3)
self.sampler = SimpleSampler(np.prod(self.all_rgbs.shape[:-1]), self.batch_size)
def __len__(self):
return len(self.all_rays)
def __getitem__(self, idx):
idx_rand = self.sampler.nextids() #torch.randint(0,len(self.all_rays),(self.batch_size,))
sample = {'rays': self.all_rays[idx_rand], 'rgbs': self.all_rgbs[idx_rand]}
return sample
================================================
FILE: dataLoader/google_objs.py
================================================
import torch, cv2
from torch.utils.data import Dataset
import json
from tqdm import tqdm
import os
from PIL import Image
from torchvision import transforms as T
import glob
from scipy.spatial.transform import Rotation as R
from .ray_utils import *
def fps_downsample(points, n_points_to_sample):
selected_points = np.zeros((n_points_to_sample, 3))
selected_idxs = []
dist = np.ones(points.shape[0]) * 100
for i in range(n_points_to_sample):
idx = np.argmax(dist).tolist()
selected_points[i] = points[idx]
print(idx)
selected_idxs.extend([idx])
dist_ = ((points - selected_points[i]) ** 2).sum(-1)
dist = np.minimum(dist, dist_).tolist()
return selected_idxs
############################# Get Spherical Path #############################
def pose_spherical_nerf(euler, radius=1.8, ep=1):
c2ws_render = np.eye(4)
c2ws_render[:3,:3] = R.from_euler('xyz', euler, degrees=True).as_matrix()
# 保留旋转矩阵的最后一列再乘个系数就能当作位置？
c2ws_render[:3,3] = c2ws_render[:3,:3] @ np.array([0.0,0.0,ep*radius])
return c2ws_render
def nerf_video_path(c2ws, theta_range=10,phi_range=20,N_views=120,radius=1.3,ep=-1):
c2ws = torch.tensor(c2ws)
mean_position = torch.mean(c2ws[:,:3, 3],dim=0).reshape(1,3).cpu().numpy()
rotvec = []
for i in range(c2ws.shape[0]):
r = R.from_matrix(c2ws[i, :3, :3])
euler_ange = r.as_euler('xyz', degrees=True).reshape(1, 3)
if i:
mask = np.abs(euler_ange - rotvec[0])>180
euler_ange[mask] += 360.0
rotvec.append(euler_ange)
# 采用欧拉角做平均的方法求旋转矩阵的平均
rotvec = np.mean(np.stack(rotvec), axis=0)
render_poses = [pose_spherical_nerf(rotvec+np.array([angle,0.0,-phi_range]), radius=radius, ep=ep) for angle in np.linspace(-theta_range,theta_range,N_views//4, endpoint=False)]
render_poses += [pose_spherical_nerf(rotvec+np.array([theta_range,0.0,angle]), radius=radius, ep=ep) for angle in np.linspace(-phi_range,phi_range,N_views//4, endpoint=False)]
render_poses += [pose_spherical_nerf(rotvec+np.array([angle,0.0,phi_range]), radius=radius, ep=ep) for angle in np.linspace(theta_range,-theta_range,N_views//4, endpoint=False)]
render_poses += [pose_spherical_nerf(rotvec+np.array([-theta_range,0.0,angle]), radius=radius, ep=ep) for angle in np.linspace(phi_range,-phi_range,N_views//4, endpoint=False)]
return render_poses
def _interpolate_trajectory(c2ws, num_views: int = 300):
"""calculate interpolate path"""
from scipy.interpolate import interp1d
from scipy.spatial.transform import Rotation, Slerp
key_rots = Rotation.from_matrix(c2ws[:, :3, :3])
key_times = list(range(len(c2ws)))
slerp = Slerp(key_times, key_rots)
interp = interp1d(key_times, c2ws[:, :3, 3], axis=0)
render_c2ws = []
for i in range(num_views):
time = float(i) / num_views * (len(c2ws) - 1)
cam_location = interp(time)
cam_rot = slerp(time).as_matrix()
c2w = np.eye(4)
c2w[:3, :3] = cam_rot
c2w[:3, 3] = cam_location
render_c2ws.append(c2w)
return np.stack(render_c2ws, axis=0)
def google_objs_path(c2ws, N_views=150):
positions = c2ws[:, :3, 3]
selected_idxs = fps_downsample(positions, 3)
selected_idxs.append(selected_idxs[0])
return _interpolate_trajectory(c2ws[selected_idxs].numpy(), N_views)
class GoogleObjsDataset(Dataset):
def __init__(self, cfg, split="train", batch_size=4096):
# self.N_vis = N_vis
self.cfg = cfg
self.root_dir = cfg.datadir
self.split = split
self.batch_size = batch_size
self.is_stack = False if "train" == split else True
self.downsample = cfg.get(f"downsample_{self.split}")
self.img_wh = (int(512 / self.downsample), int(512 / self.downsample))
self.define_transforms()
train_scene_idxs = sorted(cfg.train_scene_list)
test_scene_idxs = cfg.test_scene_list
if len(train_scene_idxs)==2:
train_scene_idxs = list(range(train_scene_idxs[0],train_scene_idxs[1]))
self.scene_idxs = train_scene_idxs if self.split=='train' else test_scene_idxs
self.train_views = cfg.train_views
self.scene_num = len(self.scene_idxs)
if 'test' == self.split:
self.test_index = train_scene_idxs.index(test_scene_idxs[0])
# self.rot = torch.tensor([[0.65561799, -0.65561799, 0.37460659],
# [0.73729737, 0.44876192, -0.50498052],
# [0.16296514, 0.60727077, 0.77760181]])
self.scene_bbox = [[-1.0, -1.0, -1.0, 0.0], [1.0, 1.0, 1.0, self.scene_num]]
# self.blender2opencv = np.array([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
self.white_bg = True
self.near_far = [0.4, 1.6]
#################################
# self.folder_path = datadir
# self.num_source_views = args.num_source_views
# self.rectify_inplane_rotation = args.rectify_inplane_rotation
self.scene_path_list = sorted(glob.glob(os.path.join(self.root_dir, "*/")))
all_rgb_files = []
# all_depth_files = []
all_pose_files = []
all_intrinsics_files = []
num_files = 250
for i, scene_idx in enumerate(self.scene_idxs):
scene_path = self.scene_path_list[scene_idx]
# print(scene_idx,scene_path)
rgb_files = [
os.path.join(scene_path, "rgb", f)
for f in sorted(os.listdir(os.path.join(scene_path, "rgb")))
]
# depth_files = [os.path.join(scene_path, 'depth', f)
# for f in sorted(os.listdir(os.path.join(scene_path, 'depth')))]
pose_files = [
f.replace("rgb", "pose").replace("png", "txt") for f in rgb_files
]
intrinsics_files = [
f.replace("rgb", "intrinsics").replace("png", "txt") for f in rgb_files
]
if (
np.min([len(rgb_files), len(pose_files), len(intrinsics_files)])
< num_files
):
print(scene_path)
continue
all_rgb_files.append(rgb_files)
# all_depth_files.append(depth_files)
all_pose_files.append(pose_files)
all_intrinsics_files.append(intrinsics_files)
index = np.arange(len(all_rgb_files))
self.all_rgb_files = np.array(all_rgb_files)[index]
# self.all_depth_files = np.array(all_depth_files)[index]
self.all_pose_files = np.array(all_pose_files)[index]
self.all_intrinsics_files = np.array(all_intrinsics_files)[index]
if self.split=='test' or self.scene_num==1:
self.read_meta()
else:
self.train_idxs = range(self.train_views)
# self.define_proj_mat()
def read_depth(self, filename):
depth = np.array(read_pfm(filename)[0], dtype=np.float32) # (800, 800)
return depth
def read_meta(self):
assert (
len(self.all_rgb_files)
== len(self.all_pose_files)
== len(self.all_intrinsics_files)
)
self.all_image_paths = []
self.all_poses = []
self.all_rays = []
self.all_rgbs = []
for idx in range(len(self.all_rgb_files)):
rgb_files = self.all_rgb_files[idx]
pose_files = self.all_pose_files[idx]
intrinsics_files = self.all_intrinsics_files[idx]
intrinsics = np.loadtxt(intrinsics_files[0])
index_3x3 = np.array([0, 1, 2, 4, 5, 6, 8, 9, 10])
self.intrinsics = intrinsics[index_3x3]
w, h = self.img_wh
self.focal = self.intrinsics[0]
self.focal *= (
self.img_wh[0] / 512
) # modify focal length to match size self.img_wh
# ray directions for all pixels, same for all images (same H, W, focal)
self.directions = get_ray_directions(
h, w, [self.focal, self.focal]
) # (h, w, 3)
self.directions = self.directions / torch.norm(
self.directions, dim=-1, keepdim=True
)
self.intrinsics = torch.tensor(
[[self.focal, 0, w / 2], [0, self.focal, h / 2], [0, 0, 1]]
).float()
self.scene_image_paths = []
self.scene_poses = []
self.scene_rays = []
self.scene_rgbs = []
# self.downsample = 1.0
img_eval_interval = (
1 # if self.N_vis < 0 else len(self.meta['frames']) // self.N_vis
)
if "train" == self.split:
cam_xyzs = []
# for i in range(len(pose_files)):
for i in range(100):
pose = np.loadtxt(pose_files[i])
cam_xyzs.append([pose[3], pose[7], pose[11]])
cam_xyzs = np.array(cam_xyzs)
idxs = fps_downsample(cam_xyzs, min(self.train_views, len(rgb_files)))
self.train_idxs = idxs
print("train idxs:", idxs)
else:
idxs = list(range(100, 200))
for i in tqdm(
idxs, desc=f"Loading data {self.split} ({len(idxs)})"
): # img_list:#
pose = np.loadtxt(pose_files[i])
pose = torch.FloatTensor(pose).view(4, 4)
self.scene_poses += [pose]
image_path = rgb_files[i]
self.scene_image_paths += [image_path]
img = Image.open(image_path)
if self.downsample != 1.0:
img = img.resize(self.img_wh, Image.LANCZOS)
img = self.transform(img) # (3, h, w)
img = img.view(3, -1).permute(1, 0) # (h*w, 3) RGBA
# img = img[:, :3] * img[:, -1:] + (1 - img[:, -1:]) # blend A to RGB
self.scene_rgbs += [img]
rays_o, rays_d = get_rays(self.directions, pose) # both (h*w, 3)
# rays_o, rays_d = rays_o@self.rot, rays_d@self.rot
self.scene_rays += [torch.cat([rays_o, rays_d], 1)] # (h*w, 6)
self.scene_poses = torch.stack(self.scene_poses)
views = 180
radius = {183: 1.3, 199: 1.7, 298: 1.5, 467: 1.1, 957: 1.9, 244: 1.2, 963: 1.2, 527: 1.2, 681:1.9,948:1.2}
self.render_path = google_objs_path(self.scene_poses, N_views=views)
name = self.scene_path_list[self.scene_idxs[0]].split('/')[-2]
np.save(f'{self.root_dir}/{name}_render_path.npy', self.render_path)
# self.render_path = nerf_video_path(self.scene_poses, N_views=views, theta_range=45,phi_range=90, radius=radius[self.scene_idxs[0]],ep=-1)
if not self.is_stack:
self.scene_rays = torch.cat(
self.scene_rays, 0
) # (len(self.meta['frames])*h*w, 3)
self.scene_rgbs = torch.cat(
self.scene_rgbs, 0
) # (len(self.meta['frames])*h*w, 3)
# self.all_depth = torch.cat(self.all_depth, 0) # (len(self.meta['frames])*h*w, 3)
else:
self.scene_rays = torch.stack(
self.scene_rays, 0
) # (len(self.meta['frames]),h*w, 3)
self.scene_rgbs = torch.stack(self.scene_rgbs, 0).reshape(
-1, *self.img_wh[::-1], 3
) # (len(self.meta['frames]),h,w,3)
# self.all_masks = torch.stack(self.all_masks, 0).reshape(-1,*self.img_wh[::-1]) # (len(self.meta['frames]),h,w,3)
######################## pre-generate and save in mem ########################
self.all_image_paths.append(self.scene_image_paths)
self.all_poses.append(self.scene_poses)
self.all_rays.append(self.scene_rays)
self.all_rgbs.append(self.scene_rgbs)
self.all_rays = torch.cat(self.all_rays)
self.all_rgbs = torch.cat(self.all_rgbs)
def get_rays(self, idx):
######################## pre-generate and save in mem ########################
return self.all_rays[idx]
def get_rgbs(self, idx):
######################## pre-generate and save in mem ########################
return self.all_rgbs[idx]
def define_transforms(self):
self.transform = T.ToTensor()
def define_proj_mat(self):
self.proj_mat = (
self.intrinsics.unsqueeze(0) @ torch.inverse(self.scene_poses)[:, :3]
)
def world2ndc(self, points, lindisp=None):
device = points.device
return (points - self.center.to(device)) / self.radius.to(device)
def update_index(self):
self.scene_idx = torch.randint(0, len(self.all_rgb_files), (1,)).item()
def __len__(self):
return 10000000 #len(self.all_rgb_files)
def __getitem__(self, idx):
#
# self.update_index()
# idx = self.scene_idx #
idx = idx % len(self.all_rgb_files)
# idx = torch.randint(len(self.all_rgb_files), (1,)).item()
######################## generate rays on the fly ########################
rgb_files = self.all_rgb_files[idx]
pose_files = self.all_pose_files[idx]
intrinsics_files = self.all_intrinsics_files[idx]
intrinsics = np.loadtxt(intrinsics_files[0])
index_3x3 = np.array([0, 1, 2, 4, 5, 6, 8, 9, 10])
intrinsics = intrinsics[index_3x3]
w, h = self.img_wh
focal = intrinsics[0]
focal *= self.img_wh[0] / 512 # modify focal length to match size self.img_wh
# ray directions for all pixels, same for all images (same H, W, focal)
directions = get_ray_directions(h, w, [focal, focal]) # (h, w, 3)
directions = directions / torch.norm(directions, dim=-1, keepdim=True)
intrinsics = torch.tensor(
[[focal, 0, w / 2], [0, focal, h / 2], [0, 0, 1]]
).float()
scene_poses = []
scene_rays = []
scene_image_paths = []
scene_rgbs = []
# downsample = 1.0
sample_views = 5
if self.scene_num>1:
ids = np.random.choice(self.train_idxs, size=sample_views)
for i in ids:
image_path = rgb_files[i]
scene_image_paths += [image_path]
img = Image.open(image_path)
idxs = torch.randint(0, w*h, (self.batch_size // sample_views,))
if self.downsample != 1.0:
img = img.resize(self.img_wh, Image.LANCZOS)
img = self.transform(img) # (3, h, w)
img = img.view(3, -1).permute(1, 0) # (h*w, 3) RGBA
# img = img[:, :3] * img[:, -1:] + (1 - img[:, -1:]) # blend A to RGB
scene_rgbs += [img[idxs]]
pose = np.loadtxt(pose_files[i])
pose = torch.FloatTensor(pose).view(4, 4)
scene_poses += [pose]
rays_o, rays_d = get_rays(directions, pose) # both (h*w, 3)
scene_rays += [torch.cat([rays_o[idxs], rays_d[idxs]], 1)] # (h*w, 6)
scene_poses = torch.stack(scene_poses)
if not self.is_stack:
scene_rays = torch.cat(scene_rays, 0) # (len(self.meta['frames])*h*w, 3)
scene_rgbs = torch.cat(scene_rgbs, 0) # (len(self.meta['frames])*h*w, 3)
else:
scene_rays = torch.stack(scene_rays, 0) # (len(self.meta['frames]),h*w, 3)
scene_rgbs = torch.stack(scene_rgbs, 0) # (len(self.meta['frames]),h*w, 3)
return {'rays': scene_rays, 'rgbs': scene_rgbs, 'idx': idx}
else:
idx_rand = torch.randint(0, len(self.all_rays), (self.batch_size,))
return {'rays': self.all_rays[idx_rand], 'rgbs': self.all_rgbs[idx_rand]}
================================================
FILE: dataLoader/image.py
================================================
import torch,imageio,cv2
from PIL import Image
Image.MAX_IMAGE_PIXELS = 1000000000
import numpy as np
import torch.nn.functional as F
from torch.utils.data import Dataset
_img_suffix = ['png','jpg','jpeg','bmp','tif']
def load(path):
suffix = path.split('.')[-1]
if suffix in _img_suffix:
img = np.array(Image.open(path))#.convert('L')
scale = 256.**(1+np.log2(np.max(img))//8)-1
return img/scale
elif 'exr' == suffix:
return imageio.imread(path)
elif 'npy' == suffix:
return np.load(path)
def srgb_to_linear(img):
limit = 0.04045
return np.where(img > limit, np.power((img + 0.055) / 1.055, 2.4), img / 12.92)
class ImageDataset(Dataset):
def __init__(self, cfg, batchsize, split='train', continue_sampling=False, tolinear=False, HW=-1, perscent=1.0, delete_region=None,mask=None):
datadir = cfg.datadir
self.batchsize = batchsize
self.continue_sampling = continue_sampling
img = load(datadir).astype(np.float32)
if HW>0:
img = cv2.resize(img,[HW,HW])
if tolinear:
img = srgb_to_linear(img)
self.img = torch.from_numpy(img)
H,W = self.img.shape[:2]
y, x = torch.meshgrid(torch.arange(0, H), torch.arange(0, W), indexing='ij')
self.coordiante = torch.stack((x,y),-1).float()+0.5
n_channel = self.img.shape[-1]
self.image = self.img
self.img, self.coordiante = self.img.reshape(H*W,-1), self.coordiante.reshape(H*W,2)
# if continue_sampling:
# coordiante_tmp = self.coordiante.view(1,1,-1,2)/torch.tensor([W,H])*2-1.0
# self.img = F.grid_sample(self.img.view(1,H,W,-1).permute(0,3,1,2),coordiante_tmp, mode='bilinear', align_corners=True).reshape(self.img.shape[-1],-1).t()
if 'train'==split:
self.mask = torch.ones_like(y)>0
if mask is not None:
self.mask = mask>0
print(torch.sum(mask)/1.0/HW/HW)
elif delete_region is not None:
if isinstance(delete_region[0], list):
for item in delete_region:
t_l_x,t_l_y,width,height = item
self.mask[t_l_y:t_l_y+height,t_l_x:t_l_x+width] = False
else:
t_l_x,t_l_y,width,height = delete_region
self.mask[t_l_y:t_l_y+height,t_l_x:t_l_x+width] = False
else:
index = torch.randperm(len(self.img))[:int(len(self.img)*perscent)]
self.mask[:] = False
self.mask.view(-1)[index] = True
self.mask = self.mask.view(-1)
self.image, self.coordiante = self.img[self.mask], self.coordiante[self.mask]
else:
self.image = self.img
self.HW = [H,W]
self.scene_bbox = [[0., 0.], [W, H]]
cfg.aabb = self.scene_bbox
#
def __len__(self):
return 10000
def __getitem__(self, idx):
H,W = self.HW
device = self.image.device
idx = torch.randint(0,len(self.image),(self.batchsize,), device=device)
if self.continue_sampling:
coordinate = self.coordiante[idx] + torch.rand((self.batchsize,2))-0.5
coordinate_tmp = (coordinate.view(1,1,self.batchsize,2))/torch.tensor([W,H],device=device)*2-1.0
rgb = F.grid_sample(self.img.view(1,H,W,-1).permute(0,3,1,2),coordinate_tmp, mode='bilinear',
align_corners=False, padding_mode='border').reshape(self.img.shape[-1],-1).t()
sample = {'rgb': rgb,
'xy': coordinate}
else:
sample = {'rgb': self.image[idx],
'xy': self.coordiante[idx]}
return sample
================================================
FILE: dataLoader/image_set.py
================================================
import torch,cv2
import numpy as np
import torch.nn.functional as F
from torch.utils.data import Dataset
def srgb_to_linear(img):
limit = 0.04045
return torch.where(img > limit, torch.pow((img + 0.055) / 1.055, 2.4), img / 12.92)
def load(path, HW=512):
suffix = path.split('.')[-1]
if 'npy' == suffix:
img = np.load(path)
# img = 0.3*img[...,:1] + 0.59*img[...,1:2] + 0.11*img[...,2:]
if img.shape[-2]!=HW:
for i in range(img.shape[0]):
img[i] = cv2.resize(img[i],[HW,HW])
return img
class ImageSetDataset(Dataset):
def __init__(self, cfg, batchsize, split='train', continue_sampling=False, HW=512, N=10, tolinear=True):
datadir = cfg.datadir
self.batchsize = batchsize
self.continue_sampling = continue_sampling
imgs = load(datadir,HW=HW)[:N]
self.imgs = torch.from_numpy(imgs).float()/255
if tolinear:
self.imgs = srgb_to_linear(self.imgs)
D,H,W = self.imgs.shape[:3]
y, x = torch.meshgrid(torch.arange(0, H), torch.arange(0, W), indexing='ij')
self.coordinate = torch.stack((x,y),-1).float()+0.5
self.imgs, self.coordinate = self.imgs.reshape(D,H*W,-1), self.coordinate.reshape(H*W,2)
self.DHW = [D,H,W]
self.scene_bbox = [[0., 0., 0.], [W, H, D]]
cfg.aabb = self.scene_bbox
# self.down_scale = 512.0/H
def __len__(self):
return 1000000
def __getitem__(self, idx):
D,H,W = self.DHW
pix_idx = torch.randint(0,H*W,(self.batchsize,))
img_idx = torch.randint(0,D,(self.batchsize,))
if self.continue_sampling:
coordinate = self.coordinate[pix_idx] + torch.rand((self.batchsize,2)) - 0.5
coordinate = torch.cat((coordinate,img_idx.unsqueeze(-1)+0.5),dim=-1)
coordinate_tmp = (coordinate.view(1,1,1,self.batchsize,3))/torch.tensor([W,H,D])*2-1.0
rgb = F.grid_sample(self.imgs.view(1,D,H,W,-1).permute(0,4,1,2,3),coordinate_tmp, mode='bilinear',
align_corners=False, padding_mode='border').reshape(self.imgs.shape[-1],-1).t()
# coordinate[:,:2] *= self.down_scale
sample = {'rgb': rgb,
'xy': coordinate}
else:
sample = {'rgb': self.imgs[img_idx,pix_idx],
'xy': torch.cat((self.coordinate[pix_idx],img_idx.unsqueeze(-1)+0.5),dim=-1)}
# 'xy': torch.cat((self.coordiante[pix_idx],img_idx.expand_as(pix_idx).unsqueeze(-1)),dim=-1)}
return sample
================================================
FILE: dataLoader/llff.py
================================================
import torch
from torch.utils.data import Dataset
import glob
import numpy as np
import os
from PIL import Image
from torchvision import transforms as T
from .ray_utils import *
def normalize(v):
"""Normalize a vector."""
return v / np.linalg.norm(v)
def average_poses(poses):
"""
Calculate the average pose, which is then used to center all poses
using @center_poses. Its computation is as follows:
1. Compute the center: the average of pose centers.
2. Compute the z axis: the normalized average z axis.
3. Compute axis y': the average y axis.
4. Compute x' = y' cross product z, then normalize it as the x axis.
5. Compute the y axis: z cross product x.
Note that at step 3, we cannot directly use y' as y axis since it's
not necessarily orthogonal to z axis. We need to pass from x to y.
Inputs:
poses: (N_images, 3, 4)
Outputs:
pose_avg: (3, 4) the average pose
"""
# 1. Compute the center
center = poses[..., 3].mean(0) # (3)
# 2. Compute the z axis
z = normalize(poses[..., 2].mean(0)) # (3)
# 3. Compute axis y' (no need to normalize as it's not the final output)
y_ = poses[..., 1].mean(0) # (3)
# 4. Compute the x axis
x = normalize(np.cross(z, y_)) # (3)
# 5. Compute the y axis (as z and x are normalized, y is already of norm 1)
y = np.cross(x, z) # (3)
pose_avg = np.stack([x, y, z, center], 1) # (3, 4)
return pose_avg
def center_poses(poses, blender2opencv):
"""
Center the poses so that we can use NDC.
See https://github.com/bmild/nerf/issues/34
Inputs:
poses: (N_images, 3, 4)
Outputs:
poses_centered: (N_images, 3, 4) the centered poses
pose_avg: (3, 4) the average pose
"""
poses = poses @ blender2opencv
pose_avg = average_poses(poses) # (3, 4)
pose_avg_homo = np.eye(4)
pose_avg_homo[:3] = pose_avg # convert to homogeneous coordinate for faster computation
pose_avg_homo = pose_avg_homo
# by simply adding 0, 0, 0, 1 as the last row
last_row = np.tile(np.array([0, 0, 0, 1]), (len(poses), 1, 1)) # (N_images, 1, 4)
poses_homo = \
np.concatenate([poses, last_row], 1) # (N_images, 4, 4) homogeneous coordinate
poses_centered = np.linalg.inv(pose_avg_homo) @ poses_homo # (N_images, 4, 4)
# poses_centered = poses_centered @ blender2opencv
poses_centered = poses_centered[:, :3] # (N_images, 3, 4)
return poses_centered, pose_avg_homo
def viewmatrix(z, up, pos):
vec2 = normalize(z)
vec1_avg = up
vec0 = normalize(np.cross(vec1_avg, vec2))
vec1 = normalize(np.cross(vec2, vec0))
m = np.eye(4)
m[:3] = np.stack([-vec0, vec1, vec2, pos], 1)
return m
def render_path_spiral(c2w, up, rads, focal, zdelta, zrate, N_rots=2, N=120):
render_poses = []
rads = np.array(list(rads) + [1.])
for theta in np.linspace(0., 2. * np.pi * N_rots, N + 1)[:-1]:
c = np.dot(c2w[:3, :4], np.array([np.cos(theta), -np.sin(theta), -np.sin(theta * zrate), 1.]) * rads)
z = normalize(c - np.dot(c2w[:3, :4], np.array([0, 0, -focal, 1.])))
render_poses.append(viewmatrix(z, up, c))
return render_poses
def get_spiral(c2ws_all, near_fars, rads_scale=1.0, N_views=120):
# center pose
c2w = average_poses(c2ws_all)
# Get average pose
up = normalize(c2ws_all[:, :3, 1].sum(0))
# Find a reasonable "focus depth" for this dataset
dt = 0.75
close_depth, inf_depth = near_fars.min() * 0.9, near_fars.max() * 5.0
focal = 1.0 / (((1.0 - dt) / close_depth + dt / inf_depth))
# Get radii for spiral path
zdelta = near_fars.min() * .2
tt = c2ws_all[:, :3, 3]
rads = np.percentile(np.abs(tt), 90, 0) * rads_scale
render_poses = render_path_spiral(c2w, up, rads, focal, zdelta, zrate=.5, N=N_views)
return np.stack(render_poses)
class LLFFDataset(Dataset):
def __init__(self, cfg , split='train', hold_every=8):
"""
spheric_poses: whether the images are taken in a spheric inward-facing manner
default: False (forward-facing)
val_num: number of val images (used for multigpu training, validate same image for all gpus)
"""
self.root_dir = cfg.datadir
self.split = split
self.hold_every = hold_every
self.is_stack = False if 'train' == split else True
self.downsample = cfg.get(f'downsample_{self.split}')
self.define_transforms()
self.blender2opencv = np.eye(4)#np.array([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
self.read_meta()
self.white_bg = False
# self.near_far = [np.min(self.near_fars[:,0]),np.max(self.near_fars[:,1])]
self.near_far = [0.0, 1.0]
self.scene_bbox = [[-1.5, -1.67, -1.0], [1.5, 1.67, 1.0]]
# self.scene_bbox = torch.tensor([[-1.67, -1.5, -1.0], [1.67, 1.5, 1.0]])
# self.center = torch.mean(self.scene_bbox, dim=0).float().view(1, 1, 3)
# self.invradius = 1.0 / (self.scene_bbox[1] - self.center).float().view(1, 1, 3)
def read_meta(self):
print(self.root_dir)
poses_bounds = np.load(os.path.join(self.root_dir, 'poses_bounds.npy')) # (N_images, 17)
self.image_paths = sorted(glob.glob(os.path.join(self.root_dir, 'images_4/*')))
# load full resolution image then resize
if self.split in ['train', 'test']:
assert len(poses_bounds) == len(self.image_paths), \
'Mismatch between number of images and number of poses! Please rerun COLMAP!'
poses = poses_bounds[:, :15].reshape(-1, 3, 5) # (N_images, 3, 5)
self.near_fars = poses_bounds[:, -2:] # (N_images, 2)
hwf = poses[:, :, -1]
# Step 1: rescale focal length according to training resolution
H, W, self.focal = poses[0, :, -1] # original intrinsics, same for all images
self.img_wh = np.array([int(W / self.downsample), int(H / self.downsample)])
self.focal = [self.focal * self.img_wh[0] / W, self.focal * self.img_wh[1] / H]
# Step 2: correct poses
# Original poses has rotation in form "down right back", change to "right up back"
# See https://github.com/bmild/nerf/issues/34
poses = np.concatenate([poses[..., 1:2], -poses[..., :1], poses[..., 2:4]], -1)
# (N_images, 3, 4) exclude H, W, focal
self.poses, self.pose_avg = center_poses(poses, self.blender2opencv)
# Step 3: correct scale so that the nearest depth is at a little more than 1.0
# See https://github.com/bmild/nerf/issues/34
near_original = self.near_fars.min()
scale_factor = near_original * 0.75 # 0.75 is the default parameter
# the nearest depth is at 1/0.75=1.33
self.near_fars /= scale_factor
self.poses[..., 3] /= scale_factor
# build rendering path
N_views, N_rots = 120, 2
tt = self.poses[:, :3, 3] # ptstocam(poses[:3,3,:].T, c2w).T
up = normalize(self.poses[:, :3, 1].sum(0))
rads = np.percentile(np.abs(tt), 90, 0)
self.render_path = get_spiral(self.poses, self.near_fars, N_views=N_views)
# distances_from_center = np.linalg.norm(self.poses[..., 3], axis=1)
# val_idx = np.argmin(distances_from_center) # choose val image as the closest to
# center image
# ray directions for all pixels, same for all images (same H, W, focal)
W, H = self.img_wh
self.directions = get_ray_directions_blender(H, W, self.focal) # (H, W, 3)
average_pose = average_poses(self.poses)
dists = np.sum(np.square(average_pose[:3, 3] - self.poses[:, :3, 3]), -1)
i_test = np.arange(0, self.poses.shape[0], self.hold_every) # [np.argmin(dists)]
img_list = i_test if self.split != 'train' else list(set(np.arange(len(self.poses))) - set(i_test))
# use first N_images-1 to train, the LAST is val
self.all_rays = []
self.all_rgbs = []
for i in img_list:
image_path = self.image_paths[i]
c2w = torch.FloatTensor(self.poses[i])
img = Image.open(image_path).convert('RGB')
if self.downsample != 1.0:
img = img.resize(self.img_wh, Image.LANCZOS)
img = self.transform(img) # (3, h, w)
img = img.view(3, -1).permute(1, 0) # (h*w, 3) RGB
self.all_rgbs += [img]
rays_o, rays_d = get_rays(self.directions, c2w) # both (h*w, 3)
rays_o, rays_d = ndc_rays_blender(H, W, self.focal[0], 1.0, rays_o, rays_d)
# viewdir = rays_d / torch.norm(rays_d, dim=-1, keepdim=True)
self.all_rays += [torch.cat([rays_o, rays_d], 1)] # (h*w, 6)
if not self.is_stack:
self.all_rays = torch.cat(self.all_rays, 0) # (len(self.meta['frames])*h*w, 3)
self.all_rgbs = torch.cat(self.all_rgbs, 0) # (len(self.meta['frames])*h*w,3)
else:
self.all_rays = torch.stack(self.all_rays, 0) # (len(self.meta['frames]),h,w, 3)
self.all_rgbs = torch.stack(self.all_rgbs, 0).reshape(-1,*self.img_wh[::-1], 3) # (len(self.meta['frames]),h,w,3)
def define_transforms(self):
self.transform = T.ToTensor()
def __len__(self):
return len(self.all_rgbs)
def __getitem__(self, idx):
sample = {'rays': self.all_rays[idx],
'rgbs': self.all_rgbs[idx]}
return sample
================================================
FILE: dataLoader/nsvf.py
================================================
import torch
from torch.utils.data import Dataset
from tqdm import tqdm
import os
from PIL import Image
from torchvision import transforms as T
from .ray_utils import *
trans_t = lambda t : torch.Tensor([
[1,0,0,0],
[0,1,0,0],
[0,0,1,t],
[0,0,0,1]]).float()
rot_phi = lambda phi : torch.Tensor([
[1,0,0,0],
[0,np.cos(phi),-np.sin(phi),0],
[0,np.sin(phi), np.cos(phi),0],
[0,0,0,1]]).float()
rot_theta = lambda th : torch.Tensor([
[np.cos(th),0,-np.sin(th),0],
[0,1,0,0],
[np.sin(th),0, np.cos(th),0],
[0,0,0,1]]).float()
def pose_spherical(theta, phi, radius):
c2w = trans_t(radius)
c2w = rot_phi(phi/180.*np.pi) @ c2w
c2w = rot_theta(theta/180.*np.pi) @ c2w
c2w = torch.Tensor(np.array([[-1,0,0,0],[0,0,1,0],[0,1,0,0],[0,0,0,1]])) @ c2w
return c2w
class NSVF(Dataset):
"""NSVF Generic Dataset."""
def __init__(self, datadir, split='train', downsample=1.0, wh=[800,800], is_stack=False):
self.root_dir = datadir
self.split = split
self.is_stack = is_stack
self.downsample = downsample
self.img_wh = (int(wh[0]/downsample),int(wh[1]/downsample))
self.define_transforms()
self.white_bg = True
self.near_far = [0.5,6.0]
self.scene_bbox = torch.from_numpy(np.loadtxt(f'{self.root_dir}/bbox.txt')).float()[:6].view(2,3)
self.blender2opencv = np.array([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
self.read_meta()
self.define_proj_mat()
self.center = torch.mean(self.scene_bbox, axis=0).float().view(1, 1, 3)
self.radius = (self.scene_bbox[1] - self.center).float().view(1, 1, 3)
def bbox2corners(self):
corners = self.scene_bbox.unsqueeze(0).repeat(4,1,1)
for i in range(3):
corners[i,[0,1],i] = corners[i,[1,0],i]
return corners.view(-1,3)
def read_meta(self):
with open(os.path.join(self.root_dir, "intrinsics.txt")) as f:
focal = float(f.readline().split()[0])
self.intrinsics = np.array([[focal,0,400.0],[0,focal,400.0],[0,0,1]])
self.intrinsics[:2] *= (np.array(self.img_wh)/np.array([800,800])).reshape(2,1)
pose_files = sorted(os.listdir(os.path.join(self.root_dir, 'pose')))
img_files = sorted(os.listdir(os.path.join(self.root_dir, 'rgb')))
if self.split == 'train':
pose_files = [x for x in pose_files if x.startswith('0_')]
img_files = [x for x in img_files if x.startswith('0_')]
elif self.split == 'val':
pose_files = [x for x in pose_files if x.startswith('1_')]
img_files = [x for x in img_files if x.startswith('1_')]
elif self.split == 'test':
test_pose_files = [x for x in pose_files if x.startswith('2_')]
test_img_files = [x for x in img_files if x.startswith('2_')]
if len(test_pose_files) == 0:
test_pose_files = [x for x in pose_files if x.startswith('1_')]
test_img_files = [x for x in img_files if x.startswith('1_')]
pose_files = test_pose_files
img_files = test_img_files
# ray directions for all pixels, same for all images (same H, W, focal)
self.directions = get_ray_directions(self.img_wh[1], self.img_wh[0], [self.intrinsics[0,0],self.intrinsics[1,1]], center=self.intrinsics[:2,2]) # (h, w, 3)
self.directions = self.directions / torch.norm(self.directions, dim=-1, keepdim=True)
frames = 200
self.render_path = torch.stack([pose_spherical(angle, -30.0, 4.0) for angle in np.linspace(-180,180,frames+1)[:-1]], 0)
self.poses = []
self.all_rays = []
self.all_rgbs = []
assert len(img_files) == len(pose_files)
for img_fname, pose_fname in tqdm(zip(img_files, pose_files), desc=f'Loading data {self.split} ({len(img_files)})'):
image_path = os.path.join(self.root_dir, 'rgb', img_fname)
img = Image.open(image_path)
if self.downsample!=1.0:
img = img.resize(self.img_wh, Image.LANCZOS)
img = self.transform(img) # (4, h, w)
img = img.view(img.shape[0], -1).permute(1, 0) # (h*w, 4) RGBA
if img.shape[-1]==4:
img = img[:, :3] * img[:, -1:] + (1 - img[:, -1:]) # blend A to RGB
self.all_rgbs += [img]
c2w = np.loadtxt(os.path.join(self.root_dir, 'pose', pose_fname)) #@ self.blender2opencv
c2w = torch.FloatTensor(c2w)
self.poses.append(c2w) # C2W
rays_o, rays_d = get_rays(self.directions, c2w) # both (h*w, 3)
self.all_rays += [torch.cat([rays_o, rays_d], 1)] # (h*w, 8)
# w2c = torch.inverse(c2w)
#
self.poses = torch.stack(self.poses)
if 'train' == self.split:
if self.is_stack:
self.all_rays = torch.stack(self.all_rays, 0).reshape(-1,*self.img_wh[::-1], 6) # (len(self.meta['frames])*h*w, 3)
self.all_rgbs = torch.stack(self.all_rgbs, 0).reshape(-1,*self.img_wh[::-1], 3) # (len(self.meta['frames])*h*w, 3)
else:
self.all_rays = torch.cat(self.all_rays, 0) # (len(self.meta['frames])*h*w, 3)
self.all_rgbs = torch.cat(self.all_rgbs, 0) # (len(self.meta['frames])*h*w, 3)
else:
self.all_rays = torch.stack(self.all_rays, 0) # (len(self.meta['frames]),h*w, 3)
self.all_rgbs = torch.stack(self.all_rgbs, 0).reshape(-1,*self.img_wh[::-1], 3) # (len(self.meta['frames]),h,w,3)
def define_transforms(self):
self.transform = T.ToTensor()
def define_proj_mat(self):
self.proj_mat = torch.from_numpy(self.intrinsics[:3,:3]).unsqueeze(0).float() @ torch.inverse(self.poses)[:,:3]
def world2ndc(self, points):
device = points.device
return (points - self.center.to(device)) / self.radius.to(device)
def __len__(self):
if self.split == 'train':
return len(self.all_rays)
return len(self.all_rgbs)
def __getitem__(self, idx):
if self.split == 'train': # use data in the buffers
sample = {'rays': self.all_rays[idx],
'rgbs': self.all_rgbs[idx]}
else: # create data for each image separately
img = self.all_rgbs[idx]
rays = self.all_rays[idx]
sample = {'rays': rays,
'rgbs': img}
return sample
================================================
FILE: dataLoader/ray_utils.py
================================================
import torch, re, json
import numpy as np
from torch import searchsorted
from kornia import create_meshgrid
def load_json(path):
with open(path, 'r') as f:
return json.load(f)
# from utils import index_point_feature
def depth2dist(z_vals, cos_angle):
# z_vals: [N_ray N_sample]
device = z_vals.device
dists = z_vals[..., 1:] - z_vals[..., :-1]
dists = torch.cat([dists, torch.Tensor([1e10]).to(device).expand(dists[..., :1].shape)], -1) # [N_rays, N_samples]
dists = dists * cos_angle.unsqueeze(-1)
return dists
def ndc2dist(ndc_pts, cos_angle):
dists = torch.norm(ndc_pts[:, 1:] - ndc_pts[:, :-1], dim=-1)
dists = torch.cat([dists, 1e10 * cos_angle.unsqueeze(-1)], -1) # [N_rays, N_samples]
return dists
def get_ray_directions(H, W, focal, center=None):
"""
Get ray directions for all pixels in camera coordinate.
Reference: https://www.scratchapixel.com/lessons/3d-basic-rendering/
ray-tracing-generating-camera-rays/standard-coordinate-systems
Inputs:
H, W, focal: image height, width and focal length
Outputs:
directions: (H, W, 3), the direction of the rays in camera coordinate
"""
grid = create_meshgrid(H, W, normalized_coordinates=False)[0] + 0.5
i, j = grid.unbind(-1)
# the direction here is without +0.5 pixel centering as calibration is not so accurate
# see https://github.com/bmild/nerf/issues/24
cent = center if center is not None else [W / 2, H / 2]
directions = torch.stack([(i - cent[0]) / focal[0], (j - cent[1]) / focal[1], torch.ones_like(i)], -1) # (H, W, 3)
return directions
def get_ray_directions_blender(H, W, focal, center=None):
"""
Get ray directions for all pixels in camera coordinate.
Reference: https://www.scratchapixel.com/lessons/3d-basic-rendering/
ray-tracing-generating-camera-rays/standard-coordinate-systems
Inputs:
H, W, focal: image height, width and focal length
Outputs:
directions: (H, W, 3), the direction of the rays in camera coordinate
"""
grid = create_meshgrid(H, W, normalized_coordinates=False)[0]+0.5
i, j = grid.unbind(-1)
# the direction here is without +0.5 pixel centering as calibration is not so accurate
# see https://github.com/bmild/nerf/issues/24
cent = center if center is not None else [W / 2, H / 2]
directions = torch.stack([(i - cent[0]) / focal[0], -(j - cent[1]) / focal[1], -torch.ones_like(i)],
-1) # (H, W, 3)
return directions
def get_rays(directions, c2w):
"""
Get ray origin and normalized directions in world coordinate for all pixels in one image.
Reference: https://www.scratchapixel.com/lessons/3d-basic-rendering/
ray-tracing-generating-camera-rays/standard-coordinate-systems
Inputs:
directions: (H, W, 3) precomputed ray directions in camera coordinate
c2w: (3, 4) transformation matrix from camera coordinate to world coordinate
Outputs:
rays_o: (H*W, 3), the origin of the rays in world coordinate
rays_d: (H*W, 3), the normalized direction of the rays in world coordinate
"""
# Rotate ray directions from camera coordinate to the world coordinate
rays_d = directions @ c2w[:3, :3].T # (H, W, 3)
# rays_d = rays_d / torch.norm(rays_d, dim=-1, keepdim=True)
# The origin of all rays is the camera origin in world coordinate
rays_o = c2w[:3, 3].expand(rays_d.shape) # (H, W, 3)
rays_d = rays_d.view(-1, 3)
rays_o = rays_o.view(-1, 3)
return rays_o, rays_d
def ndc_rays_blender(H, W, focal, near, rays_o, rays_d):
# Shift ray origins to near plane
t = -(near + rays_o[..., 2]) / rays_d[..., 2]
rays_o = rays_o + t[..., None] * rays_d
# Projection
o0 = -1. / (W / (2. * focal)) * rays_o[..., 0] / rays_o[..., 2]
o1 = -1. / (H / (2. * focal)) * rays_o[..., 1] / rays_o[..., 2]
o2 = 1. + 2. * near / rays_o[..., 2]
d0 = -1. / (W / (2. * focal)) * (rays_d[..., 0] / rays_d[..., 2] - rays_o[..., 0] / rays_o[..., 2])
d1 = -1. / (H / (2. * focal)) * (rays_d[..., 1] / rays_d[..., 2] - rays_o[..., 1] / rays_o[..., 2])
d2 = -2. * near / rays_o[..., 2]
rays_o = torch.stack([o0, o1, o2], -1)
rays_d = torch.stack([d0, d1, d2], -1)
return rays_o, rays_d
def ndc_rays(H, W, focal, near, rays_o, rays_d):
# Shift ray origins to near plane
t = (near - rays_o[..., 2]) / rays_d[..., 2]
rays_o = rays_o + t[..., None] * rays_d
# Projection
o0 = 1. / (W / (2. * focal)) * rays_o[..., 0] / rays_o[..., 2]
o1 = 1. / (H / (2. * focal)) * rays_o[..., 1] / rays_o[..., 2]
o2 = 1. - 2. * near / rays_o[..., 2]
d0 = 1. / (W / (2. * focal)) * (rays_d[..., 0] / rays_d[..., 2] - rays_o[..., 0] / rays_o[..., 2])
d1 = 1. / (H / (2. * focal)) * (rays_d[..., 1] / rays_d[..., 2] - rays_o[..., 1] / rays_o[..., 2])
d2 = 2. * near / rays_o[..., 2]
rays_o = torch.stack([o0, o1, o2], -1)
rays_d = torch.stack([d0, d1, d2], -1)
return rays_o, rays_d
# Hierarchical sampling (section 5.2)
def sample_pdf(bins, weights, N_samples, det=False, pytest=False):
device = weights.device
# Get pdf
weights = weights + 1e-5 # prevent nans
pdf = weights / torch.sum(weights, -1, keepdim=True)
cdf = torch.cumsum(pdf, -1)
cdf = torch.cat([torch.zeros_like(cdf[..., :1]), cdf], -1) # (batch, len(bins))
# Take uniform samples
if det:
u = torch.linspace(0., 1., steps=N_samples, device=device)
u = u.expand(list(cdf.shape[:-1]) + [N_samples])
else:
u = torch.rand(list(cdf.shape[:-1]) + [N_samples], device=device)
# Pytest, overwrite u with numpy's fixed random numbers
if pytest:
np.random.seed(0)
new_shape = list(cdf.shape[:-1]) + [N_samples]
if det:
u = np.linspace(0., 1., N_samples)
u = np.broadcast_to(u, new_shape)
else:
u = np.random.rand(*new_shape)
u = torch.Tensor(u)
# Invert CDF
u = u.contiguous()
inds = searchsorted(cdf.detach(), u, right=True)
below = torch.max(torch.zeros_like(inds - 1), inds - 1)
above = torch.min((cdf.shape[-1] - 1) * torch.ones_like(inds), inds)
inds_g = torch.stack([below, above], -1) # (batch, N_samples, 2)
matched_shape = [inds_g.shape[0], inds_g.shape[1], cdf.shape[-1]]
cdf_g = torch.gather(cdf.unsqueeze(1).expand(matched_shape), 2, inds_g)
bins_g = torch.gather(bins.unsqueeze(1).expand(matched_shape), 2, inds_g)
denom = (cdf_g[..., 1] - cdf_g[..., 0])
denom = torch.where(denom < 1e-5, torch.ones_like(denom), denom)
t = (u - cdf_g[..., 0]) / denom
samples = bins_g[..., 0] + t * (bins_g[..., 1] - bins_g[..., 0])
return samples
def dda(rays_o, rays_d, bbox_3D):
inv_ray_d = 1.0 / (rays_d + 1e-6)
t_min = (bbox_3D[:1] - rays_o) * inv_ray_d # N_rays 3
t_max = (bbox_3D[1:] - rays_o) * inv_ray_d
t = torch.stack((t_min, t_max)) # 2 N_rays 3
t_min = torch.max(torch.min(t, dim=0)[0], dim=-1, keepdim=True)[0]
t_max = torch.min(torch.max(t, dim=0)[0], dim=-1, keepdim=True)[0]
return t_min, t_max
def ray_marcher(rays,
N_samples=64,
lindisp=False,
perturb=0,
bbox_3D=None):
"""
sample points along the rays
Inputs:
rays: ()
Returns:
"""
# Decompose the inputs
N_rays = rays.shape[0]
rays_o, rays_d = rays[:, 0:3], rays[:, 3:6] # both (N_rays, 3)
near, far = rays[:, 6:7], rays[:, 7:8] # both (N_rays, 1)
if bbox_3D is not None:
# cal aabb boundles
near, far = dda(rays_o, rays_d, bbox_3D)
# Sample depth points
z_steps = torch.linspace(0, 1, N_samples, device=rays.device) # (N_samples)
if not lindisp: # use linear sampling in depth space
z_vals = near * (1 - z_steps) + far * z_steps
else: # use linear sampling in disparity space
z_vals = 1 / (1 / near * (1 - z_steps) + 1 / far * z_steps)
z_vals = z_vals.expand(N_rays, N_samples)
if perturb > 0: # perturb sampling depths (z_vals)
z_vals_mid = 0.5 * (z_vals[:, :-1] + z_vals[:, 1:]) # (N_rays, N_samples-1) interval mid points
# get intervals between samples
upper = torch.cat([z_vals_mid, z_vals[:, -1:]], -1)
lower = torch.cat([z_vals[:, :1], z_vals_mid], -1)
perturb_rand = perturb * torch.rand(z_vals.shape, device=rays.device)
z_vals = lower + (upper - lower) * perturb_rand
xyz_coarse_sampled = rays_o.unsqueeze(1) + \
rays_d.unsqueeze(1) * z_vals.unsqueeze(2) # (N_rays, N_samples, 3)
return xyz_coarse_sampled, rays_o, rays_d, z_vals
def read_pfm(filename):
file = open(filename, 'rb')
color = None
width = None
height = None
scale = None
endian = None
header = file.readline().decode('utf-8').rstrip()
if header == 'PF':
color = True
elif header == 'Pf':
color = False
else:
raise Exception('Not a PFM file.')
dim_match = re.match(r'^(\d+)\s(\d+)\s$', file.readline().decode('utf-8'))
if dim_match:
width, height = map(int, dim_match.groups())
else:
raise Exception('Malformed PFM header.')
scale = float(file.readline().rstrip())
if scale < 0: # little-endian
endian = '<'
scale = -scale
else:
endian = '>' # big-endian
data = np.fromfile(file, endian + 'f')
shape = (height, width, 3) if color else (height, width)
data = np.reshape(data, shape)
data = np.flipud(data)
file.close()
return data, scale
class SimpleSampler:
def __init__(self, total, batch):
self.total = total
self.batch = batch
self.curr = total
self.ids = None
def nextids(self):
self.curr += self.batch
if self.curr + self.batch > self.total:
self.ids = torch.LongTensor(np.random.permutation(self.total))
self.curr = 0
return self.ids[self.curr:self.curr + self.batch]
def ndc_bbox(all_rays):
near_min = torch.min(all_rays[...,:3].view(-1,3),dim=0)[0]
near_max = torch.max(all_rays[..., :3].view(-1, 3), dim=0)[0]
far_min = torch.min((all_rays[...,:3]+all_rays[...,3:6]).view(-1,3),dim=0)[0]
far_max = torch.max((all_rays[...,:3]+all_rays[...,3:6]).view(-1, 3), dim=0)[0]
print(f'===> ndc bbox near_min:{near_min} near_max:{near_max} far_min:{far_min} far_max:{far_max}')
return torch.stack((torch.minimum(near_min,far_min),torch.maximum(near_max,far_max)))
def pose_from_json(meta, transpose):
c2ws = []
for frame in meta['frames']:
c2ws.append(np.array(frame['transform_matrix'])@transpose)
return np.stack(c2ws)
def rotation_matrix_from_vectors(vec1, vec2):
""" Find the rotation matrix that aligns vec1 to vec2
:param vec1: A 3d "source" vector
:param vec2: A 3d "destination" vector
:return mat: A transform matrix (3x3) which when applied to vec1, aligns it with vec2.
"""
a, b = (vec1 / np.linalg.norm(vec1)).reshape(3), (vec2 / np.linalg.norm(vec2)).reshape(3)
v = np.cross(a, b)
if any(v): # if not all zeros then
c = np.dot(a, b)
s = np.linalg.norm(v)
kmat = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])
return np.eye(3) + kmat + kmat.dot(kmat) * ((1 - c) / (s ** 2))
else:
return np.eye(3) # cross of all zeros only occurs on identical directions
def normalize(x):
return x / np.linalg.norm(x)
def rotation_up(poses):
up = normalize(np.linalg.lstsq(poses[:, :3, 1],np.ones((poses.shape[0],1)))[0])
rot = rotation_matrix_from_vectors(up,np.array([0.,1.,0.]))
return rot
def search_orientation(points):
from scipy.spatial.transform import Rotation as R
bbox_sizes,rot_mats,bboxs = [],[],[]
for y_angle in np.linspace(-45,45,15):
rotvec = np.array([0,y_angle,0])/180*np.pi
rot = R.from_rotvec(rotvec).as_matrix()
point_orientation = rot@points
bbox = np.max(point_orientation,axis=1) - np.min(point_orientation,axis=1)
bbox_sizes.append(np.prod(bbox))
rot_mats.append(rot)
bboxs.append(bbox)
rot = rot_mats[np.argmin(bbox_sizes)]
bbox = bboxs[np.argmin(bbox_sizes)]
return rot,bbox
def load_point_txt(path):
points = []
with open(path, "r") as f:
for line in f:
if line[0] == "#":
continue
els = line.split(" ")
err = float(els[7])
if err > 0.25:
continue
points.append([float(els[1]), float(els[2]), float(els[3])])
points = np.stack(points)
# mask out outerliner with a sphere cover 95 percet poitns
center = np.mean(points, axis=0)
dist = np.sqrt(np.sum((points - center) ** 2, axis=1))
radius_idx = np.argsort(dist)[:int(points.shape[0] * 0.95)]
points = points[radius_idx]
return points
p34_to_44 = lambda p: np.concatenate([p, np.tile(np.reshape(np.eye(4)[-1, :], [1, 1, 4]), [p.shape[0], 1, 1])], 1)
def orientation(poses, point_path=None):
if point_path is not None:
points = load_point_txt(point_path)
else:
points = poses[:,:3,-1]
# np.savetxt('points_1.txt', points)
# np.savetxt('poses_center_1.txt', poses[:, :3, -1])
rot_up = rotation_up(poses)
poses_new = rot_up[None] @ poses[:, :3]
points = rot_up[:3, :3] @ points.T
center = np.mean(poses_new[:,:3,-1], axis=0, keepdims=True)
poses_new[:, :3, -1] -= center
points -= center.T
# scale = 1.5
rot, bbox = search_orientation(poses_new[:, :3, -1].T)
poses_new[:, :3, :4] = rot[None] @ poses_new[:, :3, :4]
points = rot @ points
aabb = np.stack((np.min(points,axis=1),np.max(points,axis=1)))
poses_new[:, :3, -1] /= np.min(aabb[1]-aabb[0])
points /= np.min(aabb[1] - aabb[0])
aabb = (aabb/np.min(aabb[1]-aabb[0])*1.5).tolist()
np.savetxt('points.txt', points.T)
np.savetxt('poses_center.txt', poses_new[:, :3, -1])
return p34_to_44(poses_new), aabb
def spherify_poses(poses, radus=1):
p34_to_44 = lambda p: np.concatenate([p, np.tile(np.reshape(np.eye(4)[-1, :], [1, 1, 4]), [p.shape[0], 1, 1])], 1)
rays_d = poses[:, :3, 2:3]
rays_o = poses[:, :3, 3:4]
def min_line_dist(rays_o, rays_d):
A_i = np.eye(3) - rays_d * np.transpose(rays_d, [0, 2, 1])
b_i = -A_i @ rays_o
pt_mindist = np.squeeze(-np.linalg.inv((np.transpose(A_i, [0, 2, 1]) @ A_i).mean(0)) @ (b_i).mean(0))
return pt_mindist
pt_mindist = min_line_dist(rays_o, rays_d)
center = pt_mindist
up = (poses[:, :3, 3] - center).mean(0)
vec0 = normalize(up)
vec1 = normalize(np.cross([.1, .2, .3], vec0))
vec2 = normalize(np.cross(vec0, vec1))
pos = center
c2w = np.stack([vec1, vec2, vec0, pos], 1)
poses_reset = np.linalg.inv(p34_to_44(c2w[None])) @ p34_to_44(poses[:, :3, :4])
rad = np.sqrt(np.mean(np.sum(np.square(poses_reset[:, :3, 3]), -1)))*radus
# print(poses_reset,center,rad)
scale = 1. / rad
poses_reset[:, :3, 3] *= scale
# bds *= sc
rad *= scale
# print ('========>',poses_reset)
centroid = np.mean(poses_reset[:, :3, 3], 0)
zh = centroid[2]
radcircle = np.sqrt(rad ** 2 - zh ** 2)
render_poses = []
for th in np.linspace(0., 2. * np.pi, 120):
camorigin = np.array([radcircle * np.cos(th), radcircle * np.sin(th), zh])
up = np.array([0, 0, -1.])
vec2 = normalize(camorigin)
vec0 = normalize(np.cross(vec2, up))
vec1 = normalize(np.cross(vec2, vec0))
pos = camorigin
p = np.stack([vec0, vec1, vec2, pos], 1)
render_poses.append(p)
render_poses = np.stack(render_poses, 0)
# render_poses = np.concatenate([render_poses, np.broadcast_to(poses[0, :3, -1:], render_poses[:, :3, -1:].shape)], -1)
# poses_reset = np.concatenate(
# [poses_reset[:, :3, :4], np.broadcast_to(poses[0, :3, -1:], poses_reset[:, :3, -1:].shape)], -1)
return poses_reset, render_poses, scale
================================================
FILE: dataLoader/sdf.py
================================================
import torch
import numpy as np
from torch.utils.data import Dataset
def N_to_reso(avg_reso, bbox):
xyz_min, xyz_max = bbox
dim = len(xyz_min)
n_voxels = avg_reso**dim
voxel_size = ((xyz_max - xyz_min).prod() / n_voxels).pow(1 / dim)
return torch.ceil((xyz_max - xyz_min) / voxel_size).long().tolist()
def load(path, split, dtype='points'):
if 'grid' == dtype:
sdf = torch.from_numpy(np.load(path).astype(np.float32))
D, H, W = sdf.shape
z, y, x = torch.meshgrid(torch.arange(0, D), torch.arange(0, H), torch.arange(0, W), indexing='ij')
coordiante = torch.stack((x,y,z),-1).reshape(D*H*W,3)#*2-1 # normalize to [-1,1]
sdf = sdf.reshape(D*H*W,-1)
DHW = [D,H,W]
elif 'points' == dtype:
DHW = [640] * 3
sdf_dict = np.load(path, allow_pickle=True).item()
sdf = torch.from_numpy(sdf_dict[f'sdfs_{split}'].astype(np.float32)).reshape(-1,1)
coordiante = torch.from_numpy(sdf_dict[f'points_{split}'].astype(np.float32))
aabb = [[-1,-1,-1],[1,1,1]]
coordiante = ((coordiante + 1) / 2 * (torch.tensor(DHW[::-1]))).reshape(-1,3)
DHW = DHW[::-1]
return coordiante, sdf, DHW
class SDFDataset(Dataset):
def __init__(self, cfg, split='train'):
datadir = cfg.datadir
self.coordiante, self.sdf, self.DHW = load(datadir, split)
[D, H, W] = self.DHW
self.scene_bbox = [[0., 0., 0.], [W, H, D]]
cfg.aabb = self.scene_bbox
def __len__(self):
return len(self.sdf)
def __getitem__(self, idx):
sample = {'rgb': self.sdf[idx],
'xy': self.coordiante[idx]}
return sample
================================================
FILE: dataLoader/tankstemple.py
================================================
import torch
from torch.utils.data import Dataset
from tqdm import tqdm
import os
from PIL import Image
from torchvision import transforms as T
from .ray_utils import *
def circle(radius=3.5, h=0.0, axis='z', t0=0, r=1):
if axis == 'z':
return lambda t: [radius * np.cos(r * t + t0), radius * np.sin(r * t + t0), h]
elif axis == 'y':
return lambda t: [radius * np.cos(r * t + t0), h, radius * np.sin(r * t + t0)]
else:
return lambda t: [h, radius * np.cos(r * t + t0), radius * np.sin(r * t + t0)]
def cross(x, y, axis=0):
T = torch if isinstance(x, torch.Tensor) else np
return T.cross(x, y, axis)
def normalize(x, axis=-1, order=2):
if isinstance(x, torch.Tensor):
l2 = x.norm(p=order, dim=axis, keepdim=True)
return x / (l2 + 1e-8), l2
else:
l2 = np.linalg.norm(x, order, axis)
l2 = np.expand_dims(l2, axis)
l2[l2 == 0] = 1
return x / l2,
def cat(x, axis=1):
if isinstance(x[0], torch.Tensor):
return torch.cat(x, dim=axis)
return np.concatenate(x, axis=axis)
def look_at_rotation(camera_position, at=None, up=None, inverse=False, cv=False):
"""
This function takes a vector 'camera_position' which specifies the location
of the camera in world coordinates and two vectors `at` and `up` which
indicate the position of the object and the up directions of the world
coordinate system respectively. The object is assumed to be centered at
the origin.
The output is a rotation matrix representing the transformation
from world coordinates -> view coordinates.
Input:
camera_position: 3
at: 1 x 3 or N x 3 (0, 0, 0) in default
up: 1 x 3 or N x 3 (0, 1, 0) in default
"""
if at is None:
at = torch.zeros_like(camera_position)
else:
at = torch.tensor(at).type_as(camera_position)
if up is None:
up = torch.zeros_like(camera_position)
up[2] = -1
else:
up = torch.tensor(up).type_as(camera_position)
z_axis = normalize(at - camera_position)[0]
x_axis = normalize(cross(up, z_axis))[0]
y_axis = normalize(cross(z_axis, x_axis))[0]
R = cat([x_axis[:, None], y_axis[:, None], z_axis[:, None]], axis=1)
return R
def gen_path(pos_gen, at=(0, 0, 0), up=(0, -1, 0), frames=180):
c2ws = []
for t in range(frames):
c2w = torch.eye(4)
cam_pos = torch.tensor(pos_gen(t * (360.0 / frames) / 180 * np.pi))
cam_rot = look_at_rotation(cam_pos, at=at, up=up, inverse=False, cv=True)
c2w[:3, 3], c2w[:3, :3] = cam_pos, cam_rot
c2ws.append(c2w)
return torch.stack(c2ws)
class TanksTempleDataset(Dataset):
"""NSVF Generic Dataset."""
def __init__(self, cfg, split='train'):
self.root_dir = cfg.datadir
self.split = split
self.is_stack = False if 'train'==split else True
self.downsample = cfg.get(f'downsample_{self.split}')
self.img_wh = (int(1920 / self.downsample), int(1080 / self.downsample))
self.define_transforms()
self.white_bg = True
self.near_far = [0.01, 6.0]
self.scene_bbox = (torch.from_numpy(np.loadtxt(f'{self.root_dir}/bbox.txt')).float()[:6].view(2, 3) * 1.2).tolist()
self.blender2opencv = np.array([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
self.read_meta()
self.define_proj_mat()
# self.center = torch.mean(self.scene_bbox, axis=0).float().view(1, 1, 3)
# self.radius = (self.scene_bbox[1] - self.center).float().view(1, 1, 3)
def bbox2corners(self):
corners = self.scene_bbox.unsqueeze(0).repeat(4, 1, 1)
for i in range(3):
corners[i, [0, 1], i] = corners[i, [1, 0], i]
return corners.view(-1, 3)
def read_meta(self):
self.intrinsics = np.loadtxt(os.path.join(self.root_dir, "intrinsics.txt"))
self.intrinsics[:2] *= (np.array(self.img_wh) / np.array([1920, 1080])).reshape(2, 1)
pose_files = sorted(os.listdir(os.path.join(self.root_dir, 'pose')))
img_files = sorted(os.listdir(os.path.join(self.root_dir, 'rgb')))
if self.split == 'train':
pose_files = [x for x in pose_files if x.startswith('0_')]
img_files = [x for x in img_files if x.startswith('0_')]
elif self.split == 'val':
pose_files = [x for x in pose_files if x.startswith('1_')]
img_files = [x for x in img_files if x.startswith('1_')]
elif self.split == 'test':
test_pose_files = [x for x in pose_files if x.startswith('2_')]
test_img_files = [x for x in img_files if x.startswith('2_')]
if len(test_pose_files) == 0:
test_pose_files = [x for x in pose_files if x.startswith('1_')]
test_img_files = [x for x in img_files if x.startswith('1_')]
pose_files = test_pose_files
img_files = test_img_files
# ray directions for all pixels, same for all images (same H, W, focal)
self.directions = get_ray_directions(self.img_wh[1], self.img_wh[0],
[self.intrinsics[0, 0], self.intrinsics[1, 1]],
center=self.intrinsics[:2, 2]) # (h, w, 3)
self.directions = self.directions / torch.norm(self.directions, dim=-1, keepdim=True)
self.poses = []
ray_per_frame = self.img_wh[0]*self.img_wh[1]
self.all_rays = torch.empty(len(pose_files),ray_per_frame,6)
self.all_rgbs = torch.empty(len(pose_files),ray_per_frame,3)
assert len(img_files) == len(pose_files)
for i in tqdm(range(len(pose_files)),desc=f'Loading data {self.split} ({len(img_files)})'):
img_fname, pose_fname = img_files[i], pose_files[i]
image_path = os.path.join(self.root_dir, 'rgb', img_fname)
img = Image.open(image_path)
if self.downsample != 1.0:
img = img.resize(self.img_wh, Image.LANCZOS)
img = self.transform(img) # (4, h, w)
img = img.view(img.shape[0], -1).permute(1, 0) # (h*w, 4) RGBA
if img.shape[-1] == 4:
img = img[:, :3] * img[:, -1:] + (1 - img[:, -1:]) # blend A to RGB
# self.all_rgbs.append(img)
c2w = np.loadtxt(os.path.join(self.root_dir, 'pose', pose_fname)) # @ cam_trans
c2w = torch.FloatTensor(c2w)
self.poses.append(c2w) # C2W
rays_o, rays_d = get_rays(self.directions, c2w) # both (h*w, 3)
self.all_rays[i] = torch.cat([rays_o, rays_d], 1) # (h*w, 6)
self.all_rgbs[i] = img
self.poses = torch.stack(self.poses)
frames = 200
scene_bbox = torch.tensor(self.scene_bbox).float()
center = torch.mean(scene_bbox, dim=0)
radius = torch.norm(scene_bbox[1] - center) * 1.2
up = torch.mean(self.poses[:, :3, 1], dim=0).tolist()
pos_gen = circle(radius=radius, h=-0.2 * up[1], axis='y')
self.render_path = gen_path(pos_gen, up=up, frames=frames)
self.render_path[:, :3, 3] += center
if 'train' == self.split:
if not self.is_stack:
# self.all_rays = torch.stack(self.all_rays, 0).reshape(-1, *self.img_wh[::-1], 6) # (len(self.meta['frames])*h*w, 3)
# self.all_rgbs = torch.stack(self.all_rgbs, 0).reshape(-1, *self.img_wh[::-1], 3) # (len(self.meta['frames])*h*w, 3)
# else:
self.all_rays = self.all_rays.reshape(-1,6) # (len(self.meta['frames])*h*w, 3)
self.all_rgbs = self.all_rgbs.reshape(-1,3) # (len(self.meta['frames])*h*w, 3)
else:
# self.all_rays = torch.stack(self.all_rays, 0) # (len(self.meta['frames]),h*w, 3)
self.all_rgbs = self.all_rgbs.reshape(-1, *self.img_wh[::-1], 3) # (len(self.meta['frames]),h,w,3)
def define_transforms(self):
self.transform = T.ToTensor()
def define_proj_mat(self):
self.proj_mat = torch.from_numpy(self.intrinsics[:3, :3]).unsqueeze(0).float() @ torch.inverse(self.poses)[:, :3]
def world2ndc(self, points):
device = points.device
return (points - self.center.to(device)) / self.radius.to(device)
def __len__(self):
if self.split == 'train':
return len(self.all_rays)
return len(self.all_rgbs)
def __getitem__(self, idx):
if self.split == 'train': # use data in the buffers
sample = {'rays': self.all_rays[idx],
'rgbs': self.all_rgbs[idx]}
else: # create data for each image separately
img = self.all_rgbs[idx]
rays = self.all_rays[idx]
sample = {'rays': rays,
'rgbs': img}
return sample
================================================
FILE: dataLoader/your_own_data.py
================================================
import torch,cv2
from torch.utils.data import Dataset
import json
from tqdm import tqdm
import os
from PIL import Image
from torchvision import transforms as T
from .ray_utils import *
class YourOwnDataset(Dataset):
def __init__(self, datadir, split='train', downsample=1.0, is_stack=False, N_vis=-1):
self.N_vis = N_vis
self.root_dir = datadir
self.split = split
self.is_stack = is_stack
self.downsample = downsample
self.define_transforms()
self.scene_bbox = torch.tensor([[-1.5, -1.5, -1.5], [1.5, 1.5, 1.5]])
self.blender2opencv = np.array([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
self.read_meta()
self.define_proj_mat()
self.white_bg = True
self.near_far = [0.1,100.0]
self.center = torch.mean(self.scene_bbox, axis=0).float().view(1, 1, 3)
self.radius = (self.scene_bbox[1] - self.center).float().view(1, 1, 3)
self.downsample=downsample
def read_depth(self, filename):
depth = np.array(read_pfm(filename)[0], dtype=np.float32) # (800, 800)
return depth
def read_meta(self):
with open(os.path.join(self.root_dir, f"transforms_{self.split}.json"), 'r') as f:
self.meta = json.load(f)
w, h = int(self.meta['w']/self.downsample), int(self.meta['h']/self.downsample)
self.img_wh = [w,h]
self.focal_x = 0.5 * w / np.tan(0.5 * self.meta['camera_angle_x']) # original focal length
self.focal_y = 0.5 * h / np.tan(0.5 * self.meta['camera_angle_y']) # original focal length
self.cx, self.cy = self.meta['cx'],self.meta['cy']
# ray directions for all pixels, same for all images (same H, W, focal)
self.directions = get_ray_directions(h, w, [self.focal_x,self.focal_y], center=[self.cx, self.cy]) # (h, w, 3)
self.directions = self.directions / torch.norm(self.directions, dim=-1, keepdim=True)
self.intrinsics = torch.tensor([[self.focal_x,0,self.cx],[0,self.focal_y,self.cy],[0,0,1]]).float()
self.image_paths = []
self.poses = []
self.all_rays = []
self.all_rgbs = []
self.all_masks = []
self.all_depth = []
img_eval_interval = 1 if self.N_vis < 0 else len(self.meta['frames']) // self.N_vis
idxs = list(range(0, len(self.meta['frames']), img_eval_interval))
for i in tqdm(idxs, desc=f'Loading data {self.split} ({len(idxs)})'):#img_list:#
frame = self.meta['frames'][i]
pose = np.array(frame['transform_matrix']) @ self.blender2opencv
c2w = torch.FloatTensor(pose)
self.poses += [c2w]
image_path = os.path.join(self.root_dir, f"{frame['file_path']}.png")
self.image_paths += [image_path]
img = Image.open(image_path)
if self.downsample!=1.0:
img = img.resize(self.img_wh, Image.LANCZOS)
img = self.transform(img) # (4, h, w)
img = img.view(-1, w*h).permute(1, 0) # (h*w, 4) RGBA
if img.shape[-1]==4:
img = img[:, :3] * img[:, -1:] + (1 - img[:, -1:]) # blend A to RGB
self.all_rgbs += [img]
rays_o, rays_d = get_rays(self.directions, c2w) # both (h*w, 3)
self.all_rays += [torch.cat([rays_o, rays_d], 1)] # (h*w, 6)
self.poses = torch.stack(self.poses)
if not self.is_stack:
self.all_rays = torch.cat(self.all_rays, 0) # (len(self.meta['frames])*h*w, 3)
self.all_rgbs = torch.cat(self.all_rgbs, 0) # (len(self.meta['frames])*h*w, 3)
# self.all_depth = torch.cat(self.all_depth, 0) # (len(self.meta['frames])*h*w, 3)
else:
self.all_rays = torch.stack(self.all_rays, 0) # (len(self.meta['frames]),h*w, 3)
self.all_rgbs = torch.stack(self.all_rgbs, 0).reshape(-1,*self.img_wh[::-1], 3) # (len(self.meta['frames]),h,w,3)
# self.all_masks = torch.stack(self.all_masks, 0).reshape(-1,*self.img_wh[::-1]) # (len(self.meta['frames]),h,w,3)
def define_transforms(self):
self.transform = T.ToTensor()
def define_proj_mat(self):
self.proj_mat = self.intrinsics.unsqueeze(0) @ torch.inverse(self.poses)[:,:3]
def world2ndc(self,points,lindisp=None):
device = points.device
return (points - self.center.to(device)) / self.radius.to(device)
def __len__(self):
return len(self.all_rgbs)
def __getitem__(self, idx):
if self.split == 'train': # use data in the buffers
sample = {'rays': self.all_rays[idx],
'rgbs': self.all_rgbs[idx]}
else: # create data for each image separately
img = self.all_rgbs[idx]
rays = self.all_rays[idx]
mask = self.all_masks[idx] # for quantity evaluation
sample = {'rays': rays,
'rgbs': img}
return sample
================================================
FILE: models/FactorFields.py
================================================
import torch, math
import torch.nn
import torch.nn.functional as F
import numpy as np
import time, skimage
from utils import N_to_reso, N_to_vm_reso
# import BasisCoding
def grid_mapping(positions, freq_bands, aabb, basis_mapping='sawtooth'):
aabbSize = max(aabb[1] - aabb[0])
scale = aabbSize[..., None] / freq_bands
if basis_mapping == 'triangle':
pts_local = (positions - aabb[0]).unsqueeze(-1) % scale
pts_local_int = ((positions - aabb[0]).unsqueeze(-1) // scale) % 2
pts_local = pts_local / (scale / 2) - 1
pts_local = torch.where(pts_local_int == 1, -pts_local, pts_local)
elif basis_mapping == 'sawtooth':
pts_local = (positions - aabb[0])[..., None] % scale
pts_local = pts_local / (scale / 2) - 1
pts_local = pts_local.clamp(-1., 1.)
elif basis_mapping == 'sinc':
pts_local = torch.sin((positions - aabb[0])[..., None] / (scale / np.pi) - np.pi / 2)
elif basis_mapping == 'trigonometric':
pts_local = (positions - aabb[0])[..., None] / scale * 2 * np.pi
pts_local = torch.cat((torch.sin(pts_local), torch.cos(pts_local)), dim=-1)
elif basis_mapping == 'x':
pts_local = (positions - aabb[0]).unsqueeze(-1) / scale
# elif basis_mapping=='hash':
# pts_local = (positions - aabb[0])/max(aabbSize)
return pts_local
def dct_dict(n_atoms_fre, size, n_selete, dim=2):
"""
Create a dictionary using the Discrete Cosine Transform (DCT) basis. If n_atoms is
not a perfect square, the returned dictionary will have ceil(sqrt(n_atoms))**2 atoms
:param n_atoms:
Number of atoms in dict
:param size:
Size of first patch dim
:return:
DCT dictionary, shape (size*size, ceil(sqrt(n_atoms))**2)
"""
# todo flip arguments to match random_dictionary
p = n_atoms_fre # int(math.ceil(math.sqrt(n_atoms)))
dct = np.zeros((p, size))
for k in range(p):
basis = np.cos(np.arange(size) * k * math.pi / p)
if k > 0:
basis = basis - np.mean(basis)
dct[k] = basis
kron = np.kron(dct, dct)
if 3 == dim:
kron = np.kron(kron, dct)
if n_selete < kron.shape[0]:
idx = [x[0] for x in np.array_split(np.arange(kron.shape[0]), n_selete)]
kron = kron[idx]
for col in range(kron.shape[0]):
norm = np.linalg.norm(kron[col]) or 1
kron[col] /= norm
kron = torch.FloatTensor(kron)
return kron
def positional_encoding(positions, freqs):
freq_bands = (2 ** torch.arange(freqs).float()).to(positions.device) # (F,)
pts = (positions[..., None] * freq_bands).reshape(
positions.shape[:-1] + (freqs * positions.shape[-1],)) # (..., DF)
pts = torch.cat([torch.sin(pts), torch.cos(pts)], dim=-1)
return pts
def raw2alpha(sigma, dist):
# sigma, dist [N_rays, N_samples]
alpha = 1. - torch.exp(-sigma * dist)
T = torch.cumprod(torch.cat([torch.ones_like(alpha[..., :1]), 1. - alpha + 1e-10], -1), -1)
weights = alpha * T[..., :-1] # [N_rays, N_samples]
return alpha, weights, T[..., -1:]
class AlphaGridMask(torch.nn.Module):
def __init__(self, device, aabb, alpha_volume):
super(AlphaGridMask, self).__init__()
self.device = device
self.aabb = aabb.to(self.device)
self.aabbSize = self.aabb[1] - self.aabb[0]
self.invgridSize = 1.0 / self.aabbSize * 2
self.alpha_volume = alpha_volume.view(1, 1, *alpha_volume.shape[-3:])
self.gridSize = torch.LongTensor([alpha_volume.shape[-1], alpha_volume.shape[-2], alpha_volume.shape[-3]]).to(
self.device)
def sample_alpha(self, xyz_sampled):
xyz_sampled = self.normalize_coord(xyz_sampled)
alpha_vals = F.grid_sample(self.alpha_volume, xyz_sampled.view(1, -1, 1, 1, 3), align_corners=True).view(-1)
return alpha_vals
def normalize_coord(self, xyz_sampled):
return (xyz_sampled - self.aabb[0]) * self.invgridSize - 1
class MLPMixer(torch.nn.Module):
def __init__(self,
in_dim,
out_dim=16,
num_layers=2,
hidden_dim=64, pe=0, with_dropout=False):
super().__init__()
self.with_dropout = with_dropout
self.in_dim = in_dim + 2 * in_dim * pe
self.num_layers = num_layers
self.hidden_dim = hidden_dim
self.pe = pe
backbone = []
for l in range(num_layers):
if l == 0:
layer_in_dim = self.in_dim
else:
layer_in_dim = self.hidden_dim
if l == num_layers - 1:
layer_out_dim, bias = out_dim, False
else:
layer_out_dim, bias = self.hidden_dim, True
backbone.append(torch.nn.Linear(layer_in_dim, layer_out_dim, bias=bias))
self.backbone = torch.nn.ModuleList(backbone)
# torch.nn.init.constant_(backbone[0].weight.data, 1.0/self.in_dim)
def forward(self, x, is_train=False):
# x: [B, 3]
h = x
if self.pe > 0:
h = torch.cat([h, positional_encoding(h, self.pe)], dim=-1)
if self.with_dropout and is_train:
h = F.dropout(h, p=0.1)
for l in range(self.num_layers):
h = self.backbone[l](h)
if l != self.num_layers - 1: # l!=0 and
h = F.relu(h, inplace=True)
# h = torch.sin(h)
# sigma, feat = h[...,0], h[...,1:]
return h
class MLPRender_Fea(torch.nn.Module):
def __init__(self, inChanel, num_layers=3, hidden_dim=64, viewpe=6, feape=2):
super(MLPRender_Fea, self).__init__()
self.in_mlpC = 3 + inChanel + 2 * viewpe * 3 + 2 * feape * inChanel
self.num_layers = num_layers
self.viewpe = viewpe
self.feape = feape
mlp = []
for l in range(num_layers):
if l == 0:
in_dim = self.in_mlpC
else:
in_dim = hidden_dim
if l == num_layers - 1:
out_dim, bias = 3, False # 3 rgb
else:
out_dim, bias = hidden_dim, True
mlp.append(torch.nn.Linear(in_dim, out_dim, bias=bias))
self.mlp = torch.nn.ModuleList(mlp)
# torch.nn.init.constant_(self.mlp[-1].bias, 0)
def forward(self, viewdirs, features):
indata = [features, viewdirs]
if self.feape > 0:
indata += [positional_encoding(features, self.feape)]
if self.viewpe > 0:
indata += [positional_encoding(viewdirs, self.viewpe)]
h = torch.cat(indata, dim=-1)
for l in range(self.num_layers):
h = self.mlp[l](h)
if l != self.num_layers - 1:
h = F.relu(h, inplace=True)
rgb = torch.sigmoid(h)
return rgb
class FactorFields(torch.nn.Module):
def __init__(self, cfg, device):
super(FactorFields, self).__init__()
self.cfg = cfg
self.device = device
self.matMode = [[0, 1], [0, 2], [1, 2]]
self.vecMode = [2, 1, 0]
self.n_scene, self.scene_idx = 1, 0
self.alphaMask = None
self.coeff_type, self.basis_type = cfg.model.coeff_type, cfg.model.basis_type
self.setup_params(self.cfg.dataset.aabb)
if self.cfg.model.coeff_type != 'none':
self.coeffs = self.init_coef()
if self.cfg.model.basis_type != 'none':
self.basises = self.init_basis()
out_dim = cfg.model.out_dim
if 'vm' in self.coeff_type:
in_dim = sum(cfg.model.basis_dims) * 3
elif 'x' in self.cfg.model.basis_type:
in_dim = len(
cfg.model.basis_dims) * 2 * self.in_dim if self.cfg.model.basis_mapping == 'trigonometric' else len(
cfg.model.basis_dims) * self.in_dim
else:
in_dim = sum(cfg.model.basis_dims)
self.linear_mat = MLPMixer(in_dim, out_dim, num_layers=cfg.model.num_layers, hidden_dim=cfg.model.hidden_dim,
with_dropout=cfg.model.with_dropout).to(device)
if 'reconstruction' in cfg.defaults.mode:
# self.cur_volumeSize = N_to_reso(cfg.training.volume_resoInit, self.aabb)
# self.update_renderParams(self.cur_volumeSize)
view_pe, fea_pe = cfg.renderer.view_pe, cfg.renderer.fea_pe
num_layers, hidden_dim = cfg.renderer.num_layers, cfg.renderer.hidden_dim
self.renderModule = MLPRender_Fea(inChanel=out_dim - 1, num_layers=num_layers, hidden_dim=hidden_dim,
viewpe=view_pe, feape=fea_pe).to(device)
self.is_unbound = self.cfg.dataset.is_unbound
if self.is_unbound:
self.bg_len = 0.2
self.inward_aabb = torch.tensor([[-1.0, -1.0, -1.0], [1.0, 1.0, 1.0]]).to(device)
self.aabb = self.inward_aabb * (1 + self.bg_len)
else:
self.inward_aabb = self.aabb
# self.freq_bands = torch.FloatTensor(cfg.model.freq_bands).to(device)
self.cur_volumeSize = N_to_reso(cfg.training.volume_resoInit ** self.in_dim, self.aabb)
self.update_renderParams(self.cur_volumeSize)
print('=====> total parameters: ', self.n_parameters())
def setup_params(self, aabb):
self.in_dim = len(aabb[0]) - 1 if (
'images' == self.cfg.defaults.mode or 'reconstructions' == self.cfg.defaults.mode) else len(aabb[0])
self.aabb = torch.FloatTensor(aabb)[:, :self.in_dim].to(self.device)
self.basis_dims = self.cfg.model.basis_dims
if 'reconstruction' not in self.cfg.defaults.mode:
self.basis_reso = self.cfg.model.basis_resos if 'image' in self.cfg.defaults.mode else np.round(
np.array(self.cfg.model.basis_resos) * (min(aabb[1][:self.in_dim]) + 1) / 1024.0).astype('int').tolist()
self.T_basis = self.cfg.model.T_basis if self.cfg.model.T_basis>0 else sum(np.power(np.array(self.basis_reso), self.in_dim) * np.array(self.cfg.model.basis_dims))
self.T_coeff = self.cfg.model.T_coeff if self.cfg.model.T_coeff>0 else self.cfg.model.total_params - self.T_basis
self.T_coeff = self.T_coeff if self.T_coeff > 0 else 8 ** self.in_dim * sum(self.basis_dims)
if 'image' == self.cfg.defaults.mode:
self.freq_bands = max(aabb[1][:self.in_dim]) / torch.FloatTensor(self.basis_reso).to(self.device)
else:
self.freq_bands = torch.FloatTensor(self.cfg.model.freq_bands).to(self.device)
self.coeff_reso = N_to_reso(self.T_coeff // sum(self.basis_dims), self.aabb[:, :self.in_dim])[::-1] # DHW
self.n_scene = 1 # int(aabb[1][-1]) if 'images' == self.cfg.defaults.mode else 1
if 'sdf' == self.cfg.defaults.mode:
self.freq_bands *= 0.5
elif 'images' == self.cfg.defaults.mode:
self.coeff_reso = [aabb[1][-1]] + self.coeff_reso
self.aabb = torch.FloatTensor(aabb).to(self.device)
if 'vec' in self.coeff_type or 'cp' in self.coeff_type or 'vm' in self.coeff_type:
self.coeff_reso = aabb[1]
else:
self.coeff_reso = N_to_reso(self.cfg.model.coeff_reso ** self.in_dim, self.aabb[:, :self.in_dim])[::-1]
self.T_coeff = sum(self.cfg.model.basis_dims) * np.prod(self.coeff_reso)
self.T_basis = self.cfg.model.total_params - self.T_coeff
scale = self.T_basis / sum(
np.power(np.array(self.cfg.model.basis_resos), self.in_dim) * np.array(self.cfg.model.basis_dims))
scale = np.power(scale, 1.0 / self.in_dim)
self.basis_reso = self.cfg.model.basis_resos if (
'vec' in self.basis_type or 'cp' in self.basis_type) else np.round(
np.array(self.cfg.model.basis_resos) * scale).astype('int').tolist()
# self.freq_bands = self.cfg.dataset.scene_reso / torch.FloatTensor(self.basis_resos).to(self.device)
self.freq_bands = torch.FloatTensor(self.cfg.model.freq_bands).to(self.device) \
if (
'reconstructions' == self.cfg.defaults.mode or 'x' in self.basis_type or 'vec' in self.basis_type or 'cp' in self.basis_type) else torch.FloatTensor(
self.cfg.model.freq_bands).to(self.device) * (self.cfg.dataset.scene_reso / float(
max(self.basis_reso)) / max(self.cfg.model.freq_bands))
self.n_scene = int(aabb[1][-1]) if 'reconstructions' == self.cfg.defaults.mode else 1
# print(self.coeff_reso,self.basis_reso,self.freq_bands)
def init_coef(self):
n_scene = self.n_scene if 'reconstructions' == self.cfg.defaults.mode or 'images' == self.cfg.defaults.mode else 1
if 'hash' in self.coeff_type or 'grid' in self.coeff_type:
coeffs = [
self.cfg.model.coef_init * torch.ones((1, sum(self.basis_dims), *self.coeff_reso), device=self.device)
for _ in range(n_scene)]
coeffs = torch.nn.ParameterList(coeffs)
elif 'cp' in self.coeff_type or 'vm' in self.coeff_type:
coeffs = []
for i in range(len(self.coeff_reso)):
coeffs.append(self.cfg.model.coef_init * torch.ones(
(1, sum(self.basis_dims), max(256, self.coeff_reso[i]), n_scene), device=self.device))
coeffs = torch.nn.ParameterList(coeffs)
elif 'vec' in self.coeff_type:
coeffs = self.cfg.model.coef_init * torch.ones(
(1, sum(self.basis_dims), max(256, max(self.coeff_reso)), n_scene), device=self.device)
coeffs = torch.nn.ParameterList([coeffs])
elif 'mlp' in self.coeff_type:
coeffs = torch.nn.ParameterList(
[MLPMixer(self.in_dim, sum(self.basis_dims), num_layers=2, hidden_dim=64, pe=4).to(self.device) for _ in
range(n_scene)])
return coeffs
def init_basis(self):
if 'hash' in self.basis_type:
import tinycudann as tcnn
n_levels = len(self.basis_reso)
if 'reconstruction' not in self.cfg.defaults.mode:
base_resolution_low, base_resolution_high = 32, int(max(self.aabb[1]).item()) // 2
per_level_scale = torch.pow(self.freq_bands[-1] / self.freq_bands[0], 1.0 / n_levels).item()
else:
base_resolution_low = torch.round(self.basis_reso[0] * self.freq_bands[0]).long().item()
per_level_scale = torch.pow(self.freq_bands[0] / self.freq_bands[-1], 1.0 / n_levels).item()
base_resolution_high = torch.round(
self.basis_reso[n_levels // 2] * self.freq_bands[n_levels // 2]).long().item()
log2_hashmap_size = np.round(np.log2(self.T_basis / 3 / np.mean(self.basis_dims)))
# per_level_scale = torch.pow((self.basis_reso[-1] * self.freq_bands[-1])/(self.basis_reso[0] * self.freq_bands[0]),1.0/n_levels).item()
# per_level_scale = torch.pow(self.freq_bands[0] / self.freq_bands[-1], 1.0 / n_levels).item()
basises = []
if len(self.basis_dims) == 1 or sum(self.basis_dims) > 32:
encoding_config_low = {
"otype": "HashGrid",
"n_levels": n_levels,
"n_features_per_level": sum(self.basis_dims) // n_levels,
"log2_hashmap_size": log2_hashmap_size,
"base_resolution": base_resolution_low,
"per_level_scale": per_level_scale # 1.25992 #1.38191 #
}
basises.append(tcnn.Encoding(
n_input_dims=self.in_dim,
encoding_config=encoding_config_low))
else:
encoding_config_low = {
"otype": "HashGrid",
"n_levels": n_levels // 2,
"n_features_per_level": min(16, self.basis_dims[0]),
"log2_hashmap_size": log2_hashmap_size,
"base_resolution": base_resolution_low,
"per_level_scale": per_level_scale # 1.25992 #1.38191 #
}
encoding_config_high = {
"otype": "HashGrid",
"n_levels": n_levels // 2,
"n_features_per_level": self.basis_dims[-1],
"log2_hashmap_size": log2_hashmap_size,
"base_resolution": base_resolution_high,
"per_level_scale": per_level_scale # 1.25992 #1.38191 #
}
basises = []
basises.append(tcnn.Encoding(
n_input_dims=self.in_dim,
encoding_config=encoding_config_low))
if self.basis_dims[0] == 32:
basises.append(tcnn.Encoding(
n_input_dims=self.in_dim,
encoding_config=encoding_config_low))
basises.append(tcnn.Encoding(
n_input_dims=self.in_dim,
encoding_config=encoding_config_high))
return torch.nn.ParameterList(basises)
else:
basises, coeffs, n_params_basis = [], [], 0
# in_dim = self.in_dim if 'images' != self.cfg.defaults.mode else self.in_dim - 1
for i, (basis_dim, reso) in enumerate(zip(self.basis_dims, self.basis_reso)):
# reso_cur = N_to_reso(reso, aabb)[::-1]
if 'mlp' in self.basis_type:
basises.append(MLPMixer(self.in_dim, basis_dim, num_layers=2, \
hidden_dim=64, pe=4).to(self.device))
elif 'grid' in self.basis_type:
basises.append(torch.nn.Parameter(dct_dict(int(np.power(basis_dim, 1. / self.in_dim) + 1), reso,
n_selete=basis_dim, dim=self.in_dim).reshape(
[1, basis_dim] + [reso] * self.in_dim).to(self.device)))
# basises.append(torch.nn.Parameter(torch.ones([1, basis_dim] + [reso] * self.in_dim).to(self.device)))
elif 'vm' in self.basis_type:
reso_level = N_to_vm_reso(reso ** self.in_dim, self.aabb[:, :self.in_dim])
for i in range(len(self.matMode)):
mat_id_0, mat_id_1 = self.matMode[i]
basises.append(torch.nn.Parameter(
0.1 * torch.randn((1, basis_dim, reso_level[mat_id_1], reso_level[mat_id_0]),
device=self.device)))
elif 'cp' in self.basis_type:
for _ in range(self.in_dim - 1):
basises.append(torch.nn.Parameter(
0.1 * torch.randn((1, basis_dim, max(reso, 128), 1), device=self.device)))
elif 'x' in self.basis_type:
continue
return torch.nn.ParameterList(basises)
def get_coeff(self, xyz_sampled):
N_points, dim = xyz_sampled.shape
in_dim = self.in_dim
pts = self.normalize_coord(xyz_sampled).view([1, -1] + [1] * (dim - 1) + [dim])
if self.coeff_type in 'hash':
coeffs = self.coeffs(pts * 0.5 + 0.5).float()
elif 'grid' in self.coeff_type:
coeffs = F.grid_sample(self.coeffs[self.scene_idx], pts, mode=self.cfg.model.coef_mode, align_corners=False,
padding_mode='border').view(-1, N_points).t()
elif 'vec' in self.coeff_type:
pts = pts.view(1, -1, 1, in_dim)
idx = (self.scene_idx + 0.5) / self.n_scene * 2 - 1
pts = torch.stack((torch.ones_like(pts[..., 0]) * idx, pts[..., 0]), dim=-1)
coeffs = F.grid_sample(self.coeffs[0], pts, mode=self.cfg.model.coef_mode,
align_corners=False, padding_mode='border').view(-1, N_points).t()
elif 'cp' in self.coeff_type:
pts = pts.squeeze(2)
idx = (self.scene_idx + 0.5) / self.n_scene * 2 - 1
pts = torch.stack((torch.ones_like(pts) * idx, pts), dim=-2)
coeffs = F.grid_sample(self.coeffs[0], pts[..., 0], mode=self.cfg.model.coef_mode,
align_corners=False, padding_mode='border').view(-1, N_points).t()
for i in range(1, in_dim):
coeffs = coeffs * F.grid_sample(self.coeffs[i], pts[..., i], mode=self.cfg.model.coef_mode,
align_corners=False, padding_mode='border').view(-1, N_points).t()
elif 'vm' in self.coeff_type:
pts = pts.squeeze(2)
idx = (self.scene_idx + 0.5) / self.n_scene * 2 - 1
pts = torch.stack((torch.ones_like(pts) * idx, pts), dim=-2)
coeffs = []
for i in range(in_dim):
coeffs.append(F.grid_sample(self.coeffs[i], pts[..., self.vecMode[i]], mode=self.cfg.model.coef_mode,
align_corners=False, padding_mode='border').view(-1, N_points).t())
coeffs = torch.cat(coeffs, dim=-1)
elif 'mlp' in self.coeff_type:
coeffs = self.coeffs[self.scene_idx](pts.view(N_points, in_dim))
elif 'hash' in self.coeff_type:
coeffs = self.coeffs[self.scene_idx]((pts.view(N_points, in_dim) + 1) / 2)
return coeffs
def get_basis(self, x):
N_points = x.shape[0]
if 'images' == self.cfg.defaults.mode:
x = x[..., :-1]
if 'hash' in self.basis_type:
x = (x - self.aabb[0]) / torch.max(self.aabb[1] - self.aabb[0])
if len(self.basises) == 1:
basises = self.basises[0](x).float()
if len(self.basises) == 2:
basises = torch.cat((self.basises[0](x), self.basises[1](x)), dim=-1).float()
elif len(self.basises) == 3:
basises = torch.cat((self.basises[0](x), self.basises[1](x), self.basises[2](x)), dim=-1).float()
else:
freq_len = len(self.freq_bands)
xyz = grid_mapping(x, self.freq_bands, self.aabb[:, :self.in_dim], self.cfg.model.basis_mapping).view(1, *(
[1] * (self.in_dim - 1)), -1, self.in_dim, freq_len)
basises = []
for i in range(freq_len):
if 'mlp' in self.basis_type:
basises.append(self.basises[i](xyz[..., i].view(-1, self.in_dim)))
elif 'grid' in self.basis_type:
basises.append(
F.grid_sample(self.basises[i], xyz[..., i], mode=self.cfg.model.basis_mode,
align_corners=True).view(-1, N_points).T)
elif 'vm' in self.basis_type:
coordinate_mat = torch.stack((xyz[..., self.matMode[0], i], xyz[..., self.matMode[1], i],
xyz[..., self.matMode[2], i])).view(3, -1, 1, 2)
for idx_mat in range(self.in_dim):
basises.append(F.grid_sample(self.basises[i * self.in_dim + idx_mat], coordinate_mat[[idx_mat]],
align_corners=True).view(-1, x.shape[0]).t())
elif 'cp' in self.basis_type:
for idx_axis in range(self.in_dim - 1):
coordinate_vec = torch.stack(
(torch.zeros_like(xyz[..., idx_axis + 1, i]), xyz[..., idx_axis + 1, i]), dim=-1).squeeze(2)
if 0 == idx_axis:
basises_level = F.grid_sample(self.basises[i * (self.in_dim - 1) + idx_axis],
coordinate_vec,
align_corners=True).view(-1, x.shape[0]).t()
else:
basises_level = basises_level * F.grid_sample(
self.basises[i * (self.in_dim - 1) + idx_axis], coordinate_vec,
align_corners=True).view(-1, x.shape[0]).t()
basises.append(basises_level)
elif 'x' in self.basis_type:
basises.append(xyz[..., i].view(x.shape[0], -1))
if isinstance(basises, list):
basises = torch.cat(basises, dim=-1)
if 'vm' in self.basis_type: # switch order
basises = basises.view(x.shape[0], freq_len, -1).permute(0, 2, 1).reshape(x.shape[0], -1)
return basises
@torch.no_grad()
def normalize_basis(self):
for basis in self.basises:
basis.data = basis.data / torch.norm(basis.data, dim=(2, 3), keepdim=True)
def get_coding(self, x):
if self.cfg.model.coeff_type != 'none' and self.cfg.model.basis_type != 'none':
coeff = self.get_coeff(x)
basises = self.get_basis(x)
return basises * coeff, coeff
elif self.cfg.model.coeff_type != 'none':
coeff = self.get_coeff(x)
return coeff, coeff
elif self.cfg.model.basis_type != 'none':
basises = self.get_basis(x)
return basises, basises
def n_parameters(self):
total = sum(p.numel() for p in self.parameters())
if 'fix' in self.cfg.model.basis_type:
total -= self.T_basis
return total
def get_optparam_groups(self, lr_small=0.001, lr_large=0.02):
grad_vars = []
if self.cfg.training.linear_mat:
grad_vars += [{'params': self.linear_mat.parameters(), 'lr': lr_small}]
if 'none' != self.coeff_type and self.cfg.training.coeff:
grad_vars += [{'params': self.coeffs.parameters(), 'lr': lr_large}]
if 'fix' not in self.cfg.model.basis_type and 'none' != self.cfg.model.basis_type and self.cfg.training.basis:
grad_vars += [{'params': self.basises.parameters(), 'lr': lr_large}]
if 'reconstruction' in self.cfg.defaults.mode and self.cfg.training.renderModule:
grad_vars += [{'params': self.renderModule.parameters(), 'lr': lr_small}]
return grad_vars
def set_optimizable(self, items, statue):
for item in items:
if item == 'basis' and self.cfg.model.basis_type != 'none':
for item in self.basises:
item.requires_grad = statue
elif item == 'coeff' and self.cfg.model.coeff_type != 'none':
for item in self.basises:
item.requires_grad = statue
elif item == 'proj':
self.linear_mat.requires_grad = statue
elif item == 'renderer':
self.renderModule.requires_grad = statue
def TV_loss(self, reg):
total = 0
for idx in range(len(self.basises)):
total = total + reg(self.basises[idx]) * 1e-2
return total
def sample_point_ndc(self, rays_o, rays_d, is_train=True, N_samples=-1):
N_samples = N_samples if N_samples > 0 else self.nSamples
near, far = self.cfg.dataset.near_far
interpx = torch.linspace(near, far, N_samples).unsqueeze(0).to(rays_o)
if is_train:
interpx += torch.rand_like(interpx).to(rays_o) * ((far - near) / N_samples)
rays_pts = rays_o[..., None, :] + rays_d[..., None, :] * interpx[..., None]
mask_outbbox = ((self.aabb[0, :self.in_dim] > rays_pts) | (rays_pts > self.aabb[1, :self.in_dim])).any(dim=-1)
return rays_pts, interpx, ~mask_outbbox
def sample_point(self, rays_o, rays_d, is_train=True, N_samples=-1):
N_samples = N_samples if N_samples > 0 else self.nSamples
vec = torch.where(rays_d == 0, torch.full_like(rays_d, 1e-6), rays_d)
rate_a = (self.aabb[1, :self.in_dim] - rays_o) / vec
rate_b = (self.aabb[0, :self.in_dim] - rays_o) / vec
t_min = torch.minimum(rate_a, rate_b).amax(-1).clamp(min=0.05, max=1e3)
rng = torch.arange(N_samples)[None].float()
if is_train:
rng = rng.repeat(rays_d.shape[-2], 1)
rng += torch.rand_like(rng[:, [0]])
step = self.stepSize * rng.to(rays_o.device)
interpx = (t_min[..., None] + step)
rays_pts = rays_o[..., None, :] + rays_d[..., None, :] * interpx[..., None]
mask_outbbox = ((self.aabb[0, :self.in_dim] > rays_pts) | (rays_pts > self.aabb[1, :self.in_dim])).any(dim=-1)
return rays_pts, interpx, ~mask_outbbox
def sample_point_unbound(self, rays_o, rays_d, is_train=True, N_samples=-1):
N_samples = N_samples if N_samples > 0 else self.nSamples
N_inner, N_outer = 3 * N_samples // 4, N_samples // 4
b_inner = torch.linspace(0, 2, N_inner + 1).to(self.device)
b_outer = 2 / torch.linspace(1, 1 / 16, N_outer + 1).to(self.device)
if is_train:
rng = torch.rand((N_inner + N_outer), device=self.device)
interpx = torch.cat([
b_inner[1:] * rng[:N_inner] + b_inner[:-1] * (1 - rng[:N_inner]),
b_outer[1:] * rng[N_inner:] + b_outer[:-1] * (1 - rng[N_inner:]),
])[None]
else:
interpx = torch.cat([
(b_inner[1:] + b_inner[:-1]) * 0.5,
(b_outer[1:] + b_outer[:-1]) * 0.5,
])[None]
rays_pts = rays_o[:, None, :] + rays_d[:, None, :] * interpx[..., None]
norm = rays_pts.abs().amax(dim=-1, keepdim=True)
inner_mask = (norm <= 1)
rays_pts = torch.where(
inner_mask,
rays_pts,
rays_pts / norm * ((1 + self.bg_len) - self.bg_len / norm)
)
return rays_pts, interpx, inner_mask.squeeze(-1)
def normalize_coord(self, xyz_sampled):
invaabbSize = 2.0 / (self.aabb[1] - self.aabb[0])
return (xyz_sampled - self.aabb[0]) * invaabbSize - 1
def basis2density(self, density_features):
if self.cfg.renderer.fea2denseAct == "softplus":
return F.softplus(density_features + self.cfg.renderer.density_shift)
elif self.cfg.renderer.fea2denseAct == "relu":
return F.relu(density_features + self.cfg.renderer.density_shift)
@torch.no_grad()
def cal_mean_coef(self, state_dict):
if 'grid' in self.coeff_type or 'mlp' in self.coeff_type:
key_list = []
for item in state_dict.keys():
if 'coeffs.0' in item:
key_list.append(item)
for key in key_list:
average = torch.zeros_like(state_dict[key])
for i in range(self.n_scene):
item = key.replace('0', f'{i}', 1)
average += state_dict[item]
state_dict.pop(item, None)
average /= self.n_scene
state_dict[key] = average
elif 'vec' in self.coeff_type:
state_dict['coeffs.0'] = torch.mean(state_dict['coeffs.0'], dim=-1, keepdim=True)
elif 'cp' in self.coeff_type or 'vm' in self.coeff_type:
for i in range(3):
state_dict[f'coeffs.{i}'] = torch.mean(state_dict[f'coeffs.{i}'], dim=-1, keepdim=True)
return state_dict
def save(self, path):
ckpt = {'state_dict': self.state_dict(), 'cfg': self.cfg}
if self.alphaMask is not None:
alpha_volume = self.alphaMask.alpha_volume.bool().cpu().numpy()
ckpt.update({'alphaMask.shape': alpha_volume.shape})
ckpt.update({'alphaMask.mask': np.packbits(alpha_volume.reshape(-1))})
ckpt.update({'alphaMask.aabb': self.alphaMask.aabb.cpu()})
# average the coeff for saving if batch training
if 'reconstruction' in self.cfg.defaults.mode:
ckpt['state_dict'] = self.cal_mean_coef(ckpt['state_dict'])
torch.save(ckpt, path)
def load(self, ckpt):
if 'alphaMask.aabb' in ckpt.keys():
length = np.prod(ckpt['alphaMask.shape'])
alpha_volume = torch.from_numpy(
np.unpackbits(ckpt['alphaMask.mask'])[:length].reshape(ckpt['alphaMask.shape']))
self.alphaMask = AlphaGridMask(self.device, ckpt['alphaMask.aabb'].to(self.device),
alpha_volume.float().to(self.device))
self.load_state_dict(ckpt['state_dict'])
volumeSize = N_to_reso(self.cfg.training.volume_resoFinal ** self.in_dim, self.aabb)
self.update_renderParams(volumeSize)
def update_renderParams(self, gridSize):
self.aabbSize = self.aabb[1] - self.aabb[0]
self.gridSize = torch.LongTensor(gridSize).to(self.device)
units = self.aabbSize / (self.gridSize - 1)
self.stepSize = torch.mean(units) * self.cfg.renderer.step_ratio
aabbDiag = torch.sqrt(torch.sum(torch.square(self.aabbSize)))
self.nSamples = int((aabbDiag / self.stepSize).item()) + 1
@torch.no_grad()
def upsample_volume_grid(self, res_target):
self.update_renderParams(res_target)
if self.cfg.dataset.dataset_name == 'google_objs' and self.n_scene == 1 and self.cfg.model.coeff_type == 'grid':
coeffs = [
F.interpolate(self.coeffs[0].data, size=None, scale_factor=1.3, align_corners=True, mode='trilinear')]
self.coeffs = torch.nn.ParameterList(coeffs)
def compute_alpha(self, xyz_locs, length=1):
if self.alphaMask is not None:
alphas = self.alphaMask.sample_alpha(xyz_locs)
alpha_mask = alphas > 0
else:
alpha_mask = torch.ones_like(xyz_locs[:, 0], dtype=bool)
sigma = torch.zeros(xyz_locs.shape[:-1], device=xyz_locs.device)
if alpha_mask.any():
feats, _ = self.get_coding(xyz_locs[alpha_mask])
validsigma = self.linear_mat(feats, is_train=False)[..., 0]
sigma[alpha_mask] = self.basis2density(validsigma)
alpha = 1 - torch.exp(-sigma * length).view(xyz_locs.shape[:-1])
return alpha
@torch.no_grad()
def getDenseAlpha(self, gridSize=None, times=16):
gridSize = self.gridSize.tolist() if gridSize is None else gridSize
aabbSize = self.inward_aabb[1] - self.inward_aabb[0]
units = aabbSize / (torch.LongTensor(gridSize).to(self.device) - 1)
units_half = 1.0 / (torch.LongTensor(gridSize) - 1) * 0.5
stepSize = torch.mean(units)
samples = torch.stack(torch.meshgrid(
[torch.linspace(units_half[0], 1 - units_half[0], gridSize[0]),
torch.linspace(units_half[1], 1 - units_half[1], gridSize[1]),
torch.linspace(units_half[2], 1 - units_half[2], gridSize[2])], indexing='ij'
), -1).to(self.device)
dense_xyz = self.inward_aabb[0] * (1 - samples) + self.inward_aabb[1] * samples
dense_xyz = dense_xyz.transpose(0, 2).contiguous()
alpha = torch.zeros_like(dense_xyz[..., 0])
for _ in range(times):
for i in range(gridSize[2]):
shiftment = (torch.rand(dense_xyz[i].shape) * 2 - 1).to(self.device) * (
units / 2 * 1.2) if times > 1 else 0.0
alpha[i] += self.compute_alpha((dense_xyz[i] + shiftment).view(-1, 3),
stepSize * self.cfg.renderer.distance_scale).view(
(gridSize[1], gridSize[0]))
return alpha / times, dense_xyz
@torch.no_grad()
def updateAlphaMask(self, gridSize=(200, 200, 200), is_update_alphaMask=False):
alpha, dense_xyz = self.getDenseAlpha(gridSize)
total_voxels = gridSize[0] * gridSize[1] * gridSize[2]
ks = 3
alpha = alpha.clamp(0, 1)[None, None]
alpha = F.max_pool3d(alpha, kernel_size=ks, padding=ks // 2, stride=1).view(gridSize[::-1])
# filter floaters
min_size = np.mean(alpha.shape[-3:]).item()
alphaMask_thres = self.cfg.renderer.alphaMask_thres if is_update_alphaMask else 0.08
if self.is_unbound:
alphaMask_thres = 0.04
alpha = (alpha >= alphaMask_thres).float()
else:
alpha = skimage.morphology.remove_small_objects(alpha.cpu().numpy() >= alphaMask_thres, min_size=min_size,
connectivity=1)
alpha = torch.FloatTensor(alpha).to(self.device)
if is_update_alphaMask:
self.alphaMask = AlphaGridMask(self.device, self.inward_aabb, alpha)
valid_xyz = dense_xyz[alpha > 0.5]
xyz_min = valid_xyz.amin(0)
xyz_max = valid_xyz.amax(0)
if not self.is_unbound:
pad = (xyz_max - xyz_min) / 20
xyz_min -= pad
xyz_max += pad
new_aabb = torch.stack((xyz_min, xyz_max))
total = torch.sum(alpha)
return new_aabb
@torch.no_grad()
def shrink(self, new_aabb):
self.setup_params(new_aabb.tolist())
if self.cfg.model.coeff_type != 'none':
del self.coeffs
self.coeffs = self.init_coef()
if self.cfg.model.basis_type != 'none':
del self.basises
self.basises = self.init_basis()
self.aabb = self.inward_aabb = new_aabb
self.cfg.dataset.aabb = self.aabb.tolist()
self.update_renderParams(self.gridSize.tolist())
@torch.no_grad()
def filtering_rays(self, all_rays, all_rgbs, N_samples=256, chunk=10240 * 5, bbox_only=False):
tt = time.time()
N = torch.tensor(all_rays.shape[:-1]).prod()
mask_filtered = []
length_current = 0
idx_chunks = torch.split(torch.arange(N), chunk)
for idx_chunk in idx_chunks:
rays_chunk = all_rays[idx_chunk].to(self.device)
rays_o, rays_d = rays_chunk[..., :3], rays_chunk[..., 3:6]
if bbox_only:
vec = torch.where(rays_d == 0, torch.full_like(rays_d, 1e-6), rays_d)
rate_a = (self.aabb[1] - rays_o) / vec
rate_b = (self.aabb[0] - rays_o) / vec
t_min = torch.minimum(rate_a, rate_b).amax(-1) # .clamp(min=near, max=far)
t_max = torch.maximum(rate_a, rate_b).amin(-1) # .clamp(min=near, max=far)
mask_inbbox = t_max > t_min
else:
xyz_sampled, _, _ = self.sample_point(rays_o, rays_d, N_samples=N_samples, is_train=False)
mask_inbbox = (self.alphaMask.sample_alpha(xyz_sampled).view(xyz_sampled.shape[:-1]) > 0).any(-1)
# mask_filtered.append(mask_inbbox.cpu())
length = torch.sum(mask_inbbox)
all_rays[length_current:length_current + length], all_rgbs[length_current:length_current + length] = \
rays_chunk[mask_inbbox].cpu(), all_rgbs[idx_chunk][mask_inbbox.cpu()]
length_current += length
return all_rays[:length_current], all_rgbs[:length_current]
def forward(self, rays_chunk, white_bg=True, is_train=False, ndc_ray=False, N_samples=-1):
# sample points
viewdirs = rays_chunk[:, 3:6]
if self.is_unbound:
xyz_sampled, z_vals, inner_mask = self.sample_point_unbound(rays_chunk[:, :3], viewdirs, is_train=is_train,
N_samples=N_samples)
dists = torch.cat((z_vals[:, 1:] - z_vals[:, :-1], z_vals[:, -1:] - z_vals[:, -2:-1]), dim=-1)
elif ndc_ray:
xyz_sampled, z_vals, inner_mask = self.sample_point_ndc(rays_chunk[:, :3], viewdirs, is_train=is_train,
N_samples=N_samples)
dists = torch.cat((z_vals[:, 1:] - z_vals[:, :-1], torch.zeros_like(z_vals[:, :1])), dim=-1)
rays_norm = torch.norm(viewdirs, dim=-1, keepdim=True)
dists = dists * rays_norm
viewdirs = viewdirs / rays_norm
else:
xyz_sampled, z_vals, inner_mask = self.sample_point(rays_chunk[:, :3], viewdirs, is_train=is_train,
N_samples=N_samples)
dists = torch.cat((z_vals[:, 1:] - z_vals[:, :-1], torch.zeros_like(z_vals[:, :1])), dim=-1)
viewdirs = viewdirs.view(-1, 1, 3).expand(xyz_sampled.shape)
ray_valid = torch.ones_like(xyz_sampled[..., 0]).bool() if self.is_unbound else inner_mask
if self.alphaMask is not None:
alpha_inner_valid = self.alphaMask.sample_alpha(xyz_sampled[inner_mask]) > 0.5
ray_valid[inner_mask.clone()] = alpha_inner_valid
sigma = torch.zeros(xyz_sampled.shape[:-1], device=xyz_sampled.device)
rgb = torch.zeros((*xyz_sampled.shape[:2], 3), device=xyz_sampled.device)
coeffs = torch.zeros((1, sum(self.cfg.model.basis_dims)), device=xyz_sampled.device)
if ray_valid.any():
feats, coeffs = self.get_coding(xyz_sampled[ray_valid])
feat = self.linear_mat(feats, is_train=is_train)
sigma[ray_valid] = self.basis2density(feat[..., 0])
alpha, weight, bg_weight = raw2alpha(sigma, dists * self.cfg.renderer.distance_scale)
app_mask = weight > self.cfg.renderer.rayMarch_weight_thres
ray_valid_new = torch.logical_and(ray_valid, app_mask)
app_mask = ray_valid_new[ray_valid]
if app_mask.any():
valid_rgbs = self.renderModule(viewdirs[ray_valid_new], feat[app_mask, 1:])
rgb[ray_valid_new] = valid_rgbs
acc_map = torch.sum(weight, -1)
rgb_map = torch.sum(weight[..., None] * rgb, -2)
if white_bg or (is_train and torch.rand((1,)) < 0.5):
rgb_map = rgb_map + (1. - acc_map[..., None])
rgb_map = rgb_map.clamp(0, 1)
with torch.no_grad():
depth_map = torch.sum(weight * z_vals, -1)
return rgb_map, depth_map, coeffs
================================================
FILE: models/__init__.py
================================================
================================================
FILE: models/sh.py
================================================
import torch
################## sh function ##################
C0 = 0.28209479177387814
C1 = 0.4886025119029199
C2 = [
1.0925484305920792,
-1.0925484305920792,
0.31539156525252005,
-1.0925484305920792,
0.5462742152960396
]
C3 = [
-0.5900435899266435,
2.890611442640554,
-0.4570457994644658,
0.3731763325901154,
-0.4570457994644658,
1.445305721320277,
-0.5900435899266435
]
C4 = [
2.5033429417967046,
-1.7701307697799304,
0.9461746957575601,
-0.6690465435572892,
0.10578554691520431,
-0.6690465435572892,
0.47308734787878004,
-1.7701307697799304,
0.6258357354491761,
]
def eval_sh(deg, sh, dirs):
"""
Evaluate spherical harmonics at unit directions
using hardcoded SH polynomials.
Works with torch/np/jnp.
... Can be 0 or more batch dimensions.
:param deg: int SH max degree. Currently, 0-4 supported
:param sh: torch.Tensor SH coeffs (..., C, (max degree + 1) ** 2)
:param dirs: torch.Tensor unit directions (..., 3)
:return: (..., C)
"""
assert deg <= 4 and deg >= 0
assert (deg + 1) ** 2 == sh.shape[-1]
C = sh.shape[-2]
result = C0 * sh[..., 0]
if deg > 0:
x, y, z = dirs[..., 0:1], dirs[..., 1:2], dirs[..., 2:3]
result = (result -
C1 * y * sh[..., 1] +
C1 * z * sh[..., 2] -
C1 * x * sh[..., 3])
if deg > 1:
xx, yy, zz = x * x, y * y, z * z
xy, yz, xz = x * y, y * z, x * z
result = (result +
C2[0] * xy * sh[..., 4] +
C2[1] * yz * sh[..., 5] +
C2[2] * (2.0 * zz - xx - yy) * sh[..., 6] +
C2[3] * xz * sh[..., 7] +
C2[4] * (xx - yy) * sh[..., 8])
if deg > 2:
result = (result +
C3[0] * y * (3 * xx - yy) * sh[..., 9] +
C3[1] * xy * z * sh[..., 10] +
C3[2] * y * (4 * zz - xx - yy)* sh[..., 11] +
C3[3] * z * (2 * zz - 3 * xx - 3 * yy) * sh[..., 12] +
C3[4] * x * (4 * zz - xx - yy) * sh[..., 13] +
C3[5] * z * (xx - yy) * sh[..., 14] +
C3[6] * x * (xx - 3 * yy) * sh[..., 15])
if deg > 3:
result = (result + C4[0] * xy * (xx - yy) * sh[..., 16] +
C4[1] * yz * (3 * xx - yy) * sh[..., 17] +
C4[2] * xy * (7 * zz - 1) * sh[..., 18] +
C4[3] * yz * (7 * zz - 3) * sh[..., 19] +
C4[4] * (zz * (35 * zz - 30) + 3) * sh[..., 20] +
C4[5] * xz * (7 * zz - 3) * sh[..., 21] +
C4[6] * (xx - yy) * (7 * zz - 1) * sh[..., 22] +
C4[7] * xz * (xx - 3 * yy) * sh[..., 23] +
C4[8] * (xx * (xx - 3 * yy) - yy * (3 * xx - yy)) * sh[..., 24])
return result
def eval_sh_bases(deg, dirs):
"""
Evaluate spherical harmonics bases at unit directions,
without taking linear combination.
At each point, the final result may the be
obtained through simple multiplication.
:param deg: int SH max degree. Currently, 0-4 supported
:param dirs: torch.Tensor (..., 3) unit directions
:return: torch.Tensor (..., (deg+1) ** 2)
"""
assert deg <= 4 and deg >= 0
result = torch.empty((*dirs.shape[:-1], (deg + 1) ** 2), dtype=dirs.dtype, device=dirs.device)
result[..., 0] = C0
if deg > 0:
x, y, z = dirs.unbind(-1)
result[..., 1] = -C1 * y;
result[..., 2] = C1 * z;
result[..., 3] = -C1 * x;
if deg > 1:
xx, yy, zz = x * x, y * y, z * z
xy, yz, xz = x * y, y * z, x * z
result[..., 4] = C2[0] * xy;
result[..., 5] = C2[1] * yz;
result[..., 6] = C2[2] * (2.0 * zz - xx - yy);
result[..., 7] = C2[3] * xz;
result[..., 8] = C2[4] * (xx - yy);
if deg > 2:
result[..., 9] = C3[0] * y * (3 * xx - yy);
result[..., 10] = C3[1] * xy * z;
result[..., 11] = C3[2] * y * (4 * zz - xx - yy);
result[..., 12] = C3[3] * z * (2 * zz - 3 * xx - 3 * yy);
result[..., 13] = C3[4] * x * (4 * zz - xx - yy);
result[..., 14] = C3[5] * z * (xx - yy);
result[..., 15] = C3[6] * x * (xx - 3 * yy);
if deg > 3:
result[..., 16] = C4[0] * xy * (xx - yy);
result[..., 17] = C4[1] * yz * (3 * xx - yy);
result[..., 18] = C4[2] * xy * (7 * zz - 1);
result[..., 19] = C4[3] * yz * (7 * zz - 3);
result[..., 20] = C4[4] * (zz * (35 * zz - 30) + 3);
result[..., 21] = C4[5] * xz * (7 * zz - 3);
result[..., 22] = C4[6] * (xx - yy) * (7 * zz - 1);
result[..., 23] = C4[7] * xz * (xx - 3 * yy);
result[..., 24] = C4[8] * (xx * (xx - 3 * yy) - yy * (3 * xx - yy));
return result
================================================
FILE: renderer.py
================================================
import torch,os,imageio,sys
from tqdm.auto import tqdm
from dataLoader.ray_utils import get_rays
from utils import *
from dataLoader.ray_utils import ndc_rays_blender
def render_ray(rays, factor_fields, chunk=4096, N_samples=-1, ndc_ray=False, white_bg=True, is_train=False, device='cuda'):
rgbs, alphas, depth_maps, weights, coeffs = [], [], [], [], []
N_rays_all = rays.shape[0]
for chunk_idx in range(N_rays_all // chunk + int(N_rays_all % chunk > 0)):
rays_chunk = rays[chunk_idx * chunk:(chunk_idx + 1) * chunk].to(device)
if is_train:
rgb_map, depth_map, coeff = factor_fields(rays_chunk, is_train=is_train, white_bg=white_bg, ndc_ray=ndc_ray, N_samples=N_samples)
coeffs.append(coeff)
else:
rgb_map, depth_map, _ = factor_fields(rays_chunk, is_train=is_train, white_bg=white_bg, ndc_ray=ndc_ray, N_samples=N_samples)
rgbs.append(rgb_map)
depth_maps.append(depth_map)
if is_train:
return torch.cat(rgbs), torch.cat(depth_maps), torch.cat(coeffs)
else:
return torch.cat(rgbs), torch.cat(depth_maps)
@torch.no_grad()
def evaluation(test_dataset,factor_fields, renderer, savePath=None, N_vis=5, prtx='', N_samples=-1,
white_bg=False, ndc_ray=False, compute_extra_metrics=True, device='cuda'):
PSNRs, rgb_maps, depth_maps = [], [], []
ssims,l_alex,l_vgg=[],[],[]
os.makedirs(savePath, exist_ok=True)
os.makedirs(savePath+"/rgbd", exist_ok=True)
try:
tqdm._instances.clear()
except Exception:
pass
torch.cuda.empty_cache()
img_eval_interval = 1 if N_vis < 0 else max(test_dataset.all_rays.shape[0] // N_vis,1)
idxs = list(range(0, test_dataset.all_rays.shape[0], img_eval_interval))
for idx, samples in tqdm(enumerate(test_dataset.all_rays[0::img_eval_interval]), file=sys.stdout):
W, H = test_dataset.img_wh
rays = samples.view(-1,samples.shape[-1])
rgb_map, depth_map = renderer(rays, factor_fields, chunk=1024, N_samples=N_samples,
ndc_ray=ndc_ray, white_bg = white_bg, device=device)
rgb_map = rgb_map.clamp(0.0, 1.0)
rgb_map, depth_map = rgb_map.reshape(H, W, 3).cpu(), depth_map.reshape(H, W).cpu()
depth_map, _ = visualize_depth_numpy(depth_map.numpy())
if len(test_dataset.all_rgbs):
gt_rgb = test_dataset.all_rgbs[idxs[idx]].view(H, W, 3)
loss = torch.mean((rgb_map - gt_rgb) ** 2)
PSNRs.append(-10.0 * np.log(loss.item()) / np.log(10.0))
if compute_extra_metrics:
ssim = rgb_ssim(rgb_map, gt_rgb, 1)
l_a = rgb_lpips(gt_rgb.numpy(), rgb_map.numpy(), 'alex', factor_fields.device)
l_v = rgb_lpips(gt_rgb.numpy(), rgb_map.numpy(), 'vgg', factor_fields.device)
ssims.append(ssim)
l_alex.append(l_a)
l_vgg.append(l_v)
rgb_map = (rgb_map.numpy() * 255).astype('uint8')
gt_rgb = (gt_rgb.numpy() * 255).astype('uint8')
# rgb_map = np.concatenate((rgb_map, depth_map), axis=1)
rgb_maps.append(rgb_map)
depth_maps.append(depth_map)
if savePath is not None:
imageio.imwrite(f'{savePath}/{prtx}{idx:03d}.png', rgb_map)
rgb_map = np.concatenate((rgb_map, gt_rgb, depth_map), axis=1)
imageio.imwrite(f'{savePath}/rgbd/{prtx}{idx:03d}.png', rgb_map)
torch.cuda.empty_cache()
imageio.mimwrite(f'{savePath}/{prtx}video.mp4', np.stack(rgb_maps), fps=30, quality=10)
imageio.mimwrite(f'{savePath}/{prtx}depthvideo.mp4', np.stack(depth_maps), fps=30, quality=10)
n_params = factor_fields.n_parameters()
if PSNRs:
psnr = np.mean(np.asarray(PSNRs))
if compute_extra_metrics:
ssim = np.mean(np.asarray(ssims))
l_a = np.mean(np.asarray(l_alex))
l_v = np.mean(np.asarray(l_vgg))
np.savetxt(f'{savePath}/{prtx}mean.txt', np.asarray([psnr, ssim, l_a, l_v, n_params]))
print(f"PSNR={psnr} SSIM={ssim} {l_a} {l_v} ")
else:
np.savetxt(f'{savePath}/{prtx}mean.txt', np.asarray([psnr,n_params]))
return PSNRs
@torch.no_grad()
def evaluation_path(test_dataset,factor_fields, c2ws, renderer, savePath=None, prtx='', N_samples=-1,
white_bg=False, ndc_ray=False, compute_extra_metrics=True, device='cuda'):
PSNRs, rgb_maps, depth_maps = [], [], []
ssims,l_alex,l_vgg=[],[],[]
os.makedirs(savePath, exist_ok=True)
os.makedirs(savePath+"/rgbd", exist_ok=True)
try:
tqdm._instances.clear()
except Exception:
pass
near_far = test_dataset.near_far
for idx, c2w in tqdm(enumerate(c2ws)):
W, H = test_dataset.img_wh
c2w = torch.FloatTensor(c2w)
rays_o, rays_d = get_rays(test_dataset.directions, c2w) # both (h*w, 3)
if ndc_ray:
rays_o, rays_d = ndc_rays_blender(H, W, test_dataset.focal[0], 1.0, rays_o, rays_d)
rays = torch.cat([rays_o, rays_d], 1) # (h*w, 6)
rgb_map, depth_map = renderer(rays, factor_fields, chunk=8192, N_samples=N_samples,
ndc_ray=ndc_ray, white_bg = white_bg, device=device)
rgb_map = rgb_map.clamp(0.0, 1.0)
rgb_map, depth_map = rgb_map.reshape(H, W, 3).cpu(), depth_map.reshape(H, W).cpu()
depth_map, _ = visualize_depth_numpy(depth_map.numpy(),near_far)
rgb_map = (rgb_map.numpy() * 255).astype('uint8')
# rgb_map = np.concatenate((rgb_map, depth_map), axis=1)
rgb_maps.append(rgb_map)
depth_maps.append(depth_map)
if savePath is not None:
imageio.imwrite(f'{savePath}/{prtx}{idx:03d}.png', rgb_map)
rgb_map = np.concatenate((rgb_map, depth_map), axis=1)
imageio.imwrite(f'{savePath}/rgbd/{prtx}{idx:03d}.png', rgb_map)
imageio.mimwrite(f'{savePath}/{prtx}video.mp4', np.stack(rgb_maps), fps=30, quality=8)
imageio.mimwrite(f'{savePath}/{prtx}depthvideo.mp4', np.stack(depth_maps), fps=30, quality=8)
if PSNRs:
psnr = np.mean(np.asarray(PSNRs))
if compute_extra_metrics:
ssim = np.mean(np.asarray(ssims))
l_a = np.mean(np.asarray(l_alex))
l_v = np.mean(np.asarray(l_vgg))
np.savetxt(f'{savePath}/{prtx}mean.txt', np.asarray([psnr, ssim, l_a, l_v]))
else:
np.savetxt(f'{savePath}/{prtx}mean.txt', np.asarray([psnr]))
return PSNRs
================================================
FILE: requirements.txt
================================================
imageio==2.26.0
kornia==0.6.10
lpips==0.1.4
matplotlib==3.7.1
omegaconf==2.3.0
opencv_python==4.7.0.72
plyfile==0.7.4
scikit-image
tqdm==4.65.0
trimesh==3.20.2
jupyterlab
================================================
FILE: run_batch.py
================================================
import os
import threading, queue
import numpy as np
import time
if __name__ == '__main__':
################ per scene NeRF ################
commands = {
'-grid': '', \
# '-DVGO-like': 'model.basis_type=none model.coeff_reso=80', \
# '-noC': 'model.coeff_type=none', \
# '-SL':'model.basis_dims=[18] model.basis_resos=[70] model.freq_bands=[8.]', \
# '-CP': f'model.coeff_type=vec model.basis_type=cp model.freq_bands=[1.,1.,1.,1.,1.,1.] model.basis_resos=[512,512,512,512,512,512] model.basis_dims=[32,32,32,32,32,32]', \
# '-iNGP-like': 'model.basis_type=hash model.coeff_type=none', \
# '-hash': f'model.basis_type=hash model.coef_init=1.0', \
# '-sinc': f'model.basis_mapping=sinc', \
# '-tria': f'model.basis_mapping=triangle', \
# '-vm': f'model.coeff_type=vm model.basis_type=vm', \
# '-mlpB': 'model.basis_type=mlp', \
# '-mlpC': 'model.coeff_type=mlp', \
# '-occNet': f'model.basis_type=x model.coeff_type=none model.basis_mapping=x model.num_layers=8 model.hidden_dim=256 ', \
# '-nerf': f'model.basis_type=x model.coeff_type=none model.basis_mapping=trigonometric ' \
# f'model.num_layers=8 model.hidden_dim=256 ' \
# f'model.freq_bands=[1.,2.,4.,8.,16.,32.,64,128,256.,512.] model.basis_dims=[1,1,1,1,1,1,1,1,1,1] model.basis_resos=[1024,512,256,128,64,32,16,8,4,2]', \
# '-iNGP-like-sl': 'model.basis_type=hash model.coeff_type=none model.basis_dims=[16] model.freq_bands=[8.] model.basis_resos=[64] ', \
# '-hash-sl': f'model.basis_type=hash model.coef_init=1.0 model.basis_dims=[16] model.freq_bands=[8.] model.basis_resos=[64] ', \
# '-DCT':'model.basis_type=fix-grid', \
}
################ per scene NeRF ################
###### uncomment the following five lines if you want to train on all scenes #########
cmds = []
for name in commands.keys(): #
# for scene in ['ship', 'mic', 'chair', 'lego', 'drums', 'ficus', 'hotdog', 'materials']:#
# cmd = f'python train_per_scene.py configs/nerf.yaml defaults.expname={scene}{name} dataset.datadir=./data/nerf_synthetic/{scene} {commands[name]}'
# cmds.append(cmd)
for scene in ['Ignatius','Truck']:#
if scene != 'Ignatius':
cmd = f'python train_per_scene.py configs/nerf.yaml defaults.expname={scene}{name} dataset.datadir=./data/TanksAndTemple/{scene} {commands[name]} ' \
f' dataset.dataset_name=tankstemple '
cmds.append(cmd)
cmd = f'python train_per_scene.py configs/nerf.yaml defaults.expname={scene}{name} dataset.datadir=./data/TanksAndTemple/{scene} {commands[name]} ' \
f' dataset.dataset_name=tankstemple exportation.render_only=1 exportation.render_path=1 exportation.render_test=0 ' \
f' defaults.ckpt=/mnt/qb/home/geiger/zyu30/Projects/Anpei/Code/factor-fields/logs/{scene}-grid/{scene}-grid.th '
cmds.append(cmd)
################ generalization NeRF ################
commands = {
# '-grid': '', \
# '-DVGO-like': 'model.basis_type=none model.coeff_reso=48',
# '-SL':'model.basis_dims=[72] model.basis_resos=[48] model.freq_bands=[6.]', \
# '-CP': f'model.coeff_type=vec model.basis_type=cp model.freq_bands=[1.,1.,1.,1.,1.,1.] model.basis_resos=[512,512,512,512,512,512] model.basis_dims=[32,32,32,32,32,32]', \
# '-hash': f'model.basis_type=hash model.coef_init=1.0 ', \
# '-sinc': f'model.basis_mapping=sinc', \
# '-tria': f'model.basis_mapping=triangle', \
# '-vm': f'model.coeff_type=vm model.basis_type=vm', \
# '-mlpB': 'model.basis_type=mlp', \
# '-mlpC': 'model.coeff_type=mlp', \
# '-hash-sl': f'model.basis_type=hash model.coef_init=1.0 model.basis_dims=[16] model.freq_bands=[8.] model.basis_resos=[64] ', \
# '-DCT':'model.basis_type=fix-grid', \
}
# for name in commands.keys(): #
# config = commands[name]
# config = f'python train_across_scene2.py configs/nerf_set.yaml defaults.expname=google-obj{name} {config} ' \
# f'training.volume_resoFinal=128 dataset.datadir=./data/google_scanned_objects/'
# cmds.append(config)
# # =========> fine tuning <================
# views = 5
# for name in commands.keys(): #
# for scene in [183,199,298,467,957,244,963,527]:#
# cmd = f'python train_across_scene.py configs/nerf_ft.yaml defaults.expname=google_objs_{name}_{scene}_{views}_views ' \
# f'dataset.datadir=/home/anpei/Dataset/google_scanned_objects/ {commands[name]} ' \
# f'dataset.train_views={views} ' \
# f'dataset.train_scene_list=[{scene}] ' \
# f'dataset.test_scene_list=[{scene}] ' \
# f'defaults.ckpt=/home/anpei/Code/NeuBasis/log/google-obj{name}//google-obj{name}.th '
# cmds.append(cmd)
# for scene in ['Ignatius','Barn','Truck','Family','Caterpillar']:#'Ignatius','Barn','Truck','Family','Caterpillar'
# cmds.append(f'python train_basis.py configs/tnt.yaml defaults.expname=tnt_{scene} ' \
# f'dataset.datadir=./data/TanksAndTemple/{scene}'
# )
# cmds = []
# for scene in ['room']:#,'hall','kitchen','living_room','room2','sofa','meeting_room','room','salon2'
# cmds.append(f'python train_basis.py configs/colmap_new.yaml defaults.expname=indoor_{scene} ' \
# f'dataset.datadir=./data/indoor/real/{scene}'
# # f'defaults.ckpt=/home/anpei/code/NeuBasis2/log/basis_ship/basis_ship.th exporation.render_only=True'
# )
# cmds = []
# cmds.append(f'python 2D_regression.py configs/image.yaml defaults.expname=NeRF model.basis_type=x model.coeff_type=none model.basis_mapping=trigonometric ' \
# f'model.num_layers=8 model.hidden_dim=256 ' \
# f'model.freq_bands=[1.,2.,4.,8.,16.,32.,64,128,256.,512.] model.basis_dims=[1,1,1,1,1,1,1,1,1,1] model.basis_resos=[1024,512,256,128,64,32,16,8,4,2]')
# cmds.append(f'python 2D_regression.py configs/image.yaml defaults.expname=NeuBasis-grid')
# cmds.append(f'python 2D_regression.py defaults.expname=NeuBasis-mlpB model.basis_type=mlp')
# cmds.append(f'python 2D_regression.py defaults.expname=NeuBasis-mlpC model.coeff_type=mlp')
# cmds.append(f'python 2D_regression.py configs/image.yaml defaults.expname=DVGO-like model.basis_type=none')
# cmds.append(f'python 2D_regression.py configs/image.yaml defaults.expname=NeuBasis-noC model.coeff_type=none')
# cmds.append(f'python 2D_regression.py configs/image.yaml defaults.expname=NeuBasis-sinc model.basis_mapping=sinc')
# cmds.append(f'python 2D_regression.py configs/image.yaml defaults.expname=NeuBasis-tria model.basis_mapping=triangle')
# cmds.append(f'python 2D_regression.py configs/image.yaml defaults.expname=NeuBasis-SL model.basis_dims=[144] model.basis_resos=[14] model.freq_bands=[73.14]')
# cmds.append(f'python 2D_regression.py configs/image.yaml defaults.expname=NeuBasis-DCT model.basis_type=fix-grid')
# cmds.append(f'python 2D_regression.py configs/image.yaml defaults.expname=NeuBasis-CP model.coeff_type=vec model.basis_type=cp \
# model.freq_bands=[1.,1.,1.,1.,1.,1.] model.basis_resos=[1024,1024,1024,1024,1024,1024] model.basis_dims=[64,64,64,32,32,32]')
# cmds.append(f'python 2D_regression.py configs/image.yaml defaults.expname=iNGP-like model.basis_type=hash model.coeff_type=none')
# cmds.append(f'python 2D_regression.py configs/image.yaml defaults.expname=NeuBasis-hash model.basis_type=hash model.coef_init=0.1 basis_dims=[16,16,16,16,16,16]')
#setting available gpus
gpu_idx = [0]
gpus_que = queue.Queue(len(gpu_idx))
for i in gpu_idx:
gpus_que.put(i)
# os.makedirs(f"log/{expFolder}", exist_ok=True)
def run_program(gpu, cmd):
cmd = f'{cmd} '
print(cmd)
os.system(cmd)
gpus_que.put(gpu)
ths = []
for i in range(len(cmds)):
gpu = gpus_que.get()
t = threading.Thread(target=run_program, args=(gpu, cmds[i]), daemon=True)
t.start()
ths.append(t)
for th in ths:
th.join()
# import os
# import numpy as np
# root = f'/mnt/qb/home/geiger/zyu30/Projects/Anpei/Code/factor-fields/logs/'
# # root = '/cluster/home/anchen/root/Code/NeuBasis/log/'
# scores = []
# # for scene in ['ship', 'mic', 'chair', 'lego', 'drums', 'ficus', 'hotdog', 'materials']:
# for scene in ['Caterpillar','Family','Ignatius','Truck']:
# scores.append(np.loadtxt(f'{root}/{scene}-grid/imgs_test_all/mean.txt'))
# # os.system(f'cp {root}/{scene}-grid/imgs_test_all/video.mp4 /mnt/qb/home/geiger/zyu30/Projects/Anpei/Code/factor-fields/logs/video/{scene}.mp4')
# os.system(f'cp {root}/{scene}-grid/{scene}-grid/imgs_path_all/video.mp4 /mnt/qb/home/geiger/zyu30/Projects/Anpei/Code/factor-fields/logs/video/{scene}.mp4')
# # print(np.mean(np.stack(scores),axis=0))
================================================
FILE: scripts/2D_regression.ipynb
================================================
{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"id": "f969c229-5a8a-44b6-91a3-bba55968b202",
"metadata": {},
"outputs": [],
"source": [
"import torch,imageio,sys,time\n",
"import numpy as np\n",
"from tqdm import tqdm\n",
"import matplotlib.pyplot as plt\n",
"from omegaconf import OmegaConf\n",
"\n",
"sys.path.append('..')\n",
"from models.FactorFields import FactorFields \n",
"\n",
"from utils import SimpleSampler\n",
"from dataLoader import dataset_dict\n",
"from torch.utils.data import DataLoader\n",
"imageio.plugins.freeimage.download()\n",
"\n",
"device = 'cuda'\n",
"torch.cuda.set_device(0)\n",
"\n",
"%load_ext autoreload\n",
"%autoreload 2"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "b9acd258-ab23-489a-93a0-7bc799bbab24",
"metadata": {},
"outputs": [],
"source": [
"def PSNR(a,b):\n",
" if type(a).__module__ == np.__name__:\n",
" mse = np.mean((a-b)**2)\n",
" else:\n",
" mse = torch.mean((a-b)**2).item()\n",
" psnr = -10.0 * np.log(mse) / np.log(10.0)\n",
" return psnr\n",
"\n",
"@torch.no_grad()\n",
"def eval_img(reso, chunk=10240):\n",
" y = torch.arange(0, reso[0])\n",
" x = torch.arange(0, reso[1])\n",
" yy, xx = torch.meshgrid((y, x), indexing='ij')\n",
" res = []\n",
" \n",
" coordiantes = torch.stack((xx,yy),dim=-1).reshape(-1,2) + 0.5 #/(torch.FloatTensor(reso[::-1])-1)*2-1\n",
" coordiantes = torch.split(coordiantes,chunk,dim=0)\n",
" for coordiante in tqdm(coordiantes):\n",
"\n",
" feats,_ = model.get_coding(coordiante.to(model.device))\n",
" y_recon = model.linear_mat(feats)\n",
" \n",
" res.append(y_recon.cpu())\n",
" return torch.cat(res).reshape(reso[0],reso[1],-1)\n",
"\n",
"def srgb_to_linear(img):\n",
"\tlimit = 0.04045\n",
"\treturn np.where(img > limit, np.power((img + 0.055) / 1.055, 2.4), img / 12.92)"
]
},
{
"cell_type": "markdown",
"id": "ee3a42f2-3314-4cf4-ba64-a015cd84d96f",
"metadata": {},
"source": [
"# Data loader\n",
"### please install the av package with \"pip install av\" if raise a error \"pyav: pip install imageio[pyav]\""
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "9d2a3772-7637-4087-8890-d5e15122ffa3",
"metadata": {},
"outputs": [],
"source": [
"base_conf = OmegaConf.load('../configs/defaults.yaml')\n",
"second_conf = OmegaConf.load('../configs/image.yaml')\n",
"cfg = OmegaConf.merge(base_conf, second_conf)\n",
"\n",
"dataset = dataset_dict[cfg.dataset.dataset_name]\n",
"tolinear = False if cfg.dataset.datadir.endswith('exr') else True\n",
"train_dataset = dataset(cfg.dataset, cfg.training.batch_size, split='train', tolinear=tolinear)\n",
"train_loader = DataLoader(train_dataset,\n",
" num_workers=8,\n",
" persistent_workers=True,\n",
" batch_size=None,\n",
" pin_memory=True)\n",
"\n",
"cfg.model.out_dim = train_dataset.img.shape[-1]\n",
"batch_size = cfg.training.batch_size\n",
"n_iter = cfg.training.n_iters\n",
"\n",
"H,W = train_dataset.HW\n",
"cfg.dataset.aabb = train_dataset.scene_bbox"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "634c9519-4be6-47b8-b030-616bdc4c6237",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"=====> total parameters: 1429376\n",
"FactorFields(\n",
" (coeffs): ParameterList( (0): Parameter containing: [torch.cuda.FloatTensor of size 1x144x63x63 (GPU 0)])\n",
" (basises): ParameterList(\n",
" (0): Parameter containing: [torch.cuda.FloatTensor of size 1x32x32x32 (GPU 0)]\n",
" (1): Parameter containing: [torch.cuda.FloatTensor of size 1x32x51x51 (GPU 0)]\n",
" (2): Parameter containing: [torch.cuda.FloatTensor of size 1x32x70x70 (GPU 0)]\n",
" (3): Parameter containing: [torch.cuda.FloatTensor of size 1x16x89x89 (GPU 0)]\n",
" (4): Parameter containing: [torch.cuda.FloatTensor of size 1x16x108x108 (GPU 0)]\n",
" (5): Parameter containing: [torch.cuda.FloatTensor of size 1x16x128x128 (GPU 0)]\n",
" )\n",
" (linear_mat): MLPMixer(\n",
" (backbone): ModuleList(\n",
" (0): Linear(in_features=144, out_features=64, bias=True)\n",
" (1): Linear(in_features=64, out_features=4, bias=False)\n",
" )\n",
" )\n",
")\n",
"total parameters: 1429376\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"Iteration 09999: loss_dist = 0.00000031 psnr = 65.022: 100%|██████████| 10000/10000 [01:32<00:00, 108.54it/s]\n",
"100%|████████████████████████████████████████████████████████████████████| 103/103 [00:00<00:00, 1117.18it/s]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"65.70337395134173 92.13329100608826\n"
]
},
{
"data": {
"text/plain": [
"╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮\n", "│ in <module>:78 │\n", "│ │\n", "│ 75 │ for i in range(model.coeffs[0].shape[1]): │\n", "│ 76 │ │ imageio.mimwrite('/mnt/qb/home/geiger/zyu30/Projects/Anpei/FactorFields/video/im │\n", "│ 77 │ │ │\n", "│ ❱ 78 │ │ os.makedirs(f'/mnt/qb/home/geiger/zyu30/Projects/Anpei/FactorFields/video/image/ │\n", "│ 79 │ │ stream = ffmpeg.input('/mnt/qb/home/geiger/zyu30/Projects/Anpei/FactorFields/vid │\n", "│ 80 │ │ stream = ffmpeg.output(stream, f'/mnt/qb/home/geiger/zyu30/Projects/Anpei/Factor │\n", "│ 81 │ │ ffmpeg.run(stream,overwrite_output=True) │\n", "╰──────────────────────────────────────────────────────────────────────────────────────────────────╯\n", "NameError: name 'os' is not defined\n", "\n" ], "text/plain": [ "\u001b[31m╭─\u001b[0m\u001b[31m──────────────────────────────\u001b[0m\u001b[31m \u001b[0m\u001b[1;31mTraceback \u001b[0m\u001b[1;2;31m(most recent call last)\u001b[0m\u001b[31m \u001b[0m\u001b[31m───────────────────────────────\u001b[0m\u001b[31m─╮\u001b[0m\n", "\u001b[31m│\u001b[0m in \u001b[92m