er 不是二 我故意的！别讲我二。留言二的那个，你别跑！
人脸识别demo完成，搞一搞yolov5+deepsort
环境还是上个环境，装了就没问题
直接上百度盘链接
链接：https://pan.baidu.com/s/16q3wOgZ4zGqkEkAzVmk2Mw?
提取码：yolo?
代码，模型，数据都有，下载下来直接run? mian.py就行
import numpy as np
import tracker
from detector import Detector
import cv2
if __name__ == '__main__':
# 根据视频尺寸，填充一个polygon，供撞线计算使用
mask_image_temp = np.zeros((1080, 1920), dtype=np.uint8)
# 初始化2个撞线polygon
list_pts_blue = [[204, 305], [227, 431], [605, 522], [1101, 464], [1900, 601], [1902, 495], [1125, 379], [604, 437],
[299, 375], [267, 289]]
ndarray_pts_blue = np.array(list_pts_blue, np.int32)
polygon_blue_value_1 = cv2.fillPoly(mask_image_temp, [ndarray_pts_blue], color=1)
polygon_blue_value_1 = polygon_blue_value_1[:, :, np.newaxis]
# 填充第二个polygon
mask_image_temp = np.zeros((1080, 1920), dtype=np.uint8)
list_pts_yellow = [[181, 305], [207, 442], [603, 544], [1107, 485], [1898, 625], [1893, 701], [1101, 568],
[594, 637], [118, 483], [109, 303]]
ndarray_pts_yellow = np.array(list_pts_yellow, np.int32)
polygon_yellow_value_2 = cv2.fillPoly(mask_image_temp, [ndarray_pts_yellow], color=2)
polygon_yellow_value_2 = polygon_yellow_value_2[:, :, np.newaxis]
# 撞线检测用mask，包含2个polygon，（值范围 0、1、2），供撞线计算使用
polygon_mask_blue_and_yellow = polygon_blue_value_1 + polygon_yellow_value_2
# 缩小尺寸，1920x1080->960x540
polygon_mask_blue_and_yellow = cv2.resize(polygon_mask_blue_and_yellow, (960, 540))
# 蓝 色盘 b,g,r
blue_color_plate = [255, 0, 0]
# 蓝 polygon图片
blue_image = np.array(polygon_blue_value_1 * blue_color_plate, np.uint8)
# 黄 色盘
yellow_color_plate = [0, 255, 255]
# 黄 polygon图片
yellow_image = np.array(polygon_yellow_value_2 * yellow_color_plate, np.uint8)
# 彩色图片（值范围 0-255）
color_polygons_image = blue_image + yellow_image
# 缩小尺寸，1920x1080->960x540
color_polygons_image = cv2.resize(color_polygons_image, (960, 540))
# list 与蓝色polygon重叠
list_overlapping_blue_polygon = []
# list 与黄色polygon重叠
list_overlapping_yellow_polygon = []
# 进入数量
down_count = 0
# 离开数量
up_count = 0
font_draw_number = cv2.FONT_HERSHEY_SIMPLEX
draw_text_postion = (int(960 * 0.01), int(540 * 0.05))
# 初始化 yolov5
detector = Detector()
# 打开视频
capture = cv2.VideoCapture('./video/test.mp4')
# capture = cv2.VideoCapture('/mnt/datasets/datasets/towncentre/TownCentreXVID.avi')
while True:
# 读取每帧图片
_, im = capture.read()
if im is None:
break
# 缩小尺寸，1920x1080->960x540
im = cv2.resize(im, (960, 540))
list_bboxs = []
bboxes = detector.detect(im)
# 如果画面中 有bbox
if len(bboxes) > 0:
list_bboxs = tracker.update(bboxes, im)
# 画框
# 撞线检测点，(x1，y1)，y方向偏移比例 0.0~1.0
output_image_frame = tracker.draw_bboxes(im, list_bboxs, line_thickness=None)
pass
else:
# 如果画面中 没有bbox
output_image_frame = im
pass
# 输出图片
output_image_frame = cv2.add(output_image_frame, color_polygons_image)
if len(list_bboxs) > 0:
# ----------------------判断撞线----------------------
for item_bbox in list_bboxs:
x1, y1, x2, y2, label, track_id = item_bbox
# 撞线检测点，(x1，y1)，y方向偏移比例 0.0~1.0
y1_offset = int(y1 + ((y2 - y1) * 0.6))
# 撞线的点
y = y1_offset
x = x1
if polygon_mask_blue_and_yellow[y, x] == 1:
# 如果撞 蓝polygon
if track_id not in list_overlapping_blue_polygon:
list_overlapping_blue_polygon.append(track_id)
pass
# 判断 黄polygon list 里是否有此 track_id
# 有此 track_id，则 认为是 外出方向
if track_id in list_overlapping_yellow_polygon:
# 外出+1
up_count += 1
print(f'类别: {label} | id: {track_id} | 上行撞线 | 上行撞线总数: {up_count} | 上行id列表: {list_overlapping_yellow_polygon}')
# 删除 黄polygon list 中的此id
list_overlapping_yellow_polygon.remove(track_id)
pass
else:
# 无此 track_id，不做其他操作
pass
elif polygon_mask_blue_and_yellow[y, x] == 2:
# 如果撞 黄polygon
if track_id not in list_overlapping_yellow_polygon:
list_overlapping_yellow_polygon.append(track_id)
pass
# 判断 蓝polygon list 里是否有此 track_id
# 有此 track_id，则 认为是 进入方向
if track_id in list_overlapping_blue_polygon:
# 进入+1
down_count += 1
print(f'类别: {label} | id: {track_id} | 下行撞线 | 下行撞线总数: {down_count} | 下行id列表: {list_overlapping_blue_polygon}')
# 删除 蓝polygon list 中的此id
list_overlapping_blue_polygon.remove(track_id)
pass
else:
# 无此 track_id，不做其他操作
pass
pass
else:
pass
pass
pass
# ----------------------清除无用id----------------------
list_overlapping_all = list_overlapping_yellow_polygon + list_overlapping_blue_polygon
for id1 in list_overlapping_all:
is_found = False
for _, _, _, _, _, bbox_id in list_bboxs:
if bbox_id == id1:
is_found = True
break
pass
pass
if not is_found:
# 如果没找到，删除id
if id1 in list_overlapping_yellow_polygon:
list_overlapping_yellow_polygon.remove(id1)
pass
if id1 in list_overlapping_blue_polygon:
list_overlapping_blue_polygon.remove(id1)
pass
pass
list_overlapping_all.clear()
pass
# 清空list
list_bboxs.clear()
pass
else:
# 如果图像中没有任何的bbox，则清空list
list_overlapping_blue_polygon.clear()
list_overlapping_yellow_polygon.clear()
pass
pass
text_draw = 'DOWN: ' + str(down_count) + \
' , UP: ' + str(up_count)
output_image_frame = cv2.putText(img=output_image_frame, text=text_draw,
org=draw_text_postion,
fontFace=font_draw_number,
fontScale=1, color=(255, 255, 255), thickness=2)
cv2.imshow('demo', output_image_frame)
cv2.waitKey(1)
pass
pass
capture.release()
cv2.destroyAllWindows()
撞线对我目前来说没什么大用，但是考虑到后面可以做个电子围栏还是可以留下的。
detect.py
import torch
import numpy as np
from models.experimental import attempt_load
from utils.datasets import letterbox
from utils.general import non_max_suppression, scale_coords
from utils.torch_utils import select_device
class Detector:
def __init__(self):
self.img_size = 640
self.threshold = 0.3
self.stride = 1
self.weights = './weights/yolov5m.pt'
self.device = '0' if torch.cuda.is_available() else 'cpu'
self.device = select_device(self.device)
model = attempt_load(self.weights, map_location=self.device)
model.to(self.device).eval()
model.half()
self.m = model
self.names = model.module.names if hasattr(
model, 'module') else model.names
def preprocess(self, img):
img0 = img.copy()
img = letterbox(img, new_shape=self.img_size)[0]
img = img[:, :, ::-1].transpose(2, 0, 1)
img = np.ascontiguousarray(img)
img = torch.from_numpy(img).to(self.device)
img = img.half()
img /= 255.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
return img0, img
def detect(self, im):
im0, img = self.preprocess(im)
pred = self.m(img, augment=False)[0]
pred = pred.float()
pred = non_max_suppression(pred, self.threshold, 0.4)
boxes = []
for det in pred:
if det is not None and len(det):
det[:, :4] = scale_coords(
img.shape[2:], det[:, :4], im0.shape).round()
for *x, conf, cls_id in det:
lbl = self.names[int(cls_id)]
if lbl not in ['person', 'bicycle', 'car', 'motorcycle', 'bus', 'truck']:
continue
pass
x1, y1 = int(x[0]), int(x[1])
x2, y2 = int(x[2]), int(x[3])
boxes.append(
(x1, y1, x2, y2, lbl, conf))
return boxes
if lbl not in ['person', 'bicycle', 'car', 'motorcycle', 'bus', 'truck'] 我只需要检测人就行 留个person其余全删。
准备工作到此结束。后面就开始结合人脸识别
