YOLOで来場者数をカウント(rev2)

2025-01-14 raspi python yolo

前回の改良版で、カメラを横切った人をカウントします。 左から右に横切るとentered、右から左に横切るとexited、の数値が加算されます。

理論的には、これらを差し引きすると部屋の中にいる人数を推定できるはずですが、実際は同時に複数人が入り混じると誤計測も発生するため、あまり厳密な用途には向かなそうです。

import cv2
from ultralytics import YOLO

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

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)]

n_people_entered = 0
n_people_exited = 0
edge_left = 80
edge_right = 560

# Dictionary to track people states
people = {}

while True:
    ret, image = cam.read()
    if not ret:
        break

    results = model.track(image)
    font = cv2.FONT_HERSHEY_SIMPLEX

    # Draw boundaries
    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)

    # Display count
    m = f"Entered: {n_people_entered} | Exited: {n_people_exited}"
    cv2.putText(image, m, (10, 450), font, 1.5, (0, 0, 0), 3, cv2.LINE_AA)

    for box in results[0].boxes:
        if box.cls != 0:  # Ignore non-person objects
            continue

        # Get unique ID of the detection
        if box.id is None: continue
        detection_id = int(box.id)

        x = int((box.xyxy[0][2] + box.xyxy[0][0]) / 2)  # Calculate box center x-coordinate
        in_left_edge = x < edge_left
        in_right_edge = x > edge_right

        # Initialize state for new detections
        if detection_id not in people:
            people[detection_id] = {"left_crossed": False, "right_crossed": False, "has_entered": False, "has_exited": False}

        # Entry logic: left to right
        if not people[detection_id]["has_entered"] and people[detection_id]["left_crossed"] and in_right_edge:
            n_people_entered += 1
            people[detection_id]["has_entered"] = True
            people[detection_id]["left_crossed"] = False
            people[detection_id]["right_crossed"] = True
            cv2.rectangle(image, (edge_right, 0), (640, 640), (0, 255, 0), 4)

        # Exit logic: right to left
        elif not people[detection_id]["has_exited"] and people[detection_id]["right_crossed"] and in_left_edge:
            n_people_exited += 1
            people[detection_id]["has_exited"] = True
            people[detection_id]["right_crossed"] = False
            people[detection_id]["left_crossed"] = True
            cv2.rectangle(image, (0, 0), (edge_left, 480), (255, 0, 0), 4)

        # Set flags for initial crossing
        elif in_right_edge and not people[detection_id]["right_crossed"]:
            people[detection_id]["right_crossed"] = True
        elif in_left_edge and not people[detection_id]["left_crossed"]:
            people[detection_id]["left_crossed"] = True

        # Draw bounding box
        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[detection_id % len(colors)], 2)

    # Reset state for people who have crossed both boundaries (complete cycle)
    for detection_id in list(people.keys()):
        if people[detection_id]["has_entered"] and people[detection_id]["has_exited"]:
            # Reset state after the cycle is complete (exit after entry)
            people[detection_id]["has_entered"] = False
            people[detection_id]["has_exited"] = False
            people[detection_id]["left_crossed"] = False
            people[detection_id]["right_crossed"] = False

    # Show the updated frame
    cv2.imshow('yolo11n', image)
    k = cv2.waitKey(1)
    if k != -1:  # Exit on key press
        break

cam.release()
cv2.destroyAllWindows()