MD-DAN/test.py at main · wangnaa/MD-DAN · 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
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import scipy.io as sio
import load_data
import network
import os
from skimage.measure import compare_psnr, compare_ssim

######### hyper-parameters
batch_size = 1
num_chanel_Res = 16
num_blocks_Res = 9

Mask = sio.loadmat('Mask_8.mat')['Mask']

save_path = 'snapshot'
if not os.path.exists(save_path):
    os.mkdir(save_path)

train_path = "sliceT1.txt"
test_path = "T1_mat/Test.txt"

######### data_loader and network
test_loader = torch.utils.data.DataLoader(load_data.Test_List(test_path),batch_size=1,shuffle=False, num_workers=1,drop_last=True)
# modelG_0 = network.Generator_cat(num_blocks_Res,num_chanel_Res).cuda()
modelG_0 = torch.load(save_path+'/modelFFA_res_0.pk')
modelG_1 = torch.load(save_path+'/modelFFA_res_1.pk')
modelG_2 = torch.load(save_path+'/modelFFA_res_2.pk')
modelG_3 = torch.load(save_path+'/modelFFA_res_3.pk')

modelupdata_z = network.updata_z(Mask).cuda()


with torch.no_grad():
    psnr_under = 0
    ssim_under = 0
    psnr_x2_0 = 0
    ssim_x2_0 = 0
    psnr_z_0 = 0
    ssim_z_0 = 0
    psnr_x2_1 = 0
    ssim_x2_1 = 0
    psnr_z_1 = 0
    ssim_z_1 = 0
    psnr_x2_2 = 0
    ssim_x2_2 = 0
    psnr_z_2 = 0
    ssim_z_2 = 0
    psnr_x2_3 = 0
    ssim_x2_3 = 0
    psnr_z_3 = 0
    ssim_z_3 = 0

    for T1_test, T2_test, under_T2_test, under_T2_test_K in test_loader:
        T1_test = T1_test.view(-1, 1, 240, 240)
        T1_test = T1_test.cuda()
        T2_test = T2_test.view(-1, 1, 240, 240)
        T2_test = T2_test.cuda()
        T2_test_ = T2_test.cpu().numpy().squeeze()
        under_T2_test = under_T2_test.view(-1, 1, 240, 240)
        under_T2_test = under_T2_test.cuda()
        under_T2_test_K = under_T2_test_K.view(-1, 2, 240, 240)

        under_T2_test_ = under_T2_test.cpu().numpy().squeeze()
        psnr_under += compare_psnr(T2_test_, under_T2_test_, data_range=T2_test_.max())
        ssim_under += compare_ssim(T2_test_, under_T2_test_, data_range=T2_test_.max())

        x2_0 = modelG_0(T1_test, under_T2_test)
        x2_0_ = x2_0.cpu()
        x2_0_ = x2_0_.data.numpy().squeeze()
        psnr_x2_0 += compare_psnr(T2_test_, x2_0_, data_range=T2_test_.max())
        ssim_x2_0 += compare_ssim(T2_test_, x2_0_, data_range=T2_test_.max())

        z_0 = modelupdata_z(x2_0, under_T2_test_K)
        z_0_ = z_0.cpu()
        z_0_ = z_0_.data.numpy().squeeze()
        psnr_z_0 += compare_psnr(T2_test_, z_0_, data_range=T2_test_.max())
        ssim_z_0 += compare_ssim(T2_test_, z_0_, data_range=T2_test_.max())

        x2_1 = modelG_1(T1_test, z_0)
        x2_1_ = x2_1.cpu()
        x2_1_ = x2_1_.data.numpy().squeeze()
        psnr_x2_1 += compare_psnr(T2_test_, x2_1_, data_range=T2_test_.max())
        ssim_x2_1 += compare_ssim(T2_test_, x2_1_, data_range=T2_test_.max())

        z_1 = modelupdata_z(x2_1, under_T2_test_K)
        z_1_ = z_1.cpu()
        z_1_ = z_1_.data.numpy().squeeze()
        psnr_z_1 += compare_psnr(T2_test_, z_1_, data_range=T2_test_.max())
        ssim_z_1 += compare_ssim(T2_test_, z_1_, data_range=T2_test_.max())

        x2_2 = modelG_2(T1_test, z_1)
        x2_2_ = x2_2.cpu()
        x2_2_ = x2_2_.data.numpy().squeeze()
        psnr_x2_2 += compare_psnr(T2_test_, x2_2_, data_range=T2_test_.max())
        ssim_x2_2 += compare_ssim(T2_test_, x2_2_, data_range=T2_test_.max())

        z_2 = modelupdata_z(x2_2, under_T2_test_K)
        z_2_ = z_2.cpu()
        z_2_ = z_2_.data.numpy().squeeze()
        psnr_z_2 += compare_psnr(T2_test_, z_2_, data_range=T2_test_.max())
        ssim_z_2 += compare_ssim(T2_test_, z_2_, data_range=T2_test_.max())

        x2_3 = modelG_3(T1_test, z_2)
        x2_3_ = x2_3.cpu()
        x2_3_ = x2_3_.data.numpy().squeeze()
        psnr_x2_3 += compare_psnr(T2_test_, x2_3_, data_range=T2_test_.max())
        ssim_x2_3 += compare_ssim(T2_test_, x2_3_, data_range=T2_test_.max())

        z_3 = modelupdata_z(x2_3, under_T2_test_K)
        z_3_ = z_3.cpu()
        z_3_ = z_3_.data.numpy().squeeze()
        psnr_z_3 += compare_psnr(T2_test_, z_3_, data_range=T2_test_.max())
        ssim_z_3 += compare_ssim(T2_test_, z_3_, data_range=T2_test_.max())


    psnr_under /= len(test_loader)
    ssim_under /= len(test_loader)

    psnr_x2_0 /= len(test_loader)
    ssim_x2_0 /= len(test_loader)
    psnr_z_0 /= len(test_loader)
    ssim_z_0 /= len(test_loader)

    psnr_x2_1 /= len(test_loader)
    ssim_x2_1 /= len(test_loader)
    psnr_z_1 /= len(test_loader)
    ssim_z_1 /= len(test_loader)
    psnr_x2_2 /= len(test_loader)
    ssim_x2_2 /= len(test_loader)
    psnr_z_2 /= len(test_loader)
    ssim_z_2 /= len(test_loader)
    psnr_x2_3 /= len(test_loader)
    ssim_x2_3 /= len(test_loader)
    psnr_z_3 /= len(test_loader)
    ssim_z_3 /= len(test_loader)

print('psnr_under:{:.4f},\tssim_under:{:.4f}'.format(psnr_under.item(), ssim_under.item()))
print('psnr_x2_0:{:.4f},\tssim_x2_0:{:.4f},psnr_z_0:{:.4f},\tssim_z_0:{:.4f}'.format(psnr_x2_0.item(),ssim_x2_0.item(),psnr_z_0.item(),ssim_z_0.item()))
print('psnr_x2_1:{:.4f},\tssim_x2_1:{:.4f},psnr_z_1:{:.4f},\tssim_z_1:{:.4f}'.format(psnr_x2_1.item(),ssim_x2_1.item(),psnr_z_1.item(),ssim_z_1.item()))
print('psnr_x2_2:{:.4f},\tssim_x2_2:{:.4f},psnr_z_2:{:.4f},\tssim_z_2:{:.4f}'.format(psnr_x2_2.item(),ssim_x2_2.item(),psnr_z_2.item(),ssim_z_2.item()))
print('psnr_x2_3:{:.4f},\tssim_x2_3:{:.4f},psnr_z_3:{:.4f},\tssim_z_3:{:.4f}'.format(psnr_x2_3.item(),ssim_x2_3.item(),psnr_z_3.item(),ssim_z_3.item()))