ELectron的node环境中使用onnxruntime-node跑onnx模型推理图片
我是首次接触模型这块的内容,onnxruntime用js调用分为两块一个是onnxruntime-web和onnxruntime-node,api都是相同的.注意的是两者都是支持gpu加速的但是不支持cuda的使用推理模型的步骤:
- 创建tensor(张量,数学不错的话应该知道,就是有形状的数据)
- 将tensor放在创建的模型实列上推理
- 得到模型的输出结果,处理结果数据
在推理图片之前我们需要知道一些前期的工作
模型的输入内容 一般模型都是自己公司开发的或者有第三方提供的都会有个pipeline.json文件文件中就说明了,模型需要怎么处理数据这里我也提供一个模型可视化网站,导入模型就可以看到一些相关的信息
创建tensor主要做的就是把图片处理成模型想要的形状的数据,处理图片的流程就是,
- 改变图片的尺寸
- 将图片转换成rgb模式(如果是rgb图片则不用转换)
- 归一化
- 转换图片格式NCWH格式(具体看模型需求)
概念解释:
通道:比如rgb就是三通道的r(red)、g(green)、b(blue),但是不一定图片都是三通道的比如png的图片还有透明度,还有cmyk的图片,但在onnx模型中基本都是要rgb的
归一化:就是减去rgb的均值除以标准差,一般在模型的pipeline.json文件中就有表明
NCWH格式:n(表示图片的数量一般都是1)、c(表示图片的通道数rgb则为3)、w和h(代表着图片的宽高)
上面这个模型要的要的就是图片数量是1,通道是3,图片宽高是300x300的tensor
下面看一下代码上具体实现吧
js
//导入相关依赖包
const ort = require('onnxruntime-node');//推理模型
const sharp = require('sharp');//用来处理图片
/**
* @description - 图片处理
* @param {string} imagePath 图片的路径
* @returns {array} 处理后图片的原始数据
*/
const preprocessImage = async (imagePath) => {
// 使用 sharp 从文件中读取图像并调整大小
const { data, info } = await sharp(imagePath)
.resize(300, 300) //改变图片的尺寸
.toFormat('jpg') // 将图片转成rgb
.raw() // 输出原始 RGB 数据
.toBuffer({ resolveWithObject: true }); //得到图片的buffer原始数据
// 将数据转换为 Float32Array 并进行归一化
const mean = [123.675, 116.28, 103.53]; // RGB均值从pipeline.json中获取
const std = [1, 1, 1]; // 标准差,保持为1从pipeline.json中获取
const normalizedData = new Float32Array(info.width * info.height * 3);//创建float32位的数组,这个数据类型也是看模型需要什么,上面的模型可是化的图片中的input中的tensor要的是float32
for (let i = 0; i < info.width * info.height; i++) {//归一化
for (let channel = 0; channel < 3; channel++) {
normalizedData[i * 3 + channel] =
(data[i * 3 + channel] - mean[channel]) / std[channel];
}
}
// 归一化后转换为 NCHW 格式
const nchwData = new Float32Array(1 * 3 * 300 * 300);
for (let c = 0; c < 3; c++) {
for (let h = 0; h < 300; h++) {
for (let w = 0; w < 300; w++) {
nchwData[c * 300 * 300 + h * 300 + w] =
normalizedData[h * 300 * 3 + w * 3 + c];
}
}
}
return nchwData;
};
/**
* @description - 推理图片
* @param {string} imagePath 图片路径
*/
const runInference = async (imagePath) => {
try {
const normalizedImage = await preprocessImage(imagePath);//获取处理后的图片数据
const session = await ort.InferenceSession.create('./end2end.onnx'); //创建模型会话 替换为你的 ONNX 模型路径
console.log('Session loaded', session);//打印session可以看到模型输入的输入输出的属性名称
const inputTensor = new ort.Tensor(//创建tensor
'float32',
normalizedImage,
[1, 3, 300, 300]
); // 适应 ONNX 模型的输入格式
const feeds = { input: inputTensor }; // 根据你的模型输入名称设置
const output = await session.run(feeds);//推理模型得到输出结果打印输出结果可以看到输出结果中有
let boxes = getOutput(output);
let originalBoxes = getBoxesPosition_contain(
{ info: { width: 640, height: 427 } },
{ info: { width: 300, height: 300 } },
boxes
);
rectangle('./demo.jpg', originalBoxes);
} catch (error) {
console.error('Error during inference:', error);
}
};
看一下我这边的输出结果,cupData就是推理图片的结果,我们最主要关注的就是dets中的dims属性,因为它告诉了我们cupData的数据形状。比如我dims的意思就是200条识别结果,每条识别结果是由5个数据组成一个的,然后将cpuData数据从前开始每5个分割一下得到一个二维的数组;下面是我对我的模型数据结果的后续处理(仅供参考,每个模型要求的处理都是不一样,比如在我这个模型里面每5个一条数据,那么这五个数据的每一项依次表示的识别框的左上角和右下角在图片的坐标,最后的代表的置信度,这组数据的下标就是对应的识别物参考表的id)具体的还是要看你自己模型
js
function getOutput(output) {
// 提取 `dets` 和 `labels` 的数据
const dets = output.dets.data;
const labels = output.labels.data;
// 确认数据的长度匹配
const numDetections = output.dets.dims[1]; // 300 个检测框
// 组合边界框与标签信息
const results = [];
for (let i = 0; i < numDetections; i++) {
const x1 = Math.round(dets[i * 5 + 0]); // 边界框的 x1 坐标
const y1 = Math.round(dets[i * 5 + 1]); // 边界框的 y1 坐标
const x2 = Math.round(dets[i * 5 + 2]); // 边界框的 x2 坐标
const y2 = Math.round(dets[i * 5 + 3]); // 边界框的 y2 坐标
const confidence = dets[i * 5 + 4]; // 置信度
const label = labels[i]; // 类别标签
// 将信息组合成一个对象s
results.push({
// point: { x1, y1, x2, y2 },
point: [x1, y1, x2, y2],
confidence,
classId: label,
});
}
return results.filter((result) => result.confidence > 0.4);
}
具体的pipeline.json文件解读需要科技才能访问