zukucode
主にWEB関連の情報を技術メモとして発信しています。

Tensorflow.js 学習済みモデルを使って画像分類を行う

TensorFlow.jsの学習済みモデルを使って画像分類を行う方法を紹介します。

今回はMobileNetというライブラリ(学習済みモデル)を使用します。

https://github.com/tensorflow/tfjs-models/tree/master/mobilenet

パッケージのインストール

以下のコマンドで必要なパッケージをインストールします。

なお、@tensorflow/tfjsTensorflow.js本体)のインストールは不要です。

$ npm install @tensorflow-models/mobilenet

画像の分類

以下のように、予測を行います。

const mobilenet = require('@tensorflow-models/mobilenet');

async function predict() {
  // モデルの読み込み
  const model = await mobilenet.load();
  // 分類する画像要素を取得
  const img = document.getElementById('img'); 
  // 予測
  const predictions = await model.classify(img, 5);
}

// 取得結果(predictions)
[
  {
    "className":"tabby, tabby cat",
    "probability":0.8397181630134583
  },
  {
    "className":"Egyptian cat",
    "probability":0.11327312886714935
  },
  {
    "className":"tiger cat",
    "probability":0.02227550931274891
  },
  {
    "className":"lynx, catamount",
    "probability":0.012203224934637547
  },
  {
    "className":"Cardigan, Cardigan Welsh corgi",
    "probability":0.0015334360068663955
  },
]

classifyメソッドの第一引数には、画像(img要素)だけでなく、キャンバス(canvas要素)を指定することも可能です。

第2引数には分類結果の取得件数を指定します。指定した上位X件の結果が取得されます。(この引数は省略可能で、省略時は3が設定されます)


関連記事