Введение
В этой статье в основном рассказывается об использовании YOLOv9 и OpenCV для реализации отслеживания и подсчета транспортных средств (шаги + исходный код).
Этапы реализации
Камеры видеонаблюдения могут эффективно использоваться для подсчета транспортных средств и статистики транспортных потоков в различных сценариях. Передовые методы компьютерного зрения, такие как обнаружение и отслеживание объектов, можно применять к записям наблюдения для идентификации и отслеживания транспортных средств, движущихся в поле зрения камеры.
【1】Установите Ultralytics, поскольку у него есть метод прямого использования предварительно обученной модели YoloV9.
pip install ultralytics
【2】После завершения вы можете создать функцию трекера для отслеживания объекта. Для этого мы только что создали файл Python под названием tracker.py.
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
【3】Запишите основной код для подсчета транспортных средств.
# Import the Libraries
import cv2
import pandas as pd
from ultralytics import YOLO
from tracker import *
После импорта всех необходимых библиотек пришло время импортировать модель. Нам не нужно загружать модель из какого-либо репозитория. Ultralytics отлично справляется со своей задачей, упрощая их загрузку напрямую.
model=YOLO('yolov9c.pt')
Модель yolov9c.pt будет загружена в текущий каталог. Модель была обучена на наборе данных COCO, состоящем из 80 различных категорий. Теперь укажем класс:
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)
В приведенном выше коде мы просматриваем каждый кадр видео, а затем обнаруживаем его. Затем, поскольку мы подсчитываем только транспортные средства, отфильтровываются только обнаруженные автомобили.
После этого находим центры обнаруженных транспортных средств и считаем их по мере пересечения искусственно созданной красной линии. Мы можем ясно видеть их на видеофрагменте ниже.
видеоадрес:https://mpvideo.qpic.cn/0bc3qyaeiaaae4ah2j7qj5svbbwdisdaaraa.f10002.mp4?
Мы видим, что по мере пересечения транспортным средством красной линии счетчик в левом верхнем углу видео продолжает увеличиваться.