-
Notifications
You must be signed in to change notification settings - Fork 1
/
car_counter.py
204 lines (152 loc) · 5.35 KB
/
car_counter.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
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
from ctypes import *
import math
import random
import cv2
def sample(probs):
s = sum(probs)
probs = [a / s for a in probs]
r = random.uniform(0, 1)
for i in range(len(probs)):
r = r - probs[i]
if r <= 0:
return i
return len(probs) - 1
def c_array(ctype, values):
arr = (ctype * len(values))()
arr[:] = values
return arr
class BOX(Structure):
_fields_ = [("x", c_float),
("y", c_float),
("w", c_float),
("h", c_float)]
class IMAGE(Structure):
_fields_ = [("w", c_int),
("h", c_int),
("c", c_int),
("data", POINTER(c_float))]
class METADATA(Structure):
_fields_ = [("classes", c_int),
("names", POINTER(c_char_p))]
#lib = CDLL("/home/pjreddie/documents/darknet/libdarknet.so", RTLD_GLOBAL)
lib = CDLL("libdarknet.so", RTLD_GLOBAL)
lib.network_width.argtypes = [c_void_p]
lib.network_width.restype = c_int
lib.network_height.argtypes = [c_void_p]
lib.network_height.restype = c_int
predict = lib.network_predict
predict.argtypes = [c_void_p, POINTER(c_float)]
predict.restype = POINTER(c_float)
set_gpu = lib.cuda_set_device
set_gpu.argtypes = [c_int]
make_image = lib.make_image
make_image.argtypes = [c_int, c_int, c_int]
make_image.restype = IMAGE
make_boxes = lib.make_boxes
make_boxes.argtypes = [c_void_p]
make_boxes.restype = POINTER(BOX)
free_ptrs = lib.free_ptrs
free_ptrs.argtypes = [POINTER(c_void_p), c_int]
num_boxes = lib.num_boxes
num_boxes.argtypes = [c_void_p]
num_boxes.restype = c_int
make_probs = lib.make_probs
make_probs.argtypes = [c_void_p]
make_probs.restype = POINTER(POINTER(c_float))
detect = lib.network_predict
detect.argtypes = [c_void_p, IMAGE, c_float, c_float,
c_float, POINTER(BOX), POINTER(POINTER(c_float))]
reset_rnn = lib.reset_rnn
reset_rnn.argtypes = [c_void_p]
load_net = lib.load_network
load_net.argtypes = [c_char_p, c_char_p, c_int]
load_net.restype = c_void_p
free_image = lib.free_image
free_image.argtypes = [IMAGE]
letterbox_image = lib.letterbox_image
letterbox_image.argtypes = [IMAGE, c_int, c_int]
letterbox_image.restype = IMAGE
load_meta = lib.get_metadata
lib.get_metadata.argtypes = [c_char_p]
lib.get_metadata.restype = METADATA
load_image = lib.load_image_color
load_image.argtypes = [c_char_p, c_int, c_int]
load_image.restype = IMAGE
ndarray_image = lib.ndarray_to_image
ndarray_image.argtypes = [POINTER(c_ubyte), POINTER(c_long), POINTER(c_long)]
ndarray_image.restype = IMAGE
rgbgr_image = lib.rgbgr_image
rgbgr_image.argtypes = [IMAGE]
predict_image = lib.network_predict_image
predict_image.argtypes = [c_void_p, IMAGE]
predict_image.restype = POINTER(c_float)
network_detect = lib.network_detect
network_detect.argtypes = [c_void_p, IMAGE, c_float,
c_float, c_float, POINTER(BOX), POINTER(POINTER(c_float))]
def classify(net, meta, im):
out = predict_image(net, im)
res = []
for i in range(meta.classes):
res.append((meta.names[i], out[i]))
res = sorted(res, key=lambda x: -x[1])
return res
def nparray_to_image(img):
data = img.ctypes.data_as(POINTER(c_ubyte))
image = ndarray_image(data, img.ctypes.shape, img.ctypes.strides)
return image
def detect_np(net, meta, np_img, entity_name=None, thresh=.5, hier_thresh=.5, nms=.45):
"""
Read a picture as a numpy array and run YOLO on it to yield list of detected objects
"""
im = nparray_to_image(np_img)
boxes = make_boxes(net)
probs = make_probs(net)
num = num_boxes(net)
network_detect(net, im, thresh, hier_thresh, nms, boxes, probs)
res = []
for j in range(num):
for i in range(meta.classes):
if probs[j][i] > 0:
if entity_name is None:
res.append((meta.names[i], probs[j][i], (boxes[j].x, boxes[j].y, boxes[j].w, boxes[j].h)))
elif meta.names[i] == entity_name.encode():
res.append((meta.names[i], probs[j][i], (boxes[j].x, boxes[j].y, boxes[j].w, boxes[j].h)))
res = sorted(res, key=lambda x: -x[1])
free_image(im)
free_ptrs(cast(probs, POINTER(c_void_p)), num)
if entity_name is None:
count = 0
return res, count
else:
return res, len(res)
def draw_bbox(image, box_list):
"""
Draw bounding boxes in box list
"""
img = image.copy()
for box in box_list:
x = int(box[2][0])
y = int(box[2][1])
w = int(box[2][2])
h = int(box[2][3])
# bounding box to co-ordinates
# for drawing bounding box in the image
x1 = x - int(w/2)
y1 = y + int(h/2)
x2 = x + int(w/2)
y2 = y - int(h/2)
cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 1)
return img
if __name__ == "__main__":
# Load the network, MS COCO parameters and YOLO V2 weights
net = load_net(b"cfg/yolo.cfg", b"yolov2.weights", 0)
meta = load_meta(b"cfg/coco.data")
# Read the image using OpenCV
image = cv2.imread("test_images/cars.jpeg")
# Detection using detect_np
car_box_list, count = detect_np(net, meta, image, entity_name='car')
# person_box_list, count = detect_np(net, meta, image, entity_name='person')
# Draw the bounding boxes
im_out = draw_bbox(image, car_box_list)
# Write it to output.jpeg
cv2.imwrite('test_images/output.jpeg', im_out)