nsff_pl
Neural Scene Flow Fields using pytorch-lightning. This repo reimplements the NSFF idea, but modifies several operations based on observation of NSFF results and discussions with the authors. For discussion details, please see the issues of the original repo. The code is based on my previous implementation.
The main modifications are the followings:
- Remove the blending weight in static NeRF. I adopt the addition strategy in NeRF-W.
- Remove disocclusion head. I use warped dynamic weights as an indicator of whether occlusion occurs. At the beginning of training, this indicator acts reliably as shown below:
Top: Reference image. Center: Warped images, artifacts appear at boundaries. Bottom: Estimated disocclusion.
As training goes, the disocclusion tends to get close to 1 almost everywhere, i.e. occlusion does not exist even in warping. In my opinion, this means the empty space learns to "move a little" to avoid the space occupied by dynamic objects (although the network has never been trained to do so).
- Compose static dynamic also in image warping.
Implementation details are in models/rendering.py.
The implementation is verified on several sequences, and produces visually plausible results. Qualitatively, these modifications produces better result on the kid-running scene compared to the original repo.
Full reconstruction
Left: GT. Center: this repo (PSNR=35.02). Right: pretrained model of the original repo(PSNR=30.45).
Background reconstruction
Left: this repo. Right: pretrained model of the original repo (by setting raw_blend_w to 0).
Fix-view-change-time (view 8, times from 0 to 16)
Left: this repo. Right: pretrained model of the original repo.
Fix-time-change-view (time 8, views from 0 to 16)
Left: this repo. Right: pretrained model of the original repo.
Novel view synthesis (view 8, spiral)
Time interpolation (view 8, add 10 frames between each integer time from time 0 to 29)
Left: this repo. Right: pretrained model of the original repo. The 2nd and 3rd rows are 0th frame and 29th frame to show the difference of the background.
The color of our method is more vivid and closer to the GT images both qualitatively and quantitatively (not because of gif compression). Also, even without any kind of supervision (either direct or self supervision), the network learns to separate the foreground and the background more cleanly than the original implementation, which is unexpected! Bad fg/bg separation not only means the background actually changes each frame, but also the color information is not leverage across time, so the reconstruction quality degrades, as can be shown in the original NSFF result towards the end.
Bonus - Depth
Our method also produces smoother depths, although it might not have direct impact on image quality.
Top left: static depth from this repo. Top right: full depth from this repo.
Bottom left: static depth from the original repo. Bottom right: full depth from the original repo.
More results
๐ป
Installation
Hardware
- OS: Ubuntu 18.04
- NVIDIA GPU with CUDA>=10.2 (tested with 1 RTX2080Ti)
Software
- Clone this repo by
git clone --recursive https://github.com/kwea123/nsff_pl - Python>=3.7 (installation via anaconda is recommended, use
conda create -n nsff_pl python=3.7to create a conda environment and activate it byconda activate nsff_pl) - Install core requirements by
pip install -r requirements.txt - Install
cupyviapip install cupy-cudaxxxby replacingxxxwith your cuda version.
๐
Training
Steps
Data preparation
Create a root directory (e.g. foobar), create a folder named frames and prepare your images (it is recommended to have at least 30 images) under it, so the structure looks like:
โโโ foobar
โโโ frames
โโโ 00000.png
...
โโโ 00029.png
The image names can be arbitrary, but the lexical order should be the same as time order! E.g. you can name the images as a.png, c.png, dd.png but the time order must be a -> c -> dd.
Motion Mask
In order to correctly reconstruct the camera poses, we must first filter out the dynamic areas so that feature points in these areas are not matched during estimation.
I use maskrcnn from detectron2. Only semantic masks are used, as I find flow-based masks too noisy.
Install detectron2 by python -m pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu102/torch1.8/index.html.
Modify the DYNAMIC_CATEGORIES variable in third_party/predict_mask.py to the dynamic classes in your data (only COCO classes are supported).
Next, NSFF requires depth and optical flows. We'll use some SOTA methods to perform the prediction.
Depth
The instructions and code are borrowed from DPT.
Download the model weights from here and put it in third_party/depth/weights/.
Optical Flow
The instructions and code are borrowed from RAFT.
Download raft-things.pth from google drive and put it in third_party/flow/models/.
Prediction
Thanks to owang, after preparing the images and the model weights, we can automate the whole process by a single command python preprocess.py --root_dir <path/to/foobar>.
Finally, your root directory will have all of this:
โโโ foobar
โโโ frames (original images, not used, you can delete)
โ โโโ 00000.png
โ ...
โ โโโ 00029.png
โโโ images_resized (resized images, not used, you can delete)
โ โโโ 00000.png
โ ...
โ โโโ 00029.png
โโโ images (the images to use in training)
โ โโโ 00000.png
โ ...
โ โโโ 00029.png
โโโ masks (not used but do not delete)
โ โโโ 00000.png.png
โ ...
โ โโโ 00029.png.png
โโโ database.db
โโโ sparse
โ โโโ 0
โ โโโ cameras.bin
โ โโโ images.bin
โ โโโ points3D.bin
โ โโโ project.ini
โโโ disps
โ โโโ 00000.png
โ ...
โ โโโ 00029.png
โโโ flow_fw
โ โโโ 00000.flo
โ ...
โ โโโ 00028.flo
โโโ flow_bw
โโโ 00001.flo
...
โโโ 00029.flo
Now you can start training!
Train!
Run the following command (modify the parameters according to opt.py):
python train.py \
--dataset_name monocular --root_dir $ROOT_DIR \
--img_wh 512 288 --start_end 0 30 \
--N_samples 128 --N_importance 0 --encode_t --use_viewdir \
--num_epochs 50 --batch_size 512 \
--optimizer adam --lr 5e-4 --lr_scheduler cosine \
--exp_name exp
I also implemented a hard sampling strategy to improve the quality of the hard regions. Add --hard_sampling to enable it.
Specifically, I compute the SSIM between the prediction and the GT at the end of each epoch, and use 1-SSIM as the sampling probability for the next epoch. This allows rays with larger errors to be sampled more frequently, and thus improve the result. The choice of SSIM is that it reflects more visual quality, and is less sensible to noise or small pixel displacement like PSNR.
๐
Testing
See test.ipynb for scene reconstruction, scene decomposition, fix-time-change-view, ..., etc. You can get almost everything out of this notebook. I will add more instructions inside in the future.
Use eval.py to create the whole sequence of moving views. E.g.
python eval.py \
--dataset_name monocular --root_dir $ROOT_DIR \
--N_samples 128 --N_importance 0 --img_wh 512 288 --start_end 0 30 \
--encode_t --output_transient \
--split test --video_format gif --fps 5 \
--ckpt_path kid.ckpt --scene_name kid_reconstruction
More specifically, the split argument specifies which novel view to generate:
test: test on training pose and timestest_spiral: spiral path over the whole sequence, with time gradually advances (integer time for now)test_spiralX: fix the time toXand generate spiral path around training viewX.test_fixviewX_interpY: fix the view to training poseXand interpolate the time from start to end, addingYframes between each integer timestamps.
โ ๏ธ
Other differences with the original paper
- I add entropy loss as suggested here. This allows the person to be "thin" and produces less artifact when the camera is far from the original pose.
- I explicitly zero the flow at far regions (where
z>0.95). - I add a cross entropy loss with thickness to encourage static and dynamic weights to peak at different locations (i.e. one sample points is either static or dynamic). The thickness specifies how many intervals should the peaks be separated. Empirically I found 15 to be a good value. Explanation about this loss is, for example for the "kid playing bubble scene", the kid is rotating around himself with the rotation axis almost fixed, so if no prior is added, the network learns the central part around the body as static (for example you can imagine a flag rotating around its pole, only the flag moves but the pole doesn't). This doesn't cause problem for reconstruction, but when it comes to novel view, the wrongly estimated static part causes artifacts. In order to make the network learn that the whole body is moving, I add this cross entropy loss to force the static peak to be at least
thickness//2far from the dynamic peak. - In pose reconstruction, the original authors use the entire image to reconstruct the poses, without masking out the dynamic region. In my opinion this strategy might lead to totally wrong pose estimation in some cases, so I opt to reconstruct the poses with dynamic region masked out. In order to set the near plane correctly, I use COLMAP combined with monodepth to get the minimum depth.
TODO
- Add COLMAP reconstruction tutorial (mask out dynamic region).
- Remove NSFF dependency for data preparation. More precisely, the original code needs quite a lot modifications to work on own data, and the depth/flow are calculated on resized images, which might reduce their accuracy.
- Add spiral path for testing.
- Add mask hard mining at the beginning of training. Or prioritized experience replay.
Acknowledgment
Thank to the authors of the NSFF paper, owang zhengqili sniklaus, for fruitful discussions and supports!
19 Nov 23, 2022
449 Dec 27, 2022
426 Jan 02, 2023
132 Dec 13, 2022
4 Mar 15, 2022
24 Nov 11, 2022
14 Oct 17, 2022
14 Nov 29, 2022
37 Jan 03, 2023
4 Feb 05, 2022
184 Dec 11, 2022
40 Dec 04, 2022
454 Jan 06, 2023
46 Dec 02, 2022
49 Dec 10, 2022
59 Dec 10, 2022
15 Nov 26, 2022
5 Aug 24, 2022
1 Jan 11, 2022
3 Jan 24, 2022