はやし雑記

はやしです

トップ > 技術系 > 顔検出 (CenterFace) をブラウザのWASMで動かした

顔検出 (CenterFace) をブラウザのWASMで動かした

技術系

前回の続きです。

blog.hayashikun.com

この記事では、前回の記事で実装したCenterFaceをwasmで動くようにして、それをブラウザで動かしたことについて書きます。

以下のように、Webカメラの映像から顔を検出して、顔の上に緑の四角を描画するというのをやりました。

https://github.com/hayashikun/wasabi/blob/master/web-wasm.gif?raw=true

今回実装したものは、こちらにデプロイしているので、PCやスマホから試すことができます。

wasbi

試してみるとわかると思いますが、結構軽快に動作しています。

リポジトリは前回と同じです。

wasmとして動かしているのが web-wasm です。前回同様にRustで書いています。 github.com

web のフロントエンドが web です。Typescriptで書いていて、webpackでバンドルしました。 github.com

Rust実装

今回、Rustで実装したのはこれだけです。

use console_error_panic_hook;
use serde::{Deserialize, Serialize};
use tract_onnx::prelude::Tensor;
use tract_onnx::prelude::tract_ndarray::Array4;
use wasm_bindgen::prelude::*;

use wasabi::center_face::{CenterFace, Face};

#[wasm_bindgen]
pub struct App {
    width: u32,
    height: u32,
    cf: CenterFace,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct DetectResult {
    faces: Vec<Face>,
}

#[wasm_bindgen]
impl App {
    pub fn new(width: u32, height: u32) -> App {
        console_error_panic_hook::set_once();

        App {
            width,
            height,
            cf: CenterFace::new(width, height).unwrap(),
        }
    }

    pub fn detect(&self, array: Vec<u8>) -> String {
        let w = self.width as usize;
        let h = self.height as usize;
        let image: Tensor = Array4::from_shape_fn(
            (1, 3, h, w),
            |(_, c, y, x)| {
                array[(y * w + x) * 4 + c] as f32
            },
        ).into();
        let faces: Vec<Face> = self.cf.detect(image).unwrap();
        let result = DetectResult { faces };

        serde_json::to_string(&result).unwrap_or("{}".to_string())
    }
}

#[wasm_bindgen] を付けた App は、JS/TSから呼び出すことができます。 App::new の際に、CenterFace::new をしています。

App.detect では、入力された画像情報を Tensor に変換して、それを CenterFace.detect に与えています。 後述しますが、フロント側ではcanvas要素のピクセルデータをJS/TSで取得しています。 canvas要素のピクセル情報は、一次元配列で rgba となっているので、array[(y * w + x) * 4 + c]としています。

推論の結果は、一旦JSON文字列に変換して返しています。

依存関係は以下のとおりです。

[dependencies]
tract-onnx = "0.16.5"
wasabi = { path = ".." }
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
wasm-bindgen = "0.2.80"
console_error_panic_hook = "0.1.7"

wasabi が前回実装した CenterFace を含むcrateです。 wasm-bindgen はJSから関数やstructを扱えるようにしてくれるcrateです。 console_error_panic_hook は、エラーメッセージをわかりやすくしてくれるものです。

フロントエンド

HTML / Typescript 実装

TSで App をimportして初期化するのはとても簡単です。

import * as wasabi from "../../web-wasm/pkg";

let app = wasabi.App.new(width, height)

wasmとtypescriptの連携は、後述するWasmPackがやってくれるので、複雑なことは何もせずにstructや関数を取り扱うことができます。

HTMLの主要部分は以下のようになっています。

<style>
    #video {
        display: none;
    }

    #hidden-canvas {
        display: none;
    }
</style>

<video id="video" autoplay="autoplay"></video>
<canvas id="hidden-canvas"></canvas>
<canvas id="display-canvas"></canvas>

今回はvideo 要素とcanvas要素を2つ置いています。ただし、表示するのは1つのcanvas要素のみです。

index.ts は以下の通りです。

import * as wasabi from "../../web-wasm/pkg";

const width = 32 * 4 * 2
const height = 32 * 3 * 2
const scale = 3

let app = wasabi.App.new(width, height)

let video = document.getElementById("video") as HTMLVideoElement;
navigator.mediaDevices.getUserMedia({
    audio: false,
    video: {
        width: {ideal: width * scale},
        height: {ideal: height * scale}
    }
}).then(function (stream) {
    video.srcObject = stream;
});

let hCanvas = document.getElementById("hidden-canvas") as HTMLCanvasElement;
hCanvas.width = width
hCanvas.height = height
let hCtx = hCanvas.getContext('2d')

let dCanvas = document.getElementById("display-canvas") as HTMLCanvasElement;
dCanvas.width = width * scale
dCanvas.height = height * scale
let dCtx = dCanvas.getContext('2d')

updateCanvas();

function updateCanvas() {
    hCtx!.drawImage(video, 0, 0, hCanvas.width, hCanvas.height);
    dCtx!.drawImage(video, 0, 0, dCanvas.width, dCanvas.height);
    let data = hCtx!.getImageData(0, 0, width, height).data
    let faces: [any] | undefined = JSON.parse(app.detect(Uint8Array.from(data))).faces;
    if (faces) {
        faces.forEach((f) => {
            dCtx!.beginPath();
            dCtx!.rect(
                f.x1 * scale, f.y1 * scale,
                (f.x2 - f.x1) * scale, (f.y2 - f.y1) * scale
            );
            dCtx!.fillStyle = "#00FF00AA";
            dCtx!.fill();
            dCtx!.closePath();
        })
    }

    requestAnimationFrame(updateCanvas)
}

まず、videoにWebカメラの映像を投影します。

そして、videoのイメージをcanvas#display-canvas に表示します。 canvas#display-canvas は表示するので、実際の大きさにします。

canvas#hidden-canvas にもvideoのイメージを写します。 canvas#hidden-canvas は canvas#display-canvas よりも小さくします。 そして、canvas#hidden-canvasの縦横のサイズはそれぞれ32の倍数で、wasabi.App.new に渡した値にします。

なぜこのようなことをしているかというと、Webカメラからの映像は、推論にそのままかけるには大きすぎて時間がかかり、fpsが下がるからです。 逆に、推論にかける大きさにしたcanvas#hidden-canvas は、解像度が低く、見た目が悪いからです。

推論の結果の座標値は拡大させて、canvas#display-canvas上に緑のマスクを表示しています。

今は推論させる画像サイズは256x192にしていますが、やはり若干のカクつきはあります。 顔が一つだけで、それなりに大きければ128x96でもなんとか検出できます。 また、全てのフレームで推論させていますが、多分半分くらい間引いてもそんなに違和感はないと思います。

wasm-pack-plugin

Wasmとそのbindingはwebpackでwasm-pack-pluginを使用して作ります。

webpack.config.js は以下のようになります。

const path = require("path");
const HtmlWebpackPlugin = require("html-webpack-plugin");
const WasmPackPlugin = require("@wasm-tool/wasm-pack-plugin");

module.exports = {
    resolve: {
        extensions: [".ts", ".tsx", ".js", ".jsx", ".wasm"]
    },
    mode: "development",
    module: {
        rules: [
            {
                test: /\.tsx?$/,
                loader: "ts-loader",
                options: {
                    transpileOnly: true
                }
            }
        ]
    },
    plugins: [
        new HtmlWebpackPlugin({
            template: path.join(__dirname, "src/index.html")
        }),
        new WasmPackPlugin({
            crateDirectory: path.join(__dirname, "../web-wasm")
        })
    ],
    experiments: {
        asyncWebAssembly: true,
    },
};

WasmPackPlugin にcrateのpathを指定すれば、buildなど含めて全てやってくれます。

なので、自らwab-wasmのcrateをbuildする必要なく、webpack serveとすればdevサーバーが起動してbuildまでやってくれます。 webpack --mode production とすれば、全てを同梱したdistを吐いてくれるので、それを適当におけば簡単にdeployできます。

Web Frontendで機械学習モデルを動かすことについて

かなり前に、onnx-js や tensorflow.js を使って、Web Frontendで機械学習モデルを動かすというのをやっていました。

blog.hayashikun.com

blog.hayashikun.com

これらと比べても、今回やったようにrustでonnxを含んだwasmを作って、それをWeb Frontendで動かすというのはそこまで大変ではなかったです。 動作については、onnx-js （今は onnxruntime-webになった）や tensorflow.js で同様のことをやって比べてみるというのをやっていないのでなんとも言えないですが、少なくとも今回のモデルで比較的軽快に動いていたので、動作としても問題無いのでは無いかと思います。

以上！ 次回はいよいよproxy-wasmで動かします。