opencv调用yolov3模型来进行图像检测

之前使用了opencv来调用ssd的模型来检测物体，今天学了一下用opencv调用yolov3的模型来检测物体，二者在预测图形的部分，代码流程差不多，反正就是加载模型然后预测输出，但是对于输出结果的处理，二者就有区别，闲话不多说，进入正题：

yolov3模型以及网络参数：链接：https://pan.baidu.com/s/1dce1q11ZMGwyIT3OLafJQw
提取码：bj6m

代码如下：

'''author:nike hu'''
import cv2
import torch

global objName
global imageW, imageH

def get_class(filename):
    global objName
    objName = []
    with open(filename) as f:
        class_name = f.readlines()
        for i in class_name:
            objName.append(i.strip())
    return objName


def test(image, model, model_layer, thred=0.5):
    global objName, imageW, imageH
    imageH = image.shape[0] # 输入照片的高
    imageW = image.shape[1] # 宽
    show_image = image # 这个用来显示图片
    model = cv2.dnn.readNetFromDarknet(model_layer, model)
    # 加载yolov3模型,第一个参数是网络的每一层的信息，第二个参数是训练好的模型，我们现在使用的是官方训练的模型
    image = cv2.dnn.blobFromImage(image, 0.007843, (416, 416), (0, 0, 0), True, False)
    '''
    函数cv2.dnn.blobFromImage(image[, scalefactor[, size[, mean[, swapRB[, crop[, ddepth]]]]]])
作用：对图像进行预处理，包括减均值，比例缩放，裁剪，交换通道等，返回一个4通道的blob(blob可以简单理解为一个N维的数组，用于神经网络的输入)
参数：image:输入图像（1、3或者4通道）
    可选参数
    scalefactor:图像各通道数值的缩放比例
    size:输出图像的空间尺寸,如size=(200,300)表示高h=300,宽w=200
    mean:用于各通道减去的值，以降低光照的影响(e.g. image为bgr3通道的图像，mean=[104.0, 177.0, 123.0],表示b通道的值-104，g-177,r-123)
    swapRB:交换RB通道，默认为False.(cv2.imread读取的是彩图是bgr通道)
    crop:图像裁剪,默认为False.当值为True时，先按比例缩放，然后从中心裁剪成size尺寸
    ddepth:输出的图像深度，可选CV_32F 或者 CV_8U.
    总之，这个函数就是用来处理图片的，将图片的维度转化为神经网络需要的维度。
    '''
    model.setInput(image) # 感觉这一步就是将数据放到神经网络的输入层
    outinfo = model.getUnconnectedOutLayersNames()
    # 这里会输出一个列表，['yolo_82', 'yolo_94', 'yolo_106']，82这些应该是对应的多少层输出的数据，分别对应13x13, 26x26, 52x52这三个尺度的特征图
    out = model.forward(outinfo)
    '''这里跟ssd有差别，ssd不加括号的参数就能检测，但是这里如果不加上面那个参数，最后输出维度(8112, 85)，但是8112！=（13*13 + 26*26 + 52*52）*3，
    而且这样预测，预测不出来啥玩意，神经网络向前推进,最后输出维度(8112, 85),其中的8112代表预测框，85分别对应x,y,w,h，置信度，以及各个类别对应的可能性，
    我们使用的是coco数据集，coco数据集有80个类别'''
    out = torch.cat((torch.from_numpy(out[0]), torch.from_numpy(out[1]), torch.from_numpy(out[2])), dim=0) # 将三个输出尺寸的维度进行合并
    out[:, :4] = change_to_conner(out[:, :4]) # 将中心点转化为左上角和右下角
    # thred_id = (out[:, 4] > thred).unsqueeze(-1).expand_as(out) # 首先根据置信度来筛选
    # pred_box = out[thred_id].view(-1, 85) # 开始删选
    # print(pred_box[:, 4], torch.argmax(pred_box[:, 5:], dim=1))
    pred_box = out
    boxes_id = cv2.dnn.NMSBoxes(pred_box[:, :4].tolist(), pred_box[:, 4].tolist(), thred, 0.3)
    # nms筛选，指的注意的是第一个参数和第二个参数都得是list,而且第一个参数数值要转化为正常的尺度，这个返回的是输入的list的符合要求的下标，维度为[n, 1]
    if boxes_id == (): # 这里代表没有预测结果
        return
    boxes_id = torch.tensor(boxes_id).long().t().squeeze(0) # 转化为tensor类型，并且将数据类型转化为Long，还要将维度转化为[n],不然后面一步会出错
    pred_box = pred_box[boxes_id] # 根据上面返回的下标选择合适的边框
    classes_id = torch.argmax(pred_box[:, 5:], dim=1) # 找到预测可能性最大的下标
    pred_num = pred_box.numpy() # 转化为Numpy数据，分别对应x,y,w,h,置信度，类别概率
    for i, box in enumerate(pred_num):
        score = round(box[4], 2) # 保留小数点后两位
        class_id = classes_id[i] # 这里预测的类别数据为浮点型，转化一下
        x1 = int(box[0]) # 左上角的x坐标，注意转化为整形
        y1 = int(box[1]) # 左上角的y坐标
        x2 = int(box[2]) # 右下角的x坐标
        y2 = int(box[3]) # 右下角的y坐标
        cv2.rectangle(show_image, (x1, y1), (x2, y2), (0, 255, 255), 2) # 画边框
        str_show = str(score) + '  ' + objName[class_id]
        # 要做图像上写的字符，注意，opencv只能写上英文，输入中文，要迂回一下，有兴趣者看我相关文章,类别下标是从0开始的
        cv2.putText(show_image, str_show, (x1 + 15, y1 + 15), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2, 8)
        # 将英文字符串画在图上
    cv2.imshow('detect_image', show_image)
    cv2.waitKey()

# 接下来将x,y,w,h转化为左上角和右下角,预测的值在0-1之间，是根据图片尺寸进行等比例缩小了的，相当于x = x(真)/imagew，现在要将x->x(真）
def change_to_conner(image_array):
    global imageW, imageH
    change_array = torch.zeros_like(image_array)
    change_array[:,0] = (image_array[:,0] - image_array[:,2] / 2) * imageW # 左上角x
    change_array[:,1] = (image_array[:,1] - image_array[:,3] / 2) * imageH # 左上角y
    change_array[:,2] = (image_array[:,0] + image_array[:,2] / 2) * imageW# 右下角x
    change_array[:,3] = (image_array[:,1] + image_array[:,3] / 2) * imageH # 右下角y
    change_array = change_array.long()
    return change_array


if __name__ == '__main__':
    get_class('coco.names')
    image = cv2.imread('images/2.png')
    if image.shape[1] > 1600:
        image = cv2.resize(image, None, fx=0.3, fy=0.3)
    model_path = './yolov3.weights'
    model_layer = './yolov3.cfg'
    test(image, model_path, model_layer)

大家如果对cv2.dnn.NMSBoxes()这个函数还有疑惑，点击了解，大家就可以对这个函数进一步理解了。

大家看看检测效果：
大家可以看见在第一张图中小车的预测没有预测准确，可能跟模型的训练程度和ssd本身的网络设计有关，但是有些奇怪的是，官方给的预测效果，ssd在速度上是比不上yolov3的，但是我的实验结果确相反，我是在六年前的笔记本上用cpu跑的这段代码，预测一张图片，ssd大概在0.3s左右，但是yolov3却在3s左右，可能跟我使用cpu跑有关，后面再去试试用yolov3官方给的接口看看效率吧。

2020 4.23