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

2024-07-24 14:42:09 浏览数 (3)

1. 准备NanoDet-PyTorch工程

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

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

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

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

在这里插入图片描述在这里插入图片描述
2. 创建python虚拟环境
代码语言:javascript复制
python -m venv Virtual-NanoDet
代码语言:javascript复制
 source myvenv/bin/activate
3. pip安装依赖包

如:

代码语言:javascript复制
pip install cmake
pip install numpy matplotlib pandas scipy opencv-python imutils -i https://pypi.tuna.tsinghua.edu.cn/simple
在这里插入图片描述在这里插入图片描述

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

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

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

代码语言:javascript复制
## 运行程序(先进入自建的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模型(注意路径):

代码语言:javascript复制
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)。

以上。

0 人点赞