python yolov8 依据track 判断是否拥堵

思路

根据视频获取每帧，获取每个tarck 的xyxy 计算这个tarck 中心点 ，和下个帧的tarck 对比计算同方向车辆
获取最大车辆
获取70%车辆的通过帧数，
每计算一帧当前帧都会增加1
当前tarck的帧数除以70%的通过帧数 会得到占比,在1到3之间可算拥堵

代码

import time

import cv2
import numpy as np
import math
from ultralytics import RTDETR, YOLO
from scipy.spatial import distance
def run():
    """Test training the YOLO model from scratch."""
    model = YOLO("car.pt")
    video_path = "car4.mp4"
    cap = cv2.VideoCapture(video_path)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    index = 0
    cache = {}

    in_count = 0
    out_count = 0
    # denslty_times = {}
    max_cars = 0
    max_frame = 0
    car_5_p = 0.5 # 50%-60% 平均车辆通过帧数
    # 拥堵率
    congestion = 0
    # Loop through the video frames
    while cap.isOpened():
        # Read a frame from the video
        success, frame = cap.read()
        index += 1
        if success:
            st = time.time()
            max_8_rate = []
            # if index % 2 ==0 :
            #     continue
            # Run YOLOv8 tracking on the frame, persisting tracks between frames
            # results = model.track(frame, persist=True, device="cuda:0")
            # results = model.track(frame, persist=True, tracker="botsort.yaml")
            results = model.track(frame, device="cuda:0",  persist=True, tracker="bytetrack.yaml")
            et = time.time()
            print(f'track time {(et - st) * 1000} ms')
            # Visualize the results on the frame

            et1 = time.time()
            # print(f'plot time {(et1 - et) * 1000} ms')
            for result in results:
                boxes = result.boxes

                if boxes.is_track:
                    for box in boxes:
                        id = int(box.id)
                        xyxy = box.xyxy[0]
                        points = [
                            (xyxy[0] , xyxy[1] ),
                            (xyxy[2] , xyxy[3] )
                         ]
                        center = np.mean(points, axis=0)
                        if id in cache.keys():
                            old = cache[id]
                            # temp = [
                            #     old,
                            #     center
                            # ]
                            # diff = np.mean(temp, axis=0)
                            old_center_point = old['data']
                            move = center - old['data']
                            y = center[1] - old_center_point[1]
                            x = center[0] - old_center_point[0]
                            # y = 1 - 0
                            # x = 1 - 0
                            angle = cel_angle(center[0], center[1], old_center_point[0],
                                              old_center_point[1])


                                # angle += 360
                            # print(f'tacker: {id} angle: {angle} ')

                            # if angle > 0 and angle < 90:
                            #     print(f'id {id} 上 move 东南')
                            # elif angle > 90 and angle < 180:
                            #     print(f'id {id} 右上 move 西南')
                            # elif angle > 180 and angle < 270:
                            #     print(f'id {id} 右 move 西北')
                            # elif angle > 270 and angle < 360:
                            #     print(f'id {id}  move 东北')
                            # elif angle == 0:
                            #     print(f'id {id}  move 东')
                            # elif angle == 90:
                            #     print(f'id {id}  move 南')
                            # elif angle == 180 :
                            #     print(f'id {id}  move 西')
                            # elif angle == 270:
                            #     print(f'id {id}  move 北')
                            # else :
                            #     print(f'id {id}  move 0')

                            current_frame_count =cache[id]['frame'] + 1
                            cache[id] = {
                                'data': center,
                                'clear': False,
                                'frame': current_frame_count
                            }
                            if current_frame_count < max_frame:
                                # 计算 当前帧 在平均帧的占比
                                ccp = current_frame_count / car_5_p
                                if ccp > 1 and ccp < 3:
                                    max_8_rate.append(ccp)
                                    # print(f'car {id} current {current_frame_count} pass p {ccp * 100}%')
                        else:
                            cache[id] = {
                                'data':center,
                                'clear':False,
                                'frame': 1
                            }
                            in_count += 1

                    # 清除过期 车辆
                    for key in list(cache.keys()):
                        car_cache = cache[key]
                        if car_cache['clear']:
                            del cache[key]
                            out_count += 1
                            # print(f'remove key {key}')
                        if car_cache['frame'] > max_frame:
                            max_frame = car_cache['frame']
                        car_cache['clear'] = True

                    # 替换最大车辆
                    if len(cache) > max_cars:
                        max_cars = len(cache)

                    # 计算 小于 70%的车的帧数
                    car_p = len(cache) / max_cars
                    if  0.7 < car_p and car_p < 1:
                        # 计算所有车平均帧数
                        temp_sum = 0
                        for key in list(cache.keys()):
                            car_cache = cache[key]
                            temp_sum += car_cache['frame']
                        car_5_p = temp_sum / len(cache)

                    print(boxes.id.tolist())

            congestion = np.mean(max_8_rate)
            # print(f'car sum:{len(cache)} max frame:{max_frame} 50%:{car_5_p} frame  Congestion {congestion}')
            # Display the annotated frame
            annotated_frame = results[0].plot()
            cv2.imshow("YOLOv8 Tracking", annotated_frame)


            et2 = time.time()
            # print(f'show time {(et2 - et1) * 1000} ms')

            # Break the loop if 'q' is pressed
            if cv2.waitKey(1) & 0xFF == ord("q"):
                break
        else:
            # Break the loop if the end of the video is reached
            break

    # Release the video capture object and close the display window
    cap.release()
    cv2.destroyAllWindows()


def cel_angle(x1,y1,x2,y2):
    angle = 0.0
    dx = x2 - x1
    dy = y2 - y1
    if x2 == x1:
        angle = math.pi / 2.0
        if y2 == y1:
            angle = 0.0
        elif y2 < y1:
            angle = 3.0 * math.pi / 2.0
    elif x2 > x1 and y2 > y1:
        angle = math.atan(dx / dy)
    elif x2 > x1 and y2 < y1:
        angle = math.pi / 2 + math.atan(-dy / dx)
    elif x2 < x1 and y2 < y1:
        angle = math.pi + math.atan(dx / dy)
    elif x2 < x1 and y2 > y1:
        angle = 3.0 * math.pi / 2.0 + math.atan(dy / -dx)
    return (angle * 180 / math.pi)
if __name__ == '__main__':
    run()