AMS-Net: Adaptive Multi-Scale Network for Image Compressive Sensing

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 (zkyhitsz@gmail.com).

Parapare

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

Train Model

TrainSet

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

Generate custom training set

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 .

Train

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

Test model

Test

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

Reconstrction results

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()

RedamancyAY / AMS-Net