-
Notifications
You must be signed in to change notification settings - Fork 1
/
demo.py
123 lines (101 loc) · 3.57 KB
/
demo.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
#!/usr/env/bin python3
# -*- coding: utf-8 -*-
import argparse
import numpy as np
import sys
import subprocess
import os
import yaml
import chainer
from chainer import cuda, optimizers, serializers
from chainer import training
import cv2
import matplotlib as mpl
mpl.rcParams['axes.xmargin'] = 0
mpl.rcParams['axes.ymargin'] = 0
import matplotlib.pyplot as plt
subprocess.call(['sh', 'setup.sh'])
from config_utils import *
from models.yolov2_base import *
from utils.image_loader import load_img_like_darknet
from utils.postprocess import visualize_with_label, clip_bbox
from utils.postprocess import select_bbox_by_class, select_bbox_by_obj
chainer.cuda.set_max_workspace_size(chainer.cuda.get_max_workspace_size())
os.environ["CHAINER_TYPE_CHECK"] = "0"
chainer.global_config.autotune = True
chainer.global_config.type_check = False
from collections import OrderedDict
yaml.add_constructor(yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG,
lambda loader, node: OrderedDict(loader.construct_pairs(node)))
def demo_yolov2():
"""Demo yolov2."""
config, args = parse_args()
h = args.height
w = args.width
config['model']['height'] = h
config['model']['width'] = w
model = get_model(config["model"])
if args.gpu != -1:
model.to_gpu(args.gpu)
img_path = args.img_path
thresh = args.thresh
nms_thresh = args.nms_thresh
label_names = open(args.name, 'r').read().split('\n')[:-1]
orig_img, img, delta = load_img_like_darknet(img_path, h, w)
orig_shape = orig_img.shape[:2]
img_shape = img.shape[2:]
# Dummy: create graph
dummy = np.zeros((1, 3, h, w), dtype='f') + 1
if args.gpu != -1:
dummy = chainer.cuda.to_gpu(dummy, device=args.gpu)
model.inference(dummy, img_shape)
# Inference
if args.gpu != -1:
start, stop = create_timer()
img = chainer.cuda.to_gpu(img, device=args.gpu)
else:
start = time.time()
bbox_pred, conf, prob = model.inference(img, img_shape)#img_shape)
bbox_pred = bbox_pred.reshape(-1, 4)
conf = conf.reshape(-1)
prob = prob.reshape(-1, model.n_classes)
bbox_pred = chainer.cuda.to_cpu(bbox_pred)
conf = chainer.cuda.to_cpu(conf)
prob = chainer.cuda.to_cpu(prob)
if args.gpu != -1:
print_timer(start, stop, sentence="Inference time(gpu)")
else:
print("Inference time(cpu):", time.time() - start)
# Post processing
start = time.time()
if args.nms == 'class':
bbox_pred, prob, cls_inds, index = \
select_bbox_by_class(bbox_pred, conf, prob, thresh, nms_thresh)
else:
bbox_pred, prob, cls_inds, index = \
select_bbox_by_obj(bbox_pred, conf, prob, thresh, nms_thresh)
bbox_pred[:, 0] -= bbox_pred[:, 2] / 2 # left_x
bbox_pred[:, 1] -= bbox_pred[:, 3] / 2 # top_y
bbox_pred[:, 2] += bbox_pred[:, 0] # right_x
bbox_pred[:, 3] += bbox_pred[:, 1] # bottom_y
bbox_pred[:, ::2] -= delta[1]
bbox_pred[:, 1::2] -= delta[0]
bbox_pred = clip_bbox(bbox_pred, orig_shape)
pred_names = [label_names[i] for i in cls_inds]
print("Post Processing: ", time.time() - start)
print(bbox_pred)
print(prob)
print(pred_names)
# Visualize
index = index[:, 0] if args.nms == 'class' else None
ax = visualize_with_label(orig_img, bbox_pred, prob, pred_names,
index=index)
if args.save:
ax.axis('off')
plt.savefig("{}.png".format(args.save),
bbox_inches="tight", pad_inches=0.0, transparent=True)
plt.show()
def main():
demo_yolov2()
if __name__ == '__main__':
main()