QuantizationAwareDeepOptics - GitLite

KO | EN

gitlite — search

#android #hacktoberfest #self-hosted #go #cloud-native #java #openai #http #golang #editor

QuantizationAwareDeepOptics

(★ 16)

[CVPR 2022] Code for the paper "Quantization-aware Deep Optics for Diffractive Snapshot Hyperspectral Imaging".

File Explorer

ZIP Download

- __init__.py
- res_block_u_net.py
- __init__.py
- camera.py
- demosaicing.py
- diffractive_optical_element.py
- noise.py
- propagation.py
- sensor.py
- sensor_bayer_pattern.py
- sensor_srfs.py
- util.py
- __init__.py
- hyperspectral.py
- - __init__.py
  - data_utils.py
  - dataset_loader.py
  - file_io.py
- __init__.py
- hyper_spec_drawer.py
- mask_generator.py
constants.py
envrionment.yaml
evaluator.py
log.py
loss.py
main.py
metrics.py
model.py
README.md
summary.py
trainer.py

# Use via CDN

jsDelivr

jsDelivr serves any public GitHub repository as a CDN with zero setup. Pick a version and a file to get a ready-to-paste link and snippet.

Version

File

Link

Example

Command Glossary

Commands referenced in this DOCs, explained below.

🔍

`conda env`

View Details ▼

Manage conda environments.

conda env create {{[-f|--file]}} {{path/to/file}}

Create an environment from an environment file (YAML, TXT, etc.):

conda env remove {{[-n|--name]}} {{environment_name}}

Delete an environment and everything in it:

conda env update {{[-f|--file]}} {{path/to/file}} --prune

Update an environment based on an environment file:

🔍

`python`

View Details ▼

Python language interpreter.

python

Start a REPL (interactive shell):

python {{path/to/file.py}}

Execute a specific Python file:

python -i {{path/to/file.py}}

Execute a specific Python file and start a REPL:

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

This repository provides the code of the paper *"Quantization-aware Deep Optics for Diffractive Snapshot Hyperspectral Imaging", CVPR 2022*.

## Environment

Firstly, create a virtual Python 3.8 environment in Anaconda with the necessary dependencies from the **environment.yaml** file provided in the code.

```
conda env create -f ./environment.yaml
```

Then, activate the created environment and continue to train or test.

## Train

### Dataset Preparation

To train the QDO model for hyperspectral imaging, the dataset should be downloaded to your computer in advance.
(e.g., [CAVE](https://www.cs.columbia.edu/CAVE/databases/multispectral/), [ICVL](http://icvl.cs.bgu.ac.il/hyperspectral/), or [Harvard](http://vision.seas.harvard.edu/hyperspec/index.html).)

Then, edit the ```DATASET_PATH``` dictionary in **util/data/dataset_loader.py** to indicate the name and path to your dataset. Here is an example:
```
DATASET_PATH = {
  "dataset_name1": "/PATH1/TO/YOUR/DATSET1",
  "dataset_name2": "/PATH2/TO/YOUR/DATSET2",
  "dataset_name3": "/PATH2/TO/YOUR/DATSET3"
}
```
And there should be three directories in your dataset path: [train, validation, test] to indicate which part should be used for training, validation, and testing.

### Argument Configuration

After the dataset is prepared, configure the ```dataset_name``` and ```dataset_loader_func_name``` of ```controlled_training_args``` dictionary in **tasks/hyperspectral.py**.

The ```dataset_loader_func_name``` can be any function provided in **util/data/dataset_loader.py** or any function you implement using TensorFlow Dataset. 
(You should also put your customized dataset loader function in *util/data/dataset_loader.py* and set the ```dataset_loader_func_name``` to your customized function name. So that the model can automatically import and use it.)

The python file in tasks could be duplicated and renamed to store different task configurations, including dataset, training options, loss, network arguments, etc.

Current **tasks/hyperspectral.py** has already provided an example configuration for training using the ICVL dataset.

### Start Training

After the configuration, the training can be started with the following commands:
```bash
python main.py --task hyperspectral \
               --doe_material SK1300 --quantization_level 4 \
               --alpha_blending --adaptive_quantization \
               --sensor_distance_mm 50 --scene_depth_m 1 \
               --tag QDOA4LevelHyperspctralTraining \
```
This example shows a 4-level QDO+A model using SK1300 as the DOE material. Its sensor distance is 50mm, and the scene depth is 1m.

```--task``` argument is the name of the python file in the **tasks** package (without ".py").

```--doe_material``` argument indicates the material refractive index used for DOE simulation. Supported options: SK1300, BK7, NOA61.

```--quantization_level``` argument indicates the level count for the DOE quantization.

```--alpha_blending``` argument (optional) indicates whether to use alpha-blending for quantization-aware training. If this option is not given, the model will use STE.

```--adaptive_quantization``` argument (optional) defines whether to use adaptive quantization (for the QDO+A model) or not (for the QDO model).

```--sensor_distance_mm``` argument defines the distance between the DOE to the sensor plane in millimeters.

```--scene_depth_m``` argument defines the distance between the scene to the DOE in meters.

```--tag``` argument is a label that makes it easier to manage checkpoints and log files.

When the training starts, the trainer will save checkpoints and current task arguments into **./checkpoint/** as 2 JSON files named **controlled_model_args.json** and **controlled_training_args.json**. These files are important for the trainer to continue training and necessary for the evaluator to test the model. Visualization summary results, including DOE height maps, PSFs, and encoded images, will also be saved to the *./logs/* directory. Tensorboard can be used for viewing these results.

The checkpoint of a QDO+A 4-level model trained on ICVL is provided [here](https://mega.nz/file/FkNkWBpA#H3PXA0DIuDVl3G2xcTeVPp5Yx6lc7A03tyoqNaRGL8k). You can put extracted files into **./checkpoint/<TAG_NAME>** for later testing. The **<TAG_NAME>** can be any value you want as the directory name.

## Test

After training, evaluation can be performed using the following commands:
```bash
# For QDO+A/QDO model, using the test set
python evaluator.py --checkpoint_dir CHEKPOINT_DIR  \
                    --tag_name TAG_NAME \
                    --tag_vars TAG_VARS
                    
# For the conventional DO model, using the test set
python evaluator.py --checkpoint_dir CHEKPOINT_DIR \
                    --tag_name TAG_NAME \
                    --tag_vars TAG_VARS \
                    --test_q TEST_QUANTIZATION_LEVEL
                    
# For QDO+A/QDO model, using the real RGB data
python evaluator.py --checkpoint_dir CHEKPOINT_DIR \
                    --tag_name TAG_NAME \
                    --tag_vars TAG_VARS \
                    --real_data_dir REAL_DATA_DIR
```

```--checkpoint_dir``` argument is the name of the subdirectory in **./checkpoint/**.

```--tag_name``` argument is the inner tag name indicating the sub-directory in ```-checkpoint_dir``` given above.

```--tag_vars``` argument (optional) is the string value to insert in the %s placeholder of tag_name. Do not set this argument if there is no "%s" in your ```--tag_name``` argument.

```--test_q``` argument (optional) indicates the quantization level used during the test. Only the conventional DO model test needs this argument.

```--real_data_dir``` argument (optional) is the path to the directory storing real captured RGB images (PNG files).

The evaluator will output test results into **./eval-res/**, including .csv files with test metrics, visualized RGB images, and corresponding hyperspectral .mat files.

# Citation

😊 Please consider citing our paper if you find this repository helpful in your research.
```
@inproceedings{li2022quantization,
  title={Quantization-Aware Deep Optics for Diffractive Snapshot Hyperspectral Imaging},
  author={Li, Lingen and Wang, Lizhi and Song, Weitao and Zhang, Lei and Xiong, Zhiwei and Huang, Hua},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  pages={19780--19789},
  year={2022}
}
```