Pytorch_ADMM-CSNet/test.py at main · yangyan92/Pytorch_ADMM-CSNet · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
"""
    ADMM_-CSNET test example (v1) with MR slices
    By Xing Li, Yan Yang, Jian Sun, Huibin Li, Zongben Xu

    Please cite the below paper for the code:

    Yan Yang, Jian Sun, Huibin Li, Zongben Xu. ADMM-CSNet: A Deep Learning Approach for Image Compressive Sensing,
    TPAMI(2019).
"""
from __future__ import print_function, division
import os
import argparse
from network.CSNet_Layers import ADMMCSNetLayer
from utils.dataset import get_test_data
import torch.utils.data as data
from utils.my_loss import MyLoss
from utils.metric import complex_psnr
import gc
from scipy.io import loadmat
from utils.fftc import *
from os.path import join

os.environ["CUDA_VISIBLE_DEVICES"] = "0"

if __name__ == '__main__':

    ###############################################################################
    # parameters
    ###############################################################################
    parser = argparse.ArgumentParser(description=' main ')
    parser.add_argument('--data_dir', default='data/', type=str,
                        help='directory of data')
    parser.add_argument('--batch_size', default=1, type=int, help='batch size')
    parser.add_argument('--outf', type=str, default='csnet_model', help='path of log files')
    args = parser.parse_args()

    ###############################################################################
    # load data info
    ###############################################################################
    test = get_test_data(args.data_dir)
    test_loader = data.DataLoader(dataset=test, batch_size=args.batch_size, shuffle=False, num_workers=4,
                                  pin_memory=False)

    ###############################################################################
    # mask
    ###############################################################################
    dir = 'data/mask_20'
    data = loadmat(join(dir, os.listdir(dir)[0]))
    mask_data = data['mask']
    mask = ifftshift(torch.Tensor(mask_data)).cuda()

    ###############################################################################
    # Build model
    ###############################################################################
    print('Loading model ...\n')
    model = ADMMCSNetLayer(mask).cuda()
    model.load_state_dict(torch.load(os.path.join(args.outf, 'cs_net_sample0.2.pth')))
    model.eval()

    ###############################################################################
    # loss
    ###############################################################################
    criterion = MyLoss().cuda()

    ###############################################################################
    # test
    ###############################################################################
    test_err = 0
    test_psnr = 0
    test_batches = 0
    for batch , (label,num) in enumerate(test_loader):
        gc.collect()
        with torch.no_grad():
            full_kspace = torch.fft.fft2(label.cuda())
            test_output = model(full_kspace)
            test_loss_normal = criterion(test_output, label.cuda())
            test_err += test_loss_normal.item()
            test_batches += 1
            test_psnr_value = complex_psnr(abs(test_output).cpu().numpy(), abs(label).cpu().numpy(),
                                           peak='normalized')
            test_psnr += test_psnr_value
    test_err /= test_batches
    test_psnr /= test_batches
    print("test_loss ", test_err)
    print("test_psnr ", test_psnr)