-
Notifications
You must be signed in to change notification settings - Fork 10
/
Copy pathtest.py
executable file
·111 lines (77 loc) · 3.25 KB
/
test.py
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
from scadec.unet_bn import Unet_bn
from scadec.train import Trainer_bn
from scadec import image_util
from scadec import util
import scipy.io as spio
import numpy as np
import os
####################################################
#### PREPARE WORKSPACE ###
####################################################
# here indicating the GPU you want to use. if you don't have GPU, just leave it.
gpu_vis = '3'
os.environ['CUDA_VISIBLE_DEVICES'] = gpu_vis; # 0,1,2,3
# here specify the path of the model you want to load
gpu_ind = '0'
model_path = 'gpu' + gpu_ind + '/models/60099_cpkt/models/final/model.cpkt'
data_channels = 2
truth_channels = 1
####################################################
#### FUNCTIONS ###
####################################################
# make the data a 4D vector
def preprocess(data, channels):
nx = data.shape[1]
ny = data.shape[2]
return data.reshape((-1, nx, ny, channels))
####################################################
#### lOAD MODEL ###
####################################################
# set up args for the unet, should be exactly the same as the loading model
kwargs = {
"layers": 5,
"conv_times": 2,
"features_root": 64,
"filter_size": 3,
"pool_size": 2,
"summaries": True
}
net = Unet_bn(img_channels=data_channels, truth_channels=truth_channels, cost="mean_squared_error", **kwargs)
####################################################
#### lOAD TRAIN ###
####################################################
#preparing training data
data_mat = spio.loadmat('train_np/obhatGausWeak128_40.mat', squeeze_me=True)
truths_mat = spio.loadmat('train_np/obGausWeak128_40.mat', squeeze_me=True)
data = data_mat['obhatGausWeak128']
data = preprocess(data, data_channels) # 4 dimension -> 3 dimension if you do data[:,:,:,1]
truths = preprocess(truths_mat['obGausWeak128'], truth_channels)
data_provider = image_util.SimpleDataProvider(data, truths)
####################################################
#### lOAD TEST ###
####################################################
vdata_mat = spio.loadmat('test_np_noise/obhatGausWeak{}Noise128.mat'.format(level), squeeze_me=True)
vtruths_mat = spio.loadmat('valid_np/obGausN1S128val.mat', squeeze_me=True)
vdata = vdata_mat['obhatGausWeak128']
vdata = preprocess(vdata, data_channels)
vtruths = preprocess(vtruths_mat['obGausN1S128val'], truth_channels)
valid_provider = image_util.SimpleDataProvider(vdata, vtruths)
####################################################
#### PREDICT ###
####################################################
predicts = []
valid_x, valid_y = valid_provider('full')
num = valid_x.shape[0]
for i in range(num):
print('')
print('')
print('************* {} *************'.format(i))
print('')
print('')
x_train, y_train = data_provider(23)
x_input = valid_x[i:i+1,:,:,:]
x_input = np.concatenate((x_input, x_train), axis=0)
predict = net.predict(model_path, x_input, 1, True)
predicts.append(predict[0:1,:,:])
predicts = np.concatenate(predicts, axis=0)
util.save_mat(predicts, 'test{}Noise.mat'.format(level))