hy clear Blog

【TensorFlow.js】ブラウザで物体検出をする(React)

2024/07/19

2024/07/19

📰 アフィリエイト広告を利用しています

はじめに

TensorFlow.jsで物体検出をブラウザで行います。今回は公式のモデルを使用して、画像の物体検知をしました。
作成したサイトは以下のURLで確認できます。最終的な目標は独自のモデルで動画の物体検知とトラッキングを行うことです。できるのかな?

デモサイト

https://dev.hy-clear.com/webapp/object_detection

参考にした公式サンプル

https://github.com/tensorflow/tfjs-models/tree/master/coco-ssd/demo

※ ニューラルネットワーク周りに関する知識がほぼないのでおかしなことを書いている可能性があります。自己責任でお願いします。

準備

ReactとTypescriptで作ります。設定は省略します。
まず、TensorFlow.jsとモデルをinstallします。

npm install @tensorflow/tfjs @tensorflow-models/coco-ssd

img でsrcに指定した画像を推論し、結果を canvas に設定する。
画像はお好きなもので

App.tsx
import { useRef } from "react"

function App() {
  const imgRef = useRef<HTMLImageElement>(null)

  const predict = async () => {}
  return (
    <>
      <div>
        <img ref={imgRef} src={"test_image.jpg"} />
        <button onClick={()=> predict()}>Predict</button>
        <canvas id="canvas"></canvas>
      </div>
    </>
  )
}

export default App

モデルのロード

必要なライブラリをインポートします。tfjs-backend はインポートしないとError: No backend found in registry.がでます。

App.tsx
import '@tensorflow/tfjs-backend-cpu';
import '@tensorflow/tfjs-backend-webgl';

import * as cocoSsd from "@tensorflow-models/coco-ssd"

モデルをロードする処理。モデルは使いまわしたいので変数に代入しておく。

App.tsx
  var model: cocoSsd.ObjectDetection | null = null

  const loadModel = async () => {
    model = await cocoSsd.load()
  }

  const predict = async () => {
    if (model === null) {
      await loadModel()
    }
  }

推論する

img要素から画像を取得し、推論した結果を取得します。

const img = imgRef.current as HTMLImageElemen
//推論にかかる時間を測定する
const startTime = performance.now()
const predictions = await model!.detect(img)
// かかった時間を出力
console.log(performance.now() - startTime)

Canvasに表示する

推論した結果を反映した画像をCanvasに表示します。
まず、画像サイズに合わせたCanvasを準備します。

// canvasの要素を取得する
const canvas = document.getElementById("canvas") as HTMLCanvasElement;

// canvas要素のサイズを画像と合わせる
canvas.width = img.width
canvas.height = img.height

// canvasの二次元描画コンテキストを取得
const context = canvas.getContext("2d") as CanvasRenderingContext2D;
// canvasのサイズに合わせて画像サイズを設定
var scaleFactor = Math.min(canvas.width / img.width, canvas.height / img.height);
var scaledWidth = img.width * scaleFactor;
var scaledHeight = img.height * scaleFactor;
context.drawImage(img, 0, 0, scaledWidth, scaledHeight)

これで画像と同じサイズのcanvas要素が作成でき、画像が表示できます。

参考:CanvasRenderingContext2D

https://developer.mozilla.org/ja/docs/Web/API/CanvasRenderingContext2D

結果を表示する

canvasに結果を表示していきます。
predictions に結果の配列があります。配列の要素は以下です。

  • bbox検出した位置情報
  • score検出したスコア
  • class検出した物体名

for (let i = 0; i < predictions.length; i++) {
  context.beginPath();
  context.rect(...predictions[i].bbox)
  context.lineWidth = 1;
  context.strokeStyle = "red";
  context.fillStyle = "red"
  context.stroke();
  context.fillText(
    predictions[i].score.toFixed(3) + " " + predictions[i].class, predictions[i].bbox[0],
    predictions[i].bbox[1] > 10 ? predictions[i].bbox[1] - 5 : 10
  )

最後に

公式のサンプル通りでは割と簡単に表示できました。
独自のモデルで試してはいるのですが、難しそうです。

すべてのコード
App.tsx
import { useRef } from "react"
import '@tensorflow/tfjs-backend-cpu';
import '@tensorflow/tfjs-backend-webgl';

import * as cocoSsd from "@tensorflow-models/coco-ssd"

function App() {
  const imgRef = useRef<HTMLImageElement>(null)
  var model: cocoSsd.ObjectDetection | null = null

  const loadModel = async () => {
    model = await cocoSsd.load()
  }

  const predict = async () => {
    if (model === null) {
      await loadModel()
    }
    const img = imgRef.current as HTMLImageElement
    //推論にかかる時間を測定する
    const startTime = performance.now()
    const predictions = await model!.detect(img)
    // かかった時間を出力
    console.log(performance.now() - startTime)

    // canvasの要素を取得する
    const canvas = document.getElementById("canvas") as HTMLCanvasElement;

    // canvas要素のサイズを画像と合わせる
    canvas.width = img.width
    canvas.height = img.height

    // canvasの二次元描画コンテキストを取得
    const context = canvas.getContext("2d") as CanvasRenderingContext2D;
    // canvasのサイズに合わせて画像サイズを設定
    var scaleFactor = Math.min(canvas.width / img.width, canvas.height / img.height);
    var scaledWidth = img.width * scaleFactor;
    var scaledHeight = img.height * scaleFactor;
    context.drawImage(img, 0, 0, scaledWidth, scaledHeight)

    for (let i = 0; i < predictions.length; i++) {
      context.beginPath();
      context.rect(...predictions[i].bbox)
      context.lineWidth = 1;
      context.strokeStyle = "red";
      context.fillStyle = "red"
      context.stroke();
      context.fillText(
        predictions[i].score.toFixed(3) + " " + predictions[i].class, predictions[i].bbox[0],
        predictions[i].bbox[1] > 10 ? predictions[i].bbox[1] - 5 : 10
      )
    }
  }
  return (
    <>
      <div>
        <img ref={imgRef} src={"test_image.jpg"} />
        <div>
          <button onClick={()=> predict()}>Predict</button>
        </div>
        <canvas id="canvas"></canvas>
      </div>
    </>
  )
}

export default App