This repository is the pytorch
code for paper "From Patch to Pixel: A Transformer-based Hierarchical Framework for Compressive Image Sensing"
.
1) Datasets
Training set: BSDS500
, testing sets: McM18
, LIVE29
, General100
and OST300
.
2)Project structure
(TCS-Net)
|-dataset
| |-train
| |-BSDS500 (.jpg)
| |-test
| |-McM18
| |-LIVE29
| |-General100
| |-OST300
|-reconstructed_images
| |-McM18
| |-grey
| |-... (Testing results .png)
| |-rgb
| |-... (Testing results .png)
| |-... (Testing sets)
| |-Res_(...).txt
|-models
| |-__init__.py
| |-net.py
| |-modules.py
|-trained_models
| |-1
| |-4
| |-... (Sampling rates)
|-config
| |-__init__.py
| |-config.py
| |-loader.py
|-test.py
|-train.py
|-train.sh
3) Competting methods
Methods | Sources | Year |
---|---|---|
Conf. Comput. Vis. Pattern Recog. | 2016 | |
Proc. Adv. Neural Inf. Process. Syst. | 2017 | |
Proc. Adv. Neural Inf. Process. Syst. | 2017 | |
Conf. Comput. Vis. Pattern Recog. | 2018 | |
Proc. Int. Conf. Mach. Learn. | 2019 | |
Trans. Image Process. | 2020 | |
Trans. Image Process. | 2021 | |
CSformer | arXiv | 2022 |
4) Performance demonstrates
Visual comparisons of reconstruction images (original images are drawn from dataset LIVE29
):
1) Re-training TCS-Net.
- Put the
BSDS500
andVOC2012
images into./dataset/train/
. - e.g., If you want to train TCS-Net at sampling rate
τ = 0.1
withGPU No.0
, please run the following command. The train set will be automatically packaged and our model will be trained with its default parameters (please make sure you have enough GPU RAM):
python train.py --rate 0.1 --GPU 0
- You can also run our shell script directly as well, it will automatically train the model under all sampling rates, i.e.,
τ ∈ {0.01, 0.04, 0.1, 0.25}
:
sh train.sh
- The trained models (.pth) will save in the
trained_models
folder.
2) Testing TCS-Net.
-
We provide the trained models so that you can put them under
TCS-Net/trained_models/
and use them for testing directly; all trained TCS-Net models can be found in this GoogleDrive link; Please note that thefolder's names
are the100 times of sampling rates
, e.g., the folder named10
includes trained models atsampling rate = 0.1
. -
Put the testing folders into
./dataset/test/
. -
e.g., if you want to test TCS-Net at sampling rate τ = 0.1 with GPU No.0, please run:
python test.py --rate 0.1 --GPU 0
- After that, the reconstructed images, PSNR and SSIM results will be saved to
./reconstructed_images/
.
We appreciate your reading and attention. For more details about TCS-Net, please refer to our paper.