【SDK案例系列 07】基于 MindX SDK + Pytorch YoLoV4的目标检测

源码下载：

https://gitee.com/open-ascend/atlas_mindxsdk_samples/blob/master/contrib/cv/object_detection/image_yolov4

一、安装昇腾驱动

先安装昇腾驱动，昇腾驱动请参考各个产品安装手册，安装完成后npu-smi info 显示安装成功

[root@localhost ~]#
[root@localhost ~]# npu-smi info
+-------------------------------------------------------------------------------------------------+
| npu-smi 22.0.2                   Version: 22.0.2                                                |
+------------------+--------------+---------------------------------------------------------------+
| NPU    Name      | Health       | Power(W)             Temp(C)           Hugepages-Usage(page)  |
| Chip   Device    | Bus-Id       | AICore(%)            Memory-Usage(MB)                         |
+==================+==============+===============================================================+
| 1      310       | OK           | 12.8                 45                0   / 0                |
| 0      0         | 0000:05:00.0 | 0                    2621  / 8192                             |
+==================+==============+===============================================================+

二、安装MindX SDK > mxVision

（1）MindX SDK需要通过官网获取。

（2）mxVision说明手册：

https://www.hiascend.com/document/detail/zh/mind-sdk/30rc3/quickstart/visionquickstart/visionquickstart_0000.html

（3）安装MindX SDK

./Ascend-mindxsdk-mxvision_3.0.RC2_linux-aarch64.run --install --install-path=/usr/local/sdk_home

–install-path为指定安装的路径

（4）安装成功后会提示如下信息

Installing collected packages:mindx
Successfully installed mindx-3.0.RC2

（5）安装成功后在对应目录下查看，能看到mxVision

[root@localhost sdk_home]#
[root@localhost sdk_home]# pwd
/usr/local/sdk_home
[root@localhost sdk_home]# ls
mxVision mxVision-3.0.RC2
[root@localhost sdk_home]#
[root@localhost sdk_home]#

（6）MindX SDK使用中需要用到OSD功能，安装后需要执行以下命令，生成om文件

bash /usr/local/sdk_home/mxVision/operators/opencvosd/generate_osd_om.sh

执行成功后，显示如下效果

[root@localhost ~]# bash /usr/local/sdk_home/mxVision/operators/opencvosd/generate_osd_om.sh
ASCEND_HOME is set to /usr/local/Ascend by user
Set ASCEND_VERSION to the default value:ascend-toolkit/latest
ATC start working now,please wait for a moment.
ATC run success, welcome to the next use.

The model has been successfully converted to om,please get it under /usr/local/sdk_home/mxVision/operators/opencvosd.
[root@localhost ~]# 

（9）安装完MindX SDK后，需要配置环境变量

.bashrc文件添加以下环境变量

# 安装mxVision时配置
. /usr/local/sdk_home/mxVision/set_env.sh

用户也可以通过修改~/.bashrc文件方式设置永久环境变量，操作如下：

a) 以运行用户在任意目录下执行vi ~/.bashrc命令，打开.bashrc文件，在文件最后一行后面添加上述内容。

b) 执行:wq!命令保存文件并退出。

c) 执行source ~/.bashrc命令使其立即生效。

三、ATC模型转换

1、获取权重文件。

（1）从源码包中获取训练后的权重文件yolov4.pth。

（2）开源 https://github.com/Tianxiaomo/pytorch-YOLOv4 提供经过训练的YOLOv4权重文件，可直接下载使用。

2、导出onnx文件。

（1）克隆代码仓至服务器任意路径下，并进入代码仓目录

git clone https://github.com/Tianxiaomo/pytorch-YOLOv4
cd pytorch-YOLOv4

（2）修改demo_pytorch2onnx.py脚本

def transform_to_onnx(weight_file, batch_size, n_classes, IN_IMAGE_H, IN_IMAGE_W):

   model = Yolov4(n_classes=n_classes, inference=False) # 修改为False

    pretrained_dict = torch.load(weight_file, map_location=torch.device('cpu'))# 修改为cpu
    model.load_state_dict(pretrained_dict)

    input_names = ["input"]
    output_names = ['feature_map_1', 'feature_map_2','feature_map_3']# 修改为3个feature_map

    dynamic = False
    if batch_size <= 0:
        dynamic = True

...

def main(weight_file, image_path, batch_size, n_classes, IN_IMAGE_H, IN_IMAGE_W):

    if batch_size <= 0:
        onnx_path_demo = transform_to_onnx(weight_file, batch_size, n_classes, IN_IMAGE_H, IN_IMAGE_W)
    else:
        # Transform to onnx as specified batch size
        transform_to_onnx(weight_file, batch_size, n_classes, IN_IMAGE_H, IN_IMAGE_W)
        # Transform to onnx for demo
        onnx_path_demo = transform_to_onnx(weight_file, 1, n_classes, IN_IMAGE_H, IN_IMAGE_W)

     # 注释掉图片推理部分
    # session = onnxruntime.InferenceSession(onnx_path_demo)
    # session = onnx.load(onnx_path)
    # print("The model expects input shape: ", session.get_inputs()[0].shape)

    # image_src = cv2.imread(image_path)
    # detect(session, image_src)

（3）将yolov4.pth权重文件移到当前目录下。

demo_pytorch2onnx.py脚本将.pth文件转换为.onnx文件，执行如下命令。

python3 demo_pytorch2onnx.py yolov4.pth data/dog.jpg 1 80 608 608

运行成功后，在当前目录生成模型文件yolov4_1_3_608_608_static.onnx。

（4）修改模型。

dy_resize.py文件在image_yolov4/data/models/yolov4目录下。

该模型含有动态Resize算子（上采样），通过计算维度变化，改为静态算子，不影响模型的精度。
运行脚本：

python3 dy_resize.py yolov4_1_3_608_608_static.onnx

生成的yolov4_1_3_608_608_static_dbs.onnx可用ATC工具转换为离线om模型。

3、把生成的yolov4_1_3_608_608_static_dbs.onnx放在yolov4_demo/models/yolov4目录下

[root@localhost yolov4]#
[root@localhost yolov4]# ls
aipp.config  atc_310.sh  atc_310P3.sh  coco2014.names  dy_resize.py  yolov4_1_3_608_608_static_dbs.onnx  yolov4_bs1_fp16.cfg  yolov4.pth
[root@localhost yolov4]# 

4、执行模型转换命令

Ascend310芯片模型转换命令如下：

atc \
    --mode=0 \
    --framework=5 \
    --model=./yolov4_1_3_608_608_static_dbs.onnx \
    --output=./yolov4_bs1_with_aipp \
    --input_format=NCHW \
    --input_shape="input:1,3,608,608" \
    --out_nodes="Conv_495:0;Conv_518:0;Conv_541:0" \
    --log=error \
    --soc_version=Ascend310 \
    --insert_op_conf=aipp.config

Ascend310P3芯片模型转换命令如下：

atc \
    --mode=0 \
    --framework=5 \
    --model=./yolov4_1_3_608_608_static_dbs.onnx \
    --output=./yolov4_bs1_with_aipp \
    --input_format=NCHW \
    --input_shape="input:1,3,608,608" \
    --out_nodes="Conv_495:0;Conv_518:0;Conv_541:0" \
    --log=error \
    --soc_version=Ascend310P3 \
    --insert_op_conf=aipp.config

参数说明：

–model：待转换的ONNX模型。

–framework：5代表ONNX模型。

–output：输出的om模型。

–input_format：输入数据的格式。

–input_shape：输入数据的shape。

–insert_op_conf=./aipp.config：AIPP插入节点，通过config文件配置算子信息，功能包括图片色域转换、裁剪、归一化，主要用于处理原图输入数据，常与DVPP配合使用，详见下文数据预处理。

详细ATC命令转换学习请参考：

https://support.huawei.com/enterprise/zh/doc/EDOC1100234054?idPath=23710424%7C251366513%7C22892968%7C251168373

3、模型转换后，会在目录下生成yolov4_bs1_with_aipp.om

[root@localhost yolov4]#
[root@localhost yolov4]# ls
aipp.config  atc_310.sh  atc_310P3.sh  coco2014.names  dy_resize.py  yolov4_1_3_608_608_static_dbs.onnx  yolov4_bs1_with_aipp.om  yolov4_bs1_fp16.cfg  yolov4.pth
[root@localhost yolov4]# 

四、使用image_yolov4

1、修改run_cpp.sh & run_python.sh中MX_SDK_HOME为MindX SDK安装目录

export MX_SDK_HOME=/usr/local/sdk_home/mxVision

2、执行run_cpp.sh 或者 run_python.sh

bash run_cpp.sh

bash run_python.sh

3、目标检测结果与test.jpg一致

目标检测结果：
Results:{"MxpiObject":[{"classVec":[{"classId":16,"className":"dog","confidence":0.99178299999999997,"headerVec":[]}],"x0":179.345978,"x1":843.84118699999999,"y0":155.08517499999999,"y1":539.24890100000005}]}

五、image_yolov4详解

1、技术流程图

视频解码：调用DVPP解码能力，转换为 YUV 格式图像数据。

图像缩放：调用DVPP，将图像缩放到一定尺寸大小。

目标检测：YoLoV4模型针对图像进行目标检测。

模型后处理：针对推理结果进行后处理文字转换。

数据序列化：将stream结果组装成json字符串输出。

2、pipeline详解

{
    "classification": {
        "stream_config": {  ##设置业务流在哪个芯片上处理
            "deviceId": "0"
        },
        "mxpi_imagedecoder0": {  ##图像解码（纯硬件）
            "factory": "mxpi_imagedecoder",
            "next": "mxpi_imageresize0"
        },
        "mxpi_imageresize0": {  ##图像缩放（纯硬件）
            "factory": "mxpi_imageresize",
            "next": "mxpi_modelinfer0"
        },
        "mxpi_modelinfer0": {  ##模型推理
            "props": {
                "modelPath": "data/models/yolov4/yolov4_bs1_with_aipp.om",  ##模型路径
                "postProcessConfigPath": "data/models/yolov4/yolov4_bs1_fp16.cfg",
                "labelPath": "data/models/yolov4/coco2014.names",
                "postProcessLibPath": "libMpYOLOv3PostProcessor.so"
            },
            "factory": "mxpi_modelinfer",
            "next": "mxpi_dataserialize0"
        },
        "mxpi_dataserialize0": {  ##数据序列化
            "props": {
                "outputDataKeys": "mxpi_modelinfer0"
            },
            "factory": "mxpi_dataserialize",
            "next": "appsink0"
        },
        "appsrc0": {
            "props": {
                "blocksize": "409600"
            },
            "factory": "appsrc",
            "next": "mxpi_imagedecoder0"
        },
        "appsink0": {  ##输出推理结果
            "props": {
                "blocksize": "4096000"
            },
            "factory": "appsink"
        }
    }
}

3、C++源码详解

int main(int argc, char* argv[])
{
    // 读取pipeline配置文件
    std::string pipelineConfigPath = "data/pipeline/Sample.pipeline";
    std::string pipelineConfig = ReadPipelineConfig(pipelineConfigPath);
    if (pipelineConfig == "") {
        LogError << "Read pipeline failed.";
        return APP_ERR_COMM_INIT_FAIL;
    }
    // 初始化 Stream manager 资源
    MxStream::MxStreamManager mxStreamManager;
    APP_ERROR ret = mxStreamManager.InitManager();
    if (ret != APP_ERR_OK) {
        LogError << GetError(ret) << "Failed to init Stream manager.";
        return ret;
    }
    // 根据指定的pipeline配置创建Stream
    ret = mxStreamManager.CreateMultipleStreams(pipelineConfig);
    if (ret != APP_ERR_OK) {
        LogError << GetError(ret) << "Failed to create Stream.";
        return ret;
    }
    // 读取测试图片
    MxStream::MxstDataInput dataBuffer;
    ret = ReadFile("data/test.jpg", dataBuffer);
    if (ret != APP_ERR_OK) {
        LogError << GetError(ret) << "Failed to read image file.";
        return ret;
    }
    std::string streamName = "classification";
    int inPluginId = 0;
    // 发送测试图片到Stream进行推理
    ret = mxStreamManager.SendData(streamName, inPluginId, dataBuffer);
    if (ret != APP_ERR_OK) {
        LogError << GetError(ret) << "Failed to send data to stream.";
        delete dataBuffer.dataPtr;
        dataBuffer.dataPtr = nullptr;
        return ret;
    }
    // 获取推理结果
    MxStream::MxstDataOutput* output = mxStreamManager.GetResult(streamName, inPluginId);
    if (output == nullptr) {
        LogError << "Failed to get pipeline output.";
        delete dataBuffer.dataPtr;
        dataBuffer.dataPtr = nullptr;
        return ret;
    }
    // 打印推理结果
    std::string result = std::string((char *)output->dataPtr, output->dataSize);
    LogInfo << "Results:" << result;

    // 销毁Stream
    mxStreamManager.DestroyAllStreams();
    delete dataBuffer.dataPtr;
    dataBuffer.dataPtr = nullptr;

    delete output;
    return 0;
}

4、Python源码详解

if __name__ == '__main__':
    # 初始化 Stream manager 资源
    streamManagerApi = StreamManagerApi()
    ret = streamManagerApi.InitManager()
    if ret != 0:
        print("Failed to init Stream manager, ret=%s" % str(ret))
        exit()

    # 根据指定的pipeline配置创建Stream
    with open("data/pipeline/Sample.pipeline", 'rb') as f:
        pipelineStr = f.read()
    ret = streamManagerApi.CreateMultipleStreams(pipelineStr)
    if ret != 0:
        print("Failed to create Stream, ret=%s" % str(ret))
        exit()

    # 读取测试图片
    dataInput = MxDataInput()
    with open("data/test.jpg", 'rb') as f:
        dataInput.data = f.read()

    # 发送测试图片到Stream进行推理
    streamName = b'classification'
    inPluginId = 0
    uniqueId = streamManagerApi.SendDataWithUniqueId(streamName, inPluginId, dataInput)
    if uniqueId < 0:
        print("Failed to send data to stream.")
        exit()

    # 获取推理结果
    inferResult = streamManagerApi.GetResultWithUniqueId(streamName, uniqueId, 3000)
    if inferResult.errorCode != 0:
        print("GetResultWithUniqueId error. errorCode=%d, errorMsg=%s" % (
            inferResult.errorCode, inferResult.data.decode()))
        exit()

    # 打印推理结果
    print(inferResult.data.decode())

    # 销毁Stream
    streamManagerApi.DestroyAllStreams()