使用 ONNX Runtime Web 在 Web 应用中对图像进行分类

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

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

下图展示了使用此模板构建的网站外观。它循环遍历示例图像列表,使用 SqueezeNet 模型调用推理会话,然后返回推理结果的分数和标签。

模板输出示例

Image of browser inferencing on sample images.

目录

设备上推理

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

SqueezeNet 机器学习模型

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

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

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

模板

此模板的目标是为您的加速机器学习 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(红、绿、蓝)值放入数组中。为此,我们将通过每个像素的 4 个 RGBA 通道循环遍历 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];
}

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

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;

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

<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 Static Web Apps

现在我们已经构建了网站,可以将其部署到 Azure Static Web Apps。请查阅此处的文档,了解如何使用 Azure 进行部署。

TypeScript Notebook

我们已经讲解了如何使用此模板,这里还有一个额外内容!在模板的 notebook 文件夹下有一个 notebook,其中包含此代码,供您进行实验和尝试您可能需要的更改。这样,如果您想尝试不同的模型或图像,可以很容易地做到。要使用 TypeScript Jupyter notebook,请下载 VS Code Jupyter notebooks 扩展。

更多资源