Python OpenCV 实现yolo目标检测

YOLO 是现今非常流行的目标检测框架。源代码是用C 写的。这里我们利用opencv 调用训练好的yolo 模型来实现一个demo。

#首先导入相应的模块
import cv2 as cv
import argparse
import sys
import numpy as np
import os.path

# 初始化变量
confThreshold = 0.5  #置信度阈值
nmsThreshold = 0.4  # 非极大值一直阈值
inpWidth = 416   # 网络输入图像的宽
inpHeight = 416 # 网络输入图像的高
# 构建参数解析器
parser = argparse.ArgumentParser(description='Object Detection using YOLO in OPENCV')
parser.add_argument('--image', help='Path to image file.')
parser.add_argument('--video', help='Path to video file.')
args = parser.parse_args()
# 载入类别名称

classesFile = "coco.names";
classes = None
with open(classesFile, 'rt') as f:
    classes = f.read().rstrip('\n').split('\n')

#定义 模型的配置和权重文件.
modelConfiguration = "yolov3.cfg";
modelWeights = "yolov3.weights";
net = cv.dnn.readNetFromDarknet(modelConfiguration, modelWeights)
net.setPreferableBackend(cv.dnn.DNN_BACKEND_OPENCV)
net.setPreferableTarget(cv.dnn.DNN_TARGET_CPU)

# 获得输出层的名字
def getOutputsNames(net):
    # 获得网络中所有层的名字
    layersNames = net.getLayerNames()
    # Get the names of the output layers, i.e. the layers with unconnected outputs
    return [layersNames[i[0] - 1] for i in net.getUnconnectedOutLayers()]

# 画预测的包围框
def drawPred(classId, conf, left, top, right, bottom):
    cv.rectangle(frame, (left, top), (right, bottom), (255, 178, 50),3)
     label = '%.2f'%conf
     #得到类别名称的标签和它的置信度得分
     if classes:
         assert(classId<len(classes))
         label = '%s:%s' %(classes[classId], label)
     # 在包围框上面显示标注
     labelsize,baseline = cv.getTextSize(label, cv.FONT_HERSHEY_SIMPLEX, 0.5, 1)
     top = max(top, labelSize[1])
	cv.rectangle(frame, (left, top - round(1.5*labelSize[1])), (left + round(1.5*labelSize[0]), top +baseLine), (255, 255, 255), cv.FILLED)
    cv.putText(frame, label, (left, top), cv.FONT_HERSHEY_SIMPLEX, 0.75, (0,0,0), 1)

# 使用nms 处理置信度得分低的框
def postprocess(frame, outs):
    frameHeight = frame.shape[0]
    frameWidth = frame.shape[1]

    classIds = []
    confidences = []
    boxes = []
    # Scan through all the bounding boxes output from the network and keep only the
    # ones with high confidence scores. Assign the box's class label as the class with the highest score.
    classIds = []
    confidences = []
    boxes = []
    for out in outs:
        for detection in out:
            scores = detection[5:]
            classId = np.argmax(scores)
            confidence = scores[classId]
            if confidence > confThreshold:
                center_x = int(detection[0] * frameWidth)
                center_y = int(detection[1] * frameHeight)
                width = int(detection[2] * frameWidth)
                height = int(detection[3] * frameHeight)
                left = int(center_x - width / 2)
                top = int(center_y - height / 2)
                classIds.append(classId)
                confidences.append(float(confidence))
                boxes.append([left, top, width, height])

    # 运行非极大值抑制来缓解具有低置信度得分的冗余的覆盖包围框
    indices = cv.dnn.NMSBoxes(boxes, confidences, confThreshold, nmsThreshold)
    for i in indices:
        i = i[0]
        box = boxes[i]
        left = box[0]
        top = box[1]
        width = box[2]
        height = box[3]
        drawPred(classIds[i], confidences[i], left, top, left + width, top + height)

# 处理输入
winName = 'object detection in OpenCV'
cv.namedWindow(winName, cv.WINDOW_NORMAL)

outputFile = "yolo_out_py.avi"
if (args.image):
    # 打开图像文件
    if not os.path.isfile(args.image):
        print("Input image file ", args.image, " doesn't exist")
        sys.exit(1)
    cap = cv.VideoCapture(args.image)
    outputFile = args.image[:-4]+'_yolo_out_py.jpg'
elif (args.video):
    # 打开视频文件
    if not os.path.isfile(args.video):
        print("Input video file ", args.video, " doesn't exist")
        sys.exit(1)
    cap = cv.VideoCapture(args.video)
    outputFile = args.video[:-4]+'_yolo_out_py.avi'
else:
    # 网络摄像头输入
    cap = cv.VideoCapture(0)
# 获取Get the video writer initialized to save the output video
if (not args.image):
    vid_writer = cv.VideoWriter(outputFile, cv.VideoWriter_fourcc('M','J','P','G'), 30, (round(cap.get(cv.CAP_PROP_FRAME_WIDTH)),round(cap.get(cv.CAP_PROP_FRAME_HEIGHT))))

while cv.waitKey(1) < 0:
    # 从视频中获取帧
    hasFrame, frame = cap.read()    
    # 如果到达了视频的末尾，则停止程序
    if not hasFrame:
        print("Done processing !!!")
        print("Output file is stored as ", outputFile)
        cv.waitKey(3000)
        break
    # 从一帧创建一个4D 的结构
    blob = cv.dnn.blobFromImage(frame, scalefactor=1/255, size=(inpWidth, inpHeight), mean=[0,0,0], swapRB=1, crop=False)
    # 设置网络的输入
    net.setInput(blob)
    # 运行网络的前向传播
    outs = net.forward(getOutputsNames(net))

    # 去除掉具有低置信度得分的包围框
    postprocess(frame, outs)
    # 放置有效率的信息。 函数 getPerfProfile  返回推断的整体的时间 和每一层的时间
    t, _ = net.getPerfProfile()
    label = 'Inference time: %.2f ms' % (t * 1000.0 / cv.getTickFrequency())
    cv.putText(frame, label, (0, 15), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255))

    # 将检测结果写入视频帧
    if (args.image):
        cv.imwrite(outputFile, frame.astype(np.uint8));
    else:
        vid_writer.write(frame.astype(np.uint8))

    cv.imshow(winName, frame)

这样一个opencv yolo 物体检测的模块就写好了。