-
Notifications
You must be signed in to change notification settings - Fork 6
/
Copy patheval.py
130 lines (108 loc) · 4.68 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
import torch
import torch.nn as nn
from data import *
import argparse
from utils.vocapi_evaluator import VOCAPIEvaluator
from utils.cocoapi_evaluator import COCOAPIEvaluator
parser = argparse.ArgumentParser(description='YOLO-v2 Detector Evaluation')
parser.add_argument('-v', '--version', default='yolo_v2',
help='yolo_v2, yolo_v3, slim_yolo_v2, tiny_yolo_v3.')
parser.add_argument('-d', '--dataset', default='voc',
help='voc, coco-val, coco-test.')
parser.add_argument('--trained_model', type=str,
default='weights_yolo_v2/yolo_v2_72.2.pth',
help='Trained state_dict file path to open')
parser.add_argument('-size', '--input_size', default=416, type=int,
help='input_size')
parser.add_argument('--cuda', action='store_true', default=False,
help='Use cuda')
args = parser.parse_args()
def voc_test(model, device, input_size):
evaluator = VOCAPIEvaluator(data_root=VOC_ROOT,
img_size=input_size,
device=device,
transform=BaseTransform(input_size),
labelmap=VOC_CLASSES,
display=True
)
# VOC evaluation
evaluator.evaluate(model)
def coco_test(model, device, input_size, test=False):
if test:
# test-dev
print('test on test-dev 2017')
evaluator = COCOAPIEvaluator(
data_dir=coco_root,
img_size=input_size,
device=device,
testset=True,
transform=BaseTransform(input_size)
)
else:
# eval
evaluator = COCOAPIEvaluator(
data_dir=coco_root,
img_size=input_size,
device=device,
testset=False,
transform=BaseTransform(input_size)
)
# COCO evaluation
evaluator.evaluate(model)
if __name__ == '__main__':
# dataset
if args.dataset == 'voc':
print('eval on voc ...')
num_classes = 20
elif args.dataset == 'coco-val':
print('eval on coco-val ...')
num_classes = 80
elif args.dataset == 'coco-test':
print('eval on coco-test-dev ...')
num_classes = 80
else:
print('unknow dataset !! we only support voc, coco-val, coco-test !!!')
exit(0)
# cuda
if args.cuda:
print('use cuda')
torch.backends.cudnn.benchmark = True
device = torch.device("cuda")
else:
device = torch.device("cpu")
# input size
input_size = [args.input_size, args.input_size]
# load net
if args.version == 'yolo_v2':
from models.yolo_v2 import myYOLOv2
anchor_size = ANCHOR_SIZE if args.dataset == 'voc' else ANCHOR_SIZE_COCO
net = myYOLOv2(device, input_size=input_size, num_classes=num_classes, anchor_size=anchor_size)
elif args.version == 'yolo_v3':
from models.yolo_v3 import myYOLOv3
anchor_size = MULTI_ANCHOR_SIZE if args.dataset == 'voc' else MULTI_ANCHOR_SIZE_COCO
net = myYOLOv3(device, input_size=input_size, num_classes=num_classes, anchor_size=anchor_size)
elif args.version == 'yolo_v3_spp':
from models.yolo_v3_spp import myYOLOv3Spp
anchor_size = MULTI_ANCHOR_SIZE if args.dataset == 'voc' else MULTI_ANCHOR_SIZE_COCO
net = myYOLOv3Spp(device, input_size=input_size, num_classes=num_classes, anchor_size=anchor_size)
elif args.version == 'slim_yolo_v2':
from models.slim_yolo_v2 import SlimYOLOv2
anchor_size = ANCHOR_SIZE if args.dataset == 'voc' else ANCHOR_SIZE_COCO
net = SlimYOLOv2(device, input_size=input_size, num_classes=num_classes, anchor_size=anchor_size)
elif args.version == 'tiny_yolo_v3':
from models.tiny_yolo_v3 import YOLOv3tiny
anchor_size = TINY_MULTI_ANCHOR_SIZE if args.dataset == 'voc' else TINY_MULTI_ANCHOR_SIZE_COCO
net = YOLOv3tiny(device, input_size=input_size, num_classes=num_classes, anchor_size=anchor_size)
# load net
net.load_state_dict(torch.load(args.trained_model, map_location='cuda'))
net.eval()
print('Finished loading model!')
net = net.to(device)
# evaluation
with torch.no_grad():
if args.dataset == 'voc':
voc_test(net, device, input_size)
elif args.dataset == 'coco-val':
coco_test(net, device, input_size, test=False)
elif args.dataset == 'coco-test':
coco_test(net, device, input_size, test=True)