实战 | 基于YOLOv9和OpenCV实现车辆跟踪计数（步骤 + 源码）

pip install ultralytics

import math

class CustomTracker:
    def __init__(self):
        # Store the center positions of the objects
        self.custom_center_points = {}
        # Keep the count of the IDs
        # each time a new object id detected, the count will increase by one
        self.custom_id_count = 0

    def custom_update(self, custom_objects_rect):
        # Objects boxes and ids
        custom_objects_bbs_ids = []

        # Get center point of new object
        for custom_rect in custom_objects_rect:
            x, y, w, h = custom_rect
            cx = (x + x + w) // 2
            cy = (y + y + h) // 2

            # Find out if that object was detected already
            same_object_detected = False
            for custom_id, pt in self.custom_center_points.items():
                dist = math.hypot(cx - pt[0], cy - pt[1])

                if dist < 35:
                    self.custom_center_points[custom_id] = (cx, cy)
                    custom_objects_bbs_ids.append([x, y, w, h, custom_id])
                    same_object_detected = True
                    break

            # New object is detected we assign the ID to that object
            if same_object_detected is False:
                self.custom_center_points[self.custom_id_count] = (cx, cy)
                custom_objects_bbs_ids.append([x, y, w, h, self.custom_id_count])
                self.custom_id_count += 1

        # Clean the dictionary by center points to remove IDS not used anymore
        new_custom_center_points = {}
        for custom_obj_bb_id in custom_objects_bbs_ids:
            _, _, _, _, custom_object_id = custom_obj_bb_id
            center = self.custom_center_points[custom_object_id]
            new_custom_center_points[custom_object_id] = center

        # Update dictionary with IDs not used removed
        self.custom_center_points = new_custom_center_points.copy()
        return custom_objects_bbs_ids

# Import the Libraries
import cv2
import pandas as pd
from ultralytics import YOLO
from tracker import *

model=YOLO('yolov9c.pt')

class_list = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
              'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter',
              'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe',
              'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite',
              'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup',
              'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog',
              'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop',
              'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book',
              'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush']

tracker=CustomTracker()
count=0

cap = cv2.VideoCapture('traffictrim.mp4')

# Get video properties
fps = int(cap.get(cv2.CAP_PROP_FPS))
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

# Create VideoWriter object to save the modified frames
output_video_path = 'output_video.mp4'
fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # You can use other codecs like 'XVID' based on your system
out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))

# Looping over each frame and Performing the Detection

down = {}
counter_down = set()
while True:
    ret, frame = cap.read()
    if not ret:
        break
    count += 1

    results = model.predict(frame)

    a = results[0].boxes.data
    a = a.detach().cpu().numpy()
    px = pd.DataFrame(a).astype("float")
    # print(px)

    list = []

    for index, row in px.iterrows():
        #        print(row)
        x1 = int(row[0])
        y1 = int(row[1])
        x2 = int(row[2])
        y2 = int(row[3])
        d = int(row[5])
        c = class_list[d]
        if 'car' in c:
            list.append([x1, y1, x2, y2])

    bbox_id = tracker.custom_update(list)
    # print(bbox_id)
    for bbox in bbox_id:
        x3, y3, x4, y4, id = bbox
        cx = int(x3 + x4) // 2
        cy = int(y3 + y4) // 2
        # cv2.circle(frame,(cx,cy),4,(0,0,255),-1) #draw ceter points of bounding box
        # cv2.rectangle(frame, (x3, y3), (x4, y4), (0, 255, 0), 2)  # Draw bounding box
        # cv2.putText(frame,str(id),(cx,cy),cv2.FONT_HERSHEY_COMPLEX,0.8,(0,255,255),2)

        y = 308
        offset = 7

        ''' condition for red line '''
        if y < (cy + offset) and y > (cy - offset):
            ''' this if condition is putting the id and the circle on the object when the center of the object touched the red line.'''

            down[id] = cy  # cy is current position. saving the ids of the cars which are touching the red line first.
            # This will tell us the travelling direction of the car.
            if id in down:
                cv2.circle(frame, (cx, cy), 4, (0, 0, 255), -1)
                #cv2.putText(frame, str(id), (cx, cy), cv2.FONT_HERSHEY_COMPLEX, 0.8, (0, 255, 255), 2)
                counter_down.add(id)

                # # line
    text_color = (255, 255, 255)  # white color for text
    red_color = (0, 0, 255)  # (B, G, R)

    # print(down)
    cv2.line(frame, (282, 308), (1004, 308), red_color, 3)  # starting cordinates and end of line cordinates
    cv2.putText(frame, ('red line'), (280, 308), cv2.FONT_HERSHEY_SIMPLEX, 0.5, text_color, 1, cv2.LINE_AA)


    downwards = (len(counter_down))
    cv2.putText(frame, ('Vehicle Counter - ') + str(downwards), (60, 40), cv2.FONT_HERSHEY_SIMPLEX, 0.5, red_color, 1,
                cv2.LINE_AA)

    cv2.line(frame,(282,308),(1004,308),red_color,3)  #  starting cordinates and end of line cordinates
    cv2.putText(frame,('red line'),(280,308),cv2.FONT_HERSHEY_SIMPLEX, 0.5, text_color, 1, cv2.LINE_AA)
    
    # This will write the Output Video to the location specified above
    out.write(frame)