search

computational-imaging / neural-holography

Code and data for Neural Holography
http://www.computationalimaging.org/publications/neuralholography/
Other
160 stars 45 forks source link
camera-in-the-loop double-phase gerchberg-saxton gradient-descent holography holonet

readme

Neural Holography with Camera-in-the-loop Training

Project Page | Paper

Yifan Peng, Suyeon Choi, Nitish Padmanaban, Gordon Wetzstein

This repository contains the scripts associated with the SIGGRAPH Asia 2020 paper "Neural Holography with Camera-in-the-loop Training"

Update 20201203:

We just released the second part of our scripts.

It contains all the code that can reproduce our work, including the camera-in-the-loop optimization, the parameterized wave propagation model / Holonet training code. Also, we are publishing our hardware SDK incorporation code for automated pipelines. Please have a look and feel free to modify and use for your work!

The specific updates can be found in following sections:

Getting Started

Our code requires PyTorch >1.7.0, as it uses Complex64 type Tensors. You can implement it in previous versions of PyTorch with the complex number operations implemented in utils/utils.py.

You can set up a conda environment with all dependencies like so:

For Windows:

conda env create -f environment_windows.yml
conda activate neural-holography

For Linux: (Hardware SDKs may not be compatible)

conda env create -f environment.yml
conda activate neural-holography

or you can manually set up a conda environment with (just execute setup_env.sh if you use Windows):

chmod u+x setup_env.sh
./setup_env.sh
conda activate neural-holography

You can load the submodule in utils/pytorch_prototyping folder with

git submodule init
git submodule update

To run phase generation with Holonet/U-net, download the pretrained model weights from here and place the contents in the pretrained_networks/ folder.

To run Camera-in-the-loop optimization or training, download PyCapture2 SDK and HOLOEYE SDK and place the SDKs in your environment folder. If your hardware setup (SLM, Camera, Laser) is different from ours, please modify related files in the utils/ folder according to the SDK before running the camera-in-th-loop optimization or training. Our hardware specifications can be found in the paper (Appendix B).

High-level structure

The code is organized as follows:

./utils/

Running the test

You can simply execute the following bash script with a method parameter (You can replace the parameter SGD with GS/DPAC/HOLONET/UNET.):

chmod u+x main_eval.sh
./main_eval.sh SGD

This bash script executes the phase optimization with 1) main.py for each R/G/B channel, and then executes 2) eval.py, which simulates the holographic image reconstruction for the optimized patterns with the angular spectrum method. Check the ./phases and ./recon folders after the execution.

1) Phase optimization

The SLM phase patterns can be reproduced with

SGD (Gradient Descent):

python main.py --channel=0 --method=SGD --root_path=./phases

SGD with Camera-in-the-loop optimization:

python main.py --channel=0 --method=SGD --citl=True --root_path=./phases

SGD with CITL-calibrated models:

python main.py --channel=0 --method=SGD --prop_model='MODEL' --prop_model_dir=YOUR_MODEL_PATH --root_path=./phases

HoloNet

python main.py --channel=0 --method=HOLONET --root_path=./phases --generator_dir=./pretrained_networks

GS (Gerchberg-Saxton):

python main.py --channel=0 --method=GS --root_path=./phases

DPAC (Double Phase Encoding):

python main.py --channel=0 --method=DPAC --root_path=./phases

U-net

python main.py --channel=0 --method=UNET --root_path=./phases --generator_dir=./pretrained_networks

You can set --channel=1/2 for other (green/blue) channels.

To monitor progress, the optimization code writes tensorboard summaries into a "summaries" subdirectory in the root_path.

2) Simulation/Evaluation

2.1) Ideal model simulation:

With optimized phase patterns, you can simulate the holographic image reconstruction with 

python eval.py --channel=0 --root_path=./phases/SGD_ASM --prop_model=ASM

For full-color simulation, you can set --channel=3, 0/1/2 corresponds to R/G/B, respectively.

This simulation code writes the reconstruction images in ./recon folder as default.

Feel free test other images after putting them in ./data folder!

2.2) CITL-calibrated model simulation

You can simulate those patterns with CITL-calibrated model

python eval.py --channel=0 --root_path=./phases/SGD_ASM --prop_model=MODEL

2.3) Evaluation on the physical setup

You can capture the image on the physical setup with

python eval.py --channel=0 --root_path=./phases/SGD_ASM --prop_model=CAMERA

3) Training

There are two-types of training in our work: 1) Parameterized wave propagation model (Camera-in-the-loop training) 2) Holonet.

Note that we need the camera-in-the-loop for training the wave propagation model while the Holonet can be trained offline once the CITL-calibrated model is calibrated. (You can pre-capture a bunch of phase-captured image pairs and can train the wave propagation models offline as well, though.)

You can train our wave propagation models with

python train_model.py --channel=0

You can train Holonet with

python train_holonet.py  --perfect_prop_model=True --run_id=my_first_holonet --batch_size=4 --channel=0

If you want to train it with CITL-calibrated models, set perfect_prop_model option to False and set model_path to your calibrated models.

4) Hardware (Camera, SLM, laser) Automation and Calibration

We incorporated the hardware SDKs as a pytorch module so that we can easily capture the experimental results and put them IN-THE-LOOP.

You can call the module with

camera_prop = PhysicalProp(channel, roi_res=YOUR_ROI,
                           range_row=(220, 1000), range_col=(300, 1630),
                           patterns_path=opt.calibration_path, # path of 21 x 12 calibration patterns, see Supplement.
                           show_preview=True)

Here, you may want to naively crop around the calibration patterns by setting range_row/col manually with the preview so that it can calculate the homography matrix without problems. (See Section S5 of Supplement)

Then, you can get camera-captured images by simply sending SLM phase patterns through the forward pass of the module:

captured_amp = camera_prop(slm_phase)

Citation

If you find our work useful in your research, please cite:

@article{Peng:2020:NeuralHolography,
author = {Y. Peng, S. Choi, N. Padmanaban, G. Wetzstein},
title = {Neural Holography with Camera-in-the-loop Training},
journal = {ACM Trans. Graph. (SIGGRAPH Asia)},
issue = {39},
number = {6},
year = {2020},
}

License

This project is licensed under the following license, with exception of the file "data/1.png", which is licensed under the CC-BY license.

Copyright (c) 2020, Stanford University

All rights reserved.

Redistribution and use in source and binary forms for academic and other non-commercial purposes with or without modification, are permitted provided that the following conditions are met:

THIS SOFTWARE IS PROVIDED BY THE TRUSTEES OF THE LELAND STANFORD JUNIOR UNIVERSITY "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE LELAND STANFORD JUNIOR UNIVERSITY OR ITS TRUSTEES BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

Contact

If you have any questions, please contact