代码如下:
from ultralytics import YOLO
import cv2
model = YOLO('./best.pt')
sum_child = 0
sum_youngster = 0
sum_middle_age = 0
sum_older = 0
qist = {'child': 0.0, 'youngster': 1.0, 'middle_age': 2.0, 'older': 3.0}
# 这是标签对应的序列号
people = []
cap = cv2.VideoCapture('demo.mp4')
while True:
ret, frame = cap.read()
if not ret:
break
frame = cv2.resize(frame, None, fx=0.5, fy=0.5)
result = model.track(frame, device=0, show=True)
result = result[0]
list_people = []
list_id = []
for i in range(len(result.boxes)):
box = result.boxes[i]
list_people.append(box.cls[0].item())
list_id.append(box.id[0].item())
for a in list_id:
if a not in people:
if list_people[int(a - 1.0)] == ['child']:
sum_child += 1
elif list_people[int(a - 1.0)] == qist['youngster']:
sum_youngster += 1
elif list_people[int(a - 1.0)] == qist['middle_age']:
sum_middle_age += 1
else:
sum_older += 1
people.append(a)
sum_values = [sum_child, sum_youngster, sum_middle_age, sum_older]
print(sum_values)
cap.release()
cv2.destroyAllWindows()
这里我使用的是我自己训练的yolov8目标检测的权重,这里检测四个类别:child,youngster,middle_age,older,并对视频的对于人群的人数进行统计。
结果如下: