This is the training code and test code for the paper AMS-Net: Adaptive Multi-Scale Network for Image Compressive Sensing.

If you have any question, please contact me ([email protected]).

We have put the trainset, image datasets, pre-trained model in the data directory:

• data/AMS-Net/: the pre-trained model
• data/dataset/: the image datasets
• data/trainset/: the training set to train our AMS-Net

If you have to put these files or your trained model in your custome folders, you can change the args in the next training or test cmd:

• dir_dataset: where you put your image datasets
• dir_trainset: where you put your training set
• dir_modelsave: where you put your pre-trained AMS-Net

You can download the training set from BaiduYunDrive(code: 9ene), or Google Dirve. Then put the training set in data/trainset/.

Our training set is constructed using the training set and validation set of BSDS500. We directly put all the images of the training set and validation set in the path ./data/dataset/BSDS500/.

To generate your custom training set, you can run the following code:

python trainset.py --dir_dataset "./data/dataset" --dir_trainset "./data/trainset" \
    --dataset "BSDS500" --mode "L" --crop_size 128 --crop_num 28

where the optional parameters are

• dataset: the dataset which contain $n$ images.
• mode: For the color image in the dataset, seting mode to "L" denotes that we directly convert the color image into gray image. You can set mode to "YCbCr_Y" for extracting the Y channel from YCbCr space to train the model.
• crop_size: the image size in the training set
• crop_num: the number of cropped images for each augmented iamge.

Note that the total number of images in the training set is , and each image is of size .

You can run the following code to train the model:

python model_train.py --gpu 1 --width 64 --depth 5 --T 10 --projection 1 \
    --dir_modelsave "./data/AMS-Net2"  --dir_dataset "./data/dataset" \
    --dir_trainset "./data/trainset" --trainset "BSDS500-L-n28-s128-N89600.tfrecords"

where the optional parameters are:

• gpu: set used gpu
• width, depth, T, projection: the parameters to construct the AMS-Net
• trainset: the file of trianing set that put in the dir_trainset

You can run the following code to obtain the reconstruction results:

python model_test.py --gpu 0 --dir_modelsave "./data/AMS-Net"  --dir_dataset "./data/dataset" \
    --mode "L" --width 64 --depth 5 --T 10 --projection 1 \
    --datasets "Set5" "Set11" "Set14" "BSD100" \
    --write_img 1 1 1 0

where the optional parameters are:

• gpu: set used gpu
• mode: For the color image in the dataset, seting mode to "L" denotes that we directly convert the color image into gray image. You can set mode to "YCbCr_Y" for extracting the Y channel from YCbCr space to test the model.
• width, depth, T, projection: the parameters to construct the AMS-Net
• datasets: the test sets
• write_img: whether should save reconstructed image for each test set. this arg should have the same length with the arg datasets

The test results are put in the directory result:

• result/W{width}-D{depth}-T{T}-Proj{projection}-{mode}.csv: all the psnr, ssim scores for all images in the test sets at each sampling ratio.
• result/reconstructed_imgs/W{width}-D{depth}-T{T}-Proj{projection}-{mode}/*/*.png: the reconstructed images

You can use the following code to load the reconstruction results and obtain the averge PSNR and SSIM score:

import pandas as pd
data = pd.read_csv("result/xxx.csv")
data.groupby(["dataset", "sr"]).mean()