-
Notifications
You must be signed in to change notification settings - Fork 2
/
detect.py
118 lines (96 loc) · 3.47 KB
/
detect.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
import sys
import time
from PIL import Image, ImageDraw
#from models.tiny_yolo import TinyYoloNet
from utils import *
from darknet import Darknet
namesfile=None
def detect(cfgfile, weightfile, imgfile):
m = Darknet(cfgfile)
m.print_network()
m.load_weights(weightfile)
print('Loading weights from %s... Done!' % (weightfile))
# if m.num_classes == 20:
# namesfile = 'data/voc.names'
# elif m.num_classes == 80:
# namesfile = 'data/coco.names'
# else:
# namesfile = 'data/names'
use_cuda = True
if use_cuda:
m.cuda()
img = Image.open(imgfile).convert('RGB')
sized = img.resize((m.width, m.height))
#for i in range(2):
start = time.time()
boxes = do_detect(m, sized, 0.5, 0.4, use_cuda)
finish = time.time()
#if i == 1:
print('%s: Predicted in %f seconds.' % (imgfile, (finish-start)))
class_names = load_class_names(namesfile)
plot_boxes(img, boxes, 'predictions.jpg', class_names)
def detect_cv2(cfgfile, weightfile, imgfile):
import cv2
m = Darknet(cfgfile)
m.print_network()
m.load_weights(weightfile)
print('Loading weights from %s... Done!' % (weightfile))
if m.num_classes == 20:
namesfile = 'data/voc.names'
elif m.num_classes == 80:
namesfile = 'data/coco.names'
else:
namesfile = 'data/names'
use_cuda = True
if use_cuda:
m.cuda()
img = cv2.imread(imgfile)
sized = cv2.resize(img, (m.width, m.height))
sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB)
for i in range(2):
start = time.time()
boxes = do_detect(m, sized, 0.5, 0.4, use_cuda)
finish = time.time()
if i == 1:
print('%s: Predicted in %f seconds.' % (imgfile, (finish-start)))
class_names = load_class_names(namesfile)
plot_boxes_cv2(img, boxes, savename='predictions.jpg', class_names=class_names)
def detect_skimage(cfgfile, weightfile, imgfile):
from skimage import io
from skimage.transform import resize
m = Darknet(cfgfile)
m.print_network()
m.load_weights(weightfile)
print('Loading weights from %s... Done!' % (weightfile))
if m.num_classes == 20:
namesfile = 'data/voc.names'
elif m.num_classes == 80:
namesfile = 'data/coco.names'
else:
namesfile = 'data/names'
use_cuda = True
if use_cuda:
m.cuda()
img = io.imread(imgfile)
sized = resize(img, (m.width, m.height)) * 255
for i in range(2):
start = time.time()
boxes = do_detect(m, sized, 0.5, 0.4, use_cuda)
finish = time.time()
if i == 1:
print('%s: Predicted in %f seconds.' % (imgfile, (finish-start)))
class_names = load_class_names(namesfile)
plot_boxes_cv2(img, boxes, savename='predictions.jpg', class_names=class_names)
if __name__ == '__main__':
if len(sys.argv) == 5:
cfgfile = sys.argv[1]
weightfile = sys.argv[2]
imgfile = sys.argv[3]
globals()["namesfile"] = sys.argv[4]
detect(cfgfile, weightfile, imgfile)
#detect_cv2(cfgfile, weightfile, imgfile)
#detect_skimage(cfgfile, weightfile, imgfile)
else:
print('Usage: ')
print(' python detect.py cfgfile weightfile imgfile names')
#detect('cfg/tiny-yolo-voc.cfg', 'tiny-yolo-voc.weights', 'data/person.jpg', version=1)