【AI模型】NanoDet-PyTorch环境配置及运行测试

DevFrank

发布于 2024-07-24 14:42:09

960

发布于 2024-07-24 14:42:09

文章被收录于专栏：C++开发学习交流

1. 准备NanoDet-PyTorch工程

该代码基于NanoDet项目进行小裁剪，专门用来实现Python语言、PyTorch 版本的代码，下载直接能使用，支持图片、视频文件、摄像头实时目标检测。

用于目标检测，模型小，检测速度快速，适合没GPU显卡的嵌入式设备运行，比如“树莓派”、ARM开发板、嵌入式开发板。

本文在Ubuntu18.04环境下进行测试：

首先将python的源更换为国内源：ubuntu修改python的pip源为国内源

2. 创建python虚拟环境

python -m venv Virtual-NanoDet

 source myvenv/bin/activate

3. pip安装依赖包

如：

pip install cmake
pip install numpy matplotlib pandas scipy opencv-python imutils -i https://pypi.tuna.tsinghua.edu.cn/simple

基本上按照这个流程安装这些包下来是没问题的，预计这个虚拟环境大小为2.5G；

4. 测试图片检测、视频检测、摄像头检测

文件中提供了图片和视频素材，摄像头用usb接口的就可以，下面开始运行程序：

## 运行程序(先进入自建的python venv中，再到目标文件夹中运行以下程序)

'''目标检测-图片'''
# python detect_main.py image --config ./config/nanodet-m.yml --model model/nanodet_m.pth --path  street.png

'''目标检测-视频文件'''
# python detect_main.py video --config ./config/nanodet-m.yml --model model/nanodet_m.pth --path  test.mp4

'''目标检测-摄像头'''
# python detect_main.py webcam --config ./config/nanodet-m.yml --model model/nanodet_m.pth --path  0

接下来就可以看到目标检测后的效果了：

5. 模型转换及部署

运行tools/export.py脚本可将pth转为onnx模型（注意路径）：

import os
import torch
from nanodet.model.arch import build_model
from nanodet.util import Logger, cfg, load_config, load_model_weight

def main(config, model_path, output_path, input_shape=(320, 320)):
    logger = Logger(-1, config.save_dir, False)
    model = build_model(config.model)
    checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage)
    load_model_weight(model, checkpoint, logger)
    dummy_input = torch.autograd.Variable(torch.randn(1, 3, input_shape[0], input_shape[1]))
    torch.onnx.export(model, dummy_input, output_path, verbose=True, keep_initializers_as_inputs=True, opset_version=11)
    print('finished exporting onnx ')

if __name__ == '__main__':
    cfg_path = r"../config/nanodet-m.yml"
    model_path = r"../model/nanodet_m.pth"
    out_path = r'../model/output.onnx'
    load_config(cfg, cfg_path)
    main(cfg, model_path, out_path, input_shape=(320, 320))

生成onnx通用模型后，可转为部署需要的格式，如ncnn，我们已经在ubuntu装好了ncnn，然后在ncnn/build/install/bin/下有一个onnx2ncnn脚本，执行转换程序：./onnx2ncnn output.onnx output.param output.bin

然后就可以在移动端程序中使用ncnn框架所需要的模型了（bin、param）。

以上。

本文参与腾讯云自媒体同步曝光计划，分享自作者个人站点/博客。

原始发表：2023-08-24，如有侵权请联系 cloudcommunity@tencent.com 删除

程序