【opencv有趣应用】基于MobileNet + SSD的物体检测

今天学习使用另外的模型进行物体检测

当前目标检测的算法有很多，如rcnn系列、yolo系列和ssd，前端网络如vgg、AlexNet、SqueezeNet，一种常用的方法是将前端网络设为MobileNet进行特征提取和加速，后端算法为SSD来真正进行目标检测。

之前我们学习了使用darknet上的基于coco数据集训练的yolov3模型来搞，这次我们使用在tensorflow上训练的基于coco数据集训练的MobileNet + SSD模型来搞。

import cv2
import numpy as np


def readAllCleassNames(filename):
    with open(filename, 'rt') as f:
        classNames = f.read().rstrip('\n').split('\n')
    return classNames


if __name__ == '__main__':
    # ======================= step 1: 先获得所有的类别信息
    classNames = readAllCleassNames('file/coco_classnames.txt')
    print(len(classNames))  # 一共有80个类别
    print(classNames)

    # ======================= step 2: 定义模型
    configFile = 'file/ssd_mobilenet_v3_large_coco_2020_01_14.pbtxt'
    weightFile = 'file/frozen_inference_graph.pb'

    # 定义模型
    mobileNet_SSD_Model = cv2.dnn_DetectionModel(weightFile, configFile)
    # 设置运行的背景，一般情况都是使用opencv dnn作为后台计算
    mobileNet_SSD_Model.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
    # 设置目标设备, DNN_TARGET_CPU其中表示使用CPU计算，默认是的
    mobileNet_SSD_Model.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)

    # 设置输入图像值的转换的一些参数
    mobileNet_SSD_Model.setInputSize(320, 320)  # 设置输入的数据大小
    mobileNet_SSD_Model.setInputScale(1.0 / 127.5)  # 设置数值缩放，归一化
    mobileNet_SSD_Model.setInputMean([127.5, 127.5, 127.5])  # 设置各个通道的平均值
    mobileNet_SSD_Model.setInputSwapRB(True)  # R和B通道进行交换，从BGR到RGB。

    # ======================= step 3: 得到一个图像，转换后将其输入模型。
    img = cv2.imread('images/car_person.jpeg')
    # cv2.imshow('cat', img)
    # cv2.waitKey(0)

    # 这个函数直接是一步到位，
    # 自动获取输出层的layers，
    # 自动将原始图像数据转成blob类型，
    # 自动进行前向传播计算
    # 自动根据confThreshold过滤检测结果，
    # 自动根据nmsThreshold做了NMS操作，
    # 返回类别id，各自的置信度，以及boundingbox
    classIds, confs, bbox = mobileNet_SSD_Model.detect(img, confThreshold=0.5, nmsThreshold=0.5)
    if len(classIds) != 0:
        for classID, conf, box in zip(classIds.flatten(), confs.flatten(), bbox):
            # 画出一个矩形检测框
            cv2.rectangle(img, box, color=(255, 0, 0), thickness=2)
            # 画出检测的目标名称和置信度
            cv2.putText(img, text="classN_name:{}, confidence:{}%".format(classNames[classID - 1].upper(),
                                                                          "%.2f" % (conf * 100)),
                        org=(box[0], box[1]-10), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.6,
                        color=(0, 0, 255), thickness=2)

        cv2.imshow('Output', img)