代码如下:

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,并对视频的对于人群的人数进行统计。

结果如下: 

更多推荐