基于YOLOv8的目标跟踪——汽车跟踪和计数_model.track

目标跟踪针对的是视频处理，它是目标检测的更高级的应用。目标跟踪要解决的主要问题应该是能够正确识别不同帧之间的同一个目标，而不仅仅是同一类目标。例如，在某个连续的时间段内总是出现张三这个人，目标跟踪可以在这段时间内把张三这个人标记为同一个人，从而实现跟踪其轨迹的目的。

YOLOv8不仅可以实现目标检测和目标分割，还可以实现目标跟踪。它的目标跟踪基于的是BoT-SORT和ByteTrack，而默认的是BoT-SORT算法。在这里，我们不介绍任何原理性的内容，只以一个很简单的例子讲解如何应用YOLOv8进行目标跟踪，以及它所带来的一个附加功能——计数。

我们的实例是对一段高速公路进行车辆跟踪，并分别记录从上至下（驶出）和从下至上（驶入）的各类车辆的数量。对于车辆跟踪，直接应用的是YOLOv8的函数model.track，而计数功能则需要自己实现算法。因为我们这里应用的场景比较简单，所以我们只是设置了一条水平基准线，通过前后两帧同一辆车辆的坐标位置，可以判断出其行驶的方向，当越过基准线时，就计数一次。

下面就给出完整的代码，及其详细解析。

导入模块，其中为了保存不同ID的目标，我们需要用到defaultdict：

import cv2
from ultralytics import YOLO
from collections import defaultdict

加载YOLOv8的预训练模型：

model = YOLO('yolov8l.pt')

加载待处理的视频，并获取相应的视频参数：

cap = cv2.VideoCapture("D:/track/Highway Traffic.mp4")
fps = cap.get(cv2.CAP_PROP_FPS)
size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
fNUMS = cap.get(cv2.CAP_PROP_FRAME_COUNT)

最终把结果保存为mp4格式的视频文件：

fourcc = cv2.VideoWriter_fourcc(*'mp4v')
videoWriter = cv2.VideoWriter("D:/track/counting.mp4", fourcc, fps, size)

把目标用矩形框标注：

def box_label(image, box, label='', color=(128, 128, 128), txt_color=(255, 255, 255)):
    #得到目标矩形框的左上角和右下角坐标
    p1, p2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
    #绘制矩形框
    cv2.rectangle(image, p1, p2, color, thickness=1, lineType=cv2.LINE_AA)
    if label:
        #得到要书写的文本的宽和长，用于给文本绘制背景色
        w, h = cv2.getTextSize(label, 0, fontScale=2 / 3, thickness=1)[0]  
        #确保显示的文本不会超出图片范围
        outside = p1[1] - h >= 3
        p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3
        cv2.rectangle(image, p1, p2, color, -1, cv2.LINE_AA)     #填充颜色
        #书写文本
        cv2.putText(image,
                label, (p1[0], p1[1] - 2 if outside else p1[1] + h + 2),
                0,
                2 / 3,
                txt_color,
                thickness=1,
                lineType=cv2.LINE_AA)

实现目标跟踪，并完成计数：

# track_history用于保存目标ID，以及它在各帧的目标位置坐标，这些坐标是按先后顺序存储的
track_history = defaultdict(lambda: [])
#车辆的计数变量
vehicle_in = 0
vehicle_out = 0
 
#视频帧循环
while cap.isOpened():
    #读取一帧图像
    success, frame = cap.read()
 
    if success:
        #在帧上运行YOLOv8跟踪，persist为True表示保留跟踪信息，conf为0.3表示只检测置信值大于0.3的目标
        results = model.track(frame,conf=0.3, persist=True)
        #得到该帧的各个目标的ID
        track_ids = results[0].boxes.id.int().cpu().tolist()
        #遍历该帧的所有目标
        for track_id, box in zip(track_ids, results[0].boxes.data):
            if box[-1] == 2:   #目标为小汽车
                #绘制该目标的矩形框
                box_label(frame, box, '#'+str(track_id)+' car', (167, 146, 11))
                #得到该目标矩形框的中心点坐标(x, y)
                x1, y1, x2, y2 = box[:4]
                x = (x1+x2)/2
                y = (y1+y2)/2
                #提取出该ID的以前所有帧的目标坐标，当该ID是第一次出现时，则创建该ID的字典
                track = track_history[track_id]
                track.append((float(x), float(y)))  #追加当前目标ID的坐标
                #只有当track中包括两帧以上的情况时，才能够比较前后坐标的先后位置关系
                if len(track) > 1:
                    _, h = track[-2]  #提取前一帧的目标纵坐标
                    #我们设基准线为纵坐标是size[1]-400的水平线
                    #当前一帧在基准线的上面，当前帧在基准线的下面时，说明该车是从上往下运行
                    if h < size[1]-400 and y >= size[1]-400:
                        vehicle_out +=1      #out计数加1
                    #当前一帧在基准线的下面，当前帧在基准线的上面时，说明该车是从下往上运行
                    if h > size[1]-400 and y <= size[1]-400:
                        vehicle_in +=1       #in计数加1
                      
            elif box[-1] == 5:   #目标为巴士
                box_label(frame, box, '#'+str(track_id)+' bus', (67, 161, 255))
                
                x1, y1, x2, y2 = box[:4]
                x = (x1+x2)/2
                y = (y1+y2)/2
                track = track_history[track_id]
                track.append((float(x), float(y)))  # x, y center point
                if len(track) > 1:
                    _, h = track[-2]
                    if h < size[1]-400 and y >= size[1]-400:
                        vehicle_out +=1
                    if h > size[1]-400 and y <= size[1]-400:
                        vehicle_in +=1 
               
            elif box[-1] == 7:   #目标为卡车
                box_label(frame, box, '#'+str(track_id)+' truck', (19, 222, 24))
                
                x1, y1, x2, y2 = box[:4]
                x = (x1+x2)/2
                y = (y1+y2)/2
                track = track_history[track_id]
                track.append((float(x), float(y)))  # x, y center point
                if len(track) > 1:
                    _, h = track[-2]
                    if h < size[1]-400 and y >= size[1]-400:
                        vehicle_out +=1
                    if h > size[1]-400 and y <= size[1]-400:
                        vehicle_in +=1 
              
            elif box[-1] == 3:   #目标为摩托车
                box_label(frame, box,'#'+str(track_id)+' motor', (186, 55, 2))
                
                x1, y1, x2, y2 = box[:4]
                x = (x1+x2)/2
                y = (y1+y2)/2
                track = track_history[track_id]
                track.append((float(x), float(y)))  # x, y center point
                if len(track) > 1:
                    _, h = track[-2]
                    if h < size[1]-400 and y >= size[1]-400:
                        vehicle_out +=1
                    if h > size[1]-400 and y <= size[1]-400:
                        vehicle_in +=1 
        #绘制基准线
        cv2.line(frame, (30,size[1]-400), (size[0]-30,size[1]-400), color=(25, 33, 189), thickness=2, lineType=4)
        #实时显示进、出车辆的数量
        cv2.putText(frame, 'in: '+str(vehicle_in), (595, size[1]-410),
                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
        cv2.putText(frame, 'out: '+str(vehicle_out), (573, size[1]-370),
                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
 
        cv2.putText(frame, "https://blog.csdn.net/zhaocj", (25, 50),
                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
      
        cv2.imshow("YOLOv8 Tracking", frame)  #显示标记好的当前帧图像
        videoWriter.write(frame)  #写入保存
     
        if cv2.waitKey(1) & 0xFF == ord("q"):   #'q'按下时，终止运行
            break
 
    else:  #视频播放结束时退出循环
        break
 
#释放视频捕捉对象，并关闭显示窗口
cap.release()
videoWriter.release()
cv2.destroyAllWindows()

最终的结果为：

可以看出YOLOv8很好的完成了任务，同一辆车尽管有时被识别为car，有时又被识别为truck，但它的ID始终是同一个值。如果我们仔细观察结果视频会发现，进出车辆之和要小于最大的ID数。这是因为除了没有被统计进来的车辆（如，左上角和右上角的车辆）外，主要原因还是YOLOv8并没有完整的按顺序标记车辆ID，比如没有标记#6的车辆。

对上面的代码略加修改，完全可以应用到其他场合，如下面的视频：