使用 ONNX Runtime Web 在 Web 应用程序中分类图像

在本教程中,我们将使用一个 GitHub 仓库模板来构建一个图像分类 Web 应用,该应用使用 ONNX Runtime Web。我们将使用 JavaScript 在浏览器中对计算机视觉模型进行推理。

在不常用于数据科学的语言中进行部署和推理时,最困难的部分之一是弄清楚如何进行数据处理和推理。我们已经通过 此模板 为您完成了所有繁重的工作!

下面是该模板生成的网站外观。它循环遍历一个示例图像列表,使用 SqueezeNet 模型调用推理会话,然后返回推理的得分和标签。

示例模板输出

Image of browser inferencing on sample images.

目录

设备上推理

此应用程序使用 onnxruntime-web JavaScript 库在设备(浏览器)上执行推理。

SqueezeNet 机器学习模型

我们将使用来自 ONNX 模型动物园SqueezeNet 模型。SqueezeNet 模型执行图像分类——它们将图像作为输入,并将图像中的主要对象分类到一组预定义的类别中。它们是在 ImageNet 数据集上训练的,该数据集包含来自 1000 个不同类别的图像。SqueezeNet 模型在大小和速度方面非常高效,同时提供了良好的准确性。这使得它们非常适合对大小有严格限制的平台,例如客户端推理。

如果您需要更高的模型内存和磁盘效率,可以将 ONNX 模型转换为 ORT 格式,并在应用程序中使用 ORT 模型而不是 ONNX 模型。您还可以 减小 ONNX Runtime 二进制文件的大小,使其仅包含对应用程序中特定模型的支持。

使用 NextJS (一个 ReactJS 框架) 创建静态站点,用于在浏览器中部署模型

模板

此模板的目标是为您的加速 ML Web 应用程序提供一个起点。该模板使用 NextJS 框架生成一个计算机视觉应用程序,该应用程序使用 TypeScript 编写并由 webpack 构建。让我们深入研究模板并分解代码。

utils 文件夹

Utils 文件夹中有三个文件:imageHelper.tsmodelHelper.tspredict.ts。Predict 是 Web 组件启动推理的入口点。在这里,我们导入帮助函数并调用默认函数以获取图像张量并运行模型推理。

predict.ts

// Language: typescript
// Path: react-next\utils\predict.ts
import { getImageTensorFromPath } from './imageHelper';
import { runSqueezenetModel } from './modelHelper';

export async function inferenceSqueezenet(path: string): Promise<[any,number]> {
  // 1. Convert image to tensor
  const imageTensor = await getImageTensorFromPath(path);
  // 2. Run model
  const [predictions, inferenceTime] = await runSqueezenetModel(imageTensor);
  // 3. Return predictions and the amount of time it took to inference.
  return [predictions, inferenceTime];
}

imageHelper.ts

首先,我们需要从本地文件或 URL 获取图像并将其转换为张量。imageHelper.ts 中的 getImageTensorFromPath 函数使用 JIMP 读取文件、调整大小并返回 imageDataJIMP 是一个 JavaScript 图像处理库。它有许多内置函数用于处理图像数据,例如调整大小、灰度、写入等。在此示例中,我们只需要调整大小,但在您的代码中可能需要额外的图像数据处理。

import * as Jimp from 'jimp';
import { Tensor } from 'onnxruntime-web';

export async function getImageTensorFromPath(path: string, dims: number[] =  [1, 3, 224, 224]): Promise<Tensor> {
  // 1. load the image  
  var image = await loadImagefromPath(path, dims[2], dims[3]);
  // 2. convert to tensor
  var imageTensor = imageDataToTensor(image, dims);
  // 3. return the tensor
  return imageTensor;
}

async function loadImagefromPath(path: string, width: number = 224, height: number= 224): Promise<Jimp> {
 // Use Jimp to load the image and resize it.
  var imageData = await Jimp.default.read(path).then((imageBuffer: Jimp) => {
    return imageBuffer.resize(width, height);
  });

  return imageData;
}

一旦我们有了 imageData,我们就会将其发送到 imageDataToTensor 函数,将其转换为用于推理的 ORT 张量。要在 JavaScript 中将图像转换为张量,我们需要将 RGB(红、绿、蓝)值放入数组中。为此,我们将通过每个像素的 RGBA 4 个通道循环遍历 imageBufferData。一旦我们获得了图像的 RGB 像素通道,我们就可以从 transposedData 创建 Float32Array 并除以 255 以标准化值。为什么 255 标准化像素值?归一化是一种在不扭曲差异的情况下将值更改为通用尺度的技术。255 是 RGB 值的最大值,因此除以 255 将我们的值标准化到 0 到 1 之间,而不会丢失统计差异。现在我们有了图像的 Float32Array 表示,我们可以通过传入类型、数据和维度来创建 ORT 张量。然后我们返回 inputTensor 进行推理。

function imageDataToTensor(image: Jimp, dims: number[]): Tensor {
  // 1. Get buffer data from image and create R, G, and B arrays.
  var imageBufferData = image.bitmap.data;
  const [redArray, greenArray, blueArray] = new Array(new Array<number>(), new Array<number>(), new Array<number>());

  // 2. Loop through the image buffer and extract the R, G, and B channels
  for (let i = 0; i < imageBufferData.length; i += 4) {
    redArray.push(imageBufferData[i]);
    greenArray.push(imageBufferData[i + 1]);
    blueArray.push(imageBufferData[i + 2]);
    // skip data[i + 3] to filter out the alpha channel
  }

  // 3. Concatenate RGB to transpose [224, 224, 3] -> [3, 224, 224] to a number array
  const transposedData = redArray.concat(greenArray).concat(blueArray);

  // 4. convert to float32
  let i, l = transposedData.length; // length, we need this for the loop
  // create the Float32Array size 3 * 224 * 224 for these dimensions output
  const float32Data = new Float32Array(dims[1] * dims[2] * dims[3]);
  for (i = 0; i < l; i++) {
    float32Data[i] = transposedData[i] / 255.0; // convert to float
  }
  // 5. create the tensor object from onnxruntime-web.
  const inputTensor = new Tensor("float32", float32Data, dims);
  return inputTensor;
}

modelHelper.ts

inputTensor 已准备好进行推理。让我们调用默认的 modelHelper.ts 函数并逐步了解其逻辑。首先,我们通过传入模型路径和 SessionOptions 来创建 ort.InferenceSession。对于 executionProviders,您可以使用 webgl 来使用 GPU,或使用 wasm 来使用 CPU。要了解有关可用于推理配置的 SessionOptions 的更多信息,请参阅此处的文档。

import * as ort from 'onnxruntime-web';
import _ from 'lodash';
import { imagenetClasses } from '../data/imagenet';

export async function runSqueezenetModel(preprocessedData: any): Promise<[any, number]> {
  // Create session and set options. See the docs here for more options: 
  //https://runtime.onnx.org.cn/docs/api/js/interfaces/InferenceSession.SessionOptions.html#graphOptimizationLevel
  const session = await ort.InferenceSession
                          .create('./_next/static/chunks/pages/squeezenet1_1.onnx', 
                          { executionProviders: ['webgl'], graphOptimizationLevel: 'all' });
  console.log('Inference session created');
  // Run inference and get results.
  var [results, inferenceTime] =  await runInference(session, preprocessedData);
  return [results, inferenceTime];
}

然后我们调用 runInference 函数,传入 session 和我们的输入张量 preprocessedData

async function runInference(session: ort.InferenceSession, preprocessedData: any): Promise<[any, number]> {
  // Get start time to calculate inference time.
  const start = new Date();
  // create feeds with the input name from model export and the preprocessed data.
  const feeds: Record<string, ort.Tensor> = {};
  feeds[session.inputNames[0]] = preprocessedData;
  // Run the session inference.
  const outputData = await session.run(feeds);
  // Get the end time to calculate inference time.
  const end = new Date();
  // Convert to seconds.
  const inferenceTime = (end.getTime() - start.getTime())/1000;
  // Get output results with the output name from the model export.
  const output = outputData[session.outputNames[0]];
  //Get the softmax of the output data. The softmax transforms values to be between 0 and 1
  var outputSoftmax = softmax(Array.prototype.slice.call(output.data));
  //Get the top 5 results.
  var results = imagenetClassesTopK(outputSoftmax, 5);
  console.log('results: ', results);
  return [results, inferenceTime];
}

推理完成后,我们返回前 5 个结果以及运行推理所需的时间。然后,这些信息将显示在 ImageCanvas Web 组件上。

data 文件夹

此模板中的 data 文件夹包含 imagenetClasses,用于根据推理结果索引分配标签。此外,还提供了 sample-image-urls.ts 用于测试应用程序。

ImageCanvas FSX 元素 Web 组件

ImageCanvas.tsx Web 组件包含按钮和显示元素。下面是 Web 组件的逻辑:

import { useRef, useState } from 'react';
import { IMAGE_URLS } from '../data/sample-image-urls';
import { inferenceSqueezenet } from '../utils/predict';
import styles from '../styles/Home.module.css';

interface Props {
  height: number;
  width: number;
}

const ImageCanvas = (props: Props) => {

  const canvasRef = useRef<HTMLCanvasElement>(null);
  var image: HTMLImageElement;
  const [topResultLabel, setLabel] = useState("");
  const [topResultConfidence, setConfidence] = useState("");
  const [inferenceTime, setInferenceTime] = useState("");
  
  // Load the image from the IMAGE_URLS array
  const getImage = () => {
    var sampleImageUrls: Array<{ text: string; value: string }> = IMAGE_URLS;
    var random = Math.floor(Math.random() * (9 - 0 + 1) + 0);
    return sampleImageUrls[random];
  }

  // Draw image and other  UI elements then run inference
  const displayImageAndRunInference = () => { 
    // Get the image
    image = new Image();
    var sampleImage = getImage();
    image.src = sampleImage.value;

    // Clear out previous values.
    setLabel(`Inferencing...`);
    setConfidence("");
    setInferenceTime("");

    // Draw the image on the canvas
    const canvas = canvasRef.current;
    const ctx = canvas!.getContext('2d');
    image.onload = () => {
      ctx!.drawImage(image, 0, 0, props.width, props.height);
    }
    // Run the inference
    submitInference();
  };

  const submitInference = async () => {

    // Get the image data from the canvas and submit inference.
    var [inferenceResult,inferenceTime] = await inferenceSqueezenet(image.src);

    // Get the highest confidence.
    var topResult = inferenceResult[0];

    // Update the label and confidence
    setLabel(topResult.name.toUpperCase());
    setConfidence(topResult.probability);
    setInferenceTime(`Inference speed: ${inferenceTime} seconds`);

  };

  return (
    <>
    <button
      className={styles.grid}
      onClick={displayImageAndRunInference} >
      Run Squeezenet inference
    </button>
    <br/>
    <canvas ref={canvasRef} width={props.width} height={props.height} />
    <span>{topResultLabel} {topResultConfidence}</span>
    <span>{inferenceTime}</span>
    </>
  )
};

export default ImageCanvas;

然后将此 Web 组件元素导入到 index.tsx 中。

<ImageCanvas width={240} height={240}/>

next.config.js

我们需要在 next.config.js 中添加几个插件。这是 NextJS 框架中实现的 webpack 配置。CopyPlugin 用于将 wasm 文件和模型文件夹文件复制到用于部署的 out 文件夹。

/** @type {import('next').NextConfig} */
const NodePolyfillPlugin = require("node-polyfill-webpack-plugin");
const CopyPlugin = require("copy-webpack-plugin");

module.exports = {
  reactStrictMode: true,
  //distDir: 'build',
  webpack: (config, {  }) => {

    config.resolve.extensions.push(".ts", ".tsx");
    config.resolve.fallback = { fs: false };

    config.plugins.push(
    new NodePolyfillPlugin(), 
    new CopyPlugin({
      patterns: [
        {
          from: './node_modules/onnxruntime-web/dist/ort-wasm-simd-threaded.wasm',
          to: 'static/chunks/pages',
        },             {
          from: './node_modules/onnxruntime-web/dist/ort-wasm-simd-threaded.mjs',
          to: 'static/chunks/pages',
        },          
          {
            from: './model',
            to: 'static/chunks/pages',
          },
        ],
      }),
    );

    return config;
  } 
}

package.json

由于我们希望将其部署为静态站点,我们需要将 package.json 中的构建命令更新为 next build && next export,以生成我们的静态站点输出。这会生成部署静态站点所需的所有资产,并将它们放入 out 文件夹中。

{
  "name": "ort-web-template",
  "version": "0.1.0",
  "private": true,
  "scripts": {
    "dev": "next dev",
    "build": "next build && next export",
    "start": "next start",
    "lint": "next lint"
  },
  "dependencies": {
    "fs": "^0.0.1-security",
    "jimp": "^0.16.1",
    "lodash": "^4.17.21",
    "ndarray": "^1.0.19",
    "ndarray-ops": "^1.2.2",
    "next": "^11.1.2",
    "onnxruntime-web": "^1.9.0",
    "react": "17.0.2",
    "react-dom": "17.0.2"
  },
  "devDependencies": {
    "node-polyfill-webpack-plugin": "^1.1.4",
    "copy-webpack-plugin": "^9.0.1",
    "@types/lodash": "^4.14.176",
    "@types/react": "17.0.19",
    "eslint": "7.32.0",
    "eslint-config-next": "11.1.0",
    "typescript": "4.4.2"
  }
}

在本地运行项目

我们准备好运行项目了。根据您是要启动调试、构建 out 文件夹,还是不带调试启动,运行相应的命令。

// to run with debugging
npm run dev
// to build the project
npm run build
// to run without debugging
npm run start

部署到 Azure 静态 Web 应用

现在我们已经构建了站点,可以将其部署到 Azure 静态 Web 应用。请查看此处的文档,了解如何使用 Azure 进行部署。

TypeScript 笔记本

我们已经介绍了如何使用此模板,但这里还有一个额外的好处!在模板的 notebook 文件夹中,有一个 笔记本 包含这些代码,供您试验和尝试可能需要的更改。这样,如果您有不同的模型或图像要尝试,就可以非常轻松地完成。要使用 TypeScript Jupyter 笔记本,请下载 VS Code Jupyter 笔记本扩展。

更多资源

  • 立即前往 GitHub NextJS ORT-Web Template 仓库开始使用该模板。

  • 请查看此处的发布博客。

  • 该模板使用 NextJS,这是一个用于使用 ReactJS 构建应用程序的框架。

  • 查看 ONNX Runtime Web Demo 获取更多模型。ONNX Runtime Web demo 是一个交互式演示门户,展示了在 VueJS 中运行 ONNX Runtime Web 的实际用例。它目前支持四个示例,让您快速体验 ONNX Runtime Web 的强大功能。

  • 这篇博客展示了如何将 ORT Web 与 Python 结合使用,将预训练的 AlexNet 模型部署到浏览器中。

  • 查看更多 ONNX Runtime JS 示例