利用OpenCV实现质心跟踪算法

# import the necessary packages
from scipy.spatial import distance as dist
from collections import OrderedDict
import numpy as np

class CentroidTracker():
    def __init__(self, maxDisappeared=50):

        self.nextObjectID = 0
        self.objects = OrderedDict()
        self.disappeared = OrderedDict()
        # 存储给定对象被允许标记为“消失”的最大连续帧数，直到我们需要从跟踪中注销该对象
        self.maxDisappeared = maxDisappeared

def register(self, centroid):
        # 注册对象时，我们使用下一个可用的对象ID来存储质心
        self.objects[self.nextObjectID] = centroid
        self.disappeared[self.nextObjectID] = 0
        self.nextObjectID += 1

    def deregister(self, objectID):
        # 要注销注册对象ID，我们从两个字典中都删除了该对象ID
        del self.objects[objectID]
        del self.disappeared[objectID]

 def update(self, rects):
        # 检查输入边界框矩形的列表是否为空
        if len(rects) == 0:
            # 遍历任何现有的跟踪对象并将其标记为消失
            for objectID in list(self.disappeared.keys()):
                self.disappeared[objectID] += 1
                # 如果达到给定对象被标记为丢失的最大连续帧数，请取消注册
                if self.disappeared[objectID] > self.maxDisappeared:
                    self.deregister(objectID)
            # 由于没有质心或跟踪信息要更新，请尽早返回
            return self.objects

        # 初始化当前帧的输入质心数组
        inputCentroids = np.zeros((len(rects), 2), dtype="int")
        # 在边界框矩形上循环
        for (i, (startX, startY, endX, endY)) in enumerate(rects):
            # use the bounding box coordinates to derive the centroid
            cX = int((startX + endX) / 2.0)
            cY = int((startY + endY) / 2.0)
            inputCentroids[i] = (cX, cY)

        # 如果我们当前未跟踪任何对象，请输入输入质心并注册每个质心
        if len(self.objects) == 0:
            for i in range(0, len(inputCentroids)):
                self.register(inputCentroids[i])
        # 否则，当前正在跟踪对象，因此我们需要尝试将输入质心与现有对象质心进行匹配
        else:
            # 抓取一组对象ID和相应的质心
            objectIDs = list(self.objects.keys())
            objectCentroids = list(self.objects.values())
            # 分别计算每对对象质心和输入质心之间的距离-我们的目标是将输入质心与现有对象质心匹配
            D = dist.cdist(np.array(objectCentroids), inputCentroids)
            # 为了执行此匹配，我们必须（1）在每行中找到最小值，然后（2）根据行索引的最小值对行索引进行排序，以使具有最小值的行位于索引列表的* front *处
            rows = D.min(axis=1).argsort()
            # 接下来，我们在列上执行类似的过程，方法是在每一列中找到最小值，然后使用先前计算的行索引列表进行排序
            cols = D.argmin(axis=1)[rows]
            # 为了确定是否需要更新，注册或注销对象，我们需要跟踪已经检查过的行索引和列索引
            usedRows = set()
            usedCols = set()

            # 循环遍历（行，列）索引元组的组合
            for (row, col) in zip(rows, cols):
                # 如果我们之前已经检查过行或列的值，请忽略它
                if row in usedRows or col in usedCols:
                    continue
                # 否则，获取当前行的对象ID，设置其新的质心，然后重置消失的计数器
                objectID = objectIDs[row]
                self.objects[objectID] = inputCentroids[col]
                self.disappeared[objectID] = 0
                # 表示我们已经分别检查了行索引和列索引
                usedRows.add(row)
                usedCols.add(col)
            # 计算我们尚未检查的行和列索引
            unusedRows = set(range(0, D.shape[0])).difference(usedRows)
            unusedCols = set(range(0, D.shape[1])).difference(usedCols)
            # 如果对象质心的数量等于或大于输入质心的数量
            # 我们需要检查一下其中的某些对象是否已潜在消失
            if D.shape[0] >= D.shape[1]:
                # loop over the unused row indexes
                for row in unusedRows:
                    # 抓取相应行索引的对象ID并增加消失的计数器
                    objectID = objectIDs[row]
                    self.disappeared[objectID] += 1
                    # 检查是否已将该对象标记为“消失”的连续帧数以用于注销该对象的手令
                    if self.disappeared[objectID] > self.maxDisappeared:
                        self.deregister(objectID)
            # 否则，如果输入质心的数量大于现有对象质心的数量，我们需要将每个新的输入质心注册为可跟踪对象
            else:
                for col in unusedCols:
                    self.register(inputCentroids[col])

        # return the set of trackable objects
        return self.objects

from Model.centroidtracker import CentroidTracker
import numpy as np
import imutils
import time
import cv2

# 定义最低置信度
confidence_t=0.5
# 初始化质心跟踪器和框架尺寸
ct = CentroidTracker()
(H, W) = (None, None)
# 加载检测人脸的模型
print("[INFO] loading model...")
net = cv2.dnn.readNetFromCaffe("deploy.prototxt", "res10_300x300_ssd_iter_140000_fp16.caffemodel")
# 初始化视频流并允许相机传感器预热
print("[INFO] starting video stream...")
vs = cv2.VideoCapture('11.mp4')
time.sleep(2.0)
fps = 30    #保存视频的FPS，可以适当调整
size=(600,1066)#宽高，根据frame的宽和高确定。
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
videoWriter = cv2.VideoWriter('3.mp4',fourcc,fps,size)#最后一个是保存图片的尺寸

# 循环播放图像流中的帧
while True:
    # 从视频流中读取下一帧并调整其大小
    (grabbed, frame)  = vs.read()
    if not grabbed:
        break
    frame = imutils.resize(frame, width=600)
    print(frame.shape)
    # 如果帧中尺寸为“无”，则抓住它们
    if W is None or H is None:
        (H, W) = frame.shape[:2]
    # 从帧中构造一个Blob，将其通过网络，
    # 获取输出预测，并初始化边界框矩形的列表
    blob = cv2.dnn.blobFromImage(frame, 1.0, (W, H),
                                 (104.0, 177.0, 123.0))
    net.setInput(blob)
    detections = net.forward()
    rects = []
    # 循环检测
    for i in range(0, detections.shape[2]):
        # 通过确保预测的概率大于最小阈值来过滤掉弱检测
        if detections[0, 0, i, 2] >0.5:
            # 计算对象边界框的（x，y）坐标，然后更新边界框矩形列表
            box = detections[0, 0, i, 3:7] * np.array([W, H, W, H])
            rects.append(box.astype("int"))
            # 在对象周围画一个边界框，以便我们可视化它
            (startX, startY, endX, endY) = box.astype("int")
            cv2.rectangle(frame, (startX, startY), (endX, endY),
                          (0, 255, 0), 2)
    # 使用边界框矩形的计算集更新质心跟踪器
    objects = ct.update(rects)
    # 循环跟踪对象
    for (objectID, centroid) in objects.items():
        # 在输出帧上绘制对象的ID和对象的质心
        text = "ID {}".fORMat(objectID)
        cv2.putText(frame, text, (centroid[0] - 10, centroid[1] - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
        cv2.circle(frame, (centroid[0], centroid[1]), 4, (0, 255, 0), -1)
    # 显示输出画面
    cv2.imshow("Frame", frame)
    videoWriter.write(frame)
    key = cv2.waiTKEy(1) & 0xFF
    # 如果按下“ q”键，则退出循环
    if key == ord("q"):
        break
videoWriter.release()
cv2.destroyAllwindows()
vs.release()