以身份证检测Demo 为例,记录yolox模型部署到前端的过程.

目标

将训练好的yolox模型部署到前端
其中:

1. yolox为’MEGVII旷视’发布的开源目标检测模型.
2. yolox模型使用官方开源项目YOLOX 训练的来, 原始模型是Pytorch的checkpoint格式, 该项目自带转换脚本tools/export_onnx.py可以将模型转为onnx格式发布.

前提

1. 训练好的yolox模型(onnx格式).
2. 模型转换运行环境:windows10, python3.8, 必要的包, 通过pip安装

   pip install onnx, onnx_tf, tensorflow, tensorflowjs

3. 前端开发环境:windows10, vue3.2.19, vite2, 必要的包, 通过npm安装

   npm install @tensorflow/tfjs

流程

转换模型

模型转换步骤从导出的onnx格式模型开始, 假设训练好的模型文件路径为C:\YOLOX-main\YOLOX_outputs\onnx\nano.onnx

onnx_filename = r'C:\YOLOX-main\YOLOX_outputs\onnx\nano.onnx'

从onnx转到tensorflow(SavedModel)

用python执行

import os
import onnx
from onnx_tf.backend import prepare

onnx_model = onnx.load(onnx_filename)
tf_rep = prepare(onnx_model, strict=False)
tf_pb_path = onnx_filename + '_graph.pb'
tf_rep.export_graph(tf_pb_path)

执行后得到对应的模型目录nano.onnx_graph.pb

参考:onnx-tensorflow

从tensorflow(SavedModel)转到Tensorflow.js可用的web格式

在命令行执行

tensorflowjs_converter --input_format=tf_saved_model     --output_node_names='outputnode'     C:\YOLOX-main\YOLOX_outputs\onnx\nano.onnx_graph.pb     C:\YOLOX-main\YOLOX_outputs\onnx

执行后得到两个文件model.json和group1-shard1of1.bin, 分别是图和权重数据(注意权重文件不可改名)

参考:tfjs文档

部署到前端

以下前端脚本均基于vue3, script setup 特性

用Tensorflow.js加载模型

加载图

首先将上述的model.json和group1-shard1of1.bin复制到vue3项目的public目录下.

const MODEL_URL = './model.json'
let model = null
model = await tf.loadGraphModel(MODEL_URL)

注意, 这里tf.loadGraphModel 如果是在nginx等生产环境下载入模型文件需要定义http请求头如下

const headers = new Headers()
headers.append('Content-Type', 'application/json')
const requestinit={
    method:'GET',
    headers:headers,
    mode:'same-origin'
}
model = await tf.loadGraphModel(MODEL_URL, {requestInit:requestinit})

使用0数据,测试模型

构建一个全由0组成的输入数据, 来测试模型是否正常工作, 控制台不报错即可(实际上权重的加载也是在第一次调用模型预测时完成的, 所以运行时间会比较长).

const zeros = tf.zeros([1,3,416,416])
await model.execute(zeros)

上述输入数据的shape是[1,3,416,416]是由于训练模型时用的时416*416大小的彩色图, 第二维的3表示3个颜色通道, 第0维的1表示batch中只有一张图.
实际输入数据的shape可能因yolox模型而异

优化加载过程

加载模型的时间相对较长, 为避免页面卡顿, 可以用async使模型加载异步于页面, 完整的加载过程如下

const MODEL_URL = './model.json'
let model = null

async function 载入模型(){
  if (model==null){
    //tf.setBackend('cpu')
    //tf.setBackend('webgl')
    await tf.ready()
    console.log('载入模型')
    const headers = new Headers()
    headers.append('Content-Type', 'application/json')
    const requestinit={
      method:'GET',
      headers:headers,
      mode:'same-origin'
    }
    model = await tf.loadGraphModel(MODEL_URL, {requestInit:requestinit})
    console.log('图ok')
    const zeros = tf.zeros([1,3,416,416])
    await model.execute(zeros)
    console.log('初始化ok')
    console.log(tf.engine())
  }
}

载入模型()

参考:tfjsAPI文档

模型输入预处理

async function yolox预处理(图片数据, 模型尺寸=[416,416]){
  const t图片 = tf.browser.fromPixels(图片数据)
  const 缩放 = Math.min(模型尺寸[0]/t图片.shape[0], 模型尺寸[1]/t图片.shape[1])
  const t缩放图片 = tf.image.resizeNearestNeighbor(t图片, [Math.floor(t图片.shape[0]*缩放),Math.floor(t图片.shape[1]*缩放)])
  const t输入图片 = t缩放图片.pad([[0,模型尺寸[0]-t缩放图片.shape[0]],[0,模型尺寸[1]-t缩放图片.shape[1]],[0,0]], 114)
  // 输入颜色通道前置并且转换rgb顺序为bgr顺序
  const 输入数据 = t输入图片.cast('float32').transpose([2,0,1]).reverse(0).expandDims(0)
  return {数据:输入数据, 缩放:缩放}
}

参考:yolox源码

模型输出后处理

async function yolox后处理(原始输出, 模型尺寸=[416,416], 缩放=1., 有效分阈值=0.1, 重合比阈值=0.4){
  // 输入tensorflowjs.tensor[1,3549,15], 图片宽高 = 416*416
  const 类别数 = 原始输出.shape[2] - 5
  const [原始坐标, 原始宽高, 原始加权分, 原始类别分] = 原始输出.squeeze([0]).split([2,2,1,类别数],1)
  const 模板尺度 = [8, 16, 32]
  let 坐标格数组 = []
  let 坐标加权数组 = []
  模板尺度.forEach(权重=>{
    const 坐标格 = tf.stack(tf.meshgrid(tf.range(0,Math.floor(模型尺寸[0]/权重),1), tf.range(0,Math.floor(模型尺寸[1]/权重),1)),2).reshape([-1,2])
    坐标格数组.push(坐标格)
    const 坐标加权 = tf.ones(坐标格.shape).mul(权重)
    坐标加权数组.push(坐标加权)
  })
  const 坐标格 = tf.concat(坐标格数组, 0)
  const 坐标加权 = tf.concat(坐标加权数组,0)
  const 实际宽高 = 原始宽高.exp().mul(坐标加权).div(缩放)
  const 中心坐标 = 原始坐标.add(坐标格).mul(坐标加权).div(缩放)
  const 左上坐标 = 中心坐标.sub(实际宽高.div(2))
  const 右下坐标 = 左上坐标.add(实际宽高)
  const 类别分 = 原始类别分.mul(原始加权分)
  const 有效类别 = 类别分.argMax(1)
  const 有效类别布尔索引 = tf.oneHot(有效类别, 类别数)
  const 加权类别分 = await tf.booleanMaskAsync(类别分, 有效类别布尔索引.cast('bool'))
  // 筛选
  const 原始索引 = tf.range(0,原始坐标.shape[0],1,'int32') // 行向量[N]
  const 有效分布尔索引 = 加权类别分.greater(有效分阈值)
  const 有效分 = await tf.booleanMaskAsync(加权类别分, 有效分布尔索引)
  const 有效索引 = await tf.booleanMaskAsync(原始索引, 有效分布尔索引)
  const {values:_, indices:有效分从大到小索引索引} = tf.topk(有效分,有效分.shape[0])
  const 有效分从大到小索引 = 有效索引.gather(有效分从大到小索引索引)
  let 有效输出数组 = []
  // 待筛选数据[左上坐标2, 宽高2, 有效分, 类别, 右下坐标2, 面积]
  var 待筛选数据 = tf.concat([左上坐标,实际宽高,加权类别分.reshape([-1,1]),有效类别.reshape([-1,1]),右下坐标,实际宽高.prod(1).reshape([-1,1])],1).gather(有效分从大到小索引,0)
  while(待筛选数据 && 待筛选数据.shape[0]>1){
    const 划分 = 待筛选数据.split([1,待筛选数据.shape[0]-1],0)
    const 当前框 =划分[0]
    待筛选数据 = 划分[1]
    有效输出数组.push(当前框)
    const 重合左上坐标 = 待筛选数据.slice([0,0],[待筛选数据.shape[0],2]).maximum(当前框.slice([0,0],[1,2]))
    const 重合右下坐标 = 待筛选数据.slice([0,6],[待筛选数据.shape[0],2]).minimum(当前框.slice([0,6],[1,2]))
    const 重合宽高 = 重合右下坐标.sub(重合左上坐标)
    const 重合面积 = 重合宽高.maximum(0).prod(1).reshape([-1,1])
    const 总面积 = 待筛选数据.slice([0,8],[待筛选数据.shape[0],1]).add(当前框.slice([0,8],[1,1]))
    const 面积占比 = 重合面积.div(总面积)
    const 保留的筛选数据布尔索引 = 面积占比.lessEqual(重合比阈值).reshape([-1])
    待筛选数据 = await tf.booleanMaskAsync(待筛选数据, 保留的筛选数据布尔索引, 0)
  }
  if(待筛选数据 && 待筛选数据.shape[0]>1){
    有效输出数组.push(待筛选数据)
  }
  const 有效输出 = tf.concat(有效输出数组,0)
  const 输出划分 = 有效输出.split([4,1,1,3],1)
  return {方框:输出划分[0].arraySync(), 评分:输出划分[1].arraySync(), 类别:输出划分[2].arraySync()}
}

后处理过程耗时较多, 原因是处理过程中频繁构建Tensor.

参考:yolox源码

演示

基于yolox的身份证检测Demo

后记

1. 上述yolox模型部署到前端处理一张图片需要2~4秒, 相对的python端onnx模型仅需200毫秒左右. 主要的延迟原因在于后处理过程.

2. 部署小型的目标检测模型用yolo3可能会更好, 识别效果大差不差, 代码还简单不少.