-
Notifications
You must be signed in to change notification settings - Fork 5
/
Copy patheval.py
168 lines (126 loc) · 5.99 KB
/
eval.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
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
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
import sys
from model.trainer import Trainer
sys.path.insert(0, '.')
import torch
import torch.nn.functional as F
import torch.backends.cudnn as cudnn
from torch.nn.parallel import gather
import torch.optim.lr_scheduler
import dataset.dataset as myDataLoader
import dataset.Transforms as myTransforms
from model.metric_tool import ConfuseMatrixMeter
from model.utils import BCEDiceLoss
from PIL import Image
import os, time
import numpy as np
from argparse import ArgumentParser
@torch.no_grad()
def val(args, val_loader, model):
model.eval()
salEvalVal = ConfuseMatrixMeter(n_class=2)
epoch_loss = []
total_batches = len(val_loader)
print(len(val_loader))
for iter, batched_inputs in enumerate(val_loader):
img, target = batched_inputs
img_name = val_loader.sampler.data_source.file_list[iter]
pre_img = img[:, 0:3]
post_img = img[:, 3:6]
start_time = time.time()
if args.onGPU == True:
pre_img = pre_img.cuda()
target = target.cuda()
post_img = post_img.cuda()
pre_img_var = torch.autograd.Variable(pre_img).float()
post_img_var = torch.autograd.Variable(post_img).float()
target_var = torch.autograd.Variable(target).float()
# run the mdoel
output = model(pre_img_var, post_img_var)
loss = BCEDiceLoss(output, target_var)
pred = torch.where(output > 0.5, torch.ones_like(output), torch.zeros_like(output)).long()
# torch.cuda.synchronize()
time_taken = time.time() - start_time
epoch_loss.append(loss.data.item())
# compute the confusion matrix
if args.onGPU and torch.cuda.device_count() > 1:
output = gather(pred, 0, dim=0)
f1 = salEvalVal.update_cm(pred.cpu().numpy(), target_var.cpu().numpy())
if iter % 5 == 0:
print('\r[%d/%d] F1: %3f loss: %.3f time: %.3f' % (iter, total_batches, f1, loss.data.item(), time_taken),
end='')
average_epoch_loss_val = sum(epoch_loss) / len(epoch_loss)
scores = salEvalVal.get_scores()
return average_epoch_loss_val, scores
def ValidateSegmentation(args):
os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu_id)
torch.backends.cudnn.benchmark = True
model = Trainer(args.model_type).float()
args.savedir = args.savedir + '_' + args.file_root + '_iter_' + str(args.max_steps) + '_lr_' + str(args.lr) + '/'
if args.file_root == 'LEVIR':
args.file_root = './levir_cd_256'
elif args.file_root == 'WHU':
args.file_root = './whu_cd_256'
elif args.file_root == 'CLCD':
args.file_root = './clcd_256'
elif args.file_root == 'OSCD':
args.file_root = 'oscd_256'
else:
raise TypeError('%s has not defined' % args.file_root)
if not os.path.exists(args.savedir):
os.makedirs(args.savedir)
if args.onGPU:
model = model.cuda()
mean = [0.5, 0.5, 0.5, 0.5, 0.5, 0.5]
std = [0.5, 0.5, 0.5, 0.5, 0.5, 0.5]
# compose the data with transforms
valDataset = myTransforms.Compose([
myTransforms.Normalize(mean=mean, std=std),
myTransforms.Scale(args.inWidth, args.inHeight),
myTransforms.ToTensor()
])
test_data = myDataLoader.Dataset(file_root=args.file_root, mode="test", transform=valDataset)
testLoader = torch.utils.data.DataLoader(
test_data, shuffle=False,
batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=False)
if args.onGPU:
cudnn.benchmark = True
logFileLoc = args.savedir + args.logFile
if os.path.isfile(logFileLoc):
logger = open(logFileLoc, 'a')
else:
logger = open(logFileLoc, 'w')
logger.write(
"\n%s\t%s\t%s\t%s\t%s\t%s\t%s" % ('Epoch', 'Kappa', 'IoU', 'F1', 'R', 'P', 'OA'))
logger.flush()
# load the model
model_file_name = args.savedir + 'best_model.pth'
state_dict = torch.load(model_file_name)
model.load_state_dict(state_dict)
loss_test, score_test = val(args, testLoader, model)
print("\nTest :\t Kappa (te) = %.4f\t IoU (te) = %.4f\t F1 (te) = %.4f\t R (te) = %.4f\t P (te) = %.4f OA (te) = %.4f" \
% (score_test['Kappa'], score_test['IoU'], score_test['F1'], score_test['recall'], score_test['precision'], score_test['OA']))
logger.write("\n%s\t\t%.4f\t\t%.4f\t\t%.4f\t\t%.4f\t\t%.4f\t\t%.4f" % ('Test', score_test['Kappa'], score_test['IoU'],
score_test['F1'], score_test['recall'],
score_test['precision'], score_test['OA']))
logger.flush()
logger.close()
if __name__ == '__main__':
parser = ArgumentParser()
parser.add_argument('--file_root', default="LEVIR", help='Data directory | LEVIR | WHU | CLCD | OSCD ')
parser.add_argument('--inWidth', type=int, default=256, help='Width of RGB image')
parser.add_argument('--inHeight', type=int, default=256, help='Height of RGB image')
parser.add_argument('--max_steps', type=int, default=80000, help='Max. number of iterations')
parser.add_argument('--num_workers', type=int, default=3, help='No. of parallel threads')
parser.add_argument('--batch_size', type=int, default=1, help='Batch size')
parser.add_argument('--model_type', type=str, default='small', help='select vit model type')
parser.add_argument('--lr', type=float, default=2e-4, help='Initial learning rate')
parser.add_argument('--savedir', default='./results', help='Directory to save the results')
parser.add_argument('--logFile', default='testLog.txt',
help='File that stores the training and validation logs')
parser.add_argument('--onGPU', default=True, type=lambda x: (str(x).lower() == 'true'),
help='Run on CPU or GPU. If TRUE, then GPU.')
parser.add_argument('--gpu_id', default=0, type=int, help='GPU id number')
args = parser.parse_args()
print('Called with args:')
print(args)
ValidateSegmentation(args)