-
Notifications
You must be signed in to change notification settings - Fork 3
/
detect3.py
315 lines (240 loc) · 10.4 KB
/
detect3.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
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
import torch
from PIL import Image
from PIL import ImageDraw
import numpy as np
import utils
import nets
from torchvision import transforms
import time
import os
import matplotlib.pyplot as plt
import cv2
class Detector:
def __init__(self, pnet_param="./param6/p_net.pth", rnet_param="./param6/r_net.pth", onet_param="./param6/o_net.pth",
isCuda=True):
self.isCuda = isCuda
self.pnet = nets.PNet()
self.rnet = nets.RNet()
self.onet = nets.ONet()
if self.isCuda:
self.pnet.cuda()
self.rnet.cuda()
self.onet.cuda()
self.pnet.load_state_dict(torch.load(pnet_param, map_location="cuda"))
self.rnet.load_state_dict(torch.load(rnet_param, map_location="cuda"))
self.onet.load_state_dict(torch.load(onet_param, map_location="cuda"))
self.pnet.eval() # 批归一化, 使用之前训练的Batchnormal
self.rnet.eval()
self.onet.eval()
self._image_transform = transforms.Compose(
[
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
]
)
def detect(self, image):
start_time = time.time() #获取当前时间的函数。
pnet_boxes = self.__pnet_detect(image)
if pnet_boxes.shape[0] == 0: # 防止程序格式错误
return np.array([])
end_time = time.time()
t_pnet = end_time - start_time
# return pnet_boxes
start_time = time.time()
rnet_boxes = self.__rnet_detect(image, pnet_boxes)
# print( rnet_boxes)
if rnet_boxes.shape[0] == 0:
return np.array([])
end_time = time.time()
t_rnet = end_time - start_time
start_time = time.time()
onet_boxes = self.__onet_detect(image, rnet_boxes)
if onet_boxes.shape[0] == 0:
return np.array([])
end_time = time.time()
t_onet = end_time - start_time
t_sum = t_pnet + t_rnet + t_onet
print("total:{0} pnet:{1} rnet:{2} onet:{3}".format(t_sum, t_pnet, t_rnet, t_onet))
return onet_boxes # 可以更改为p网络框进行测试
def __rnet_detect(self, image, pnet_boxes):
_img_dataset = []
_pnet_boxes = utils.convert_to_square(pnet_boxes)
for _box in _pnet_boxes:
_x1 = int(_box[0])
_y1 = int(_box[1])
_x2 = int(_box[2])
_y2 = int(_box[3])
image = np.array(image)
image = Image.fromarray(image)
img = image.crop((_x1, _y1, _x2, _y2))
img = img.resize((24, 24))
img_data = self._image_transform(img)
_img_dataset.append(img_data)
img_dataset = torch.stack(_img_dataset)
if self.isCuda:
img_dataset = img_dataset.cuda()
_cls, _offset = self.rnet(img_dataset)
cls = _cls.cpu().data.numpy()
offset = _offset.cpu().data.numpy()
boxes = []
idxs, _ = np.where(cls > 0.6) # 0.7-0.8
for idx in idxs:
_box = _pnet_boxes[idx] # 根据索引拿到裁剪框(预测框)的坐标,便于反算
_x1 = int(_box[0])
_y1 = int(_box[1])
_x2 = int(_box[2])
_y2 = int(_box[3])
ow = _x2 - _x1
oh = _y2 - _y1
x1 = _x1 + ow * offset[idx][0]
y1 = _y1 + oh * offset[idx][1]
x2 = _x2 + ow * offset[idx][2]
y2 = _y2 + oh * offset[idx][3]
# print(cls[idx][0]) # 拿到标量
boxes.append([x1, y1, x2, y2, cls[idx][0]])
return utils.nms(np.array(boxes), 0.5)
def __onet_detect(self, image, rnet_boxes):
_img_dataset = []
_rnet_boxes = utils.convert_to_square(rnet_boxes)
for _box in _rnet_boxes:
_x1 = int(_box[0])
_y1 = int(_box[1])
_x2 = int(_box[2])
_y2 = int(_box[3])
img = image.crop((_x1, _y1, _x2, _y2))
img = img.resize((48, 48))
img_data = self._image_transform(img)
_img_dataset.append(img_data)
img_dataset = torch.stack(_img_dataset)
if self.isCuda:
img_dataset = img_dataset.cuda()
_cls, _offset = self.onet(img_dataset)
cls = _cls.cpu().data.numpy()
offset = _offset.cpu().data.numpy()
boxes = []
idxs, _ = np.where(cls > 0.95) # 一般为0.99
for idx in idxs:
_box = _rnet_boxes[idx]
_x1 = int(_box[0])
_y1 = int(_box[1])
_x2 = int(_box[2])
_y2 = int(_box[3])
ow = _x2 - _x1
oh = _y2 - _y1
# 网络在输入裁剪图片的时候,坐标已经知道,无需反算预测框,直接反算预测框到真实框的位置即可。
x1 = _x1 + ow * offset[idx][0]
y1 = _y1 + oh * offset[idx][1]
x2 = _x2 + ow * offset[idx][2]
y2 = _y2 + oh * offset[idx][3]
boxes.append([x1, y1, x2, y2, cls[idx][0]])
return utils.nms(np.array(boxes), 0.7, isMin=True) # #阈值0.7。保留IOU小于0.7的框
def __pnet_detect(self, image):
boxes = []
img = image
w, h = img.size # 图片的宽高
# w, h = 960, 544
# print("------")
# print(np.shape(img))
min_side_len = min(w, h) # 得到图片最小的一边
scale = 1 #缩放比例
count = 0
while min_side_len > 12:
# print("第{}循环开始".format(count+1))
img_data = self._image_transform(img)
if self.isCuda:
img_data = img_data.cuda()
# print(img_data.shape) # torch.Size([3, 722, 1200])
img_data.unsqueeze_(0) # 加批次N, torch.Size([1, 3, 722, 1200])
_cls, _offest = self.pnet(img_data)
# print(_cls.shape) # torch.Size([1, 1, 722, 1200])
# print(_offest.shape) # torch.Size([1, 4, 722, 1200])
cls, offest = _cls[0][0].cpu().data, _offest[0].cpu().data # _cls:H, W, _offset:C H W
idxs = torch.nonzero(torch.gt(cls, 0.5), as_tuple=False) # 置信度一般为0.5,先得到布尔值,再取出索引
# print(idxs.shape) # torch.Size([N, 2])
for idx in idxs: # 拿到所有满足要求的特征图索引,开始遍历,反算到原图的真实框
boxes.append(self.__box(idx, offest, cls, scale)) # 添加符合条件的真实框, 拿到索引对应置信度的值。
scale *= 0.7
_w = int(w * scale)
_h = int(h * scale)
# print(_w, _h) # 672 380
# print(np.shape(img)) # (960, 544, 3)
img = img.resize((_w, _h)) # (960, 544, 3)
# img = np.resize(img, (_w, _h)) # # (672, 380)
# img = np.expand_dims(img, 2).repeat(3, axis=2) # 在第二维度上重复三次。
# print("========")
# print(np.shape(img)) # (672, 380, 3)
min_side_len = np.minimum(_w, _h)
count += 1
# print("第{}轮循环完成".format(count))
return utils.nms(np.array(boxes), 0.3) # 保留IOU小于0.5的框,IOU取值越小,去框越多。
def __box(self, start_index, offset, cls, scale, stride=2, side_len=12):
# 利用特征图的索引反算到原图上的预测框
_x1 = (start_index[1] * stride).float() / scale
_y1 = (start_index[0] * stride).float() / scale
_x2 = (start_index[1] * stride + side_len).float() / scale
_y2 = (start_index[0] * stride + side_len).float() / scale
ow = _x2 - _x1
oh = _y2 - _y1
# 利用预测框位置反算到真实框。
cls = cls[start_index[0], start_index[1]] # 取到对应索引位置的置信度
# print(offset.shape) # torch.Size([4, 722, 1200])
_offset = offset[:, start_index[0], start_index[1]] # C H W 取H W对应位置
# print(_offset) # tensor([ 0.0773, -0.1806, 0.3268, 0.5450])
x1 = _x1 + ow * _offset[0]
y1 = _y1 + oh * _offset[1]
x2 = _x2 + ow * _offset[2]
y2 = _y2 + oh * _offset[3]
return [x1, y1, x2, y2, cls]
if __name__ == '__main__':
with torch.no_grad() as grad:
# path = "3.mp4" # 本地视频路径
# path = r"http://vfx.mtime.cn/Video/2019/03/19/mp4/190319125415785691.mp4" # 在线视频路径
# cap = cv2.VideoCapture(path)
cap = cv2.VideoCapture(0) # 调取内置摄像头
fps = cap.get(cv2.CAP_PROP_FPS)
# print(fps)
w = int(cap.get(3)) # 获取图片的宽度
h = int(cap.get(4)) # 获取图片的高度
# print(w)
# print(h)
# 写入视频格式
# fourc = cv2.VideoWriter_fourcc(*"DVIX")
# out = cv2.VideoWriter("2.mp4", fourc, fps, (w, h))
font = cv2.FONT_HERSHEY_COMPLEX
# frame表示读出的每一张图片, ret表示这一张图片是否存在
while True:
x = time.time()
ret, frame1 = cap.read()
# 将十六进制数据转成 二进制数据
if cv2.waitKey(int(1)) & 0xFF == ord("q"): # 视频在播放的过程中按键,循环会中断)。
break
elif ret == False: # 视频播放完了,循环自动中断。
break
frame2 = Image.fromarray(frame1)
detector = Detector()
boxes = detector.detect(frame2)
print("侦测出来的框:", boxes)
for box in boxes:
# x1 = int(box[0])
# y1 = int(box[1])
# x2 = int(box[2])
# y2 = int(box[3])
# w = int(0.1 * (box[2] - box[0]))
w = 0
# h = int(0.1 * (box[3] - box[1]))
h = 0
x1 = int(box[0]) + w
y1 = int(box[1]) + h
x2 = int(box[2]) - w
y2 = int(box[3]) - h
print(box[4])
cv2.rectangle(frame1, (x1, y1), (x2, y2), [0, 0, 255], 3)
# else:
# print("矩形框没有值")
y = time.time()
print(1/(y-x))
# out.write(frame1)
# cv2.putText(frame1, "", (100, 100), font, 1, (0, 0, 255), 1, lineType=cv2.LINE_AA)
cv2.imshow("", frame1)
cap.release() # 将视频关了
cv2.destroyAllWindows()