Jiayong Peng, Zhiwei Xiong, Xin Huang, Zheng-Ping Li, Dong Liu, Feihu Xu
Photon-efficient imaging has enabled a number of applications relying on single-photon sensors that can capture a 3D image with as few as one photon per pixel. In practice, however, measurements of low photon counts are often mixed with heavy background noise, which poses a great challenge for existing computational reconstruction algorithms. In this paper, we first analyze the long-range correlations in both spatial and temporal dimensions of the measurements. Then we propose a non-local neural network for depth reconstruction by exploiting the long-range correlations. The proposed network achieves decent reconstruction fidelity even under photon counts (and signal-to-background ratio, SBR) as low as 1 photon/pixel (and 0.01), which significantly surpasses the state-of-the-art. Moreover, our non-local network trained on simulated data can be well generalized to different real-world imaging systems, which could extend the application scope of photon-efficient imaging in challenging scenarios with a strict limit on optical flux.
If you find the code useful in your research, please consider citing:
A. Environment
The code is developed using python 3.6 and Pytorch 1.0 on Ubuntu 16.04. NVIDIA 1080Ti GPU is required for training and testing.
B. Requirements
pip install -r requirements.txt
A. generate training and validation data
-
download NYU V2 data
bash ./data_gener/TrainData/download_nyu_dataset.bash
-
generate RGB-D and albedo with Matlab 2017b
matlab -nodesktop -nosplash -r ./data_gener/TrainData/ConvertRGBD
-
simulate SPAD measurements with Matlab 2017b
matlab -nodesktop -nosplash -r ./data_gener/TrainData/SimulateTrainMeasurements
B. generate testing data for Middlebury dataset
matlab -nodesktop -nosplash -r ./data_gener/TestData/middlebury/SimulateTestMeasurements
C. Indoor and outdoor real-world datasets are provided in
./data_gener/TestData/realworld
-
make train and validation list
python ./training/util/make_train_lists.py
-
edit the corresponding parameters in
./training/config.ini
-
run the main train file
python ./training/main.py
-
the training will take about 35 hours on NVIDIA 1080Ti.
-
edit the corresponding parameters in
./testing/config.ini
-
run the main test file
python ./testing/main.py
-
the test can be conducted on simulated Middlebury dataset, indoor and outdoor real-world dataset by selecting different functions in the
./testing/main.py
We provide some of the visualization results here.