YOLOで来場者数をカウント

2024-12-21 raspi python yolo

YOLOを使って得られるバウンディングボックスを利用して、画面の左端から現れて、右端まで移動した人の数をカウント(n_people)するプログラムの例です。もう少しスマートなトラッキング方法もあるような気がするのですが、、まずはバージョン1ということで、公開します。ちょっとずつブラッシュアップしましょう。

import cv2

cv2.namedWindow('yolo11n')
cam = cv2.VideoCapture(0)

from ultralytics import YOLO
model = YOLO('yolo11n-pose_ncnn_model')
colors = [(255,0,0),(0,255,0),(0,0,255),(255,255,0),(255,0,255),(0,255,255)]
people = {}
n_people = 0
edge_left = 80
edge_right = 560
while True:
    ret, image = cam.read()
    results = model.track(image)
    m = "r:%d b:%d people: %d =>" % (len(results),len(results[0].boxes),n_people)
    font = cv2.FONT_HERSHEY_SIMPLEX
    cv2.putText(image,m,(10,460),font,1,(5,5,5),1,cv2.LINE_AA)
    cv2.rectangle(image,(1,1),(edge_left-1,478),(128,128,128),2)
    cv2.rectangle(image,(edge_right-1,1),(638,638),(128,128,128),2)
    for box in results[0].boxes:
        if box.id == None or box.cls != 0: continue
        i = int(box.id)
        if i not in people: people[i] = {"x": 0, "left": False, "right": False, "ttl": 10}
        x = int((box.xyxy[0][2] + box.xyxy[0][0]) / 2)
        people[i]["x"] = x
        in_left_edge = x < edge_left
        in_right_edge = x > edge_right
        if people[i]["left"] and people[i]["right"]:
            if people[i]["ttl"] > 0:
                people[i]["ttl"] -= 1
            else:
                people[i]["left"] = False
                people[i]["right"] = False
                people[i]["ttl"] = 30
        elif people[i]["left"] and in_right_edge:
            people[i]["right"] = True
            n_people += 1
            cv2.rectangle(image,(edge_right,0),(640,640),(255,0,0),4)
        elif people[i]["right"] and in_left_edge:
            people[i]["left"] = True
            n_people -= 1
            cv2.rectangle(image,(0,0),(edge_left,480),(255,0,0),4)
        elif in_right_edge:
            people[i]["right"] = True
        elif in_left_edge:
            people[i]["left"] = True

        pt1 = (int(box.xyxy[0][0]),int(box.xyxy[0][1]))
        pt2 = (int(box.xyxy[0][2]),int(box.xyxy[0][3]))
        cv2.rectangle(image,pt1,pt2,colors[i % 6],1)
        cv2.putText(image,"%d" % i,(x,24),font,1,colors[i % 6],1,cv2.LINE_AA) 
    for k in people:     
        print(people[k])    

    cv2.imshow('yolo11n',image)
    k = cv2.waitKey(1)
    if k != -1:
        break

cam.release()
cv2.destroyAllWindows()