-
Notifications
You must be signed in to change notification settings - Fork 3
/
demo.py
48 lines (42 loc) · 1.36 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
from utils import *
from darknet import Darknet
import cv2
def demo(cfgfile, weightfile):
m = Darknet(cfgfile)
m.print_network()
m.load_weights(weightfile)
print('Loading weights from %s... Done!' % (weightfile))
namesfile = 'data/kaist_person.names'
class_names = load_class_names(namesfile)
use_cuda = True
if use_cuda:
m.cuda()
cap = cv2.VideoCapture(1)
if not cap.isOpened():
print("Unable to open camera")
exit(-1)
while True:
res, img = cap.read()
if res:
sized = cv2.resize(img, (m.width, m.height))
bboxes = do_detect(m, sized, 0.5, 0.4, use_cuda)
print('------')
draw_img = plot_boxes_cv2(img, bboxes, None, class_names)
cv2.imshow(cfgfile, draw_img)
cv2.waitKey(1)
else:
print("Unable to read image")
exit(-1)
############################################
if __name__ == '__main__':
cfgfile = 'cfg/yolov3_kaist.cfg'
weightfile = 'weights/kaist_thermal_detector.weights'
if len(sys.argv) >=1:
if len(sys.argv) == 3:
cfgfile = sys.argv[1]
weightfile = sys.argv[2]
demo(cfgfile, weightfile)
else:
print('Usage:')
print(' python demo.py [cfgfile] [weightfile]')
print(' perform detection on camera')