Yolov7をtfjsでの実装方法の解説

Yolov7をtfjsで試したのでメモ
モデルのコンバートとJSでの処理について書きました。
Non-Maximum Suppression(NMS)は未実装です。

PytorchからONNXへのConvert

PytorchのモデルからONNXに変換して動かすPythonコードは本家に置いてあります。
コンバートから試す方はこちらを参考ください。

tfjsにコンバート済みのモデルはこちらにあります。

ONNX以降のtfjsへのConvert手順

onnx-tfと、tensorflowを実行できる環境を用意します。
protobufのバージョンがコンフリクトするので環境は分けて作るのがオススメです。
なお、数少ないM1 Macユーザーはtfjsを入れるのは非常に手間です。こちらを参考ください。

実行環境ができればコマンドを叩けば直ぐに作れます。
オプションで半精度(fp16)を指定するのは精度の影響が少ない割にファイルサイズが半分になるのでオススメです。

// onnx to tensorflow
$ onnx-tf convert -i yolov7-tiny.onnx -o yolov7-tiny_tensorflow

// tensorflow to tfjs
$ tensorflowjs_converter \
--input_format tf_saved_model yolov7-tiny_tensorflow ./yolov7-tiny_tfjs_fp16 \
--quantize_float16

実装方法

画像の読み込み、前処理、推論、後処理を実装します。

画像の読み込み

// canvasに処理したい画像を読み込む
// getImageDataでUint8ClampedArrayを取り出す
async function loadImage(src) {
    const imgBlob = await fetch(src).then((resp) => resp.blob());
    const img = await createImageBitmap(imgBlob);
    const canvas = document.createElement('canvas');
    canvas.width = img.width;
    canvas.height = img.height;
    const ctx = canvas.getContext('2d');
    ctx.drawImage(img, 0, 0);
    console.log(img.width, img.height);
    return ctx.getImageData(0, 0, img.width, img.height);
}

画像の前処理

// padAndResizeToで画像をモデルの入力と同じ縦横比になるようアスペクト比でリサイズして0埋めする
const {resized, ratio, padding} = padAndResizeTo(image, 640, 640);
// 返り値のtensorを255で割ってモデルの入力になるよう転置する
const x = resized.div(255).transpose([2,0,1]).reshape([1, 3, 640, 640]);

padAndResizeToはこちらから拾ってきました。

推論

const feedDict = {"images": x };
let preds = await detModel.executeAsync(feedDict, "output");

後処理

// 推論結果をスライスしてバウンディングボックスの座標を取り出す
const data = await preds.data();
for (let i=0; i<data.length / 7; i++){
const [ batch_id,x0,y0,x1,y1,cls_id,score ] = data.slice(i*7, i*7+7);
let box = [
  (x0 - padding.left) / ratio,
  (y0 - padding.bottom) / ratio,
  (x1 - padding.left) / ratio,
  (y1 - padding.bottom) / ratio,
];
}

コードとサンプル

以上です。

その他

別で姿勢推定のトップダウンの処理を実装しました。
物体検出の後に姿勢推定処理を走らせると、M1 Macのブラウザで3fpsしか出ないのでとても悲しい。
今回の物体検出の処理はこの成果物の派生になります。

サンプル実装（カメラが起動するのでご注意ください）