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
BSDS500andVOC2012images into./dataset/train/. - e.g., If you want to train TCS-Net at sampling rate
τ = 0.1withGPU 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_modelsfolder.
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 namesare the100 times of sampling rates, e.g., the folder named10includes 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.
