程式碼

以下是使用Yolov4和DeepSort进行目标检测和跟踪的Python代码示例，需要安装相关的库和模型文件才能运行。

import cv2
import numpy as np
import time
from deep_sort import preprocessing
from deep_sort import nn_matching
from deep_sort.detection import Detection
from deep_sort.tracker import Tracker
from tools import generate_detections as gdet

# 加载yolov4模型和deepsort模型
model_path = 'models/yolov4.weights'
config_path = 'models/yolov4.cfg'
labels = open('models/coco.names').read().strip().split('\n')
net = cv2.dnn.readNetFromDarknet(config_path, model_path)
net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
deepsort = Tracker("models/deep_sort.yaml")
metric = nn_matching.NearestNeighborDistanceMetric("cosine", 0.5, 100)

# 定义输入和输出视频路径
video_path = "test.mp4"
output_path = "output.avi"

# 定义输入和输出视频对象
video = cv2.VideoCapture(video_path)
if not video.isOpened():
    print("Error: Could not open video.")
    exit()
fps = int(video.get(cv2.CAP_PROP_FPS))
width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
output = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*"MJPG"), fps, (width, height))

# 处理视频帧
while True:
    ret, frame = video.read()
    if not ret:
        break
    
    # 对视频帧进行目标检测
    blob = cv2.dnn.blobFromImage(frame, 1 / 255.0, (608, 608), swapRB=True, crop=False)
    net.setInput(blob)
    layer_names = net.getLayerNames()
    layer_names = [layer_names[i[0] - 1] for i in net.getUnconnectedOutLayers()]
    outputs = net.forward(layer_names)
    boxes = []
    confidences = []
    class_ids = []
    for output in outputs:
        for detection in output:
            scores = detection[5:]
            class_id = np.argmax(scores)
            confidence = scores[class_id]
            if confidence > 0.5 and class_id == 0:
                box = detection[0:4] * np.array([width, height, width, height])
                (centerX, centerY, w, h) = box.astype("int")
                x = int(centerX - (w / 2))
                y = int(centerY - (h / 2))
                boxes.append([x, y, int(w), int(h)])
                confidences.append(float(confidence))
                class_ids.append(class_id)
    
    # 对检测结果进行深度特征提取和目标跟踪
    features = np.array([preprocessing.extract_image_patch(frame, box) for box in boxes])
    detections = [Detection(bbox, confidence, feature) for bbox, confidence, feature in zip(boxes, confidences, features)]
    tracker_boxes = []
    if len(detections) > 0:
        # 进行目标跟踪
        tracker_boxes = np.array([d.tlwh for d in detections])
        tracker_scores = metric.matching_cascade(gdet.get_embedding(deepsort.encoder, detections), deepsort.trackers, metric.threshold)
        for i, detection in enumerate(detections):
            deepsort.update(detection, tracker_scores[i])
    
    # 在视频帧上绘制检测和跟踪结果
    for track in deepsort.tracks:
        if not track.is_confirmed() or track.time_since_update > 1:
            continue
        bbox = track.to_tlbr()
        cv2.rectangle(frame, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), (0, 255, 0), 2)
        cv2.putText(frame, str(track.track_id), (int(bbox[0]), int(bbox[1]) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
    for det in detections:
        bbox = det.to_tlbr()
        cv2.rectangle(frame, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), (0, 0, 255), 2)
    
    # 输出视频帧
    output.write(frame)
    
    # 暂停一段时间，用于视频播放
    time.sleep(1 / fps)

# 释放视频对象和输出对象
video.release()
output.release()