opencv4.4使用yolov3和yolov4模型检测目标

yolov3：

# -*- coding: utf-8 -*

import numpy as np

import cv2 as cv

import os

import time

yolo_dir = '/home/ubuntu/model/yolo4-tiny'  # YOLO文件路径

weightsPath = os.path.join(yolo_dir, 'yolov3-test.weights')  # 权重文件

configPath = os.path.join(yolo_dir, 'yolov3-test.cfg')  # 配置文件

labelsPath = os.path.join(yolo_dir, 'test.names')  # label名称

imgPath = os.path.join(yolo_dir, 'test.jpg')  # 测试图像

CONFIDENCE = 0.25  # 过滤弱检测的最小概率

THRESHOLD = 0.2  # 非最大值抑制阈值

#批量图片检测和保存位置

test_dir = '/home/ubuntu/model/yolo4-tiny/imgs'

save_dir = '/home/ubuntu/model/yolo4-tiny/imgsResultV3'

# 加载网络、配置权重

net = cv.dnn.readNetFromDarknet(configPath, weightsPath)  # #  利用下载的文件

#批量加载图像

pics = os.listdir(test_dir)

start = time.time()

minTime = 10;

maxTime = 0;

for im in pics:

# 加载图片、转为blob格式、送入网络输入层

    s = time.time()

    imgPath = os.path.join(test_dir,im)

    img = cv.imread(imgPath)

    blobImg = cv.dnn.blobFromImage(img, 1.0/255.0, (416, 416), None, True, False)   # # net需要的输入是blob格式的，用blobFromImage这个函数来转格式

    net.setInput(blobImg)  # # 调用setInput函数将图片送入输入层

    # 获取网络输出层信息（所有输出层的名字），设定并前向传播

    outInfo = net.getUnconnectedOutLayersNames()  # # 前面的yolov3架构也讲了，yolo在每个scale都有输出，outInfo是每个scale的名字信息，供net.forward使用

    layerOutputs = net.forward(outInfo)  # 得到各个输出层的、各个检测框等信息，是二维结构。

    # 拿到图片尺寸

    (H, W) = img.shape[:2]

    # 过滤layerOutputs

    # layerOutputs的第1维的元素内容: [center_x, center_y, width, height, objectness, N-class score data]

    # 过滤后的结果放入：

    boxes = [] # 所有边界框（各层结果放一起）

    confidences = [] # 所有置信度

    classIDs = [] # 所有分类ID

    # # 1）过滤掉置信度低的框框

    for out in layerOutputs:  # 各个输出层

        for detection in out:  # 各个框框

            # 拿到置信度

            scores = detection[5:]  # 各个类别的置信度

            classID = np.argmax(scores)  # 最高置信度的id即为分类id

            confidence = scores[classID]  # 拿到置信度

            # 根据置信度筛查

            if confidence > CONFIDENCE:

                box = detection[0:4] * np.array([W, H, W, H])  # 将边界框放会图片尺寸

                (centerX, centerY, width, height) = box.astype("int")

                x = int(centerX - (width / 2))

                y = int(centerY - (height / 2))

                boxes.append([x, y, int(width), int(height)])

                confidences.append(float(confidence))

                classIDs.append(classID)

    # # 2）应用非最大值抑制(non-maxima suppression，nms)进一步筛掉

    idxs = cv.dnn.NMSBoxes(boxes, confidences, CONFIDENCE, THRESHOLD) # boxes中，保留的box的索引index存入idxs

    # 得到labels列表

    with open(labelsPath, 'rt') as f:

        labels = f.read().rstrip('\n').split('\n')

    # 应用检测结果

    np.random.seed(42)

    COLORS = np.random.randint(0, 255, size=(len(labels), 3), dtype="uint8")  # 框框显示颜色，每一类有不同的颜色，每种颜色都是由RGB三个值组成的，所以size为(len(labels), 3)

    if len(idxs) > 0:

        for i in idxs.flatten():  # indxs是二维的，第0维是输出层，所以这里把它展平成1维

            (x, y) = (boxes[i][0], boxes[i][1])

            (w, h) = (boxes[i][2], boxes[i][3])

            color = [int(c) for c in COLORS[classIDs[i]]]

            cv.rectangle(img, (x, y), (x+w, y+h), color, 2)  # 线条粗细为2px

            text = "{}: {:.4f}".format(labels[classIDs[i]], confidences[i])

            cv.putText(img, text, (x, y-5), cv.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)  # cv.FONT_HERSHEY_SIMPLEX字体风格、0.5字体大小、粗细2px

    print(imgPath+"检测耗时： {:.6f} seconds".format(time.time() - s))  # # 可以打印下信息

    times = time.time() - s

    if(minTime>times):

        minTime = times

    if(maxTime<times):

        maxTime = times

    cv.imwrite(os.path.join(save_dir,im),img)

print("yolov3检测100张图总共耗时：%.6f秒" % (time.time() - start))

print("耗时最长为：%.6f秒" % (maxTime))

print("耗时最短为：%.6f秒" % (minTime))

 

 

yolov4:

 

# -*- coding: utf-8 -*

import cv2 as cv

import time

import os

 

#置信度和nms设置

confThreshold = 0.25

nmsThreshold = 0.2

 

class_names = []#初始化一个列表以存储类名

COLORS = [(0, 255, 255),(0,0,255), (255, 255, 0),(0, 255, 0), (255, 0, 0)]#颜色

color=(0,255,0)

color1=(0,0,255)

c=(0, 0, 0)

#模型参数

weights="/home/ubuntu/model/yolo4-tiny/yolov4-tiny.weights"

cfg="/home/ubuntu/model/yolo4-tiny/yolov4-tiny.cfg"

m="/home/ubuntu/model/yolo4-tiny/coco.names"

#批量图片检测和保存位置

test_dir = '/home/ubuntu/model/yolo4-tiny/imgs'

save_dir = '/home/ubuntu/model/yolo4-tiny/imgsResultV4'

# 网络设置

net = cv.dnn_DetectionModel(cfg, weights)

#GPU运行

net.setPreferableBackend(cv.dnn.DNN_BACKEND_CUDA)

net.setPreferableTarget(cv.dnn.DNN_TARGET_CUDA)

net.setInputSize(416, 416)

net.setInputScale(1.0 / 255)

net.setInputSwapRB(True)

with open(m, "r") as f:

    class_names = [cname.strip() for cname in f.readlines()]

 

#批量加载图像路径

pics = os.listdir(test_dir)

start = time.time()

minTime = 10;

maxTime = 0;

inImgTime=0;

testImgTime=0;

for im in pics:

    s = time.time()

    #加载图片

    img = os.path.join(test_dir,im)

    image = cv.imread(img)

    

    sTest = time.time()

    inImg = sTest-s

    inImgTime +=inImg

    #模型检测

    classes, confidences, boxes = net.detect(image, confThreshold, nmsThreshold)

    sTestEnd = time.time()

    testImg = sTestEnd-sTest

    testImgTime+= testImg

    count=0

    #循环画框

    for (classid, score, box) in zip(classes, confidences, boxes):

        label = "%s(%2.0f%%)" % (class_names[classid[0]], score*100)#标签置信度

        labelSize, baseLine = cv.getTextSize(label, cv.FONT_HERSHEY_SIMPLEX, 0.5, 1)

        left, top, width, height = box

        cv.rectangle(image, box, color1, 1)

        cv.rectangle(image, (left-1, top - labelSize[1]-5), (left + labelSize[0], top), color, cv.FILLED)

        cv.putText(image, label, (left, top-5), cv.FONT_HERSHEY_SIMPLEX, 0.5, c,2)

        count+=1

    times = time.time() - s

    if(minTime>times):

        minTime = times

    if(maxTime<times):

        maxTime = times

    print(img+"加载图片耗时：%.6f秒" % inImg)

    print(img+"检测图片耗时：%.6f秒" % testImg)

    print(img+"检测耗时：%.3f秒" % times)

    cv.putText(image,"Num: %s" % str(count),(25,50),cv.FONT_HERSHEY_SIMPLEX,1,[255,255,0],2)

    cv.imwrite(os.path.join(save_dir,im),image)#用原始图像名im保存转换后的图片

print("yolov4检测6张图总共耗时：%.3f秒" % (time.time() - start))

print("耗时最长为：%.3f秒" % (maxTime))

print("耗时最短为：%.3f秒" % (minTime))

print("加载图片总共耗时：%.6f秒" % (inImgTime))

print("检测图片总共耗时：%.6f秒" % (testImgTime))