基于OpenCV目标跟踪实现人员计数器

class TrackableObject:
    def __init__(self, objectID, centroid):
        # store the object ID, then initialize a list of centroids
        # using the current centroid
        self.objectID = objectID
        self.centroids = [centroid]
        # initialize a boolean used to indicate if the object has
        # already been counted or not
        self.counted = False

# import the necessary packages
from pyimagesearch.centroidtracker import CentroidTracker
from pyimagesearch.trackableobject import TrackableObject
from imutils.video import VideoStream
from imutils.video import FPS
import numpy as np
import argparse
import imutils
import time
import dlib
import cv2

# construct the argument parse and parse the arguments
ap = argparse.ArgumentParser()
ap.add_argument("-p", "--prototxt", required=True,
    help="path to Caffe 'deploy' prototxt file")
ap.add_argument("-m", "--model", required=True,
    help="path to Caffe pre-trained model")
ap.add_argument("-i", "--input", type=str,
    help="path to optional input video file")
ap.add_argument("-o", "--output", type=str,
    help="path to optional output video file")
ap.add_argument("-c", "--confidence", type=float, default=0.4,
    help="minimum probability to filter weak detections")
ap.add_argument("-s", "--skip-frames", type=int, default=30,
    help="# of skip frames between detections")
args = vars(ap.parse_args())

# initialize the list of class labels MobileNet SSD was trained to detect
CLASSES = ["background", "aeroplane", "bicycle", "bird", "boat",
    "bottle", "bus", "car", "cat", "chair", "cow", "diningtable",
    "dog", "horse", "motorbike", "person", "pottedplant", "sheep",
    "sofa", "train", "tvmonitor"]
# load our serialized model from disk
print("[INFO] loading model...")
net = cv2.dnn.readNetFromCaffe(args["prototxt"], args["model"])

# if a video path was not supplied, grab a reference to the webcam
if not args.get("input", False):
    print("[INFO] starting video stream...")
    vs = VideoStream(src=0).start()
    time.sleep(2.0)
# otherwise, grab a reference to the video file
else:
    print("[INFO] opening video file...")
    vs = cv2.VideoCapture(args["input"])

# initialize the video writer (we'll instantiate later if need be)
writer = None
# initialize the frame dimensions (we'll set them as soon as we read
# the first frame from the video)
W = None
H = None
# instantiate our centroid tracker, then initialize a list to store
# each of our dlib correlation trackers, followed by a dictionary to
# map each unique object ID to a TrackableObject
ct = CentroidTracker(maxDisappeared=40, maxDistance=50)
trackers = []
trackableObjects = {}
# initialize the total number of frames processed thus far, along
# with the total number of objects that have moved either up or down
totalFrames = 0
totalDown = 0
totalUp = 0
# start the frames per second throughput estimator
fps = FPS().start()

# loop over frames from the video stream
while True:
    # grab the next frame and handle if we are reading from either
    # VideoCapture or VideoStream
    frame = vs.read()
    frame = frame[1] if args.get("input", False) else frame
    # if we are viewing a video and we did not grab a frame then we
    # have reached the end of the video
    if args["input"] is not None and frame is None:
        break
    # resize the frame to have a maximum width of 500 pixels (the
    # less data we have, the faster we can process it), then convert
    # the frame from BGR to RGB for dlib
    frame = imutils.resize(frame, width=500)
    rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    # if the frame dimensions are empty, set them
    if W is None or H is None:
        (H, W) = frame.shape[:2]
    # if we are supposed to be writing a video to disk, initialize
    # the writer
    if args["output"] is not None and writer is None:
        fourcc = cv2.VideoWriter_fourcc(*"MJPG")
        writer = cv2.VideoWriter(args["output"], fourcc, 30,
            (W, H), True)

    # initialize the current status along with our list of bounding
    # box rectangles returned by either (1) our object detector or
    # (2) the correlation trackers
    status = "Waiting"
    rects = []
    # check to see if we should run a more computationally expensive
    # object detection method to aid our tracker
    if totalFrames % args["skip_frames"] == 0:
        # set the status and initialize our new set of object trackers
        status = "Detecting"
        trackers = []
        # convert the frame to a blob and pass the blob through the
        # network and obtain the detections
        blob = cv2.dnn.blobFromImage(frame, 0.007843, (W, H), 127.5)
        net.setInput(blob)
        detections = net.forward()

        # loop over the detections
        for i in np.arange(0, detections.shape[2]):
            # extract the confidence (i.e., probability) associated
            # with the prediction
            confidence = detections[0, 0, i, 2]
            # filter out weak detections by requiring a minimum
            # confidence
            if confidence > args["confidence"]:
                # extract the index of the class label from the
                # detections list
                idx = int(detections[0, 0, i, 1])
                # if the class label is not a person, ignore it
                if CLASSES[idx] != "person":
                    continue

                # compute the (x, y)-coordinates of the bounding box
                # for the object
                box = detections[0, 0, i, 3:7] * np.array([W, H, W, H])
                (startX, startY, endX, endY) = box.astype("int")
                # construct a dlib rectangle object from the bounding
                # box coordinates and then start the dlib correlation
                # tracker
                tracker = dlib.correlation_tracker()
                rect = dlib.rectangle(startX, startY, endX, endY)
                tracker.start_track(rgb, rect)
                # add the tracker to our list of trackers so we can
                # utilize it during skip frames
                trackers.append(tracker)

    # otherwise, we should utilize our object *trackers* rather than
    # object *detectors* to obtain a higher frame processing throughput
    else:
        # loop over the trackers
        for tracker in trackers:
            # set the status of our system to be 'tracking' rather
            # than 'waiting' or 'detecting'
            status = "Tracking"
            # update the tracker and grab the updated position
            tracker.update(rgb)
            pos = tracker.get_position()
            # unpack the position object
            startX = int(pos.left())
            startY = int(pos.top())
            endX = int(pos.right())
            endY = int(pos.bottom())
            # add the bounding box coordinates to the rectangles list
            rects.append((startX, startY, endX, endY))

    # draw a horizontal line in the center of the frame -- once an
    # object crosses this line we will determine whether they were
    # moving 'up' or 'down'
    cv2.line(frame, (0, H // 2), (W, H // 2), (0, 255, 255), 2)
    # use the centroid tracker to associate the (1) old object
    # centroids with (2) the newly computed object centroids
    objects = ct.update(rects)

    # loop over the tracked objects
    for (objectID, centroid) in objects.items():
        # check to see if a trackable object exists for the current
        # object ID
        to = trackableObjects.get(objectID, None)
        # if there is no existing trackable object, create one
        if to is None:
            to = TrackableObject(objectID, centroid)
        # otherwise, there is a trackable object so we can utilize it
        # to determine direction
        else:
            # the difference between the y-coordinate of the *current*
            # centroid and the mean of *previous* centroids will tell
            # us in which direction the object is moving (negative for
            # 'up' and positive for 'down')
            y = [c[1] for c in to.centroids]
            direction = centroid[1] - np.mean(y)
            to.centroids.append(centroid)
            # check to see if the object has been counted or not
            if not to.counted:
                # if the direction is negative (indicating the object
                # is moving up) AND the centroid is above the center
                # line, count the object
                if direction < 0 and centroid[1] < H // 2:
                    totalUp += 1
                    to.counted = True
                # if the direction is positive (indicating the object
                # is moving down) AND the centroid is below the
                # center line, count the object
                elif direction > 0 and centroid[1] > H // 2:
                    totalDown += 1
                    to.counted = True
        # store the trackable object in our dictionary
        trackableObjects[objectID] = to

        # draw both the ID of the object and the centroid of the
        # object on the output frame
        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)
    # construct a tuple of information we will be displaying on the
    # frame
    info = [
        ("Up", totalUp),
        ("Down", totalDown),
        ("Status", status),
    ]
    # loop over the info tuples and draw them on our frame
    for (i, (k, v)) in enumerate(info):
        text = "{}: {}".format(k, v)
        cv2.putText(frame, text, (10, H - ((i * 20) + 20)),
            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)

    # check to see if we should write the frame to disk
    if writer is not None:
        writer.write(frame)
    # show the output frame
    cv2.imshow("Frame", frame)
    key = cv2.waitKey(1) & 0xFF
    # if the `q` key was pressed, break from the loop
    if key == ord("q"):
        break
    # increment the total number of frames processed thus far and
    # then update the FPS counter
    totalFrames += 1
    fps.update()

# stop the timer and display FPS information
fps.stop()
print("[INFO] elapsed time: {:.2f}".format(fps.elapsed()))
print("[INFO] approx. FPS: {:.2f}".format(fps.fps()))
# check to see if we need to release the video writer pointer
if writer is not None:
    writer.release()
# if we are not using a video file, stop the camera video stream
if not args.get("input", False):
    vs.stop()
# otherwise, release the video file pointer
else:
    vs.release()
# close any open windows
cv2.destroyAllWindows()

from pyimagesearch.centroidtracker import CentroidTracker
from pyimagesearch.trackableobject import TrackableObject
from imutils.video import VideoStream
from imutils.video import FPS
import numpy as np
import argparse
import imutils
import time
import dlib
import cv2


# 构造参数解析并解析参数
ap = argparse.ArgumentParser()
ap.add_argument("-p", "--prototxt", required=True,
	help="path to Caffe 'deploy' prototxt file")
ap.add_argument("-m", "--model", required=True,
	help="path to Caffe pre-trained model")
ap.add_argument("-i", "--input", type=str,
	help="path to optional input video file")
ap.add_argument("-o", "--output", type=str,
	help="path to optional output video file")
ap.add_argument("-c", "--confidence", type=float, default=0.4,
	help="minimum probability to filter weak detections")
ap.add_argument("-s", "--skip-frames", type=int, default=30,
	help="# of skip frames between detections")
args = vars(ap.parse_args())

# 初始化类标签列表
CLASSES = ["background", "aeroplane", "bicycle", "bird", "boat",
	"bottle", "bus", "car", "cat", "chair", "cow", "diningtable",
	"dog", "horse", "motorbike", "person", "pottedplant", "sheep",
	"sofa", "train", "tvmonitor"]
# 从磁盘加载我们的序列化模型
print("[INFO] loading model...")
net = cv2.dnn.readNetFromCaffe(args["prototxt"], args["model"])

# 如果未提供视频路径，请获取网络摄像头的引用
if not args.get("input", False):
	print("[INFO] starting video stream...")
	vs = VideoStream(src=0).start()
	time.sleep(2.0)
# 否则，获取对视频文件的引用
else:
	print("[INFO] opening video file...")
	vs = cv2.VideoCapture(args["input"])

# 初始化视频写入器(如果需要，我们稍后将进行实例化)
writer = None
# 初始化帧尺寸（我们将在从视频中读取第一帧后立即设置它们）
W = None
H = None
# 实例化我们的质心跟踪器，然后初始化一个列表来存储每个dlib相关跟踪器，
# 然后是一个字典来将每个唯一的对象ID映射到TrackableObject
ct = CentroidTracker(maxDisappeared=40, maxDistance=50)
trackers = []
trackableObjects = {}
# 初始化到目前为止处理的帧总数，以及向上或向下移动的对象总数
totalFrames = 0
totalDown = 0
totalUp = 0
# 启动FPS评估器
fps = FPS().start()

# 循环视频流中的帧
while True:
	# 如果我们正在从 VideoCapture 或 VideoStream 读取数据，则抓取下一帧并处理
	frame = vs.read()
	frame = frame[1] if args.get("input", False) else frame
	# 如果我们正在观看视频并且我们没有抓取帧，那么我们已经到了视频的结尾
	if args["input"] is not None and frame is None:
		break
	# 调整帧的最大宽度为 500 像素（我们拥有的数据越少，我们处理它的速度就越快），
	# 然后将帧从 BGR 转换为 RGB 用于 dlib
	frame = imutils.resize(frame, width=500)
	rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
	# 如果帧尺寸为空，则设置它们
	if W is None or H is None:
		(H, W) = frame.shape[:2]
	# 如果我们应该将视频写入磁盘，请初始化写入器
	if args["output"] is not None and writer is None:
		fourcc = cv2.VideoWriter_fourcc(*"MJPG")
		writer = cv2.VideoWriter(args["output"], fourcc, 30,
			(W, H), True)


	# 初始化当前状态以及由（1）我们的对象检测器或（2）相关跟踪器返回的边界框矩形列表
	status = "Waiting"
	rects = []
	# 检查我们是否应该运行计算量更大的目标检测方法来帮助我们的跟踪器
	if totalFrames % args["skip_frames"] == 0:
		# 设置状态并初始化我们的新对象跟踪器集
		status = "Detecting"
		trackers = []
		# 将帧转换为 blob 并通过网络传递 blob 并获得检测结果
		blob = cv2.dnn.blobFromImage(frame, 0.007843, (W, H), 127.5)
		net.setInput(blob)
		detections = net.forward()

		# 循环检测结果
		for i in np.arange(0, detections.shape[2]):
			# 提取与预测相关的置信度(即概率)
			confidence = detections[0, 0, i, 2]
			# 通过要求最小置信度过滤掉弱检测
			if confidence > args["confidence"]:
				# 从检测列表中提取类标签的索引
				idx = int(detections[0, 0, i, 1])
				# 如果类标签不是人，则忽略它
				if CLASSES[idx] != "person":
					continue
				# 计算对象边界框的 (x, y) 坐标
				box = detections[0, 0, i, 3:7] * np.array([W, H, W, H])
				(startX, startY, endX, endY) = box.astype("int")
				# 利用边界框坐标构造一个 dlib 矩形对象，然后启动 dlib 相关跟踪器
				tracker = dlib.correlation_tracker()
				rect = dlib.rectangle(startX, startY, endX, endY)
				tracker.start_track(rgb, rect)
				# 将跟踪器添加到我们的跟踪器列表中，以便我们可以在跳帧期间使用它
				trackers.append(tracker)

	# 否则，我们应该利用目标跟踪器而不是目标检测器来获得更高的FPS
	else:
		# 遍历跟踪器
		for tracker in trackers:
			# 将系统的状态设置为“跟踪”而不是“等待”或“检测”
			status = "Tracking"
			# 更新跟踪器并获取更新的位置
			tracker.update(rgb)
			pos = tracker.get_position()
			# 解包位置对象
			startX = int(pos.left())
			startY = int(pos.top())
			endX = int(pos.right())
			endY = int(pos.bottom())
			# 将边界框坐标添加到矩形列表
			rects.append((startX, startY, endX, endY))

	# 在帧中心画一条水平线——一旦一个物体穿过这条线，我们将确定他们是在“向上”还是“向下”移动。
	cv2.line(frame, (0, H // 2), (W, H // 2), (0, 255, 255), 2)
	# 使用质心跟踪器将 (1) 旧对象质心与 (2) 新计算的对象质心相关联
	objects = ct.update(rects)

	# 循环遍历被跟踪的对象
	for (objectID, centroid) in objects.items():
		# 检查当前对象 ID 是否存在可跟踪对象
		to = trackableObjects.get(objectID, None)
		# 如果没有现有的可跟踪对象，则创建一个
		if to is None:
			to = TrackableObject(objectID, centroid)
		# 否则，有一个可追踪的物体，所以我们可以利用它来确定方向
		else:
			# *当前*质心的 y 坐标与 *previous* 质心的平均值之间的差异
			# 将告诉我们物体在哪个方向移动（“向上”为负，“向下”为正）
			y = [c[1] for c in to.centroids]
			direction = centroid[1] - np.mean(y)
			to.centroids.append(centroid)
			# 检查对象是否已被计数
			if not to.counted:
				# 如果方向为负(表示物体向上移动)且质心在中线以上，则计算物体
				if direction < 0 and centroid[1] < H // 2:
					totalUp += 1
					to.counted = True
				# 如果方向为正（表示物体正在向下移动）并且质心低于中心线，则计算物体
				elif direction > 0 and centroid[1] > H // 2:
					totalDown += 1
					to.counted = True
		# 将可跟踪对象存储在我们的字典中
		trackableObjects[objectID] = to

		# 在输出帧上绘制对象的 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)
	# 构建我们将在帧上显示的信息元组
	info = [
		("Up", totalUp),
		("Down", totalDown),
		("Status", status),
	]
	# 遍历信息元组并将它们绘制在我们的帧上
	for (i, (k, v)) in enumerate(info):
		text = "{}: {}".format(k, v)
		cv2.putText(frame, text, (10, H - ((i * 20) + 20)),
			cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)



	# 检查我们是否应该将帧写入磁盘
	if writer is not None:
		writer.write(frame)
	# 显示输出帧
	cv2.imshow("Frame", frame)
	key = cv2.waitKey(1) & 0xFF
	# 如果' q '键被按下，中断循环
	if key == ord("q"):
		break
	# 增加到目前为止处理的帧总数，然后更新 FPS 计数器
	totalFrames += 1
	fps.update()

# 停止定时器，显示FPS信息
fps.stop()
print("[INFO] elapsed time: {:.2f}".format(fps.elapsed()))
print("[INFO] approx. FPS: {:.2f}".format(fps.fps()))
# 检查我们是否需要释放视频写入器指针
if writer is not None:
	writer.release()
# 如果我们不使用视频文件，请停止摄像头视频流
if not args.get("input", False):
	vs.stop()
# 否则，释放视频文件指针
else:
	vs.release()
# 关闭所有打开的窗口
cv2.destroyAllWindows()

from scipy.spatial import distance as dist
from collections import OrderedDict
import numpy as np

class CentroidTracker:
	def __init__(self, maxDisappeared=50, maxDistance=50):
		# 初始化下一个唯一的对象ID，并使用两个有序字典来跟踪给定对象ID到其质心的映射，
		# 以及它被标记为“消失”的连续帧数
		self.nextObjectID = 0
		self.objects = OrderedDict()
		self.disappeared = OrderedDict()

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

		# 存储质心之间的最大距离以关联对象——如果距离大于这个最大距离，我们开始将对象标记为“消失”
		self.maxDistance = maxDistance

	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):
			# 使用边界框坐标推导出质心
			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

				# 如果质心之间的距离大于最大距离，则不要将两个质心关联到同一个对象
				if D[row, col] > self.maxDistance:
					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]:
				# 循环未使用的行索引
				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 self.objects

class TrackableObject:
	def __init__(self, objectID, centroid):
		# 存储对象 ID，然后使用当前质心初始化质心列表
		self.objectID = objectID
		self.centroids = [centroid]

		# 初始化一个布尔值，用于指示对象是否已被计数
		self.counted = False

python people_counter.py --prototxt mobilenet_ssd/MobileNetSSD_deploy.prototxt \
    --model mobilenet_ssd/MobileNetSSD_deploy.caffemodel \
    --input videos/example_01.mp4 --output output/output_01.avi