|
|
import tkinter.filedialog as tkinter
|
|
|
import numpy as np
|
|
|
import argparse
|
|
|
import cv2
|
|
|
from utils import FPS
|
|
|
|
|
|
|
|
|
# 图像增强
|
|
|
def enhance_image(frame):
|
|
|
# 直方图均衡化
|
|
|
frame_yuv = cv2.cvtColor(frame, cv2.COLOR_BGR2YUV)
|
|
|
frame_yuv[:, :, 0] = cv2.equalizeHist(frame_yuv[:, :, 0])
|
|
|
frame = cv2.cvtColor(frame_yuv, cv2.COLOR_YUV2BGR)
|
|
|
|
|
|
# 锐化
|
|
|
sharpen_kernel = np.array([[-1, -1, -1],
|
|
|
[-1, 9, -1],
|
|
|
[-1, -1, -1]])
|
|
|
sharpened_frame = cv2.filter2D(frame, -1, sharpen_kernel)
|
|
|
|
|
|
# 对比度增强
|
|
|
alpha = 1.5 # 控制对比度（1.0表示不变）
|
|
|
enhanced_frame = cv2.convertScaleAbs(sharpened_frame, alpha=alpha, beta=0)
|
|
|
|
|
|
# 亮度调整
|
|
|
beta = 30 # 控制亮度调整量
|
|
|
enhanced_frame = cv2.convertScaleAbs(enhanced_frame, alpha=1.0, beta=beta)
|
|
|
|
|
|
return enhanced_frame
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# 参数
|
|
|
ap = argparse.ArgumentParser()
|
|
|
ap.add_argument("-p", "--prototxt", default="mobilenet_ssd/MobileNetSSD_deploy.prototxt",
|
|
|
help="path to Caffe 'deploy' prototxt file")
|
|
|
ap.add_argument("-m", "--model", default="mobilenet_ssd/MobileNetSSD_deploy.caffemodel",
|
|
|
help="path to Caffe pre-trained model")
|
|
|
ap.add_argument("-v", "--video", default=None,
|
|
|
help="path to 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.3,
|
|
|
help="minimum probability to filter weak detections")
|
|
|
args = vars(ap.parse_args())
|
|
|
|
|
|
# SSD标签
|
|
|
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 args["video"] is None:
|
|
|
video_path = tkinter.askopenfilename(filetypes=[("视频文件", "*.mp4")])
|
|
|
print("[INFO] starting video stream...")
|
|
|
vs = cv2.VideoCapture(video_path)
|
|
|
else:
|
|
|
print("[INFO] starting video stream...")
|
|
|
vs = cv2.VideoCapture(args["video"])
|
|
|
writer = None
|
|
|
|
|
|
# 初始化目标追踪器
|
|
|
trackers = []
|
|
|
labels = []
|
|
|
fps = FPS().start()
|
|
|
|
|
|
while True:
|
|
|
# 读取一帧
|
|
|
(grabbed, frame) = vs.read()
|
|
|
|
|
|
# 是否是最后了
|
|
|
if frame is None:
|
|
|
break
|
|
|
|
|
|
# 图像增强
|
|
|
# frame = enhance_image(frame)
|
|
|
|
|
|
# 预处理操作
|
|
|
(h, w) = frame.shape[:2]
|
|
|
width = 600
|
|
|
r = width / float(w)
|
|
|
dim = (width, int(h * r))
|
|
|
frame = cv2.resize(frame, dim, interpolation=cv2.INTER_AREA)
|
|
|
rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
|
|
|
|
|
# 如果要将结果保存的话
|
|
|
if args["output"] is not None and writer is None:
|
|
|
fourcc = cv2.VideoWriter_fourcc(*"MJPG")
|
|
|
writer = cv2.VideoWriter(args["output"], fourcc, 30, (frame.shape[1], frame.shape[0]), True)
|
|
|
|
|
|
# 先检测 再追踪
|
|
|
if len(trackers) == 0:
|
|
|
# 获取blob数据
|
|
|
(h, w) = frame.shape[:2]
|
|
|
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])
|
|
|
label = CLASSES[idx]
|
|
|
|
|
|
# 只保留人的
|
|
|
if CLASSES[idx] != "person":
|
|
|
continue
|
|
|
|
|
|
# 得到BBOX
|
|
|
box = detections[0, 0, i, 3:7] * np.array([w, h, w, h])
|
|
|
(startX, startY, endX, endY) = box.astype("int")
|
|
|
|
|
|
# 使用CSRT目标追踪器
|
|
|
tracker = cv2.TrackerCSRT_create()
|
|
|
tracker.init(frame, (startX, startY, endX - startX, endY - startY))
|
|
|
|
|
|
# 保存结果
|
|
|
labels.append(label)
|
|
|
trackers.append(tracker)
|
|
|
|
|
|
# 绘图
|
|
|
cv2.rectangle(frame, (startX, startY), (endX, endY), (0, 255, 0), 2)
|
|
|
cv2.putText(frame, label, (startX, startY - 15), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 255, 0), 2)
|
|
|
|
|
|
# 如果已经有了框，就可以直接追踪了
|
|
|
else:
|
|
|
# 每一个追踪器都要进行更新
|
|
|
for (tracker, label) in zip(trackers, labels):
|
|
|
success, box = tracker.update(frame)
|
|
|
if success:
|
|
|
(startX, startY, w, h) = [int(v) for v in box]
|
|
|
endX = startX + w
|
|
|
endY = startY + h
|
|
|
|
|
|
# 画出来
|
|
|
cv2.rectangle(frame, (startX, startY), (endX, endY), (0, 255, 0), 2)
|
|
|
cv2.putText(frame, label, (startX, startY - 15), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 255, 0), 2)
|
|
|
|
|
|
# 也可以把结果保存下来
|
|
|
if writer is not None:
|
|
|
writer.write(frame)
|
|
|
|
|
|
# 显示
|
|
|
cv2.imshow("Frame", frame)
|
|
|
key = cv2.waitKey(1) & 0xFF
|
|
|
|
|
|
# 退出
|
|
|
if key == 27:
|
|
|
break
|
|
|
|
|
|
# 计算FPS
|
|
|
fps.update()
|
|
|
|
|
|
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()
|
|
|
|
|
|
cv2.destroyAllWindows()
|
|
|
vs.release()
|
