otthorn
/
nerf_plus_plus
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

first commit

Browse Source
master
Kai-46 4 years ago
commit 5fbe15ff24
58 changed files with 4470 additions and 0 deletions
Show Stats Download Patch File Download Diff File

8
.idea/.gitignore vendored
Unescape Escape View File

@ -0,0 +1,8 @@
# Default ignored files
/shelf/
/workspace.xml
# Datasource local storage ignored files
/dataSources/
/dataSources.local.xml
# Editor-based HTTP Client requests
/httpRequests/

15
.idea/deployment.xml
Unescape Escape View File

@ -0,0 +1,15 @@
<?xml version="1.0" encoding="UTF-8"?>
<project version="4">
<component name="PublishConfigData" autoUpload="Always" serverName="nerf_bg_latest_ddp">
<serverData>
<paths name="nerf_bg_latest_ddp">
<serverdata>
<mappings>
<mapping deploy="/home/zhangka2/gernot_experi/nerf_bg_latest_ddp" local="$PROJECT_DIR$" web="/" />
</mappings>
</serverdata>
</paths>
</serverData>
<option name="myAutoUpload" value="ALWAYS" />
</component>
</project>

6
.idea/inspectionProfiles/profiles_settings.xml
Unescape Escape View File

@ -0,0 +1,6 @@
<component name="InspectionProjectProfileManager">
<settings>
<option name="USE_PROJECT_PROFILE" value="false" />
<version value="1.0" />
</settings>
</component>

7
.idea/misc.xml
Unescape Escape View File

@ -0,0 +1,7 @@
<?xml version="1.0" encoding="UTF-8"?>
<project version="4">
<component name="JavaScriptSettings">
<option name="languageLevel" value="ES6" />
</component>
<component name="ProjectRootManager" version="2" project-jdk-name="Python 3.7 (nerf)" project-jdk-type="Python SDK" />
</project>

8
.idea/modules.xml
Unescape Escape View File

@ -0,0 +1,8 @@
<?xml version="1.0" encoding="UTF-8"?>
<project version="4">
<component name="ProjectModuleManager">
<modules>
<module fileurl="file://$PROJECT_DIR$/.idea/nerf_bg_latest_ddp.iml" filepath="$PROJECT_DIR$/.idea/nerf_bg_latest_ddp.iml" />
</modules>
</component>
</project>

13
.idea/nerf_bg_latest_ddp.iml
Unescape Escape View File

@ -0,0 +1,13 @@
<?xml version="1.0" encoding="UTF-8"?>
<module type="PYTHON_MODULE" version="4">
<component name="NewModuleRootManager">
<content url="file://$MODULE_DIR$" />
<orderEntry type="jdk" jdkName="Python 3.7 (nerf)" jdkType="Python SDK" />
<orderEntry type="sourceFolder" forTests="false" />
</component>
<component name="PyDocumentationSettings">
<option name="format" value="PLAIN" />
<option name="myDocStringFormat" value="Plain" />
<option name="renderExternalDocumentation" value="true" />
</component>
</module>

7
.idea/other.xml
Unescape Escape View File

@ -0,0 +1,7 @@
<?xml version="1.0" encoding="UTF-8"?>
<project version="4">
<component name="PySciProjectComponent">
<option name="PY_SCI_VIEW" value="true" />
<option name="PY_SCI_VIEW_SUGGESTED" value="true" />
</component>
</project>

14
.idea/webServers.xml
Unescape Escape View File

@ -0,0 +1,14 @@
<?xml version="1.0" encoding="UTF-8"?>
<project version="4">
<component name="WebServers">
<option name="servers">
<webServer id="1d5a6596-7d8b-45af-9960-9d4d014e6bbe" name="nerf_bg_latest_ddp">
<fileTransfer accessType="SFTP" host="isl-iam2.rr.intel.com" port="22" sshConfigId="d5ddaa0b-8e6c-4721-ad8e-298dc2859ce7" sshConfig="intel_cluster" keyPair="true">
<advancedOptions>
<advancedOptions dataProtectionLevel="Private" passiveMode="true" shareSSLContext="true" />
</advancedOptions>
</fileTransfer>
</webServer>
</option>
</component>
</project>

0
README.md
Unescape Escape View File

BIN
__pycache__/utils.cpython-37.pyc
View File

Binary file not shown.

162
conda_env_nerf.yml
Unescape Escape View File

@ -0,0 +1,162 @@
name: nerf
channels:
- pytorch
- conda-forge
- defaults
dependencies:
- _libgcc_mutex=0.1=conda_forge
- _openmp_mutex=4.5=0_gnu
- absl-py=0.9.0=py37_0
- astor=0.7.1=py_0
- boost=1.72.0=py37h9de70de_0
- boost-cpp=1.72.0=h8e57a91_0
- bzip2=1.0.8=h516909a_2
- c-ares=1.15.0=h516909a_1001
- ca-certificates=2020.4.5.1=hecc5488_0
- cairo=1.16.0=hfb77d84_1002
- certifi=2020.4.5.1=py37hc8dfbb8_0
- cffi=1.13.2=py37h8022711_0
- cloudpickle=1.3.0=py_0
- cudatoolkit=10.0.130=0
- cycler=0.10.0=py_2
- cytoolz=0.10.1=py37h516909a_0
- dask-core=2.10.1=py_0
- dbus=1.13.6=he372182_0
- decorator=4.4.1=py_0
- expat=2.2.9=he1b5a44_2
- ffmpeg=4.1.3=h167e202_0
- fontconfig=2.13.1=h86ecdb6_1001
- freetype=2.10.0=he983fc9_1
- gast=0.3.3=py_0
- gettext=0.19.8.1=hc5be6a0_1002
- giflib=5.2.1=h516909a_1
- glib=2.58.3=py37h6f030ca_1002
- gmp=6.2.0=he1b5a44_1
- gnutls=3.6.5=hd3a4fd2_1002
- graphite2=1.3.13=hf484d3e_1000
- grpcio=1.23.0=py37hb0870dc_1
- gst-plugins-base=1.14.5=h0935bb2_2
- gstreamer=1.14.5=h36ae1b5_2
- h5py=2.10.0=nompi_py37h513d04c_102
- harfbuzz=2.4.0=h9f30f68_3
- hdf5=1.10.5=nompi_h3c11f04_1104
- icu=64.2=he1b5a44_1
- ilmbase=2.4.1=h8b12597_0
- imageio=2.6.1=py37_0
- intel-openmp=2020.0=166
- jasper=1.900.1=h07fcdf6_1006
- joblib=0.14.1=py_0
- jpeg=9c=h14c3975_1001
- keras-applications=1.0.8=py_1
- keras-preprocessing=1.1.0=py_0
- kiwisolver=1.1.0=py37hc9558a2_0
- krb5=1.16.4=h2fd8d38_0
- lame=3.100=h14c3975_1001
- ld_impl_linux-64=2.33.1=h53a641e_8
- libblas=3.8.0=14_openblas
- libcblas=3.8.0=14_openblas
- libclang=9.0.1=default_hde54327_0
- libcurl=7.68.0=hda55be3_0
- libedit=3.1.20170329=hf8c457e_1001
- libffi=3.2.1=he1b5a44_1006
- libgcc-ng=9.2.0=h24d8f2e_2
- libgfortran-ng=7.3.0=hdf63c60_5
- libgomp=9.2.0=h24d8f2e_2
- libiconv=1.15=h516909a_1005
- liblapack=3.8.0=14_openblas
- liblapacke=3.8.0=14_openblas
- libllvm9=9.0.1=hc9558a2_0
- libopenblas=0.3.7=h5ec1e0e_6
- libopencv=4.2.0=py37_2
- libpng=1.6.37=hed695b0_0
- libprotobuf=3.8.0=h8b12597_0
- libssh2=1.8.2=h22169c7_2
- libstdcxx-ng=9.2.0=hdf63c60_2
- libtiff=4.1.0=hc3755c2_3
- libuuid=2.32.1=h14c3975_1000
- libwebp=1.0.2=h56121f0_5
- libxcb=1.13=h14c3975_1002
- libxkbcommon=0.10.0=he1b5a44_0
- libxml2=2.9.10=hee79883_0
- lz4-c=1.8.3=he1b5a44_1001
- markdown=3.2.1=py_0
- matplotlib-base=3.1.3=py37h250f245_0
- mkl=2020.0=166
- mock=3.0.5=py37_0
- ncurses=6.1=hf484d3e_1002
- nettle=3.4.1=h1bed415_1002
- networkx=2.4=py_0
- ninja=1.10.0=hc9558a2_0
- nspr=4.25=he1b5a44_0
- nss=3.47=he751ad9_0
- numpy=1.18.1=py37h95a1406_0
- olefile=0.46=py_0
- opencv=4.2.0=py37_2
- openh264=1.8.0=hdbcaa40_1000
- openimageio=2.1.13=hf311ebb_0
- openssl=1.1.1g=h516909a_0
- pandas=1.0.1=py37hb3f55d8_0
- patsy=0.5.1=py_0
- pcre=8.44=he1b5a44_0
- pillow=6.2.1=py37hd70f55b_1
- pip=20.0.2=py_2
- pixman=0.38.0=h516909a_1003
- protobuf=3.8.0=py37he1b5a44_2
- pthread-stubs=0.4=h14c3975_1001
- py-opencv=4.2.0=py37h5ca1d4c_2
- py-openimageio=2.1.13=py37hf311ebb_0
- pycparser=2.19=py37_1
- pyparsing=2.4.6=py_0
- python=3.7.6=h357f687_2
- python-dateutil=2.8.1=py_0
- python_abi=3.7=1_cp37m
- pytorch=1.0.1=py3.7_cuda10.0.130_cudnn7.4.2_2
- pytz=2019.3=py_0
- pywavelets=1.1.1=py37hc1659b7_0
- qt=5.12.5=hd8c4c69_1
- readline=8.0=hf8c457e_0
- scikit-image=0.16.2=py37hb3f55d8_0
- scikit-learn=0.22.1=py37hcdab131_1
- scipy=1.4.1=py37h921218d_0
- seaborn=0.10.0=py_1
- setuptools=45.2.0=py37_0
- six=1.14.0=py37_0
- sqlite=3.30.1=hcee41ef_0
- statsmodels=0.11.0=py37h516909a_0
- tensorboard=1.13.1=py37_0
- tensorboardx=2.0=py_0
- tensorflow=1.13.1=h5ece82f_5
- tensorflow-base=1.13.1=py37h5ece82f_5
- tensorflow-estimator=1.13.0=py_0
- termcolor=1.1.0=py_2
- tk=8.6.10=hed695b0_0
- toolz=0.10.0=py_0
- torchvision=0.2.2=py_3
- tornado=6.0.3=py37h516909a_4
- tqdm=4.42.1=py_0
- werkzeug=1.0.0=py_0
- wheel=0.34.2=py_1
- x264=1!152.20180806=h14c3975_0
- xorg-kbproto=1.0.7=h14c3975_1002
- xorg-libice=1.0.10=h516909a_0
- xorg-libsm=1.2.3=h84519dc_1000
- xorg-libx11=1.6.9=h516909a_0
- xorg-libxau=1.0.9=h14c3975_0
- xorg-libxdmcp=1.1.3=h516909a_0
- xorg-libxext=1.3.4=h516909a_0
- xorg-libxrender=0.9.10=h516909a_1002
- xorg-renderproto=0.11.1=h14c3975_1002
- xorg-xextproto=7.3.0=h14c3975_1002
- xorg-xproto=7.0.31=h14c3975_1007
- xz=5.2.4=h14c3975_1001
- zlib=1.2.11=h516909a_1006
- zstd=1.4.4=h3b9ef0a_1
- pip:
- configargparse==1.2.3
- future==0.18.2
- imageio-ffmpeg==0.4.1
- openexr==1.3.2
- pyexr==0.3.7
- pymcubes==0.1.0
- pyquaternion==0.9.5
prefix: /home/kz298/anaconda3/envs/nerf

48
configs/lf_data/lf_africa.txt
Unescape Escape View File

@ -0,0 +1,48 @@
### INPUT
datadir = /home/zhangka2/gernot_experi/lf_data/lf_nerf
scene = africa
expname = africa_ddp
basedir = ./logs
config = None
ckpt_path = None
no_reload = False
testskip = 1
### TRAINING
N_iters = 1250001
N_rand = 512
lrate = 0.0005
lrate_decay_factor = 0.1
lrate_decay_steps = 50000000
### CASCADE
cascade_level = 2
cascade_samples = 64,128
near_depth = 0.
far_depth = 1.
### TESTING
render_only = False
render_test = False
render_train = False
# chunk_size = 16384
chunk_size = 4096
### RENDERING
det = False
max_freq_log2 = 10
max_freq_log2_viewdirs = 4
netdepth = 8
netwidth = 256
raw_noise_std = 1.0
N_iters_perturb = 1000
inv_uniform = False
use_viewdirs = True
white_bkgd = False
### CONSOLE AND TENSORBOARD
i_img = 2000
i_print = 100
i_testset = 5000000
i_video = 5000000
i_weights = 5000

48
configs/lf_data/lf_basket.txt
Unescape Escape View File

@ -0,0 +1,48 @@
### INPUT
datadir = /home/zhangka2/gernot_experi/lf_data/lf_nerf
scene = basket
expname = basket_ddp
basedir = ./logs
config = None
ckpt_path = None
no_reload = False
testskip = 1
### TRAINING
N_iters = 1250001
N_rand = 512
lrate = 0.0005
lrate_decay_factor = 0.1
lrate_decay_steps = 50000000
### CASCADE
cascade_level = 2
cascade_samples = 64,128
near_depth = 0.
far_depth = 1.
### TESTING
render_only = False
render_test = False
render_train = False
# chunk_size = 16384
chunk_size = 4096
### RENDERING
det = False
max_freq_log2 = 10
max_freq_log2_viewdirs = 4
netdepth = 8
netwidth = 256
raw_noise_std = 1.0
N_iters_perturb = 1000
inv_uniform = False
use_viewdirs = True
white_bkgd = False
### CONSOLE AND TENSORBOARD
i_img = 2000
i_print = 100
i_testset = 5000000
i_video = 5000000
i_weights = 5000

50
configs/lf_data/lf_ship.txt
Unescape Escape View File

@ -0,0 +1,50 @@
### INPUT
datadir = /home/zhangka2/gernot_experi/lf_data/lf_nerf
scene = ship
expname = ship_ddp
basedir = ./logs
config = None
ckpt_path = None
no_reload = False
testskip = 1
### TRAINING
N_iters = 1250001
# N_rand = 512
N_rand = 1024
lrate = 0.0005
lrate_decay_factor = 0.1
lrate_decay_steps = 50000000
### CASCADE
cascade_level = 2
cascade_samples = 64,128
near_depth = 0.
far_depth = 1.
### TESTING
render_only = False
render_test = False
render_train = False
# chunk_size = 16384
# chunk_size = 4096
chunk_size = 8192
### RENDERING
det = False
max_freq_log2 = 10
max_freq_log2_viewdirs = 4
netdepth = 8
netwidth = 256
raw_noise_std = 1.0
N_iters_perturb = 1000
inv_uniform = False
use_viewdirs = True
white_bkgd = False
### CONSOLE AND TENSORBOARD
i_img = 2000
i_print = 100
i_testset = 5000000
i_video = 5000000
i_weights = 5000

48
configs/lf_data/lf_torch.txt
Unescape Escape View File

@ -0,0 +1,48 @@
### INPUT
datadir = /home/zhangka2/gernot_experi/lf_data/lf_nerf
scene = torch
expname = torch_ddp
basedir = ./logs
config = None
ckpt_path = None
no_reload = False
testskip = 1
### TRAINING
N_iters = 1250001
N_rand = 512
lrate = 0.0005
lrate_decay_factor = 0.1
lrate_decay_steps = 50000000
### CASCADE
cascade_level = 2
cascade_samples = 64,128
near_depth = 0.
far_depth = 1.
### TESTING
render_only = False
render_test = False
render_train = False
# chunk_size = 16384
chunk_size = 4096
### RENDERING
det = False
max_freq_log2 = 10
max_freq_log2_viewdirs = 4
netdepth = 8
netwidth = 256
raw_noise_std = 1.0
N_iters_perturb = 1000
inv_uniform = False
use_viewdirs = True
white_bkgd = False
### CONSOLE AND TENSORBOARD
i_img = 2000
i_print = 100
i_testset = 5000000
i_video = 5000000
i_weights = 5000

49
configs/tanks_and_temples/tat_intermediate_m60.txt
Unescape Escape View File

@ -0,0 +1,49 @@
### INPUT
datadir = /home/zhangka2/gernot_experi/gernot_data/gernot_nerf_sphere
scene = tat_intermediate_M60
expname = tat_intermediate_M60_bg_carve_latest
basedir = ./logs
config = None
ckpt_path = None
no_reload = False
testskip = 1
### TRAINING
N_iters = 1250001
# N_rand = 4096
N_rand = 2048
lrate = 0.0005
lrate_decay_factor = 0.1
lrate_decay_steps = 500000
### CASCADE
cascade_level = 2
cascade_samples = 64,128
near_depth = 0.
far_depth = 1.
### TESTING
render_only = False
render_test = False
render_train = False
chunk_size = 16384
# chunk_size = 8192
### RENDERING
det = False
max_freq_log2 = 10
max_freq_log2_viewdirs = 4
netdepth = 8
netwidth = 256
raw_noise_std = 1.0
N_iters_perturb = 1000
inv_uniform = False
use_viewdirs = True
white_bkgd = False
### CONSOLE AND TENSORBOARD
i_img = 2000
i_print = 100
i_testset = 5000000
i_video = 5000000
i_weights = 5000

49
configs/tanks_and_temples/tat_intermediate_playground.txt
Unescape Escape View File

@ -0,0 +1,49 @@
### INPUT
datadir = /home/zhangka2/gernot_experi/gernot_data/gernot_nerf_sphere
scene = tat_intermediate_Playground
expname = tat_intermediate_Playground_bg_carve_latest
basedir = ./logs
config = None
ckpt_path = None
no_reload = False
testskip = 1
### TRAINING
N_iters = 1250001
# N_rand = 4096
N_rand = 2048
lrate = 0.0005
lrate_decay_factor = 0.1
lrate_decay_steps = 500000
### CASCADE
cascade_level = 2
cascade_samples = 64,128
near_depth = 0.
far_depth = 1.
### TESTING
render_only = False
render_test = False
render_train = False
chunk_size = 16384
# chunk_size = 8192
### RENDERING
det = False
max_freq_log2 = 10
max_freq_log2_viewdirs = 4
netdepth = 8
netwidth = 256
raw_noise_std = 1.0
N_iters_perturb = 1000
inv_uniform = False
use_viewdirs = True
white_bkgd = False
### CONSOLE AND TENSORBOARD
i_img = 2000
i_print = 100
i_testset = 5000000
i_video = 5000000
i_weights = 5000

48
configs/tanks_and_temples/tat_intermediate_playground_bignet.txt
Unescape Escape View File

@ -0,0 +1,48 @@
### INPUT
datadir = /home/zhangka2/gernot_experi/gernot_data/gernot_nerf_sphere
scene = tat_intermediate_Playground
expname = tat_intermediate_Playground_ddp_bignet
basedir = ./logs
config = None
ckpt_path = None
no_reload = False
testskip = 1
### TRAINING
N_iters = 1250001
N_rand = 256
lrate = 0.0005
lrate_decay_factor = 0.1
lrate_decay_steps = 500000
### CASCADE
cascade_level = 2
cascade_samples = 64,128
near_depth = 0.
far_depth = 1.
### TESTING
render_only = False
render_test = False
render_train = False
# chunk_size = 16384
chunk_size = 4096
### RENDERING
det = False
max_freq_log2 = 10
max_freq_log2_viewdirs = 4
netdepth = 8
netwidth = 512
raw_noise_std = 1.0
N_iters_perturb = 1000
inv_uniform = False
use_viewdirs = True
white_bkgd = False
### CONSOLE AND TENSORBOARD
i_img = 2000
i_print = 100
i_testset = 5000000
i_video = 5000000
i_weights = 5000

49
configs/tanks_and_temples/tat_intermediate_train.txt
Unescape Escape View File

@ -0,0 +1,49 @@
### INPUT
datadir = /home/zhangka2/gernot_experi/gernot_data/gernot_nerf_sphere
scene = tat_intermediate_Train
expname = tat_intermediate_Train_bg_carve_latest
basedir = ./logs
config = None
ckpt_path = None
no_reload = False
testskip = 1
### TRAINING
N_iters = 1250001
# N_rand = 4096
N_rand = 2048
lrate = 0.0005
lrate_decay_factor = 0.1
lrate_decay_steps = 500000
### CASCADE
cascade_level = 2
cascade_samples = 64,128
near_depth = 0.
far_depth = 1.
### TESTING
render_only = False
render_test = False
render_train = False
chunk_size = 16384
# chunk_size = 8192
### RENDERING
det = False
max_freq_log2 = 10
max_freq_log2_viewdirs = 4
netdepth = 8
netwidth = 256
raw_noise_std = 1.0
N_iters_perturb = 1000
inv_uniform = False
use_viewdirs = True
white_bkgd = False
### CONSOLE AND TENSORBOARD
i_img = 2000
i_print = 100
i_testset = 5000000
i_video = 5000000
i_weights = 5000

53
configs/tanks_and_temples/tat_training_truck.txt
Unescape Escape View File

@ -0,0 +1,53 @@
### INPUT
datadir = /home/zhangka2/gernot_experi/gernot_data/gernot_nerf_sphere
scene = tat_training_Truck
expname = tat_training_Truck_ddp_implicit
basedir = ./logs
config = None
ckpt_path = None
no_reload = False
testskip = 1
### TRAINING
N_iters = 1250001
# N_rand = 512
N_rand = 1024
lrate = 0.0005
lrate_decay_factor = 0.1
lrate_decay_steps = 50000000
### implicit
use_implicit = True
### CASCADE
cascade_level = 2
cascade_samples = 64,128
near_depth = 0.
far_depth = 1.
### TESTING
render_only = False
render_test = False
render_train = False
# chunk_size = 16384
# chunk_size = 4096
chunk_size = 8192
### RENDERING
det = False
max_freq_log2 = 10
max_freq_log2_viewdirs = 4
netdepth = 8
netwidth = 256
raw_noise_std = 1.0
N_iters_perturb = 1000
inv_uniform = False
use_viewdirs = True
white_bkgd = False
### CONSOLE AND TENSORBOARD
i_img = 2000
i_print = 100
i_testset = 5000000
i_video = 5000000
i_weights = 5000

48
configs/tanks_and_temples/tat_training_truck_bignet.txt
Unescape Escape View File

@ -0,0 +1,48 @@
### INPUT
datadir = /home/zhangka2/gernot_experi/gernot_data/gernot_nerf_sphere
scene = tat_training_Truck
expname = tat_training_Truck_ddp_bignet
basedir = ./logs
config = None
ckpt_path = None
no_reload = False
testskip = 1
### TRAINING
N_iters = 1250001
N_rand = 256
lrate = 0.0005
lrate_decay_factor = 0.1
lrate_decay_steps = 500000
### CASCADE
cascade_level = 2
cascade_samples = 64,128
near_depth = 0.
far_depth = 1.
### TESTING
render_only = False
render_test = False
render_train = False
# chunk_size = 16384
chunk_size = 4096
### RENDERING
det = False
max_freq_log2 = 10
max_freq_log2_viewdirs = 4
netdepth = 8
netwidth = 512
raw_noise_std = 1.0
N_iters_perturb = 1000
inv_uniform = False
use_viewdirs = True
white_bkgd = False
### CONSOLE AND TENSORBOARD
i_img = 2000
i_print = 100
i_testset = 5000000
i_video = 5000000
i_weights = 5000

47
configs/tanks_and_temples/tat_training_truck_subset.txt
Unescape Escape View File

@ -0,0 +1,47 @@
### INPUT
datadir = /home/zhangka2/gernot_experi/gernot_data/gernot_nerf_sphere
scene = tat_training_Truck_subset
expname = tat_training_Truck_subset_bg_carvenew
basedir = ./logs
config = None
ckpt_path = None
no_reload = False
testskip = 1
### TRAINING
N_iters = 250001
N_rand = 2048
lrate = 0.0005
lrate_decay_factor = 0.1
lrate_decay_steps = 500000
### CASCADE
cascade_level = 2
cascade_samples = 64,64
near_depth = 0.
far_depth = 1.
### TESTING
render_only = False
render_test = False
render_train = False
chunk_size = 8192
### RENDERING
det = False
max_freq_log2 = 10
max_freq_log2_viewdirs = 4
netdepth = 8
netwidth = 256
raw_noise_std = 1.0
N_iters_perturb = 1000
inv_uniform = False
use_viewdirs = False
white_bkgd = False
### CONSOLE AND TENSORBOARD
i_img = 2000
i_print = 100
i_testset = 5000000
i_video = 5000000
i_weights = 5000

54
configs/tanks_and_temples_sparse/tat_intermediate_playground_addcarve.txt
Unescape Escape View File

@ -0,0 +1,54 @@
### INPUT
datadir = /home/zhangka2/gernot_experi/gernot_data/gernot_nerf_sphere_sparse
scene = tat_intermediate_Playground
expname = tat_intermediate_Playground_ddp_sparse_addcarve
basedir = ./logs
config = None
ckpt_path = None
no_reload = False
testskip = 1
### TRAINING
N_iters = 1250001
# N_rand = 4096
N_rand = 2048
lrate = 0.0005
lrate_decay_factor = 0.1
lrate_decay_steps = 500000
### implicit
use_implicit = True
load_min_depth = True
regularize_weight = 0.1
### CASCADE
cascade_level = 2
cascade_samples = 64,128
near_depth = 0.
far_depth = 1.
### TESTING
render_only = False
render_test = False
render_train = False
chunk_size = 16384
# chunk_size = 8192
### RENDERING
det = False
max_freq_log2 = 10
max_freq_log2_viewdirs = 4
netdepth = 8
netwidth = 256
raw_noise_std = 1.0
N_iters_perturb = 1000
inv_uniform = False
use_viewdirs = True
white_bkgd = False
### CONSOLE AND TENSORBOARD
i_img = 2000
i_print = 100
i_testset = 5000000
i_video = 5000000
i_weights = 5000

54
configs/tanks_and_temples_sparse/tat_intermediate_playground_addparam.txt
Unescape Escape View File

@ -0,0 +1,54 @@
### INPUT
datadir = /home/zhangka2/gernot_experi/gernot_data/gernot_nerf_sphere_sparse
scene = tat_intermediate_Playground
expname = tat_intermediate_Playground_ddp_sparse_addparam
basedir = ./logs
config = None
ckpt_path = None
no_reload = False
testskip = 1
### TRAINING
N_iters = 1250001
# N_rand = 4096
N_rand = 2048
lrate = 0.0005
lrate_decay_factor = 0.1
lrate_decay_steps = 500000
### implicit
use_implicit = True
load_min_depth = False
regularize_weight = 0.
### CASCADE
cascade_level = 2
cascade_samples = 64,128
near_depth = 0.
far_depth = 1.
### TESTING
render_only = False
render_test = False
render_train = False
chunk_size = 16384
# chunk_size = 8192
### RENDERING
det = False
max_freq_log2 = 10
max_freq_log2_viewdirs = 4
netdepth = 8
netwidth = 256
raw_noise_std = 1.0
N_iters_perturb = 1000
inv_uniform = False
use_viewdirs = True
white_bkgd = False
### CONSOLE AND TENSORBOARD
i_img = 2000
i_print = 100
i_testset = 5000000
i_video = 5000000
i_weights = 5000

54
configs/tanks_and_temples_sparse/tat_intermediate_playground_addregularize.txt
Unescape Escape View File

@ -0,0 +1,54 @@
### INPUT
datadir = /home/zhangka2/gernot_experi/gernot_data/gernot_nerf_sphere_sparse
scene = tat_intermediate_Playground
expname = tat_intermediate_Playground_ddp_sparse_addregularize_pretrain
basedir = ./logs
config = /home/zhangka2/gernot_experi/nerf_bg_latest_ddp/logs/tat_intermediate_Playground_ddp_sparse_addparam/model_210000.pth
ckpt_path = None
no_reload = False
testskip = 1
### TRAINING
N_iters = 1250001
# N_rand = 4096
N_rand = 2048
lrate = 0.0005
lrate_decay_factor = 0.1
lrate_decay_steps = 500000
### implicit
use_implicit = True
load_min_depth = False
regularize_weight = 0.1
### CASCADE
cascade_level = 2
cascade_samples = 64,128
near_depth = 0.
far_depth = 1.
### TESTING
render_only = False
render_test = False
render_train = False
chunk_size = 16384
# chunk_size = 8192
### RENDERING
det = False
max_freq_log2 = 10
max_freq_log2_viewdirs = 4
netdepth = 8
netwidth = 256
raw_noise_std = 1.0
N_iters_perturb = 1000
inv_uniform = False
use_viewdirs = True
white_bkgd = False
### CONSOLE AND TENSORBOARD
i_img = 2000
i_print = 100
i_testset = 5000000
i_video = 5000000
i_weights = 5000

55
configs/tanks_and_temples_sparse/tat_training_truck_addcarve.txt
Unescape Escape View File

@ -0,0 +1,55 @@
### INPUT
datadir = /home/zhangka2/gernot_experi/gernot_data/gernot_nerf_sphere_sparse
scene = tat_training_Truck
expname = tat_training_Truck_ddp_sparse_addcarve
basedir = ./logs
config = None
ckpt_path = None
no_reload = False
testskip = 1
### TRAINING
N_iters = 1250001
# N_rand = 512
N_rand = 1024
lrate = 0.0005
lrate_decay_factor = 0.1
lrate_decay_steps = 50000000
### implicit
use_implicit = True
load_min_depth = True
regularize_weight = 0.1
### CASCADE
cascade_level = 2
cascade_samples = 64,128
near_depth = 0.
far_depth = 1.
### TESTING
render_only = False
render_test = False
render_train = False
# chunk_size = 16384
# chunk_size = 4096
chunk_size = 8192
### RENDERING
det = False
max_freq_log2 = 10
max_freq_log2_viewdirs = 4
netdepth = 8
netwidth = 256
raw_noise_std = 1.0
N_iters_perturb = 1000
inv_uniform = False
use_viewdirs = True
white_bkgd = False
### CONSOLE AND TENSORBOARD
i_img = 2000
i_print = 100
i_testset = 5000000
i_video = 5000000
i_weights = 5000

55
configs/tanks_and_temples_sparse/tat_training_truck_addparam.txt
Unescape Escape View File

@ -0,0 +1,55 @@
### INPUT
datadir = /home/zhangka2/gernot_experi/gernot_data/gernot_nerf_sphere_sparse
scene = tat_training_Truck
expname = tat_training_Truck_ddp_sparse_addparam
basedir = ./logs
config = None
ckpt_path = None
no_reload = False
testskip = 1
### TRAINING
N_iters = 1250001
# N_rand = 512
N_rand = 1024
lrate = 0.0005
lrate_decay_factor = 0.1
lrate_decay_steps = 50000000
### implicit
use_implicit = True
load_min_depth = False
regularize_weight = 0.
### CASCADE
cascade_level = 2
cascade_samples = 64,128
near_depth = 0.
far_depth = 1.
### TESTING
render_only = False
render_test = False
render_train = False
# chunk_size = 16384
# chunk_size = 4096
chunk_size = 8192
### RENDERING
det = False
max_freq_log2 = 10
max_freq_log2_viewdirs = 4
netdepth = 8
netwidth = 256
raw_noise_std = 1.0
N_iters_perturb = 1000
inv_uniform = False
use_viewdirs = True
white_bkgd = False
### CONSOLE AND TENSORBOARD
i_img = 2000
i_print = 100
i_testset = 5000000
i_video = 5000000
i_weights = 5000

55
configs/tanks_and_temples_sparse/tat_training_truck_addregularize.txt
Unescape Escape View File

@ -0,0 +1,55 @@
### INPUT
datadir = /home/zhangka2/gernot_experi/gernot_data/gernot_nerf_sphere_sparse
scene = tat_training_Truck
expname = tat_training_Truck_ddp_sparse_addregularize_pretrain
basedir = ./logs
config = None
ckpt_path = /home/zhangka2/gernot_experi/nerf_bg_latest_ddp/logs/tat_training_Truck_ddp_sparse_addparam/model_245000.pth
no_reload = False
testskip = 1
### TRAINING
N_iters = 1250001
# N_rand = 512
N_rand = 1024
lrate = 0.0005
lrate_decay_factor = 0.1
lrate_decay_steps = 50000000
### implicit
use_implicit = True
load_min_depth = False
regularize_weight = 0.1
### CASCADE
cascade_level = 2
cascade_samples = 64,128
near_depth = 0.
far_depth = 1.
### TESTING
render_only = False
render_test = False
render_train = False
# chunk_size = 16384
# chunk_size = 4096
chunk_size = 8192
### RENDERING
det = False
max_freq_log2 = 10
max_freq_log2_viewdirs = 4
netdepth = 8
netwidth = 256
raw_noise_std = 1.0
N_iters_perturb = 1000
inv_uniform = False
use_viewdirs = True
white_bkgd = False
### CONSOLE AND TENSORBOARD
i_img = 2000
i_print = 100
i_testset = 5000000
i_video = 5000000
i_weights = 5000

94
data_loader_split.py
Unescape Escape View File

@ -0,0 +1,94 @@
import os
import numpy as np
import imageio
import logging
from nerf_sample_ray_split import RaySamplerSingleImage
import glob
logger = logging.getLogger(__package__)
########################################################################################################################
# camera coordinate system: x-->right, y-->down, z-->scene (opencv/colmap convention)
# poses is camera-to-world
########################################################################################################################
def find_files(dir, exts):
if os.path.isdir(dir):
# types should be ['*.png', '*.jpg']
files_grabbed = []
for ext in exts:
files_grabbed.extend(glob.glob(os.path.join(dir, ext)))
if len(files_grabbed) > 0:
files_grabbed = sorted(files_grabbed)
return files_grabbed
else:
return []
def load_data_split(basedir, scene, split, skip=1, try_load_min_depth=True):
def parse_txt(filename):
assert os.path.isfile(filename)
nums = open(filename).read().split()
return np.array([float(x) for x in nums]).reshape([4, 4]).astype(np.float32)
split_dir = '{}/{}/{}'.format(basedir, scene, split)
intrinsics_files = find_files('{}/intrinsics'.format(split_dir), exts=['*.txt'])
pose_files = find_files('{}/pose'.format(split_dir), exts=['*.txt'])
logger.info('raw intrinsics_files: {}'.format(len(intrinsics_files)))
logger.info('raw pose_files: {}'.format(len(pose_files)))
intrinsics_files = intrinsics_files[::skip]
pose_files = pose_files[::skip]
cam_cnt = len(pose_files)
# img files
img_files = find_files('{}/rgb'.format(split_dir), exts=['*.png', '*.jpg'])
if len(img_files) > 0:
logger.info('raw img_files: {}'.format(len(img_files)))
img_files = img_files[::skip]
assert(len(img_files) == cam_cnt)
else:
img_files = [None, ] * cam_cnt
# mask files
mask_files = find_files('{}/mask'.format(split_dir), exts=['*.png', '*.jpg'])
if len(mask_files) > 0:
logger.info('raw mask_files: {}'.format(len(mask_files)))
mask_files = mask_files[::skip]
assert(len(mask_files) == cam_cnt)
else:
mask_files = [None, ] * cam_cnt
# min depth files
mindepth_files = find_files('{}/min_depth'.format(split_dir), exts=['*.png', '*.jpg'])
if try_load_min_depth and len(mindepth_files) > 0:
logger.info('raw mindepth_files: {}'.format(len(mindepth_files)))
mindepth_files = mindepth_files[::skip]
assert(len(mindepth_files) == cam_cnt)
else:
mindepth_files = [None, ] * cam_cnt
# assume all images have the same size
train_imgfile = find_files('{}/{}/train/rgb'.format(basedir, scene), exts=['*.png', '*.jpg'])[0]
train_im = imageio.imread(train_imgfile)
H, W = train_im.shape[:2]
ray_samplers = []
for i in range(cam_cnt):
intrinsics = parse_txt(intrinsics_files[i])
pose = parse_txt(pose_files[i])
# read max depth
try:
max_depth = float(open('{}/max_depth.txt'.format(split_dir)).readline().strip())
except:
max_depth = None
ray_samplers.append(RaySamplerSingleImage(H=H, W=W, intrinsics=intrinsics, c2w=pose,
img_path=img_files[i],
mask_path=mask_files[i],
min_depth_path=mindepth_files[i],
max_depth=max_depth))
logger.info('Split {}, # views: {}'.format(split, cam_cnt))
return ray_samplers

131
data_verifier.py
Unescape Escape View File

@ -0,0 +1,131 @@
import cv2
import numpy as np
## pip install opencv-python=3.4.2.17 opencv-contrib-python==3.4.2.17
def skew(x):
return np.array([[0, -x[2], x[1]],
[x[2], 0, -x[0]],
[-x[1], x[0], 0]])
def two_view_geometry(intrinsics1, extrinsics1, intrinsics2, extrinsics2):
'''
:param intrinsics1: 4 by 4 matrix
:param extrinsics1: 4 by 4 W2C matrix
:param intrinsics2: 4 by 4 matrix
:param extrinsics2: 4 by 4 W2C matrix
:return:
'''
relative_pose = extrinsics2.dot(np.linalg.inv(extrinsics1))
R = relative_pose[:3, :3]
T = relative_pose[:3, 3]
tx = skew(T)
E = np.dot(tx, R)
F = np.linalg.inv(intrinsics2[:3, :3]).T.dot(E).dot(np.linalg.inv(intrinsics1[:3, :3]))
return E, F, relative_pose
def drawpointslines(img1, img2, lines1, pts2, color):
'''
draw corresponding epilines on img1 for the points in img2
'''
r, c = img1.shape
img1 = cv2.cvtColor(img1, cv2.COLOR_GRAY2BGR)
img2 = cv2.cvtColor(img2, cv2.COLOR_GRAY2BGR)
for r, pt2, cl in zip(lines1, pts2, color):
x0, y0 = map(int, [0, -r[2]/r[1]])
x1, y1 = map(int, [c, -(r[2]+r[0]*c)/r[1]])
cl = tuple(cl.tolist())
img1 = cv2.line(img1, (x0,y0), (x1,y1), cl, 1)
img2 = cv2.circle(img2, tuple(pt2), 5, cl, -1)
return img1, img2
def epipolar(coord1, F, img1, img2):
# compute epipole
pts1 = coord1.astype(int).T
color = np.random.randint(0, high=255, size=(len(pts1), 3))
# Find epilines corresponding to points in left image (first image) and
# drawing its lines on right image
lines2 = cv2.computeCorrespondEpilines(pts1.reshape(-1,1,2), 1,F)
lines2 = lines2.reshape(-1,3)
img3, img4 = drawpointslines(img2,img1,lines2,pts1,color)
## print(img3.shape)
## print(np.concatenate((img4, img3)).shape)
## cv2.imwrite('vis.png', np.concatenate((img4, img3), axis=1))
return np.concatenate((img4, img3), axis=1)
def verify_data(img1, img2, intrinsics1, extrinsics1, intrinsics2, extrinsics2):
img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
E, F, relative_pose = two_view_geometry(intrinsics1, extrinsics1,
intrinsics2, extrinsics2)
# sift = cv2.xfeatures2d.SIFT_create(nfeatures=20)
# kp1 = sift.detect(img1, mask=None)
# coord1 = np.array([[kp.pt[0], kp.pt[1]] for kp in kp1]).T
# Initiate ORB detector
orb = cv2.ORB_create()
# find the keypoints with ORB
kp1 = orb.detect(img1, None)
coord1 = np.array([[kp.pt[0], kp.pt[1]] for kp in kp1[:20]]).T
return epipolar(coord1, F, img1, img2)
if __name__ == '__main__':
from data_loader import load_data
from run_nerf import config_parser
from nerf_sample_ray import parse_camera
import os
parser = config_parser()
args = parser.parse_args()
print(args)
data = load_data(args.datadir, args.scene, testskip=1)
all_imgs = data['images']
all_cameras = data['cameras']
all_intrinsics = []
all_extrinsics = [] # W2C
for i in range(all_cameras.shape[0]):
W, H, intrinsics, extrinsics = parse_camera(all_cameras[i])
all_intrinsics.append(intrinsics)
all_extrinsics.append(np.linalg.inv(extrinsics))
#### arbitrarily select 10 pairs of images to verify pose
out_dir = os.path.join(args.basedir, args.expname, 'data_verify')
print(out_dir)
os.makedirs(out_dir, exist_ok=True)
def calc_angles(c2w_1, c2w_2):
c1 = c2w_1[:3, 3:4]
c2 = c2w_2[:3, 3:4]
c1 = c1 / np.linalg.norm(c1)
c2 = c2 / np.linalg.norm(c2)
return np.rad2deg(np.arccos(np.dot(c1.T, c2)))
images_verify = []
for i in range(10):
while True:
idx1, idx2 = np.random.choice(len(all_imgs), (2,), replace=False)
angle = calc_angles(np.linalg.inv(all_extrinsics[idx1]),
np.linalg.inv(all_extrinsics[idx2]))
if angle > 5. and angle < 10.:
break
im = verify_data(np.uint8(all_imgs[idx1]*255.), np.uint8(all_imgs[idx2]*255.),
all_intrinsics[idx1], all_extrinsics[idx1],
all_intrinsics[idx2], all_extrinsics[idx2])
cv2.imwrite(os.path.join(out_dir, '{:03d}.png'.format(i)), im)

155
ddp_model.py
Unescape Escape View File

@ -0,0 +1,155 @@
import torch
import torch.nn as nn
# import torch.nn.functional as F
# import numpy as np
from utils import TINY_NUMBER, HUGE_NUMBER
from collections import OrderedDict
from nerf_network import Embedder, MLPNet
######################################################################################
# wrapper to simplify the use of nerfnet
######################################################################################
def depth2pts_outside(ray_o, ray_d, depth):
'''
ray_o, ray_d: [..., 3]
depth: [...]; inverse of distance to sphere origin
'''
# note: d1 becomes negative if this mid point is behind camera
d1 = -torch.sum(ray_d * ray_o, dim=-1) / torch.sum(ray_d * ray_d, dim=-1)
p_mid = ray_o + d1.unsqueeze(-1) * ray_d
p_mid_norm = torch.norm(p_mid, dim=-1)
ray_d_cos = 1. / torch.norm(ray_d, dim=-1)
d2 = torch.sqrt(1. - p_mid_norm * p_mid_norm) * ray_d_cos
p_sphere = ray_o + (d1 + d2).unsqueeze(-1) * ray_d
rot_axis = torch.cross(ray_o, p_sphere, dim=-1)
rot_axis = rot_axis / torch.norm(rot_axis, dim=-1, keepdim=True)
phi = torch.asin(p_mid_norm)
theta = torch.asin(p_mid_norm * depth) # depth is inside [0, 1]
rot_angle = (phi - theta).unsqueeze(-1) # [..., 1]
# now rotate p_sphere
# Rodrigues formula: https://en.wikipedia.org/wiki/Rodrigues%27_rotation_formula
p_sphere_new = p_sphere * torch.cos(rot_angle) + \
torch.cross(rot_axis, p_sphere, dim=-1) * torch.sin(rot_angle) + \
rot_axis * torch.sum(rot_axis*p_sphere, dim=-1, keepdim=True) * (1.-torch.cos(rot_angle))
p_sphere_new = p_sphere_new / torch.norm(p_sphere_new, dim=-1, keepdim=True)
pts = torch.cat((p_sphere_new, depth.unsqueeze(-1)), dim=-1)
# now calculate conventional depth
depth_real = 1. / (depth + TINY_NUMBER) * torch.cos(theta) * ray_d_cos + d1
return pts, depth_real
class NerfNet(nn.Module):
def __init__(self, args):
'''
:param D: network depth
:param W: network width
:param input_ch: input channels for encodings of (x, y, z)
:param input_ch_viewdirs: input channels for encodings of view directions
:param skips: skip connection in network
:param use_viewdirs: if True, will use the view directions as input
'''
super().__init__()
# foreground
self.fg_embedder_position = Embedder(input_dim=3,
max_freq_log2=args.max_freq_log2 - 1,
N_freqs=args.max_freq_log2)
self.fg_embedder_viewdir = Embedder(input_dim=3,
max_freq_log2=args.max_freq_log2_viewdirs - 1,
N_freqs=args.max_freq_log2_viewdirs)
self.fg_net = MLPNet(D=args.netdepth, W=args.netwidth,
input_ch=self.fg_embedder_position.out_dim,
input_ch_viewdirs=self.fg_embedder_viewdir.out_dim,
use_viewdirs=args.use_viewdirs,
use_implicit=args.use_implicit)
# background; bg_pt is (x, y, z, 1/r)
self.bg_embedder_position = Embedder(input_dim=4,
max_freq_log2=args.max_freq_log2 - 1,
N_freqs=args.max_freq_log2)
self.bg_embedder_viewdir = Embedder(input_dim=3,
max_freq_log2=args.max_freq_log2_viewdirs - 1,
N_freqs=args.max_freq_log2_viewdirs)
self.bg_net = MLPNet(D=args.netdepth, W=args.netwidth,
input_ch=self.bg_embedder_position.out_dim,
input_ch_viewdirs=self.bg_embedder_viewdir.out_dim,
use_viewdirs=args.use_viewdirs,
use_implicit=args.use_implicit)
def forward(self, ray_o, ray_d, fg_z_max, fg_z_vals, bg_z_vals):
'''
:param ray_o, ray_d: [..., 3]
:param fg_z_max: [...,]
:param fg_z_vals, bg_z_vals: [..., N_samples]
:return
'''
# print(ray_o.shape, ray_d.shape, fg_z_max.shape, fg_z_vals.shape, bg_z_vals.shape)
ray_d_norm = torch.norm(ray_d, dim=-1, keepdim=True) # [..., 1]
viewdirs = ray_d / ray_d_norm # [..., 3]
dots_sh = list(ray_d.shape[:-1])
######### render foreground
N_samples = fg_z_vals.shape[-1]
fg_ray_o = ray_o.unsqueeze(-2).expand(dots_sh + [N_samples, 3])
fg_ray_d = ray_d.unsqueeze(-2).expand(dots_sh + [N_samples, 3])
fg_viewdirs = viewdirs.unsqueeze(-2).expand(dots_sh + [N_samples, 3])
fg_pts = fg_ray_o + fg_z_vals.unsqueeze(-1) * fg_ray_d
input = torch.cat((self.fg_embedder_position(fg_pts),
self.fg_embedder_viewdir(fg_viewdirs)), dim=-1)
fg_raw = self.fg_net(input)
# alpha blending
fg_dists = fg_z_vals[..., 1:] - fg_z_vals[..., :-1]
# account for view directions
fg_dists = ray_d_norm * torch.cat((fg_dists, fg_z_max.unsqueeze(-1) - fg_z_vals[..., -1:]), dim=-1) # [..., N_samples]
fg_alpha = 1. - torch.exp(-fg_raw['sigma'] * fg_dists) # [..., N_samples]
T = torch.cumprod(1. - fg_alpha + TINY_NUMBER, dim=-1) # [..., N_samples]
bg_lambda = T[..., -1]
T = torch.cat((torch.ones_like(T[..., 0:1]), T[..., :-1]), dim=-1) # [..., N_samples]
fg_weights = fg_alpha * T # [..., N_samples]
fg_rgb_map = torch.sum(fg_weights.unsqueeze(-1) * fg_raw['rgb'], dim=-2) # [..., 3]
fg_diffuse_rgb_map = torch.sum(fg_weights.unsqueeze(-1) * fg_raw['diffuse_rgb'], dim=-2) # [..., 3]
fg_depth_map = torch.sum(fg_weights * fg_z_vals, dim=-1) # [...,]
# render background
N_samples = bg_z_vals.shape[-1]
bg_ray_o = ray_o.unsqueeze(-2).expand(dots_sh + [N_samples, 3])
bg_ray_d = ray_d.unsqueeze(-2).expand(dots_sh + [N_samples, 3])
bg_viewdirs = viewdirs.unsqueeze(-2).expand(dots_sh + [N_samples, 3])
bg_pts, _ = depth2pts_outside(bg_ray_o, bg_ray_d, bg_z_vals) # [..., N_samples, 4]
input = torch.cat((self.bg_embedder_position(bg_pts),
self.bg_embedder_viewdir(bg_viewdirs)), dim=-1)
# near_depth: physical far; far_depth: physical near
input = torch.flip(input, dims=[-2,])
bg_z_vals = torch.flip(bg_z_vals, dims=[-1,]) # 1--->0
bg_dists = bg_z_vals[..., :-1] - bg_z_vals[..., 1:]
bg_dists = torch.cat((bg_dists, HUGE_NUMBER * torch.ones_like(bg_dists[..., 0:1])), dim=-1) # [..., N_samples]
bg_raw = self.bg_net(input)
bg_alpha = 1. - torch.exp(-bg_raw['sigma'] * bg_dists) # [..., N_samples]
# Eq. (3): T
# maths show weights, and summation of weights along a ray, are always inside [0, 1]
T = torch.cumprod(1. - bg_alpha + TINY_NUMBER, dim=-1)[..., :-1] # [..., N_samples-1]
T = torch.cat((torch.ones_like(T[..., 0:1]), T), dim=-1) # [..., N_samples]
bg_weights = bg_alpha * T # [..., N_samples]
bg_rgb_map = torch.sum(bg_weights.unsqueeze(-1) * bg_raw['rgb'], dim=-2) # [..., 3]
bg_diffuse_rgb_map = torch.sum(bg_weights.unsqueeze(-1) * bg_raw['diffuse_rgb'], dim=-2) # [..., 3]
bg_depth_map = torch.sum(bg_weights * bg_z_vals, dim=-1) # [...,]
# composite foreground and background
bg_rgb_map = bg_lambda.unsqueeze(-1) * bg_rgb_map
bg_diffuse_rgb_map = bg_lambda.unsqueeze(-1) * bg_diffuse_rgb_map
bg_depth_map = bg_lambda * bg_depth_map
rgb_map = fg_rgb_map + bg_rgb_map
diffuse_rgb_map = fg_diffuse_rgb_map + bg_diffuse_rgb_map
ret = OrderedDict([('rgb', rgb_map), # loss
('diffuse_rgb', diffuse_rgb_map), # regularize
('fg_weights', fg_weights), # importance sampling
('bg_weights', bg_weights), # importance sampling
('fg_rgb', fg_rgb_map), # below are for logging
('fg_depth', fg_depth_map),
('bg_rgb', bg_rgb_map),
('bg_depth', bg_depth_map),
('bg_lambda', bg_lambda)])
return ret

621
ddp_run_nerf.py
Unescape Escape View File

@ -0,0 +1,621 @@
import torch
import torch.nn as nn
import torch.optim
import torch.distributed
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.multiprocessing
import os
from collections import OrderedDict
from ddp_model import NerfNet
import time
from data_loader_split import load_data_split
import numpy as np
from tensorboardX import SummaryWriter
from utils import img2mse, mse2psnr, img_HWC2CHW, colorize, TINY_NUMBER
import logging
logger = logging.getLogger(__package__)
def setup_logger():
# create logger
logger = logging.getLogger(__package__)
# logger.setLevel(logging.DEBUG)
logger.setLevel(logging.INFO)
# create console handler and set level to debug
ch = logging.StreamHandler()
ch.setLevel(logging.DEBUG)
# create formatter
formatter = logging.Formatter('%(asctime)s [%(levelname)s] %(name)s: %(message)s')
# add formatter to ch
ch.setFormatter(formatter)
# add ch to logger
logger.addHandler(ch)
def intersect_sphere(ray_o, ray_d):
'''
ray_o, ray_d: [..., 3]
compute the depth of the intersection point between this ray and unit sphere
'''
# note: d1 becomes negative if this mid point is behind camera
d1 = -torch.sum(ray_d * ray_o, dim=-1) / torch.sum(ray_d * ray_d, dim=-1)
p = ray_o + d1.unsqueeze(-1) * ray_d
# consider the case where the ray does not intersect the sphere
ray_d_cos = 1. / torch.norm(ray_d, dim=-1)
d2 = torch.sqrt(1. - torch.sum(p * p, dim=-1)) * ray_d_cos
return d1 + d2
def perturb_samples(z_vals):
# get intervals between samples
mids = .5 * (z_vals[..., 1:] + z_vals[..., :-1])
upper = torch.cat([mids, z_vals[..., -1:]], dim=-1)
lower = torch.cat([z_vals[..., 0:1], mids], dim=-1)
# uniform samples in those intervals
t_rand = torch.rand_like(z_vals)
z_vals = lower + (upper - lower) * t_rand # [N_rays, N_samples]
return z_vals
def sample_pdf(bins, weights, N_samples, det=False):
'''
:param bins: tensor of shape [..., M+1], M is the number of bins
:param weights: tensor of shape [..., M]
:param N_samples: number of samples along each ray
:param det: if True, will perform deterministic sampling
:return: [..., N_samples]
'''
# Get pdf
weights = weights + TINY_NUMBER # prevent nans
pdf = weights / torch.sum(weights, dim=-1, keepdim=True) # [..., M]
cdf = torch.cumsum(pdf, dim=-1) # [..., M]
cdf = torch.cat([torch.zeros_like(cdf[..., 0:1]), cdf], dim=-1) # [..., M+1]
# Take uniform samples
dots_sh = list(weights.shape[:-1])
M = weights.shape[-1]
min_cdf = 0.00
max_cdf = 1.00 # prevent outlier samples
if det:
u = torch.linspace(min_cdf, max_cdf, N_samples, device=bins.device)
u = u.view([1]*len(dots_sh) + [N_samples]).expand(dots_sh + [N_samples,]) # [..., N_samples]
else:
sh = dots_sh + [N_samples]
u = torch.rand(*sh, device=bins.device) * (max_cdf - min_cdf) + min_cdf # [..., N_samples]
# Invert CDF
# [..., N_samples, 1] >= [..., 1, M] ----> [..., N_samples, M] ----> [..., N_samples,]
above_inds = torch.sum(u.unsqueeze(-1) >= cdf[..., :M].unsqueeze(-2), dim=-1).long()
# random sample inside each bin
below_inds = torch.clamp(above_inds-1, min=0)
inds_g = torch.stack((below_inds, above_inds), dim=-1) # [..., N_samples, 2]
cdf = cdf.unsqueeze(-2).expand(dots_sh + [N_samples, M+1]) # [..., N_samples, M+1]
cdf_g = torch.gather(input=cdf, dim=-1, index=inds_g) # [..., N_samples, 2]
bins = bins.unsqueeze(-2).expand(dots_sh + [N_samples, M+1]) # [..., N_samples, M+1]
bins_g = torch.gather(input=bins, dim=-1, index=inds_g) # [..., N_samples, 2]
# fix numeric issue
denom = cdf_g[..., 1] - cdf_g[..., 0] # [..., N_samples]
denom = torch.where(denom<TINY_NUMBER, torch.ones_like(denom), denom)
t = (u - cdf_g[..., 0]) / denom
samples = bins_g[..., 0] + t * (bins_g[..., 1] - bins_g[..., 0] + TINY_NUMBER)
return samples
def render_single_image(rank, world_size, models, ray_sampler, chunk_size):
##### parallel rendering of a single image
ray_batch = ray_sampler.get_all()
# split into ranks; make sure different processes don't overlap
rank_split_sizes = [ray_batch['ray_d'].shape[0] // world_size, ] * world_size
rank_split_sizes[-1] = ray_batch['ray_d'].shape[0] - sum(rank_split_sizes[:-1])
for key in ray_batch:
if torch.is_tensor(ray_batch[key]):
ray_batch[key] = torch.split(ray_batch[key], rank_split_sizes)[rank].to(rank)
# split into chunks and render inside each process
ray_batch_split = OrderedDict()
for key in ray_batch:
if torch.is_tensor(ray_batch[key]):
ray_batch_split[key] = torch.split(ray_batch[key], chunk_size)
# forward and backward
ret_merge_chunk = [OrderedDict() for _ in range(models['cascade_level'])]
for s in range(len(ray_batch_split['ray_d'])):
ray_o = ray_batch_split['ray_o'][s]
ray_d = ray_batch_split['ray_d'][s]
min_depth = ray_batch_split['min_depth'][s]
dots_sh = list(ray_d.shape[:-1])
for m in range(models['cascade_level']):
net = models['net_{}'.format(m)]
# sample depths
N_samples = models['cascade_samples'][m]
if m == 0:
# foreground depth
fg_far_depth = intersect_sphere(ray_o, ray_d) # [...,]
# fg_near_depth = 0.18 * torch.ones_like(fg_far_depth)
fg_near_depth = min_depth # [..., 3]
step = (fg_far_depth - fg_near_depth) / (N_samples - 1)
fg_depth = torch.stack([fg_near_depth + i * step for i in range(N_samples)], dim=-1) # [..., N_samples]
# background depth
bg_depth = torch.linspace(0., 1., N_samples).view(
[1, ] * len(dots_sh) + [N_samples,]).expand(dots_sh + [N_samples,]).to(rank)
# delete unused memory
del fg_near_depth
del step
torch.cuda.empty_cache()
else:
# sample pdf and concat with earlier samples
fg_weights = ret['fg_weights'].clone().detach()
fg_depth_mid = .5 * (fg_depth[..., 1:] + fg_depth[..., :-1]) # [..., N_samples-1]
fg_weights = fg_weights[..., 1:-1] # [..., N_samples-2]
fg_depth_samples = sample_pdf(bins=fg_depth_mid, weights=fg_weights,
N_samples=N_samples, det=True) # [..., N_samples]
fg_depth, _ = torch.sort(torch.cat((fg_depth, fg_depth_samples), dim=-1))
# sample pdf and concat with earlier samples
bg_weights = ret['bg_weights'].clone().detach()
bg_depth_mid = .5 * (bg_depth[..., 1:] + bg_depth[..., :-1])
bg_weights = bg_weights[..., 1:-1] # [..., N_samples-2]
bg_depth_samples = sample_pdf(bins=bg_depth_mid, weights=bg_weights,
N_samples=N_samples, det=True) # [..., N_samples]
bg_depth, _ = torch.sort(torch.cat((bg_depth, bg_depth_samples), dim=-1))
# delete unused memory
del fg_weights
del fg_depth_mid
del fg_depth_samples
del bg_weights
del bg_depth_mid
del bg_depth_samples
torch.cuda.empty_cache()
with torch.no_grad():
ret = net(ray_o, ray_d, fg_far_depth, fg_depth, bg_depth)
for key in ret:
if key not in ['fg_weights', 'bg_weights']:
if torch.is_tensor(ret[key]):
if key not in ret_merge_chunk[m]:
ret_merge_chunk[m][key] = [ret[key].cpu(), ]
else:
ret_merge_chunk[m][key].append(ret[key].cpu())
ret[key] = None
# clean unused memory
torch.cuda.empty_cache()
# merge results from different chunks
for m in range(len(ret_merge_chunk)):
for key in ret_merge_chunk[m]:
ret_merge_chunk[m][key] = torch.cat(ret_merge_chunk[m][key], dim=0)
# merge results from different processes
if rank == 0:
ret_merge_rank = [OrderedDict() for _ in range(len(ret_merge_chunk))]
for m in range(len(ret_merge_chunk)):
for key in ret_merge_chunk[m]:
# generate tensors to store results from other processes
sh = list(ret_merge_chunk[m][key].shape[1:])
ret_merge_rank[m][key] = [torch.zeros(*[size,]+sh, dtype=torch.float32) for size in rank_split_sizes]
torch.distributed.gather(ret_merge_chunk[m][key], ret_merge_rank[m][key])
ret_merge_rank[m][key] = torch.cat(ret_merge_rank[m][key], dim=0).reshape(
(ray_sampler.H, ray_sampler.W, -1)).squeeze()
# print(m, key, ret_merge_rank[m][key].shape)
else: # send results to main process
for m in range(len(ret_merge_chunk)):
for key in ret_merge_chunk[m]:
torch.distributed.gather(ret_merge_chunk[m][key])
# only rank 0 program returns
if rank == 0:
return ret_merge_rank
else:
return None
def log_view_to_tb(writer, global_step, log_data, gt_img, mask, prefix=''):
rgb_im = img_HWC2CHW(torch.from_numpy(gt_img))
writer.add_image(prefix + 'rgb_gt', rgb_im, global_step)
for m in range(len(log_data)):
rgb_im = img_HWC2CHW(log_data[m]['rgb'])
rgb_im = torch.clamp(rgb_im, min=0., max=1.) # just in case diffuse+specular>1
writer.add_image(prefix + 'level_{}/rgb'.format(m), rgb_im, global_step)
rgb_im = img_HWC2CHW(log_data[m]['fg_rgb'])
rgb_im = torch.clamp(rgb_im, min=0., max=1.) # just in case diffuse+specular>1
writer.add_image(prefix + 'level_{}/fg_rgb'.format(m), rgb_im, global_step)
depth = log_data[m]['fg_depth']
depth_im = img_HWC2CHW(colorize(depth, cmap_name='jet', append_cbar=True,
mask=mask))
writer.add_image(prefix + 'level_{}/fg_depth'.format(m), depth_im, global_step)
rgb_im = img_HWC2CHW(log_data[m]['bg_rgb'])
rgb_im = torch.clamp(rgb_im, min=0., max=1.) # just in case diffuse+specular>1
writer.add_image(prefix + 'level_{}/bg_rgb'.format(m), rgb_im, global_step)
depth = log_data[m]['bg_depth']
depth_im = img_HWC2CHW(colorize(depth, cmap_name='jet', append_cbar=True,
mask=mask))
writer.add_image(prefix + 'level_{}/bg_depth'.format(m), depth_im, global_step)
bg_lambda = log_data[m]['bg_lambda']
bg_lambda_im = img_HWC2CHW(colorize(bg_lambda, cmap_name='hot', append_cbar=True,
mask=mask))
writer.add_image(prefix + 'level_{}/bg_lambda'.format(m), bg_lambda_im, global_step)
def setup(rank, world_size):
os.environ['MASTER_ADDR'] = 'localhost'
# port = np.random.randint(12355, 12399)
# os.environ['MASTER_PORT'] = '{}'.format(port)
os.environ['MASTER_PORT'] = '12355'
# initialize the process group
torch.distributed.init_process_group("gloo", rank=rank, world_size=world_size)
def cleanup():
torch.distributed.destroy_process_group()
def ddp_train_nerf(rank, args):
###### set up multi-processing
setup(rank, args.world_size)
###### set up logger
logger = logging.getLogger(__package__)
setup_logger()
###### decide chunk size according to gpu memory
logger.info('gpu_mem: {}'.format(torch.cuda.get_device_properties(rank).total_memory))
if torch.cuda.get_device_properties(rank).total_memory / 1e9 > 14:
logger.info('setting batch size according to 24G gpu')
args.N_rand = 1024
args.chunk_size = 8192
else:
logger.info('setting batch size according to 12G gpu')
args.N_rand = 512
args.chunk_size = 4096
###### Create log dir and copy the config file
if rank == 0:
os.makedirs(os.path.join(args.basedir, args.expname), exist_ok=True)
f = os.path.join(args.basedir, args.expname, 'args.txt')
with open(f, 'w') as file:
for arg in sorted(vars(args)):
attr = getattr(args, arg)
file.write('{} = {}\n'.format(arg, attr))
if args.config is not None:
f = os.path.join(args.basedir, args.expname, 'config.txt')
with open(f, 'w') as file:
file.write(open(args.config, 'r').read())
torch.distributed.barrier()
ray_samplers = load_data_split(args.datadir, args.scene, split='train', try_load_min_depth=args.load_min_depth)
val_ray_samplers = load_data_split(args.datadir, args.scene, split='validation', try_load_min_depth=args.load_min_depth)
###### create network and wrap in ddp; each process should do this
# fix random seed just to make sure the network is initialized with same weights at different processes
torch.manual_seed(777)
# very important!!! otherwise it might introduce extra memory in rank=0 gpu
torch.cuda.set_device(rank)
models = OrderedDict()
models['cascade_level'] = args.cascade_level
models['cascade_samples'] = [int(x.strip()) for x in args.cascade_samples.split(',')]
for m in range(models['cascade_level']):
net = NerfNet(args).to(rank)
net = DDP(net, device_ids=[rank], output_device=rank)
optim = torch.optim.Adam(net.parameters(), lr=args.lrate)
models['net_{}'.format(m)] = net
models['optim_{}'.format(m)] = optim
start = -1
###### load pretrained weights; each process should do this
if (args.ckpt_path is not None) and (os.path.isfile(args.ckpt_path)):
ckpts = [args.ckpt_path]
else:
ckpts = [os.path.join(args.basedir, args.expname, f)
for f in sorted(os.listdir(os.path.join(args.basedir, args.expname))) if f.endswith('.pth')]
def path2iter(path):
tmp = os.path.basename(path)[:-4]
idx = tmp.rfind('_')
return int(tmp[idx + 1:])
ckpts = sorted(ckpts, key=path2iter)
logger.info('Found ckpts: {}'.format(ckpts))
if len(ckpts) > 0 and not args.no_reload:
fpath = ckpts[-1]
logger.info('Reloading from: {}'.format(fpath))
start = path2iter(fpath)
# configure map_location properly for different processes
map_location = {'cuda:%d' % 0: 'cuda:%d' % rank}
to_load = torch.load(fpath, map_location=map_location)
for m in range(models['cascade_level']):
for name in ['net_{}'.format(m), 'optim_{}'.format(m)]:
models[name].load_state_dict(to_load[name])
models[name].load_state_dict(to_load[name])
##### important!!!
# make sure different processes sample different rays
np.random.seed((rank + 1) * 777)
# make sure different processes have different perturbations in depth samples
torch.manual_seed((rank + 1) * 777)
##### only main process should do the logging
if rank == 0:
writer = SummaryWriter(os.path.join(args.basedir, 'summaries', args.expname))
# start training
what_val_to_log = 0 # helper variable for parallel rendering of a image
what_train_to_log = 0
for global_step in range(start+1, start+1+args.N_iters):
time0 = time.time()
scalars_to_log = OrderedDict()
### Start of core optimization loop
scalars_to_log['resolution'] = ray_samplers[0].resolution_level
# randomly sample rays and move to device
i = np.random.randint(low=0, high=len(ray_samplers))
ray_batch = ray_samplers[i].random_sample(args.N_rand, center_crop=False)
for key in ray_batch:
if torch.is_tensor(ray_batch[key]):
ray_batch[key] = ray_batch[key].to(rank)
# forward and backward
dots_sh = list(ray_batch['ray_d'].shape[:-1]) # number of rays
all_rets = [] # results on different cascade levels
for m in range(models['cascade_level']):
optim = models['optim_{}'.format(m)]
net = models['net_{}'.format(m)]
# sample depths
N_samples = models['cascade_samples'][m]
if m == 0:
# foreground depth
fg_far_depth = intersect_sphere(ray_batch['ray_o'], ray_batch['ray_d']) # [...,]
# fg_near_depth = 0.18 * torch.ones_like(fg_far_depth)
fg_near_depth = ray_batch['min_depth'] # [..., 3]
step = (fg_far_depth - fg_near_depth) / (N_samples - 1)
fg_depth = torch.stack([fg_near_depth + i * step for i in range(N_samples)], dim=-1) # [..., N_samples]
fg_depth = perturb_samples(fg_depth) # random perturbation during training
# background depth
bg_depth = torch.linspace(0., 1., N_samples).view(
[1, ] * len(dots_sh) + [N_samples,]).expand(dots_sh + [N_samples,]).to(rank)
bg_depth = perturb_samples(bg_depth) # random perturbation during training
else:
# sample pdf and concat with earlier samples
fg_weights = ret['fg_weights'].clone().detach()
fg_depth_mid = .5 * (fg_depth[..., 1:] + fg_depth[..., :-1]) # [..., N_samples-1]
fg_weights = fg_weights[..., 1:-1] # [..., N_samples-2]
fg_depth_samples = sample_pdf(bins=fg_depth_mid, weights=fg_weights,
N_samples=N_samples, det=False) # [..., N_samples]
fg_depth, _ = torch.sort(torch.cat((fg_depth, fg_depth_samples), dim=-1))
# sample pdf and concat with earlier samples
bg_weights = ret['bg_weights'].clone().detach()
bg_depth_mid = .5 * (bg_depth[..., 1:] + bg_depth[..., :-1])
bg_weights = bg_weights[..., 1:-1] # [..., N_samples-2]
bg_depth_samples = sample_pdf(bins=bg_depth_mid, weights=bg_weights,
N_samples=N_samples, det=False) # [..., N_samples]
bg_depth, _ = torch.sort(torch.cat((bg_depth, bg_depth_samples), dim=-1))
optim.zero_grad()
ret = net(ray_batch['ray_o'], ray_batch['ray_d'], fg_far_depth, fg_depth, bg_depth)
all_rets.append(ret)
rgb_gt = ray_batch['rgb'].to(rank)
loss = img2mse(ret['rgb'], rgb_gt)
scalars_to_log['level_{}/loss'.format(m)] = loss.item()
scalars_to_log['level_{}/pnsr'.format(m)] = mse2psnr(loss.item())
# regularize sigma with photo-consistency
diffuse_loss = img2mse(ret['diffuse_rgb'], rgb_gt)
scalars_to_log['level_{}/diffuse_loss'.format(m)] = diffuse_loss.item()
scalars_to_log['level_{}/diffuse_psnr'.format(m)] = mse2psnr(diffuse_loss.item())
loss = (1. - args.regularize_weight) * loss + args.regularize_weight * diffuse_loss
loss.backward()
optim.step()
# # clean unused memory
# torch.cuda.empty_cache()
### end of core optimization loop
dt = time.time() - time0
scalars_to_log['iter_time'] = dt
### only main process should do the logging
if rank == 0 and (global_step % args.i_print == 0 or global_step < 10):
logstr = '{} step: {} '.format(args.expname, global_step)
for k in scalars_to_log:
logstr += ' {}: {:.6f}'.format(k, scalars_to_log[k])
writer.add_scalar(k, scalars_to_log[k], global_step)
logger.info(logstr)
### each process should do this; but only main process merges the results
if global_step % args.i_img == 0 or global_step == start+1:
#### critical: make sure each process is working on the same random image
time0 = time.time()
idx = what_val_to_log % len(val_ray_samplers)
log_data = render_single_image(rank, args.world_size, models, val_ray_samplers[idx], args.chunk_size)
what_val_to_log += 1
dt = time.time() - time0
if rank == 0: # only main process should do this
logger.info('Logged a random validation view in {} seconds'.format(dt))
log_view_to_tb(writer, global_step, log_data, gt_img=val_ray_samplers[idx].get_img(), mask=None, prefix='val/')
time0 = time.time()
idx = what_train_to_log % len(ray_samplers)
log_data = render_single_image(rank, args.world_size, models, ray_samplers[idx], args.chunk_size)
what_train_to_log += 1
dt = time.time() - time0
if rank == 0: # only main process should do this
logger.info('Logged a random training view in {} seconds'.format(dt))
log_view_to_tb(writer, global_step, log_data, gt_img=ray_samplers[idx].get_img(), mask=None, prefix='train/')
log_data = None
torch.cuda.empty_cache()
if rank == 0 and (global_step % args.i_weights == 0 and global_step > 0):
# saving checkpoints and logging
fpath = os.path.join(args.basedir, args.expname, 'model_{:06d}.pth'.format(global_step))
to_save = OrderedDict()
for m in range(models['cascade_level']):
name = 'net_{}'.format(m)
to_save[name] = models[name].state_dict()
name = 'optim_{}'.format(m)
to_save[name] = models[name].state_dict()
torch.save(to_save, fpath)
# clean up for multi-processing
cleanup()
def config_parser():
import configargparse
parser = configargparse.ArgumentParser()
parser.add_argument('--config', is_config_file=True, help='config file path')
parser.add_argument("--expname", type=str, help='experiment name')
parser.add_argument("--basedir", type=str, default='./logs/', help='where to store ckpts and logs')
# dataset options
parser.add_argument("--datadir", type=str, default=None, help='input data directory')
parser.add_argument("--scene", type=str, default=None, help='scene name')
parser.add_argument("--testskip", type=int, default=8,
help='will load 1/N images from test/val sets, useful for large datasets like deepvoxels')
# model size
parser.add_argument("--netdepth", type=int, default=8, help='layers in coarse network')
parser.add_argument("--netwidth", type=int, default=256, help='channels per layer in coarse network')
parser.add_argument("--use_viewdirs", action='store_true', help='use full 5D input instead of 3D')
# checkpoints
parser.add_argument("--no_reload", action='store_true', help='do not reload weights from saved ckpt')
parser.add_argument("--ckpt_path", type=str, default=None,
help='specific weights npy file to reload for coarse network')
# batch size
parser.add_argument("--N_rand", type=int, default=32 * 32 * 2,
help='batch size (number of random rays per gradient step)')
parser.add_argument("--chunk_size", type=int, default=1024 * 8,
help='number of rays processed in parallel, decrease if running out of memory')
# iterations
parser.add_argument("--N_iters", type=int, default=250001,
help='number of iterations')
parser.add_argument("--render_splits", type=str, default='test',
help='splits to render')
# cascade training
parser.add_argument("--cascade_level", type=int, default=2,
help='number of cascade levels')
parser.add_argument("--cascade_samples", type=str, default='64,64',
help='samples at each level')
parser.add_argument("--devices", type=str, default='0,1',
help='cuda device for each level')
parser.add_argument("--bg_devices", type=str, default='0,2',
help='cuda device for the background of each level')
parser.add_argument("--world_size", type=int, default='-1',
help='number of processes')
# mixed precison training
parser.add_argument("--opt_level", type=str, default='O1',
help='mixed precison training')
parser.add_argument("--near_depth", type=float, default=0.1,
help='near depth plane')
parser.add_argument("--far_depth", type=float, default=50.,
help='far depth plane')
# learning rate options
parser.add_argument("--lrate", type=float, default=5e-4, help='learning rate')
parser.add_argument("--lrate_decay_factor", type=float, default=0.1,
help='decay learning rate by a factor every specified number of steps')
parser.add_argument("--lrate_decay_steps", type=int, default=5000,
help='decay learning rate by a factor every specified number of steps')
# rendering options
parser.add_argument("--inv_uniform", action='store_true',
help='if True, will uniformly sample inverse depths')
parser.add_argument("--det", action='store_true', help='deterministic sampling for coarse and fine samples')
parser.add_argument("--max_freq_log2", type=int, default=10,
help='log2 of max freq for positional encoding (3D location)')
parser.add_argument("--max_freq_log2_viewdirs", type=int, default=4,
help='log2 of max freq for positional encoding (2D direction)')
parser.add_argument("--N_iters_perturb", type=int, default=1000,
help='perturb and center-crop at first 1000 iterations to prevent training from getting stuck')
parser.add_argument("--raw_noise_std", type=float, default=1.,
help='std dev of noise added to regularize sigma output, 1e0 recommended')
parser.add_argument("--white_bkgd", action='store_true',
help='apply the trick to avoid fitting to white background')
# use implicit
parser.add_argument("--use_implicit", action='store_true', help='whether to use implicit regularization')
parser.add_argument("--regularize_weight", type=float, default=0.5,
help='regularizing weight of auxiliary loss')
parser.add_argument("--load_min_depth", action='store_true', help='whether to load min depth')
# no training; render only
parser.add_argument("--render_only", action='store_true',
help='do not optimize, reload weights and render out render_poses path')
parser.add_argument("--render_train", action='store_true', help='render the training set')
parser.add_argument("--render_test", action='store_true', help='render the test set instead of render_poses path')
# no training; extract mesh only
parser.add_argument("--mesh_only", action='store_true',
help='do not optimize, extract mesh from pretrained model')
parser.add_argument("--N_pts", type=int, default=256,
help='voxel resolution; N_pts * N_pts * N_pts')
parser.add_argument("--mesh_thres", type=str, default='10,20,30,40,50',
help='threshold(s) for mesh extraction; can use multiple thresholds')
# logging/saving options
parser.add_argument("--i_print", type=int, default=100, help='frequency of console printout and metric loggin')
parser.add_argument("--i_img", type=int, default=500, help='frequency of tensorboard image logging')
parser.add_argument("--i_weights", type=int, default=10000, help='frequency of weight ckpt saving')
parser.add_argument("--i_testset", type=int, default=50000, help='frequency of testset saving')
parser.add_argument("--i_video", type=int, default=50000, help='frequency of render_poses video saving')
return parser
def train():
parser = config_parser()
args = parser.parse_args()
logger.info(parser.format_values())
if args.world_size == -1:
args.world_size = torch.cuda.device_count()
logger.info('Using # gpus: {}'.format(args.world_size))
torch.multiprocessing.spawn(ddp_train_nerf,
args=(args,),
nprocs=args.world_size,
join=True)
if __name__ == '__main__':
setup_logger()
train()

391
ddp_test_nerf.py
Unescape Escape View File

@ -0,0 +1,391 @@
import torch
import torch.nn as nn
import torch.optim
import torch.distributed
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.multiprocessing
import numpy as np
import os
from collections import OrderedDict
from ddp_model import NerfNet
import time
from data_loader_split import load_data_split
from utils import mse2psnr, img_HWC2CHW, colorize, colorize_np, TINY_NUMBER, to8b
import imageio
from ddp_run_nerf import config_parser
import logging
logger = logging.getLogger(__package__)
def setup_logger():
# create logger
logger = logging.getLogger(__package__)
logger.setLevel(logging.DEBUG)
# create console handler and set level to debug
ch = logging.StreamHandler()
ch.setLevel(logging.INFO)
# create formatter
formatter = logging.Formatter('%(asctime)s [%(levelname)s] %(name)s: %(message)s')
# add formatter to ch
ch.setFormatter(formatter)
# add ch to logger
logger.addHandler(ch)
def intersect_sphere(ray_o, ray_d):
'''
ray_o, ray_d: [..., 3]
compute the depth of the intersection point between this ray and unit sphere
'''
# note: d1 becomes negative if this mid point is behind camera
d1 = -torch.sum(ray_d * ray_o, dim=-1) / torch.sum(ray_d * ray_d, dim=-1)
p = ray_o + d1.unsqueeze(-1) * ray_d
# consider the case where the ray does not intersect the sphere
ray_d_cos = 1. / torch.norm(ray_d, dim=-1)
d2 = torch.sqrt(1. - torch.sum(p * p, dim=-1)) * ray_d_cos
return d1 + d2
def perturb_samples(z_vals):
# get intervals between samples
mids = .5 * (z_vals[..., 1:] + z_vals[..., :-1])
upper = torch.cat([mids, z_vals[..., -1:]], dim=-1)
lower = torch.cat([z_vals[..., 0:1], mids], dim=-1)
# uniform samples in those intervals
t_rand = torch.rand_like(z_vals)
z_vals = lower + (upper - lower) * t_rand # [N_rays, N_samples]
return z_vals
def sample_pdf(bins, weights, N_samples, det=False):
'''
:param bins: tensor of shape [..., M+1], M is the number of bins
:param weights: tensor of shape [..., M]
:param N_samples: number of samples along each ray
:param det: if True, will perform deterministic sampling
:return: [..., N_samples]
'''
# Get pdf
weights = weights + TINY_NUMBER # prevent nans
pdf = weights / torch.sum(weights, dim=-1, keepdim=True) # [..., M]
cdf = torch.cumsum(pdf, dim=-1) # [..., M]
cdf = torch.cat([torch.zeros_like(cdf[..., 0:1]), cdf], dim=-1) # [..., M+1]
# Take uniform samples
dots_sh = list(weights.shape[:-1])
M = weights.shape[-1]
min_cdf = 0.00
max_cdf = 1.00 # prevent outlier samples
if det:
u = torch.linspace(min_cdf, max_cdf, N_samples, device=bins.device)
u = u.view([1]*len(dots_sh) + [N_samples]).expand(dots_sh + [N_samples,]) # [..., N_samples]
else:
sh = dots_sh + [N_samples]
u = torch.rand(*sh, device=bins.device) * (max_cdf - min_cdf) + min_cdf # [..., N_samples]
# Invert CDF
# [..., N_samples, 1] >= [..., 1, M] ----> [..., N_samples, M] ----> [..., N_samples,]
above_inds = torch.sum(u.unsqueeze(-1) >= cdf[..., :M].unsqueeze(-2), dim=-1).long()
# random sample inside each bin
below_inds = torch.clamp(above_inds-1, min=0)
inds_g = torch.stack((below_inds, above_inds), dim=-1) # [..., N_samples, 2]
cdf = cdf.unsqueeze(-2).expand(dots_sh + [N_samples, M+1]) # [..., N_samples, M+1]
cdf_g = torch.gather(input=cdf, dim=-1, index=inds_g) # [..., N_samples, 2]
bins = bins.unsqueeze(-2).expand(dots_sh + [N_samples, M+1]) # [..., N_samples, M+1]
bins_g = torch.gather(input=bins, dim=-1, index=inds_g) # [..., N_samples, 2]
# fix numeric issue
denom = cdf_g[..., 1] - cdf_g[..., 0] # [..., N_samples]
denom = torch.where(denom<TINY_NUMBER, torch.ones_like(denom), denom)
t = (u - cdf_g[..., 0]) / denom
samples = bins_g[..., 0] + t * (bins_g[..., 1] - bins_g[..., 0] + TINY_NUMBER)
return samples
def render_single_image(rank, world_size, models, ray_sampler, chunk_size):
##### parallel rendering of a single image
ray_batch = ray_sampler.get_all()
# split into ranks; make sure different processes don't overlap
rank_split_sizes = [ray_batch['ray_d'].shape[0] // world_size, ] * world_size
rank_split_sizes[-1] = ray_batch['ray_d'].shape[0] - sum(rank_split_sizes[:-1])
for key in ray_batch:
if torch.is_tensor(ray_batch[key]):
ray_batch[key] = torch.split(ray_batch[key], rank_split_sizes)[rank].to(rank)
# split into chunks and render inside each process
ray_batch_split = OrderedDict()
for key in ray_batch:
if torch.is_tensor(ray_batch[key]):
ray_batch_split[key] = torch.split(ray_batch[key], chunk_size)
# forward and backward
ret_merge_chunk = [OrderedDict() for _ in range(models['cascade_level'])]
for s in range(len(ray_batch_split['ray_d'])):
ray_o = ray_batch_split['ray_o'][s]
ray_d = ray_batch_split['ray_d'][s]
min_depth = ray_batch_split['min_depth'][s]
dots_sh = list(ray_d.shape[:-1])
for m in range(models['cascade_level']):
net = models['net_{}'.format(m)]
# sample depths
N_samples = models['cascade_samples'][m]
if m == 0:
# foreground depth
fg_far_depth = intersect_sphere(ray_o, ray_d) # [...,]
# fg_near_depth = 0.18 * torch.ones_like(fg_far_depth)
fg_near_depth = min_depth # [..., 3]
step = (fg_far_depth - fg_near_depth) / (N_samples - 1)
fg_depth = torch.stack([fg_near_depth + i * step for i in range(N_samples)], dim=-1) # [..., N_samples]
# background depth
bg_depth = torch.linspace(0., 1., N_samples).view(
[1, ] * len(dots_sh) + [N_samples,]).expand(dots_sh + [N_samples,]).to(rank)
# delete unused memory
del fg_near_depth
del step
torch.cuda.empty_cache()
else:
# sample pdf and concat with earlier samples
fg_weights = ret['fg_weights'].clone().detach()
fg_depth_mid = .5 * (fg_depth[..., 1:] + fg_depth[..., :-1]) # [..., N_samples-1]
fg_weights = fg_weights[..., 1:-1] # [..., N_samples-2]
fg_depth_samples = sample_pdf(bins=fg_depth_mid, weights=fg_weights,
N_samples=N_samples, det=True) # [..., N_samples]
fg_depth, _ = torch.sort(torch.cat((fg_depth, fg_depth_samples), dim=-1))
# sample pdf and concat with earlier samples
bg_weights = ret['bg_weights'].clone().detach()
bg_depth_mid = .5 * (bg_depth[..., 1:] + bg_depth[..., :-1])
bg_weights = bg_weights[..., 1:-1] # [..., N_samples-2]
bg_depth_samples = sample_pdf(bins=bg_depth_mid, weights=bg_weights,
N_samples=N_samples, det=True) # [..., N_samples]
bg_depth, _ = torch.sort(torch.cat((bg_depth, bg_depth_samples), dim=-1))
# delete unused memory
del fg_weights
del fg_depth_mid
del fg_depth_samples
del bg_weights
del bg_depth_mid
del bg_depth_samples
torch.cuda.empty_cache()
with torch.no_grad():
ret = net(ray_o, ray_d, fg_far_depth, fg_depth, bg_depth)
for key in ret:
if key not in ['fg_weights', 'bg_weights']:
if torch.is_tensor(ret[key]):
if key not in ret_merge_chunk[m]:
ret_merge_chunk[m][key] = [ret[key].cpu(), ]
else:
ret_merge_chunk[m][key].append(ret[key].cpu())
ret[key] = None
# clean unused memory
torch.cuda.empty_cache()
# merge results from different chunks
for m in range(len(ret_merge_chunk)):
for key in ret_merge_chunk[m]:
ret_merge_chunk[m][key] = torch.cat(ret_merge_chunk[m][key], dim=0)
# merge results from different processes
if rank == 0:
ret_merge_rank = [OrderedDict() for _ in range(len(ret_merge_chunk))]
for m in range(len(ret_merge_chunk)):
for key in ret_merge_chunk[m]:
# generate tensors to store results from other processes
sh = list(ret_merge_chunk[m][key].shape[1:])
ret_merge_rank[m][key] = [torch.zeros(*[size,]+sh, dtype=torch.float32) for size in rank_split_sizes]
torch.distributed.gather(ret_merge_chunk[m][key], ret_merge_rank[m][key])
ret_merge_rank[m][key] = torch.cat(ret_merge_rank[m][key], dim=0).reshape(
(ray_sampler.H, ray_sampler.W, -1)).squeeze()
# print(m, key, ret_merge_rank[m][key].shape)
else: # send results to main process
for m in range(len(ret_merge_chunk)):
for key in ret_merge_chunk[m]:
torch.distributed.gather(ret_merge_chunk[m][key])
# only rank 0 program returns
if rank == 0:
return ret_merge_rank
else:
return None
def setup(rank, world_size):
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '12355'
# initialize the process group
torch.distributed.init_process_group("gloo", rank=rank, world_size=world_size)
def cleanup():
torch.distributed.destroy_process_group()
def ddp_test_nerf(rank, args):
###### set up multi-processing
setup(rank, args.world_size)
###### set up logger
logger = logging.getLogger(__package__)
setup_logger()
###### decide chunk size according to gpu memory
if torch.cuda.get_device_properties(rank).total_memory / 1e9 > 14:
logger.info('setting batch size according to 24G gpu')
args.N_rand = 1024
args.chunk_size = 8192
else:
logger.info('setting batch size according to 12G gpu')
args.N_rand = 512
args.chunk_size = 4096
###### create network and wrap in ddp; each process should do this
# fix random seed just to make sure the network is initialized with same weights at different processes
torch.manual_seed(777)
# very important!!! otherwise it might introduce extra memory in rank=0 gpu
torch.cuda.set_device(rank)
models = OrderedDict()
models['cascade_level'] = args.cascade_level
models['cascade_samples'] = [int(x.strip()) for x in args.cascade_samples.split(',')]
for m in range(models['cascade_level']):
net = NerfNet(args).to(rank)
net = DDP(net, device_ids=[rank], output_device=rank)
optim = torch.optim.Adam(net.parameters(), lr=args.lrate)
models['net_{}'.format(m)] = net
models['optim_{}'.format(m)] = optim
start = -1
###### load pretrained weights; each process should do this
if (args.ckpt_path is not None) and (os.path.isfile(args.ckpt_path)):
ckpts = [args.ckpt_path]
else:
ckpts = [os.path.join(args.basedir, args.expname, f)
for f in sorted(os.listdir(os.path.join(args.basedir, args.expname))) if f.endswith('.pth')]
def path2iter(path):
tmp = os.path.basename(path)[:-4]
idx = tmp.rfind('_')
return int(tmp[idx + 1:])
ckpts = sorted(ckpts, key=path2iter)
logger.info('Found ckpts: {}'.format(ckpts))
if len(ckpts) > 0 and not args.no_reload:
fpath = ckpts[-1]
logger.info('Reloading from: {}'.format(fpath))
start = path2iter(fpath)
# configure map_location properly for different processes
map_location = {'cuda:%d' % 0: 'cuda:%d' % rank}
to_load = torch.load(fpath, map_location=map_location)
for m in range(models['cascade_level']):
for name in ['net_{}'.format(m), 'optim_{}'.format(m)]:
models[name].load_state_dict(to_load[name])
models[name].load_state_dict(to_load[name])
render_splits = [x.strip() for x in args.render_splits.strip().split(',')]
# start testing
for split in render_splits:
out_dir = os.path.join(args.basedir, args.expname,
'render_{}_{:06d}'.format(split, start))
if rank == 0:
os.makedirs(out_dir, exist_ok=True)
###### load data and create ray samplers; each process should do this
ray_samplers = load_data_split(args.datadir, args.scene, split, try_load_min_depth=args.load_min_depth)
for idx in range(len(ray_samplers)):
### each process should do this; but only main process merges the results
fname = '{:06d}.png'.format(idx)
if ray_samplers[idx].img_path is not None:
fname = os.path.basename(ray_samplers[idx].img_path)
if os.path.isfile(os.path.join(out_dir, fname)):
logger.info('Skipping {}'.format(fname))
continue
time0 = time.time()
ret = render_single_image(rank, args.world_size, models, ray_samplers[idx], args.chunk_size)
dt = time.time() - time0
if rank == 0: # only main process should do this
logger.info('Rendered {} in {} seconds'.format(fname, dt))
# only save last level
im = ret[-1]['rgb'].numpy()
# compute psnr if ground-truth is available
if ray_samplers[idx].img_path is not None:
gt_im = ray_samplers[idx].get_img()
psnr = mse2psnr(np.mean((gt_im - im) * (gt_im - im)))
logger.info('{}: psnr={}'.format(fname, psnr))
im = to8b(im)
imageio.imwrite(os.path.join(out_dir, fname), im)
# im = ret[-1]['diffuse_rgb'].numpy()
# im = to8b(im)
# imageio.imwrite(os.path.join(out_dir, 'diffuse_' + fname), im)
im = ret[-1]['fg_rgb'].numpy()
im = to8b(im)
imageio.imwrite(os.path.join(out_dir, 'fg_' + fname), im)
im = ret[-1]['bg_rgb'].numpy()
im = to8b(im)
imageio.imwrite(os.path.join(out_dir, 'bg_' + fname), im)
im = ret[-1]['fg_depth'].numpy()
im = colorize_np(im, cmap_name='jet', append_cbar=True)
im = to8b(im)
imageio.imwrite(os.path.join(out_dir, 'fg_depth_' + fname), im)
im = ret[-1]['bg_depth'].numpy()
im = colorize_np(im, cmap_name='jet', append_cbar=True)
im = to8b(im)
imageio.imwrite(os.path.join(out_dir, 'bg_depth_' + fname), im)
torch.cuda.empty_cache()
# clean up for multi-processing
cleanup()
def test():
parser = config_parser()
args = parser.parse_args()
logger.info(parser.format_values())
if args.world_size == -1:
args.world_size = torch.cuda.device_count()
logger.info('Using # gpus: {}'.format(args.world_size))
torch.multiprocessing.spawn(ddp_test_nerf,
args=(args,),
nprocs=args.world_size,
join=True)
if __name__ == '__main__':
setup_logger()
test()

167
nerf_network.py
Unescape Escape View File

@ -0,0 +1,167 @@
import torch
import torch.nn as nn
# import torch.nn.functional as F
# import numpy as np
from collections import OrderedDict
import logging
logger = logging.getLogger(__package__)
class Embedder(nn.Module):
def __init__(self, input_dim, max_freq_log2, N_freqs,
log_sampling=True, include_input=True,
periodic_fns=(torch.sin, torch.cos)):
'''
:param input_dim: dimension of input to be embedded
:param max_freq_log2: log2 of max freq; min freq is 1 by default
:param N_freqs: number of frequency bands
:param log_sampling: if True, frequency bands are linerly sampled in log-space
:param include_input: if True, raw input is included in the embedding
:param periodic_fns: periodic functions used to embed input
'''
super().__init__()
self.input_dim = input_dim
self.include_input = include_input
self.periodic_fns = periodic_fns
self.out_dim = 0
if self.include_input:
self.out_dim += self.input_dim
self.out_dim += self.input_dim * N_freqs * len(self.periodic_fns)
if log_sampling:
self.freq_bands = 2. ** torch.linspace(0., max_freq_log2, N_freqs)
else:
self.freq_bands = torch.linspace(2. ** 0., 2. ** max_freq_log2, N_freqs)
self.freq_bands = self.freq_bands.numpy().tolist()
def forward(self, input):
'''
:param input: tensor of shape [..., self.input_dim]
:return: tensor of shape [..., self.out_dim]
'''
assert (input.shape[-1] == self.input_dim)
out = []
if self.include_input:
out.append(input)
for i in range(len(self.freq_bands)):
freq = self.freq_bands[i]
for p_fn in self.periodic_fns:
out.append(p_fn(input * freq))
out = torch.cat(out, dim=-1)
assert (out.shape[-1] == self.out_dim)
return out
# default tensorflow initialization of linear layers
def weights_init(m):
if isinstance(m, nn.Linear):
nn.init.xavier_uniform_(m.weight.data)
if m.bias is not None:
nn.init.zeros_(m.bias.data)
class MLPNet(nn.Module):
def __init__(self, D=8, W=256, input_ch=3, input_ch_viewdirs=3, skips=[4], use_viewdirs=False
, use_implicit=False):
'''
:param D: network depth
:param W: network width
:param input_ch: input channels for encodings of (x, y, z)
:param input_ch_viewdirs: input channels for encodings of view directions
:param skips: skip connection in network
:param use_viewdirs: if True, will use the view directions as input
'''
super().__init__()
self.use_implicit = use_implicit
if self.use_implicit:
logger.info('Using implicit regularization as well!')
self.input_ch = input_ch
self.input_ch_viewdirs = input_ch_viewdirs
self.use_viewdirs = use_viewdirs
self.skips = skips
self.base_layers = []
dim = self.input_ch
for i in range(D):
self.base_layers.append(
nn.Sequential(nn.Linear(dim, W), nn.ReLU())
)
dim = W
if i in self.skips and i != (D-1): # skip connection after i^th layer
dim += input_ch
self.base_layers = nn.ModuleList(self.base_layers)
# self.base_layers.apply(weights_init) # xavier init
sigma_layers = [nn.Linear(dim, 1), ] # sigma must be positive
self.sigma_layers = nn.Sequential(*sigma_layers)
# self.sigma_layers.apply(weights_init) # xavier init
base_dim = dim
# diffuse color
diffuse_rgb_layers = []
dim = base_dim
for i in range(1):
diffuse_rgb_layers.append(nn.Linear(dim, W))
diffuse_rgb_layers.append(nn.ReLU())
dim = W
diffuse_rgb_layers.append(nn.Linear(dim, 3))
diffuse_rgb_layers.append(nn.Sigmoid())
self.diffuse_rgb_layers = nn.Sequential(*diffuse_rgb_layers)
# self.diffuse_rgb_layers.apply(weights_init)
# specular color
specular_rgb_layers = []
dim = base_dim
base_remap_layers = [nn.Linear(dim, 256), ]
self.base_remap_layers = nn.Sequential(*base_remap_layers)
# self.base_remap_layers.apply(weights_init)
dim = 256 + self.input_ch_viewdirs
for i in range(1):
specular_rgb_layers.append(nn.Linear(dim, W))
specular_rgb_layers.append(nn.ReLU())
dim = W
specular_rgb_layers.append(nn.Linear(dim, 3))
specular_rgb_layers.append(nn.Sigmoid()) # rgb values are normalized to [0, 1]
self.specular_rgb_layers = nn.Sequential(*specular_rgb_layers)
# self.specular_rgb_layers.apply(weights_init)
def forward(self, input):
'''
:param input: [..., input_ch+input_ch_viewdirs]
:return [..., 4]
'''
input_pts = input[..., :self.input_ch]
base = self.base_layers[0](input_pts)
for i in range(len(self.base_layers)-1):
if i in self.skips:
base = torch.cat((input_pts, base), dim=-1)
base = self.base_layers[i+1](base)
sigma = self.sigma_layers(base)
sigma = torch.abs(sigma)
diffuse_rgb = self.diffuse_rgb_layers(base)
base_remap = self.base_remap_layers(base)
input_viewdirs = input[..., -self.input_ch_viewdirs:]
specular_rgb = self.specular_rgb_layers(torch.cat((base_remap, input_viewdirs), dim=-1))
if self.use_implicit:
rgb = specular_rgb
else:
rgb = diffuse_rgb + specular_rgb
ret = OrderedDict([('rgb', rgb),
('diffuse_rgb', diffuse_rgb),
('sigma', sigma.squeeze(-1))])
return ret

249
nerf_sample_ray_split.py
Unescape Escape View File

@ -0,0 +1,249 @@
import numpy as np
from collections import OrderedDict
import torch
import cv2
import imageio
########################################################################################################################
# ray batch sampling
########################################################################################################################
def parse_camera(params):
H, W = params[:2]
intrinsics = params[2:18].reshape((4, 4))
c2w = params[18:34].reshape((4, 4))
return int(W), int(H), intrinsics.astype(np.float32), c2w.astype(np.float32)
def get_rays_single_image(H, W, intrinsics, c2w):
'''
:param H: image height
:param W: image width
:param intrinsics: 4 by 4 intrinsic matrix
:param c2w: 4 by 4 camera to world extrinsic matrix
:return:
'''
u, v = np.meshgrid(np.arange(W), np.arange(H))
u = u.reshape(-1).astype(dtype=np.float32) + 0.5 # add half pixel
v = v.reshape(-1).astype(dtype=np.float32) + 0.5
pixels = np.stack((u, v, np.ones_like(u)), axis=0) # (3, H*W)
rays_d = np.dot(np.linalg.inv(intrinsics[:3, :3]), pixels)
rays_d = np.dot(c2w[:3, :3], rays_d) # (3, H*W)
rays_d = rays_d.transpose((1, 0)) # (H*W, 3)
rays_o = c2w[:3, 3].reshape((1, 3))
rays_o = np.tile(rays_o, (rays_d.shape[0], 1)) # (H*W, 3)
depth = np.linalg.inv(c2w)[2, 3]
depth = depth * np.ones((rays_o.shape[0],), dtype=np.float32) # (H*W,)
return rays_o, rays_d, depth
class RaySamplerSingleImage(object):
def __init__(self, H, W, intrinsics, c2w,
img_path=None,
resolution_level=1,
mask_path=None,
min_depth_path=None,
max_depth=None):
super().__init__()
self.W_orig = W
self.H_orig = H
self.intrinsics_orig = intrinsics
self.c2w_mat = c2w
self.img_path = img_path
self.mask_path = mask_path
self.min_depth_path = min_depth_path
self.max_depth = max_depth
self.resolution_level = -1
self.set_resolution_level(resolution_level)
def set_resolution_level(self, resolution_level):
if resolution_level != self.resolution_level:
self.resolution_level = resolution_level
self.W = self.W_orig // resolution_level
self.H = self.H_orig // resolution_level
self.intrinsics = np.copy(self.intrinsics_orig)
self.intrinsics[:2, :3] /= resolution_level
# only load image at this time
if self.img_path is not None:
self.img = imageio.imread(self.img_path).astype(np.float32) / 255.
self.img = cv2.resize(self.img, (self.W, self.H), interpolation=cv2.INTER_AREA)
self.img = self.img.reshape((-1, 3))
else:
self.img = None
if self.mask_path is not None:
self.mask = imageio.imread(self.mask_path).astype(np.float32) / 255.
self.mask = cv2.resize(self.mask, (self.W, self.H), interpolation=cv2.INTER_NEAREST)
self.mask = self.mask.reshape((-1))
else:
self.mask = None
if self.min_depth_path is not None:
self.min_depth = imageio.imread(self.min_depth_path).astype(np.float32) / 255. * self.max_depth + 1e-4
self.min_depth = cv2.resize(self.min_depth, (self.W, self.H), interpolation=cv2.INTER_LINEAR)
self.min_depth = self.min_depth.reshape((-1))
else:
self.min_depth = None
self.rays_o, self.rays_d, self.depth = get_rays_single_image(self.H, self.W,
self.intrinsics, self.c2w_mat)
def get_img(self):
if self.img is not None:
return self.img.reshape((self.H, self.W, 3))
else:
return None
def get_all(self):
if self.min_depth is not None:
min_depth = self.min_depth
else:
min_depth = 1e-4 * np.ones_like(self.rays_d[..., 0])
ret = OrderedDict([
('ray_o', self.rays_o),
('ray_d', self.rays_d),
('depth', self.depth),
('rgb', self.img),
('mask', self.mask),
('min_depth', min_depth)
])
# return torch tensors
for k in ret:
if ret[k] is not None:
ret[k] = torch.from_numpy(ret[k])
return ret
def random_sample(self, N_rand, center_crop=False):
'''
:param N_rand: number of rays to be casted
:return:
'''
if center_crop:
half_H = self.H // 2
half_W = self.W // 2
quad_H = half_H // 2
quad_W = half_W // 2
# pixel coordinates
u, v = np.meshgrid(np.arange(half_W-quad_W, half_W+quad_W),
np.arange(half_H-quad_H, half_H+quad_H))
u = u.reshape(-1)
v = v.reshape(-1)
select_inds = np.random.choice(u.shape[0], size=(N_rand,), replace=False)
# Convert back to original image
select_inds = v[select_inds] * self.W + u[select_inds]
else:
# Random from one image
select_inds = np.random.choice(self.H*self.W, size=(N_rand,), replace=False)
rays_o = self.rays_o[select_inds, :] # [N_rand, 3]
rays_d = self.rays_d[select_inds, :] # [N_rand, 3]
depth = self.depth[select_inds] # [N_rand, ]
if self.img is not None:
rgb = self.img[select_inds, :] # [N_rand, 3]
else:
rgb = None
if self.mask is not None:
mask = self.mask[select_inds]
else:
mask = None
if self.min_depth is not None:
min_depth = self.min_depth[select_inds]
else:
min_depth = 1e-4 * np.ones_like(rays_d[..., 0])
ret = OrderedDict([
('ray_o', rays_o),
('ray_d', rays_d),
('depth', depth),
('rgb', rgb),
('mask', mask),
('min_depth', min_depth)
])
# return torch tensors
for k in ret:
if ret[k] is not None:
ret[k] = torch.from_numpy(ret[k])
return ret
# def random_sample_patches(self, N_patch, r_patch=16, center_crop=False):
# '''
# :param N_patch: number of patches to be sampled
# :param r_patch: patch size will be (2*r_patch+1)*(2*r_patch+1)
# :return:
# '''
# # even size patch
# # offsets to center pixels
# u, v = np.meshgrid(np.arange(-r_patch, r_patch),
# np.arange(-r_patch, r_patch))
# u = u.reshape(-1)
# v = v.reshape(-1)
# offsets = v * self.W + u
# # center pixel coordinates
# u_min = r_patch
# u_max = self.W - r_patch
# v_min = r_patch
# v_max = self.H - r_patch
# if center_crop:
# u_min = self.W // 4 + r_patch
# u_max = self.W - self.W // 4 - r_patch
# v_min = self.H // 4 + r_patch
# v_max = self.H - self.H // 4 - r_patch
# u, v = np.meshgrid(np.arange(u_min, u_max, r_patch),
# np.arange(v_min, v_max, r_patch))
# u = u.reshape(-1)
# v = v.reshape(-1)
# select_inds = np.random.choice(u.shape[0], size=(N_patch,), replace=False)
# # Convert back to original image
# select_inds = v[select_inds] * self.W + u[select_inds]
# # pick patches
# select_inds = np.stack([select_inds + shift for shift in offsets], axis=1)
# select_inds = select_inds.reshape(-1)
# rays_o = self.rays_o[select_inds, :] # [N_rand, 3]
# rays_d = self.rays_d[select_inds, :] # [N_rand, 3]
# depth = self.depth[select_inds] # [N_rand, ]
# if self.img is not None:
# rgb = self.img[select_inds, :] # [N_rand, 3]
# # ### debug
# # import imageio
# # imgs = rgb.reshape((N_patch, r_patch*2, r_patch*2, -1))
# # for kk in range(imgs.shape[0]):
# # imageio.imwrite('./debug_{}.png'.format(kk), imgs[kk])
# # ###
# else:
# rgb = None
# ret = OrderedDict([
# ('ray_o', rays_o),
# ('ray_d', rays_d),
# ('depth', depth),
# ('rgb', rgb)
# ])
# # return torch tensors
# for k in ret:
# ret[k] = torch.from_numpy(ret[k])
# return ret

181
old_scripts/data_loader.py
Unescape Escape View File

@ -0,0 +1,181 @@
import os
import numpy as np
import imageio
from collections import OrderedDict
import logging
logger = logging.getLogger(__package__)
########################################################################################################################
# camera coordinate system: x-->right, y-->down, z-->scene (opencv/colmap convention)
# poses is camera-to-world
########################################################################################################################
def load_data(basedir, scene, testskip=8):
def parse_txt(filename):
assert os.path.isfile(filename)
nums = open(filename).read().split()
return np.array([float(x) for x in nums]).reshape([4, 4]).astype(np.float32)
def dir2poses(posedir):
poses = np.stack(
[parse_txt(os.path.join(posedir, f)) for f in sorted(os.listdir(posedir)) if f.endswith('txt')], 0)
poses = poses.astype(np.float32)
return poses
def dir2intrinsics(intrinsicdir):
intrinsics = np.stack(
[parse_txt(os.path.join(intrinsicdir, f)) for f in sorted(os.listdir(intrinsicdir)) if f.endswith('txt')], 0)
intrinsics = intrinsics.astype(np.float32)
return intrinsics
intrinsics = dir2intrinsics('{}/{}/train/intrinsics'.format(basedir, scene))
testintrinsics = dir2poses('{}/{}/test/intrinsics'.format(basedir, scene))
testintrinsics = testintrinsics[::testskip]
valintrinsics = dir2poses('{}/{}/validation/intrinsics'.format(basedir, scene))
valintrinsics = valintrinsics[::testskip]
print(intrinsics.shape, testintrinsics.shape, valintrinsics.shape)
poses = dir2poses('{}/{}/train/pose'.format(basedir, scene))
testposes = dir2poses('{}/{}/test/pose'.format(basedir, scene))
testposes = testposes[::testskip]
valposes = dir2poses('{}/{}/validation/pose'.format(basedir, scene))
valposes = valposes[::testskip]
print(poses.shape, testposes.shape, valposes.shape)
imgd = '{}/{}/train/rgb'.format(basedir, scene)
imgfiles = ['{}/{}'.format(imgd, f)
for f in sorted(os.listdir(imgd)) if f.endswith('png') or f.endswith('jpg')]
imgs = [imageio.imread(f).astype(np.float32)[..., :3] / 255. for f in imgfiles]
maskd = '{}/{}/train/mask'.format(basedir, scene)
if os.path.isdir(maskd):
logger.info('Loading mask from: {}'.format(maskd))
maskfiles = ['{}/{}'.format(maskd, f)
for f in sorted(os.listdir(maskd)) if f.endswith('png') or f.endswith('jpg')]
masks = [imageio.imread(f).astype(np.float32) / 255. for f in maskfiles]
else:
masks = [None for im in imgs]
# load min_depth map
min_depthd = '{}/{}/train/min_depth'.format(basedir, scene)
if os.path.isdir(min_depthd):
logger.info('Loading min_depth from: {}'.format(min_depthd))
max_depth = float(open('{}/{}/train/max_depth.txt'.format(basedir, scene)).readline().strip())
min_depthfiles = ['{}/{}'.format(min_depthd, f)
for f in sorted(os.listdir(min_depthd)) if f.endswith('png') or f.endswith('jpg')]
min_depths = [imageio.imread(f).astype(np.float32) / 255. * max_depth + 1e-4 for f in min_depthfiles]
else:
min_depths = [None for im in imgs]
testimgd = '{}/{}/test/rgb'.format(basedir, scene)
testimgfiles = ['{}/{}'.format(testimgd, f)
for f in sorted(os.listdir(testimgd)) if f.endswith('png') or f.endswith('jpg')]
testimgs = [imageio.imread(f).astype(np.float32)[..., :3] / 255. for f in testimgfiles]
testimgfiles = testimgfiles[::testskip]
testimgs = testimgs[::testskip]
testmaskd = '{}/{}/test/mask'.format(basedir, scene)
if os.path.isdir(testmaskd):
logger.info('Loading mask from: {}'.format(testmaskd))
testmaskfiles = ['{}/{}'.format(testmaskd, f)
for f in sorted(os.listdir(testmaskd)) if f.endswith('png') or f.endswith('jpg')]
testmasks = [imageio.imread(f).astype(np.float32) / 255. for f in testmaskfiles]
else:
testmasks = [None for im in testimgs]
# load min_depth map
min_depthd = '{}/{}/test/min_depth'.format(basedir, scene)
if os.path.isdir(min_depthd):
logger.info('Loading min_depth from: {}'.format(min_depthd))
max_depth = float(open('{}/{}/test/max_depth.txt'.format(basedir, scene)).readline().strip())
min_depthfiles = ['{}/{}'.format(min_depthd, f)
for f in sorted(os.listdir(min_depthd)) if f.endswith('png') or f.endswith('jpg')]
test_min_depths = [imageio.imread(f).astype(np.float32) / 255. * max_depth + 1e-4 for f in min_depthfiles]
else:
test_min_depths = [None for im in testimgs]
valimgd = '{}/{}/validation/rgb'.format(basedir, scene)
valimgfiles = ['{}/{}'.format(valimgd, f)
for f in sorted(os.listdir(valimgd)) if f.endswith('png') or f.endswith('jpg')]
valimgs = [imageio.imread(f).astype(np.float32)[..., :3] / 255. for f in valimgfiles]
valimgfiles = valimgfiles[::testskip]
valimgs = valimgs[::testskip]
valmaskd = '{}/{}/validation/mask'.format(basedir, scene)
if os.path.isdir(valmaskd):
logger.info('Loading mask from: {}'.format(valmaskd))
valmaskfiles = ['{}/{}'.format(valmaskd, f)
for f in sorted(os.listdir(valmaskd)) if f.endswith('png') or f.endswith('jpg')]
valmasks = [imageio.imread(f).astype(np.float32) / 255. for f in valmaskfiles]
else:
valmasks = [None for im in valimgs]
# load min_depth map
min_depthd = '{}/{}/validation/min_depth'.format(basedir, scene)
if os.path.isdir(min_depthd):
logger.info('Loading min_depth from: {}'.format(min_depthd))
max_depth = float(open('{}/{}/validation/max_depth.txt'.format(basedir, scene)).readline().strip())
min_depthfiles = ['{}/{}'.format(min_depthd, f)
for f in sorted(os.listdir(min_depthd)) if f.endswith('png') or f.endswith('jpg')]
val_min_depths = [imageio.imread(f).astype(np.float32) / 255. * max_depth + 1e-4 for f in min_depthfiles]
else:
val_min_depths = [None for im in valimgs]
# data format for training/testing
print(len(imgs), len(testimgs), len(valimgs))
all_imgs = imgs + valimgs + testimgs
all_masks = masks + valmasks + testmasks
all_min_depths = min_depths + val_min_depths + test_min_depths
all_paths = imgfiles + valimgfiles + testimgfiles
counts = [0] + [len(x) for x in [imgs, valimgs, testimgs]]
counts = np.cumsum(counts)
i_split = [list(np.arange(counts[i], counts[i+1])) for i in range(3)]
intrinsics = np.concatenate([intrinsics, valintrinsics, testintrinsics], 0)
poses = np.concatenate([poses, valposes, testposes], 0)
img_sizes = np.stack([np.array(x.shape[:2]) for x in all_imgs], axis=0) # [H, W]
cnt = len(all_imgs)
all_cams = np.concatenate((img_sizes.astype(dtype=np.float32), intrinsics.reshape((cnt, -1)), poses.reshape((cnt, -1))), axis=1)
if os.path.isdir('{}/{}/camera_path/intrinsics'.format(basedir, scene)):
camera_path_intrinsics = dir2poses('{}/{}/camera_path/intrinsics'.format(basedir, scene))
camera_path_poses = dir2poses('{}/{}/camera_path/pose'.format(basedir, scene))
# assume centered principal points
# img_sizes = np.stack((camera_path_intrinsics[:, 1, 2]*2, camera_path_intrinsics[:, 0, 2]*2), axis=1) # [H, W]
# img_sizes = np.int32(img_sizes)
H = all_cams[0, 0]
W = all_cams[0, 1]
img_sizes = np.stack((np.ones_like(camera_path_intrinsics[:, 1, 2])*H, np.ones_like(camera_path_intrinsics[:, 0, 2])*W), axis=1) # [H, W]
cnt = len(camera_path_intrinsics)
render_cams = np.concatenate(
(img_sizes.astype(dtype=np.float32), camera_path_intrinsics.reshape((cnt, -1)), camera_path_poses.reshape((cnt, -1))),
axis=1)
else:
render_cams = None
print(all_cams.shape)
data = OrderedDict([('images', all_imgs),
('masks', all_masks),
('paths', all_paths),
('min_depths', all_min_depths),
('cameras', all_cams),
('i_train', i_split[0]),
('i_val', i_split[1]),
('i_test', i_split[2]),
('render_cams', render_cams)])
logger.info('Data statistics:')
logger.info('\t # of training views: {}'.format(len(data['i_train'])))
logger.info('\t # of validation views: {}'.format(len(data['i_val'])))
logger.info('\t # of test views: {}'.format(len(data['i_test'])))
if data['render_cams'] is not None:
logger.info('\t # of render cameras: {}'.format(len(data['render_cams'])))
return data

617
old_scripts/ddp_run_nerf.py
Unescape Escape View File

@ -0,0 +1,617 @@
import torch
import torch.nn as nn
import torch.optim
import torch.distributed
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.multiprocessing
import os
from collections import OrderedDict
from ddp_model import NerfNet
import time
# from data_loader import load_data
# from nerf_sample_ray import RaySamplerSingleImage
from data_loader_split import load_data_split
import numpy as np
from tensorboardX import SummaryWriter
from utils import img2mse, mse2psnr, img_HWC2CHW, colorize, TINY_NUMBER
import logging
logger = logging.getLogger(__package__)
def setup_logger():
# create logger
logger = logging.getLogger(__package__)
logger.setLevel(logging.DEBUG)
# create console handler and set level to debug
ch = logging.StreamHandler()
ch.setLevel(logging.DEBUG)
# create formatter
formatter = logging.Formatter('%(asctime)s [%(levelname)s] %(name)s: %(message)s')
# add formatter to ch
ch.setFormatter(formatter)
# add ch to logger
logger.addHandler(ch)
def intersect_sphere(ray_o, ray_d):
'''
ray_o, ray_d: [..., 3]
compute the depth of the intersection point between this ray and unit sphere
'''
# note: d1 becomes negative if this mid point is behind camera
d1 = -torch.sum(ray_d * ray_o, dim=-1) / torch.sum(ray_d * ray_d, dim=-1)
p = ray_o + d1.unsqueeze(-1) * ray_d
# consider the case where the ray does not intersect the sphere
ray_d_cos = 1. / torch.norm(ray_d, dim=-1)
d2 = torch.sqrt(1. - torch.sum(p * p, dim=-1)) * ray_d_cos
return d1 + d2
def perturb_samples(z_vals):
# get intervals between samples
mids = .5 * (z_vals[..., 1:] + z_vals[..., :-1])
upper = torch.cat([mids, z_vals[..., -1:]], dim=-1)
lower = torch.cat([z_vals[..., 0:1], mids], dim=-1)
# uniform samples in those intervals
t_rand = torch.rand_like(z_vals)
z_vals = lower + (upper - lower) * t_rand # [N_rays, N_samples]
return z_vals
def sample_pdf(bins, weights, N_samples, det=False):
'''
:param bins: tensor of shape [..., M+1], M is the number of bins
:param weights: tensor of shape [..., M]
:param N_samples: number of samples along each ray
:param det: if True, will perform deterministic sampling
:return: [..., N_samples]
'''
# Get pdf
weights = weights + TINY_NUMBER # prevent nans
pdf = weights / torch.sum(weights, dim=-1, keepdim=True) # [..., M]
cdf = torch.cumsum(pdf, dim=-1) # [..., M]
cdf = torch.cat([torch.zeros_like(cdf[..., 0:1]), cdf], dim=-1) # [..., M+1]
# Take uniform samples
dots_sh = list(weights.shape[:-1])
M = weights.shape[-1]
min_cdf = 0.00
max_cdf = 1.00 # prevent outlier samples
if det:
u = torch.linspace(min_cdf, max_cdf, N_samples, device=bins.device)
u = u.view([1]*len(dots_sh) + [N_samples]).expand(dots_sh + [N_samples,]) # [..., N_samples]
else:
sh = dots_sh + [N_samples]
u = torch.rand(*sh, device=bins.device) * (max_cdf - min_cdf) + min_cdf # [..., N_samples]
# Invert CDF
# [..., N_samples, 1] >= [..., 1, M] ----> [..., N_samples, M] ----> [..., N_samples,]
above_inds = torch.sum(u.unsqueeze(-1) >= cdf[..., :M].unsqueeze(-2), dim=-1).long()
# random sample inside each bin
below_inds = torch.clamp(above_inds-1, min=0)
inds_g = torch.stack((below_inds, above_inds), dim=-1) # [..., N_samples, 2]
cdf = cdf.unsqueeze(-2).expand(dots_sh + [N_samples, M+1]) # [..., N_samples, M+1]
cdf_g = torch.gather(input=cdf, dim=-1, index=inds_g) # [..., N_samples, 2]
bins = bins.unsqueeze(-2).expand(dots_sh + [N_samples, M+1]) # [..., N_samples, M+1]
bins_g = torch.gather(input=bins, dim=-1, index=inds_g) # [..., N_samples, 2]
# fix numeric issue
denom = cdf_g[..., 1] - cdf_g[..., 0] # [..., N_samples]
denom = torch.where(denom<TINY_NUMBER, torch.ones_like(denom), denom)
t = (u - cdf_g[..., 0]) / denom
samples = bins_g[..., 0] + t * (bins_g[..., 1] - bins_g[..., 0] + TINY_NUMBER)
return samples
def render_single_image(rank, world_size, models, ray_sampler, chunk_size):
##### parallel rendering of a single image
ray_batch = ray_sampler.get_all()
# split into ranks; make sure different processes don't overlap
rank_split_sizes = [ray_batch['ray_d'].shape[0] // world_size, ] * world_size
rank_split_sizes[-1] = ray_batch['ray_d'].shape[0] - sum(rank_split_sizes[:-1])
for key in ray_batch:
if torch.is_tensor(ray_batch[key]):
ray_batch[key] = torch.split(ray_batch[key], rank_split_sizes)[rank].to(rank)
# split into chunks and render inside each process
ray_batch_split = OrderedDict()
for key in ray_batch:
if torch.is_tensor(ray_batch[key]):
ray_batch_split[key] = torch.split(ray_batch[key], chunk_size)
# forward and backward
ret_merge_chunk = [OrderedDict() for _ in range(models['cascade_level'])]
for s in range(len(ray_batch_split['ray_d'])):
ray_o = ray_batch_split['ray_o'][s]
ray_d = ray_batch_split['ray_d'][s]
min_depth = ray_batch_split['min_depth'][s]
dots_sh = list(ray_d.shape[:-1])
for m in range(models['cascade_level']):
net = models['net_{}'.format(m)]
# sample depths
N_samples = models['cascade_samples'][m]
if m == 0:
# foreground depth
fg_far_depth = intersect_sphere(ray_o, ray_d) # [...,]
# fg_near_depth = 0.18 * torch.ones_like(fg_far_depth)
fg_near_depth = min_depth # [..., 3]
step = (fg_far_depth - fg_near_depth) / (N_samples - 1)
fg_depth = torch.stack([fg_near_depth + i * step for i in range(N_samples)], dim=-1) # [..., N_samples]
# background depth
bg_depth = torch.linspace(0., 1., N_samples).view(
[1, ] * len(dots_sh) + [N_samples,]).expand(dots_sh + [N_samples,]).to(rank)
# delete unused memory
del fg_near_depth
del step
torch.cuda.empty_cache()
else:
# sample pdf and concat with earlier samples
fg_weights = ret['fg_weights'].clone().detach()
fg_depth_mid = .5 * (fg_depth[..., 1:] + fg_depth[..., :-1]) # [..., N_samples-1]
fg_weights = fg_weights[..., 1:-1] # [..., N_samples-2]
fg_depth_samples = sample_pdf(bins=fg_depth_mid, weights=fg_weights,
N_samples=N_samples, det=True) # [..., N_samples]
fg_depth, _ = torch.sort(torch.cat((fg_depth, fg_depth_samples), dim=-1))
# sample pdf and concat with earlier samples
bg_weights = ret['bg_weights'].clone().detach()
bg_depth_mid = .5 * (bg_depth[..., 1:] + bg_depth[..., :-1])
bg_weights = bg_weights[..., 1:-1] # [..., N_samples-2]
bg_depth_samples = sample_pdf(bins=bg_depth_mid, weights=bg_weights,
N_samples=N_samples, det=True) # [..., N_samples]
bg_depth, _ = torch.sort(torch.cat((bg_depth, bg_depth_samples), dim=-1))
# delete unused memory
del fg_weights
del fg_depth_mid
del fg_depth_samples
del bg_weights
del bg_depth_mid
del bg_depth_samples
torch.cuda.empty_cache()
with torch.no_grad():
ret = net(ray_o, ray_d, fg_far_depth, fg_depth, bg_depth)
for key in ret:
if key not in ['fg_weights', 'bg_weights']:
if torch.is_tensor(ret[key]):
if key not in ret_merge_chunk[m]:
ret_merge_chunk[m][key] = [ret[key].cpu(), ]
else:
ret_merge_chunk[m][key].append(ret[key].cpu())
ret[key] = None
# clean unused memory
torch.cuda.empty_cache()
# merge results from different chunks
for m in range(len(ret_merge_chunk)):
for key in ret_merge_chunk[m]:
ret_merge_chunk[m][key] = torch.cat(ret_merge_chunk[m][key], dim=0)
# merge results from different processes
if rank == 0:
ret_merge_rank = [OrderedDict() for _ in range(len(ret_merge_chunk))]
for m in range(len(ret_merge_chunk)):
for key in ret_merge_chunk[m]:
# generate tensors to store results from other processes
sh = list(ret_merge_chunk[m][key].shape[1:])
ret_merge_rank[m][key] = [torch.zeros(*[size,]+sh, dtype=torch.float32) for size in rank_split_sizes]
torch.distributed.gather(ret_merge_chunk[m][key], ret_merge_rank[m][key])
ret_merge_rank[m][key] = torch.cat(ret_merge_rank[m][key], dim=0).reshape(
(ray_sampler.H, ray_sampler.W, -1)).squeeze()
# print(m, key, ret_merge_rank[m][key].shape)
else: # send results to main process
for m in range(len(ret_merge_chunk)):
for key in ret_merge_chunk[m]:
torch.distributed.gather(ret_merge_chunk[m][key])
# only rank 0 program returns
if rank == 0:
return ret_merge_rank
else:
return None
def log_view_to_tb(writer, global_step, log_data, gt_img, mask, prefix=''):
rgb_im = img_HWC2CHW(torch.from_numpy(gt_img))
writer.add_image(prefix + 'rgb_gt', rgb_im, global_step)
for m in range(len(log_data)):
rgb_im = img_HWC2CHW(log_data[m]['rgb'])
rgb_im = torch.clamp(rgb_im, min=0., max=1.) # just in case diffuse+specular>1
writer.add_image(prefix + 'level_{}/rgb'.format(m), rgb_im, global_step)
rgb_im = img_HWC2CHW(log_data[m]['fg_rgb'])
rgb_im = torch.clamp(rgb_im, min=0., max=1.) # just in case diffuse+specular>1
writer.add_image(prefix + 'level_{}/fg_rgb'.format(m), rgb_im, global_step)
depth = log_data[m]['fg_depth']
depth_im = img_HWC2CHW(colorize(depth, cmap_name='jet', append_cbar=True,
mask=mask))
writer.add_image(prefix + 'level_{}/fg_depth'.format(m), depth_im, global_step)
rgb_im = img_HWC2CHW(log_data[m]['bg_rgb'])
rgb_im = torch.clamp(rgb_im, min=0., max=1.) # just in case diffuse+specular>1
writer.add_image(prefix + 'level_{}/bg_rgb'.format(m), rgb_im, global_step)
depth = log_data[m]['bg_depth']
depth_im = img_HWC2CHW(colorize(depth, cmap_name='jet', append_cbar=True,
mask=mask))
writer.add_image(prefix + 'level_{}/bg_depth'.format(m), depth_im, global_step)
bg_lambda = log_data[m]['bg_lambda']
bg_lambda_im = img_HWC2CHW(colorize(bg_lambda, cmap_name='hot', append_cbar=True,
mask=mask))
writer.add_image(prefix + 'level_{}/bg_lambda'.format(m), bg_lambda_im, global_step)
def setup(rank, world_size):
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '12355'
# initialize the process group
torch.distributed.init_process_group("gloo", rank=rank, world_size=world_size)
def cleanup():
torch.distributed.destroy_process_group()
def ddp_train_nerf(rank, args):
###### set up multi-processing
setup(rank, args.world_size)
###### set up logger
logger = logging.getLogger(__package__)
setup_logger()
###### Create log dir and copy the config file
if rank == 0:
os.makedirs(os.path.join(args.basedir, args.expname), exist_ok=True)
f = os.path.join(args.basedir, args.expname, 'args.txt')
with open(f, 'w') as file:
for arg in sorted(vars(args)):
attr = getattr(args, arg)
file.write('{} = {}\n'.format(arg, attr))
if args.config is not None:
f = os.path.join(args.basedir, args.expname, 'config.txt')
with open(f, 'w') as file:
file.write(open(args.config, 'r').read())
torch.distributed.barrier()
###### load data and create ray samplers; each process should do this
# data = load_data(args.datadir, args.scene, args.testskip)
# ray_samplers = []
# for i in data['i_train']:
# ray_samplers.append(RaySamplerSingleImage(cam_params=data['cameras'][i],
# img=data['images'][i],
# img_path=data['paths'][i],
# mask=data['masks'][i],
# min_depth=data['min_depths'][i]))
#
# val_ray_samplers = []
# for i in data['i_val']:
# val_ray_samplers.append(RaySamplerSingleImage(cam_params=data['cameras'][i],
# img=data['images'][i],
# img_path=data['paths'][i],
# mask=data['masks'][i],
# min_depth=data['min_depths'][i]))
# # free memory
# del data
ray_samplers = load_data_split(args.datadir, args.scene, split='train')
val_ray_samplers = load_data_split(args.datadir, args.scene, split='validation')
###### create network and wrap in ddp; each process should do this
# fix random seed just to make sure the network is initialized with same weights at different processes
torch.manual_seed(777)
# very important!!! otherwise it might introduce extra memory in rank=0 gpu
torch.cuda.set_device(rank)
models = OrderedDict()
models['cascade_level'] = args.cascade_level
models['cascade_samples'] = [int(x.strip()) for x in args.cascade_samples.split(',')]
for m in range(models['cascade_level']):
net = NerfNet(args).to(rank)
net = DDP(net, device_ids=[rank], output_device=rank)
optim = torch.optim.Adam(net.parameters(), lr=args.lrate)
models['net_{}'.format(m)] = net
models['optim_{}'.format(m)] = optim
start = -1
###### load pretrained weights; each process should do this
if (args.ckpt_path is not None) and (os.path.isfile(args.ckpt_path)):
ckpts = [args.ckpt_path]
else:
ckpts = [os.path.join(args.basedir, args.expname, f)
for f in sorted(os.listdir(os.path.join(args.basedir, args.expname))) if f.endswith('.pth')]
def path2iter(path):
tmp = os.path.basename(path)[:-4]
idx = tmp.rfind('_')
return int(tmp[idx + 1:])
ckpts = sorted(ckpts, key=path2iter)
logger.info('Found ckpts: {}'.format(ckpts))
if len(ckpts) > 0 and not args.no_reload:
fpath = ckpts[-1]
logger.info('Reloading from: {}'.format(fpath))
start = path2iter(fpath)
# configure map_location properly for different processes
map_location = {'cuda:%d' % 0: 'cuda:%d' % rank}
to_load = torch.load(fpath, map_location=map_location)
for m in range(models['cascade_level']):
for name in ['net_{}'.format(m), 'optim_{}'.format(m)]:
models[name].load_state_dict(to_load[name])
models[name].load_state_dict(to_load[name])
##### important!!!
# make sure different processes sample different rays
np.random.seed((rank + 1) * 777)
# make sure different processes have different perturbations in depth samples
torch.manual_seed((rank + 1) * 777)
##### only main process should do the logging
if rank == 0:
writer = SummaryWriter(os.path.join(args.basedir, 'summaries', args.expname))
# start training
what_val_to_log = 0 # helper variable for parallel rendering of a image
what_train_to_log = 0
for global_step in range(start+1, start+1+args.N_iters):
time0 = time.time()
scalars_to_log = OrderedDict()
### Start of core optimization loop
scalars_to_log['resolution'] = ray_samplers[0].resolution_level
# randomly sample rays and move to device
i = np.random.randint(low=0, high=len(ray_samplers))
ray_batch = ray_samplers[i].random_sample(args.N_rand, center_crop=False)
for key in ray_batch:
if torch.is_tensor(ray_batch[key]):
ray_batch[key] = ray_batch[key].to(rank)
# forward and backward
dots_sh = list(ray_batch['ray_d'].shape[:-1]) # number of rays
all_rets = [] # results on different cascade levels
for m in range(models['cascade_level']):
optim = models['optim_{}'.format(m)]
net = models['net_{}'.format(m)]
# sample depths
N_samples = models['cascade_samples'][m]
if m == 0:
# foreground depth
fg_far_depth = intersect_sphere(ray_batch['ray_o'], ray_batch['ray_d']) # [...,]
# fg_near_depth = 0.18 * torch.ones_like(fg_far_depth)
fg_near_depth = ray_batch['min_depth'] # [..., 3]
step = (fg_far_depth - fg_near_depth) / (N_samples - 1)
fg_depth = torch.stack([fg_near_depth + i * step for i in range(N_samples)], dim=-1) # [..., N_samples]
fg_depth = perturb_samples(fg_depth) # random perturbation during training
# background depth
bg_depth = torch.linspace(0., 1., N_samples).view(
[1, ] * len(dots_sh) + [N_samples,]).expand(dots_sh + [N_samples,]).to(rank)
bg_depth = perturb_samples(bg_depth) # random perturbation during training
else:
# sample pdf and concat with earlier samples
fg_weights = ret['fg_weights'].clone().detach()
fg_depth_mid = .5 * (fg_depth[..., 1:] + fg_depth[..., :-1]) # [..., N_samples-1]
fg_weights = fg_weights[..., 1:-1] # [..., N_samples-2]
fg_depth_samples = sample_pdf(bins=fg_depth_mid, weights=fg_weights,
N_samples=N_samples, det=False) # [..., N_samples]
fg_depth, _ = torch.sort(torch.cat((fg_depth, fg_depth_samples), dim=-1))
# sample pdf and concat with earlier samples
bg_weights = ret['bg_weights'].clone().detach()
bg_depth_mid = .5 * (bg_depth[..., 1:] + bg_depth[..., :-1])
bg_weights = bg_weights[..., 1:-1] # [..., N_samples-2]
bg_depth_samples = sample_pdf(bins=bg_depth_mid, weights=bg_weights,
N_samples=N_samples, det=False) # [..., N_samples]
bg_depth, _ = torch.sort(torch.cat((bg_depth, bg_depth_samples), dim=-1))
optim.zero_grad()
ret = net(ray_batch['ray_o'], ray_batch['ray_d'], fg_far_depth, fg_depth, bg_depth)
all_rets.append(ret)
rgb_gt = ray_batch['rgb'].to(rank)
loss = img2mse(ret['rgb'], rgb_gt)
scalars_to_log['level_{}/loss'.format(m)] = loss.item()
scalars_to_log['level_{}/pnsr'.format(m)] = mse2psnr(loss.item())
# regularize sigma with photo-consistency
loss = loss + img2mse(ret['diffuse_rgb'], rgb_gt)
loss.backward()
optim.step()
# # clean unused memory
# torch.cuda.empty_cache()
### end of core optimization loop
dt = time.time() - time0
scalars_to_log['iter_time'] = dt
### only main process should do the logging
if rank == 0 and (global_step % args.i_print == 0 or global_step < 10):
logstr = '{} step: {} '.format(args.expname, global_step)
for k in scalars_to_log:
logstr += ' {}: {:.6f}'.format(k, scalars_to_log[k])
writer.add_scalar(k, scalars_to_log[k], global_step)
logger.info(logstr)
### each process should do this; but only main process merges the results
if global_step % args.i_img == 0 or global_step == start+1:
#### critical: make sure each process is working on the same random image
time0 = time.time()
idx = what_val_to_log % len(val_ray_samplers)
log_data = render_single_image(rank, args.world_size, models, val_ray_samplers[idx], args.chunk_size)
what_val_to_log += 1
dt = time.time() - time0
if rank == 0: # only main process should do this
logger.info('Logged a random validation view in {} seconds'.format(dt))
log_view_to_tb(writer, global_step, log_data, gt_img=val_ray_samplers[idx].img_orig, mask=None, prefix='val/')
time0 = time.time()
idx = what_train_to_log % len(ray_samplers)
log_data = render_single_image(rank, args.world_size, models, ray_samplers[idx], args.chunk_size)
what_train_to_log += 1
dt = time.time() - time0
if rank == 0: # only main process should do this
logger.info('Logged a random training view in {} seconds'.format(dt))
log_view_to_tb(writer, global_step, log_data, gt_img=ray_samplers[idx].img_orig, mask=None, prefix='train/')
log_data = None
torch.cuda.empty_cache()
if rank == 0 and (global_step % args.i_weights == 0 and global_step > 0):
# saving checkpoints and logging
fpath = os.path.join(args.basedir, args.expname, 'model_{:06d}.pth'.format(global_step))
to_save = OrderedDict()
for m in range(models['cascade_level']):
name = 'net_{}'.format(m)
to_save[name] = models[name].state_dict()
name = 'optim_{}'.format(m)
to_save[name] = models[name].state_dict()
torch.save(to_save, fpath)
# clean up for multi-processing
cleanup()
def config_parser():
import configargparse
parser = configargparse.ArgumentParser()
parser.add_argument('--config', is_config_file=True, help='config file path')
parser.add_argument("--expname", type=str, help='experiment name')
parser.add_argument("--basedir", type=str, default='./logs/', help='where to store ckpts and logs')
# dataset options
parser.add_argument("--datadir", type=str, default=None, help='input data directory')
parser.add_argument("--scene", type=str, default=None, help='scene name')
parser.add_argument("--testskip", type=int, default=8,
help='will load 1/N images from test/val sets, useful for large datasets like deepvoxels')
# model size
parser.add_argument("--netdepth", type=int, default=8, help='layers in coarse network')
parser.add_argument("--netwidth", type=int, default=256, help='channels per layer in coarse network')
parser.add_argument("--use_viewdirs", action='store_true', help='use full 5D input instead of 3D')
# checkpoints
parser.add_argument("--no_reload", action='store_true', help='do not reload weights from saved ckpt')
parser.add_argument("--ckpt_path", type=str, default=None,
help='specific weights npy file to reload for coarse network')
# batch size
parser.add_argument("--N_rand", type=int, default=32 * 32 * 2,
help='batch size (number of random rays per gradient step)')
parser.add_argument("--chunk_size", type=int, default=1024 * 8,
help='number of rays processed in parallel, decrease if running out of memory')
# iterations
parser.add_argument("--N_iters", type=int, default=250001,
help='number of iterations')
# cascade training
parser.add_argument("--cascade_level", type=int, default=2,
help='number of cascade levels')
parser.add_argument("--cascade_samples", type=str, default='64,64',
help='samples at each level')
parser.add_argument("--devices", type=str, default='0,1',
help='cuda device for each level')
parser.add_argument("--bg_devices", type=str, default='0,2',
help='cuda device for the background of each level')
parser.add_argument("--world_size", type=int, default='-1',
help='number of processes')
# mixed precison training
parser.add_argument("--opt_level", type=str, default='O1',
help='mixed precison training')
parser.add_argument("--near_depth", type=float, default=0.1,
help='near depth plane')
parser.add_argument("--far_depth", type=float, default=50.,
help='far depth plane')
# learning rate options
parser.add_argument("--lrate", type=float, default=5e-4, help='learning rate')
parser.add_argument("--lrate_decay_factor", type=float, default=0.1,
help='decay learning rate by a factor every specified number of steps')
parser.add_argument("--lrate_decay_steps", type=int, default=5000,
help='decay learning rate by a factor every specified number of steps')
# rendering options
parser.add_argument("--inv_uniform", action='store_true',
help='if True, will uniformly sample inverse depths')
parser.add_argument("--det", action='store_true', help='deterministic sampling for coarse and fine samples')
parser.add_argument("--max_freq_log2", type=int, default=10,
help='log2 of max freq for positional encoding (3D location)')
parser.add_argument("--max_freq_log2_viewdirs", type=int, default=4,
help='log2 of max freq for positional encoding (2D direction)')
parser.add_argument("--N_iters_perturb", type=int, default=1000,
help='perturb and center-crop at first 1000 iterations to prevent training from getting stuck')
parser.add_argument("--raw_noise_std", type=float, default=1.,
help='std dev of noise added to regularize sigma output, 1e0 recommended')
parser.add_argument("--white_bkgd", action='store_true',
help='apply the trick to avoid fitting to white background')
# no training; render only
parser.add_argument("--render_only", action='store_true',
help='do not optimize, reload weights and render out render_poses path')
parser.add_argument("--render_train", action='store_true', help='render the training set')
parser.add_argument("--render_test", action='store_true', help='render the test set instead of render_poses path')
# no training; extract mesh only
parser.add_argument("--mesh_only", action='store_true',
help='do not optimize, extract mesh from pretrained model')
parser.add_argument("--N_pts", type=int, default=256,
help='voxel resolution; N_pts * N_pts * N_pts')
parser.add_argument("--mesh_thres", type=str, default='10,20,30,40,50',
help='threshold(s) for mesh extraction; can use multiple thresholds')
# logging/saving options
parser.add_argument("--i_print", type=int, default=100, help='frequency of console printout and metric loggin')
parser.add_argument("--i_img", type=int, default=500, help='frequency of tensorboard image logging')
parser.add_argument("--i_weights", type=int, default=10000, help='frequency of weight ckpt saving')
parser.add_argument("--i_testset", type=int, default=50000, help='frequency of testset saving')
parser.add_argument("--i_video", type=int, default=50000, help='frequency of render_poses video saving')
return parser
def train():
parser = config_parser()
args = parser.parse_args()
logger.info(parser.format_values())
if args.world_size == -1:
args.world_size = torch.cuda.device_count()
logger.info('Using # gpus: {}'.format(args.world_size))
torch.multiprocessing.spawn(ddp_train_nerf,
args=(args,),
nprocs=args.world_size,
join=True)
if __name__ == '__main__':
setup_logger()
train()

231
old_scripts/nerf_sample_ray.py
Unescape Escape View File

@ -0,0 +1,231 @@
import numpy as np
from collections import OrderedDict
import torch
import cv2
########################################################################################################################
# ray batch sampling
########################################################################################################################
def parse_camera(params):
H, W = params[:2]
intrinsics = params[2:18].reshape((4, 4))
c2w = params[18:34].reshape((4, 4))
return int(W), int(H), intrinsics.astype(np.float32), c2w.astype(np.float32)
def get_rays_single_image(H, W, intrinsics, c2w):
'''
:param H: image height
:param W: image width
:param intrinsics: 4 by 4 intrinsic matrix
:param c2w: 4 by 4 camera to world extrinsic matrix
:return:
'''
u, v = np.meshgrid(np.arange(W), np.arange(H))
u = u.reshape(-1).astype(dtype=np.float32) + 0.5 # add half pixel
v = v.reshape(-1).astype(dtype=np.float32) + 0.5
pixels = np.stack((u, v, np.ones_like(u)), axis=0) # (3, H*W)
rays_d = np.dot(np.linalg.inv(intrinsics[:3, :3]), pixels)
rays_d = np.dot(c2w[:3, :3], rays_d) # (3, H*W)
rays_d = rays_d.transpose((1, 0)) # (H*W, 3)
rays_o = c2w[:3, 3].reshape((1, 3))
rays_o = np.tile(rays_o, (rays_d.shape[0], 1)) # (H*W, 3)
depth = np.linalg.inv(c2w)[2, 3]
depth = depth * np.ones((rays_o.shape[0],), dtype=np.float32) # (H*W,)
return rays_o, rays_d, depth
class RaySamplerSingleImage(object):
def __init__(self, cam_params, img_path=None, img=None, resolution_level=1, mask=None, min_depth=None):
super().__init__()
self.W_orig, self.H_orig, self.intrinsics_orig, self.c2w_mat = parse_camera(cam_params)
self.img_path = img_path
self.img_orig = img
self.mask_orig = mask
self.min_depth_orig = min_depth
self.resolution_level = -1
self.set_resolution_level(resolution_level)
def set_resolution_level(self, resolution_level):
if resolution_level != self.resolution_level:
self.resolution_level = resolution_level
self.W = self.W_orig // resolution_level
self.H = self.H_orig // resolution_level
self.intrinsics = np.copy(self.intrinsics_orig)
self.intrinsics[:2, :3] /= resolution_level
if self.img_orig is not None:
self.img = cv2.resize(self.img_orig, (self.W, self.H), interpolation=cv2.INTER_AREA)
self.img = self.img.reshape((-1, 3))
else:
self.img = None
if self.mask_orig is not None:
self.mask = cv2.resize(self.mask_orig, (self.W, self.H), interpolation=cv2.INTER_NEAREST)
self.mask = self.mask.reshape((-1))
else:
self.mask = None
if self.min_depth_orig is not None:
self.min_depth = cv2.resize(self.min_depth_orig, (self.W, self.H), interpolation=cv2.INTER_LINEAR)
self.min_depth = self.min_depth.reshape((-1))
else:
self.min_depth = None
self.rays_o, self.rays_d, self.depth = get_rays_single_image(self.H, self.W,
self.intrinsics, self.c2w_mat)
def get_all(self):
if self.min_depth is not None:
min_depth = self.min_depth
else:
min_depth = 1e-4 * np.ones_like(self.rays_d[..., 0])
ret = OrderedDict([
('ray_o', self.rays_o),
('ray_d', self.rays_d),
('depth', self.depth),
('rgb', self.img),
('mask', self.mask),
('min_depth', min_depth)
])
# return torch tensors
for k in ret:
if ret[k] is not None:
ret[k] = torch.from_numpy(ret[k])
return ret
def random_sample(self, N_rand, center_crop=False):
'''
:param N_rand: number of rays to be casted
:return:
'''
if center_crop:
half_H = self.H // 2
half_W = self.W // 2
quad_H = half_H // 2
quad_W = half_W // 2
# pixel coordinates
u, v = np.meshgrid(np.arange(half_W-quad_W, half_W+quad_W),
np.arange(half_H-quad_H, half_H+quad_H))
u = u.reshape(-1)
v = v.reshape(-1)
select_inds = np.random.choice(u.shape[0], size=(N_rand,), replace=False)
# Convert back to original image
select_inds = v[select_inds] * self.W + u[select_inds]
else:
# Random from one image
select_inds = np.random.choice(self.H*self.W, size=(N_rand,), replace=False)
rays_o = self.rays_o[select_inds, :] # [N_rand, 3]
rays_d = self.rays_d[select_inds, :] # [N_rand, 3]
depth = self.depth[select_inds] # [N_rand, ]
if self.img is not None:
rgb = self.img[select_inds, :] # [N_rand, 3]
else:
rgb = None
if self.mask is not None:
mask = self.mask[select_inds]
else:
mask = None
if self.min_depth is not None:
min_depth = self.min_depth[select_inds]
else:
min_depth = 1e-4 * np.ones_like(rays_d[..., 0])
ret = OrderedDict([
('ray_o', rays_o),
('ray_d', rays_d),
('depth', depth),
('rgb', rgb),
('mask', mask),
('min_depth', min_depth)
])
# return torch tensors
for k in ret:
if ret[k] is not None:
ret[k] = torch.from_numpy(ret[k])
return ret
# def random_sample_patches(self, N_patch, r_patch=16, center_crop=False):
# '''
# :param N_patch: number of patches to be sampled
# :param r_patch: patch size will be (2*r_patch+1)*(2*r_patch+1)
# :return:
# '''
# # even size patch
# # offsets to center pixels
# u, v = np.meshgrid(np.arange(-r_patch, r_patch),
# np.arange(-r_patch, r_patch))
# u = u.reshape(-1)
# v = v.reshape(-1)
# offsets = v * self.W + u
# # center pixel coordinates
# u_min = r_patch
# u_max = self.W - r_patch
# v_min = r_patch
# v_max = self.H - r_patch
# if center_crop:
# u_min = self.W // 4 + r_patch
# u_max = self.W - self.W // 4 - r_patch
# v_min = self.H // 4 + r_patch
# v_max = self.H - self.H // 4 - r_patch
# u, v = np.meshgrid(np.arange(u_min, u_max, r_patch),
# np.arange(v_min, v_max, r_patch))
# u = u.reshape(-1)
# v = v.reshape(-1)
# select_inds = np.random.choice(u.shape[0], size=(N_patch,), replace=False)
# # Convert back to original image
# select_inds = v[select_inds] * self.W + u[select_inds]
# # pick patches
# select_inds = np.stack([select_inds + shift for shift in offsets], axis=1)
# select_inds = select_inds.reshape(-1)
# rays_o = self.rays_o[select_inds, :] # [N_rand, 3]
# rays_d = self.rays_d[select_inds, :] # [N_rand, 3]
# depth = self.depth[select_inds] # [N_rand, ]
# if self.img is not None:
# rgb = self.img[select_inds, :] # [N_rand, 3]
# # ### debug
# # import imageio
# # imgs = rgb.reshape((N_patch, r_patch*2, r_patch*2, -1))
# # for kk in range(imgs.shape[0]):
# # imageio.imwrite('./debug_{}.png'.format(kk), imgs[kk])
# # ###
# else:
# rgb = None
# ret = OrderedDict([
# ('ray_o', rays_o),
# ('ray_d', rays_d),
# ('depth', depth),
# ('rgb', rgb)
# ])
# # return torch tensors
# for k in ret:
# ret[k] = torch.from_numpy(ret[k])
# return ret

19
render_all.sh
Unescape Escape View File

@ -0,0 +1,19 @@
#!/bin/bash
#SBATCH -p gpu
#SBATCH --gres=gpu:4
#SBATCH -c 10
#SBATCH -C pascal
#SBATCH --mem=40G
#SBATCH --time=24:00:00
#SBATCH --output=slurm_%A.out
#SBATCH --qos=high
PYTHON=/home/zhangka2/anaconda3/envs/nerf-ddp/bin/python
CODE_DIR=/home/zhangka2/gernot_experi/nerf_bg_latest_ddp
echo $CODE_DIR
#$PYTHON -u $CODE_DIR/ddp_test_nerf.py --config $CODE_DIR/configs/lf_data/lf_africa.txt
$PYTHON -u $CODE_DIR/ddp_test_nerf.py --config $CODE_DIR/configs/tanks_and_temples/tat_training_truck.txt

24
render_tat_training_truck.sh
Unescape Escape View File

@ -0,0 +1,24 @@
#!/bin/bash
#SBATCH -p q6
#SBATCH --gres=gpu:3
#SBATCH -c 8
#SBATCH -C turing
#SBATCH --mem=16G
#SBATCH --time=48:00:00
#SBATCH --output=slurm_%A.out
PYTHON=/home/zhangka2/anaconda3/envs/nerf/bin/python
CODE_DIR=/home/zhangka2/gernot_experi/nerf_bg_latest
echo $CODE_DIR
#$PYTHON -u $CODE_DIR/run_nerf.py --config $CODE_DIR/configs/tanks_and_temples/tat_training_truck_addregularize.txt
#$PYTHON -u $CODE_DIR/nerf_render_path.py --config $CODE_DIR/configs/tanks_and_temples/tat_training_truck_addregularize.txt
#$PYTHON -u $CODE_DIR/nerf_render_image.py --config $CODE_DIR/configs/tanks_and_temples/tat_training_truck_addregularize.txt
$PYTHON -u $CODE_DIR/nerf_render_path.py --config $CODE_DIR/configs/tanks_and_temples/tat_intermediate_playground.txt
$PYTHON -u $CODE_DIR/nerf_render_image.py --config $CODE_DIR/configs/tanks_and_temples/tat_intermediate_playground.txt
#$PYTHON -u $CODE_DIR/nerf_render_path.py --config $CODE_DIR/configs/tanks_and_temples/tat_training_truck_addregularize.txt
#$PYTHON -u $CODE_DIR/nerf_render_image.py --config $CODE_DIR/configs/tanks_and_temples/tat_training_truck_addregularize.txt

16
sparse_playground.sh
Unescape Escape View File

@ -0,0 +1,16 @@
#!/bin/bash
#SBATCH -p q6
#SBATCH --gres=gpu:4
#SBATCH -c 10
#SBATCH -C turing
#SBATCH --mem=60G
#SBATCH --time=48:00:00
#SBATCH --output=slurm_%A.out
PYTHON=/home/zhangka2/anaconda3/envs/nerf-ddp/bin/python
CODE_DIR=/home/zhangka2/gernot_experi/nerf_bg_latest_ddp
echo $CODE_DIR
$PYTHON -u $CODE_DIR/ddp_run_nerf.py --config $CODE_DIR/configs/tanks_and_temples_sparse/tat_intermediate_playground_addparam.txt

16
sparse_playground_addcarve.sh
Unescape Escape View File

@ -0,0 +1,16 @@
#!/bin/bash
#SBATCH -p q6
#SBATCH --gres=gpu:4
#SBATCH -c 10
#SBATCH -C turing
#SBATCH --mem=60G
#SBATCH --time=48:00:00
#SBATCH --output=slurm_%A.out
PYTHON=/home/zhangka2/anaconda3/envs/nerf-ddp/bin/python
CODE_DIR=/home/zhangka2/gernot_experi/nerf_bg_latest_ddp
echo $CODE_DIR
$PYTHON -u $CODE_DIR/ddp_run_nerf.py --config $CODE_DIR/configs/tanks_and_temples_sparse/tat_intermediate_playground_addcarve.txt

16
sparse_playground_addregularize.sh
Unescape Escape View File

@ -0,0 +1,16 @@
#!/bin/bash
#SBATCH -p q6
#SBATCH --gres=gpu:4
#SBATCH -c 10
#SBATCH -C turing
#SBATCH --mem=60G
#SBATCH --time=48:00:00
#SBATCH --output=slurm_%A.out
PYTHON=/home/zhangka2/anaconda3/envs/nerf-ddp/bin/python
CODE_DIR=/home/zhangka2/gernot_experi/nerf_bg_latest_ddp
echo $CODE_DIR
$PYTHON -u $CODE_DIR/ddp_run_nerf.py --config $CODE_DIR/configs/tanks_and_temples_sparse/tat_intermediate_playground_addregularize.txt

16
sparse_truck.sh
Unescape Escape View File

@ -0,0 +1,16 @@
#!/bin/bash
#SBATCH -p gpu
#SBATCH --gres=gpu:4
#SBATCH -c 10
#SBATCH -C turing
#SBATCH --mem=60G
#SBATCH --time=48:00:00
#SBATCH --output=slurm_%A.out
PYTHON=/home/zhangka2/anaconda3/envs/nerf-ddp/bin/python
CODE_DIR=/home/zhangka2/gernot_experi/nerf_bg_latest_ddp
echo $CODE_DIR
$PYTHON -u $CODE_DIR/ddp_run_nerf.py --config $CODE_DIR/configs/tanks_and_temples_sparse/tat_training_truck_addparam.txt

16
sparse_truck_addcarve.sh
Unescape Escape View File

@ -0,0 +1,16 @@
#!/bin/bash
#SBATCH -p gpu
#SBATCH --gres=gpu:4
#SBATCH -c 10
####SBATCH -C turing
#SBATCH --mem=60G
#SBATCH --time=48:00:00
#SBATCH --output=slurm_%A.out
PYTHON=/home/zhangka2/anaconda3/envs/nerf-ddp/bin/python
CODE_DIR=/home/zhangka2/gernot_experi/nerf_bg_latest_ddp
echo $CODE_DIR
$PYTHON -u $CODE_DIR/ddp_run_nerf.py --config $CODE_DIR/configs/tanks_and_temples_sparse/tat_training_truck_addcarve.txt

16
sparse_truck_addregularize.sh
Unescape Escape View File

@ -0,0 +1,16 @@
#!/bin/bash
#SBATCH -p gpu
#SBATCH --gres=gpu:4
#SBATCH -c 10
#SBATCH -C turing
#SBATCH --mem=60G
#SBATCH --time=48:00:00
#SBATCH --output=slurm_%A.out
PYTHON=/home/zhangka2/anaconda3/envs/nerf-ddp/bin/python
CODE_DIR=/home/zhangka2/gernot_experi/nerf_bg_latest_ddp
echo $CODE_DIR
$PYTHON -u $CODE_DIR/ddp_run_nerf.py --config $CODE_DIR/configs/tanks_and_temples_sparse/tat_training_truck_addregularize.txt

16
train_lf_africa.sh
Unescape Escape View File

@ -0,0 +1,16 @@
#!/bin/bash
#SBATCH -p gpu
#SBATCH --gres=gpu:8
#SBATCH -c 10
#SBATCH -C turing
#SBATCH --mem=80G
#SBATCH --time=24:00:00
#SBATCH --output=slurm_%A.out
#SBATCH --qos=high
PYTHON=/home/zhangka2/anaconda3/envs/nerf-ddp/bin/python
CODE_DIR=/home/zhangka2/gernot_experi/nerf_bg_latest_ddp
echo $CODE_DIR
$PYTHON -u $CODE_DIR/ddp_run_nerf.py --config $CODE_DIR/configs/lf_data/lf_africa.txt

17
train_lf_basket.sh
Unescape Escape View File

@ -0,0 +1,17 @@
#!/bin/bash
#SBATCH -p gpu
#SBATCH --gres=gpu:8
#SBATCH -c 10
#SBATCH -C turing
#SBATCH --mem=100G
#SBATCH --time=48:00:00
#SBATCH --output=slurm_%A.out
######## #SBATCH --qos=high
PYTHON=/home/zhangka2/anaconda3/envs/nerf-ddp/bin/python
CODE_DIR=/home/zhangka2/gernot_experi/nerf_bg_latest_ddp
echo $CODE_DIR
$PYTHON -u $CODE_DIR/ddp_run_nerf.py --config $CODE_DIR/configs/lf_data/lf_basket.txt

16
train_lf_ship.sh
Unescape Escape View File

@ -0,0 +1,16 @@
#!/bin/bash
#SBATCH -p q6
#SBATCH --gres=gpu:4
#SBATCH -c 10
#SBATCH -C turing
#SBATCH --mem=80G
#SBATCH --time=48:00:00
#SBATCH --output=slurm_%A.out
#SBATCH --qos=normal
PYTHON=/home/zhangka2/anaconda3/envs/nerf-ddp/bin/python
CODE_DIR=/home/zhangka2/gernot_experi/nerf_bg_latest_ddp
echo $CODE_DIR
$PYTHON -u $CODE_DIR/ddp_run_nerf.py --config $CODE_DIR/configs/lf_data/lf_ship.txt

16
train_lf_torch.sh
Unescape Escape View File

@ -0,0 +1,16 @@
#!/bin/bash
#SBATCH -p gpu
#SBATCH --gres=gpu:8
#SBATCH -c 10
#SBATCH -C turing
#SBATCH --mem=80G
#SBATCH --time=48:00:00
#SBATCH --output=slurm_%A.out
#SBATCH --qos=normal
PYTHON=/home/zhangka2/anaconda3/envs/nerf-ddp/bin/python
CODE_DIR=/home/zhangka2/gernot_experi/nerf_bg_latest_ddp
echo $CODE_DIR
$PYTHON -u $CODE_DIR/ddp_run_nerf.py --config $CODE_DIR/configs/lf_data/lf_torch.txt

18
train_tat_intermediate_m60.sh
Unescape Escape View File

@ -0,0 +1,18 @@
#!/bin/bash
#SBATCH -p q6
#SBATCH --gres=gpu:3
#SBATCH -c 8
#SBATCH -C turing
#SBATCH --mem=16G
#SBATCH --time=48:00:00
#SBATCH --output=slurm_%A.out
PYTHON=/home/zhangka2/anaconda3/envs/nerf/bin/python
CODE_DIR=/home/zhangka2/gernot_experi/nerf_bg_latest
echo $CODE_DIR
$PYTHON -u $CODE_DIR/run_nerf.py --config $CODE_DIR/configs/tanks_and_temples/tat_intermediate_m60.txt
$PYTHON -u $CODE_DIR/nerf_render_image.py --config $CODE_DIR/configs/tanks_and_temples/tat_intermediate_m60.txt

18
train_tat_intermediate_playground.sh
Unescape Escape View File

@ -0,0 +1,18 @@
#!/bin/bash
#SBATCH -p q6
#SBATCH --gres=gpu:3
#SBATCH -c 8
#SBATCH -C turing
#SBATCH --mem=16G
#SBATCH --time=48:00:00
#SBATCH --output=slurm_%A.out
PYTHON=/home/zhangka2/anaconda3/envs/nerf/bin/python
CODE_DIR=/home/zhangka2/gernot_experi/nerf_bg_latest
echo $CODE_DIR
$PYTHON -u $CODE_DIR/run_nerf.py --config $CODE_DIR/configs/tanks_and_temples/tat_intermediate_playground.txt
$PYTHON -u $CODE_DIR/nerf_render_image.py --config $CODE_DIR/configs/tanks_and_temples/tat_intermediate_playground.txt

15
train_tat_intermediate_playground_bignet.sh
Unescape Escape View File

@ -0,0 +1,15 @@
#!/bin/bash
#SBATCH -p gpu
#SBATCH --gres=gpu:8
#SBATCH -c 25
#SBATCH -C turing
#SBATCH --time=48:00:00
#SBATCH --output=slurm_%A.out
PYTHON=/home/zhangka2/anaconda3/envs/nerf-ddp/bin/python
CODE_DIR=/home/zhangka2/gernot_experi/nerf_bg_latest_ddp
echo $CODE_DIR
$PYTHON -u $CODE_DIR/ddp_run_nerf.py --config $CODE_DIR/configs/tanks_and_temples/tat_intermediate_playground_bignet.txt

18
train_tat_intermediate_train.sh
Unescape Escape View File

@ -0,0 +1,18 @@
#!/bin/bash
#SBATCH -p q6
#SBATCH --gres=gpu:3
#SBATCH -c 8
#SBATCH -C turing
#SBATCH --mem=16G
#SBATCH --time=48:00:00
#SBATCH --output=slurm_%A.out
PYTHON=/home/zhangka2/anaconda3/envs/nerf/bin/python
CODE_DIR=/home/zhangka2/gernot_experi/nerf_bg_latest
echo $CODE_DIR
$PYTHON -u $CODE_DIR/run_nerf.py --config $CODE_DIR/configs/tanks_and_temples/tat_intermediate_train.txt
$PYTHON -u $CODE_DIR/nerf_render_image.py --config $CODE_DIR/configs/tanks_and_temples/tat_intermediate_train.txt

16
train_tat_training_truck.sh
Unescape Escape View File

@ -0,0 +1,16 @@
#!/bin/bash
#SBATCH -p q6
#SBATCH --gres=gpu:4
#SBATCH -c 10
#SBATCH -C turing
#SBATCH --mem=50G
#SBATCH --time=48:00:00
#SBATCH --output=slurm_%A.out
PYTHON=/home/zhangka2/anaconda3/envs/nerf-ddp/bin/python
CODE_DIR=/home/zhangka2/gernot_experi/nerf_bg_latest_ddp
echo $CODE_DIR
$PYTHON -u $CODE_DIR/ddp_run_nerf.py --config $CODE_DIR/configs/tanks_and_temples/tat_training_truck.txt

15
train_tat_training_truck_bignet.sh
Unescape Escape View File

@ -0,0 +1,15 @@
#!/bin/bash
#SBATCH -p gpu
#SBATCH --gres=gpu:8
#SBATCH -c 25
#SBATCH -C turing
#SBATCH --time=48:00:00
#SBATCH --output=slurm_%A.out
PYTHON=/home/zhangka2/anaconda3/envs/nerf-ddp/bin/python
CODE_DIR=/home/zhangka2/gernot_experi/nerf_bg_latest_ddp
echo $CODE_DIR
$PYTHON -u $CODE_DIR/ddp_run_nerf.py --config $CODE_DIR/configs/tanks_and_temples/tat_training_truck_bignet.txt

19
train_tat_truck_sphere_subset.sh
Unescape Escape View File

@ -0,0 +1,19 @@
#!/bin/bash
#SBATCH -p gpu
#SBATCH --gres=gpu:3
#SBATCH -c 8
#SBATCH -C turing
#SBATCH --mem=16G
#SBATCH --time=48:00:00
#SBATCH --output=slurm_%A.out
#SBATCH --exclude=isl-gpu17
PYTHON=/home/zhangka2/anaconda3/envs/nerf/bin/python
CODE_DIR=/home/zhangka2/gernot_experi/nerf_bg
echo $CODE_DIR
$PYTHON $CODE_DIR/run_nerf.py --config $CODE_DIR/configs/tanks_and_temples/tat_training_truck_subset.txt
$PYTHON $CODE_DIR/nerf_render_image.py --config $CODE_DIR/configs/tanks_and_temples/tat_training_truck_subset.txt

206
utils.py
Unescape Escape View File

@ -0,0 +1,206 @@
import torch
# import torch.nn as nn
# import torch.nn.functional as F
import numpy as np
HUGE_NUMBER = 1e10
TINY_NUMBER = 1e-6 # float32 only has 7 decimal digits precision
# Misc utils
# work on tensors
# img2mse = lambda x, y: torch.mean((x - y) * (x - y))
def img2mse(x, y, mask=None):
if mask is None:
return torch.mean((x - y) * (x - y))
else:
return torch.sum((x - y) * (x - y) * mask.unsqueeze(-1)) / (torch.sum(mask) * x.shape[-1] + TINY_NUMBER)
img_HWC2CHW = lambda x: x.permute(2, 0, 1)
gray2rgb = lambda x: x.unsqueeze(2).repeat(1, 1, 3)
def normalize(x):
min = x.min()
max = x.max()
return (x - min) / ((max - min) + TINY_NUMBER)
to8b = lambda x: (255 * np.clip(x, 0, 1)).astype(np.uint8)
# gray2rgb = lambda x: np.tile(x[:,:,np.newaxis], (1, 1, 3))
mse2psnr = lambda x: -10. * np.log(x+TINY_NUMBER) / np.log(10.)
#
# def normalize(x):
# x_min, x_max = np.percentile(x, (0.5, 99.5))
# x = np.clip(x, x_min, x_max)
# x = (x - x_min) / (x_max - x_min)
# return x
########################################################################################################################
#
########################################################################################################################
from matplotlib.backends.backend_agg import FigureCanvasAgg
from matplotlib.figure import Figure
import matplotlib as mpl
from matplotlib import cm
import cv2
def get_vertical_colorbar(h, vmin, vmax, cmap_name='jet', label=None):
'''
:param w: pixels
:param h: pixels
:param vmin: min value
:param vmax: max value
:param cmap_name:
:param label
:return:
'''
fig = Figure(figsize=(1.2, 8), dpi=100)
fig.subplots_adjust(right=1.5)
canvas = FigureCanvasAgg(fig)
# Do some plotting.
ax = fig.add_subplot(111)
cmap = cm.get_cmap(cmap_name)
norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax)
tick_cnt = 6
tick_loc = np.linspace(vmin, vmax, tick_cnt)
cb1 = mpl.colorbar.ColorbarBase(ax, cmap=cmap,
norm=norm,
ticks=tick_loc,
orientation='vertical')
tick_label = ['{:3.2f}'.format(x) for x in tick_loc]
cb1.set_ticklabels(tick_label)
cb1.ax.tick_params(labelsize=18, rotation=0)
if label is not None:
cb1.set_label(label)
fig.tight_layout()
# # debug
# fig.savefig("debug3.png")
canvas.draw()
s, (width, height) = canvas.print_to_buffer()
im = np.frombuffer(s, np.uint8).reshape((height, width, 4))
im = im[:, :, :3].astype(np.float32) / 255.
if h != im.shape[0]:
w = int(im.shape[1] / im.shape[0] * h)
im = cv2.resize(im, (w, h), interpolation=cv2.INTER_AREA)
return im
# def colorize_np(x, cmap_name='jet', append_cbar=False):
# vmin = x.min()
# vmax = x.max() + TINY_NUMBER
# x = (x - vmin) / (vmax - vmin)
# # x = np.clip(x, 0., 1.)
# cmap = cm.get_cmap(cmap_name)
# x_new = cmap(x)[:, :, :3]
# cbar = get_vertical_colorbar(h=x.shape[0], vmin=vmin, vmax=vmax, cmap_name=cmap_name)
# if append_cbar:
# x_new = np.concatenate((x_new, np.zeros_like(x_new[:, :5, :]), cbar), axis=1)
# return x_new
# else:
# return x_new, cbar
# # tensor
# def colorize(x, cmap_name='jet', append_cbar=False):
# x = x.numpy()
# x, cbar = colorize_np(x, cmap_name)
# if append_cbar:
# x = np.concatenate((x, np.zeros_like(x[:, :5, :]), cbar), axis=1)
# x = torch.from_numpy(x)
# return x
def colorize_np(x, cmap_name='jet', mask=None, append_cbar=False):
if mask is not None:
# vmin, vmax = np.percentile(x[mask], (1, 99))
vmin = np.min(x[mask])
vmax = np.max(x[mask])
vmin = vmin - np.abs(vmin) * 0.01
x[np.logical_not(mask)] = vmin
x = np.clip(x, vmin, vmax)
# print(vmin, vmax)
else:
vmin = x.min()
vmax = x.max() + TINY_NUMBER
x = (x - vmin) / (vmax - vmin)
# x = np.clip(x, 0., 1.)
cmap = cm.get_cmap(cmap_name)
x_new = cmap(x)[:, :, :3]
if mask is not None:
mask = np.float32(mask[:, :, np.newaxis])
x_new = x_new * mask + np.zeros_like(x_new) * (1. - mask)
cbar = get_vertical_colorbar(h=x.shape[0], vmin=vmin, vmax=vmax, cmap_name=cmap_name)
if append_cbar:
x_new = np.concatenate((x_new, np.zeros_like(x_new[:, :5, :]), cbar), axis=1)
return x_new
else:
return x_new, cbar
# tensor
def colorize(x, cmap_name='jet', append_cbar=False, mask=None):
x = x.numpy()
if mask is not None:
mask = mask.numpy().astype(dtype=np.bool)
x, cbar = colorize_np(x, cmap_name, mask)
if append_cbar:
x = np.concatenate((x, np.zeros_like(x[:, :5, :]), cbar), axis=1)
x = torch.from_numpy(x)
return x
if __name__ == '__main__':
# # cbar = get_vertical_colorbar(h=512, vmin=0.1, vmax=5, cmap_name='jet')
# # cbar = cbar[:, :, :3]
# import imageio
#
# # imageio.imwrite('./debug.png', cbar)
#
# x = torch.rand(512, 512)
# x = colorize(x, append_cbar=True)
#
# x = np.uint8(x.numpy() * 255.)
#
# import imageio
# imageio.imwrite('./debug.png', x)
import os
import imageio
img_dir = '/phoenix/S7/kz298/latest_work/nerf/logs/dtu_scan9_3_nearfar/renderonly_train_200001'
all_imgs = []
for item in sorted(os.listdir(img_dir)):
if item[-4:] == '.png':
fpath = os.path.join(img_dir, item)
all_imgs.append(imageio.imread(fpath))
imageio.mimwrite(os.path.join(img_dir, 'video.mp4'), all_imgs, fps=3, quality=8)
Write Preview
Loading…
Cancel
Save