TF.js 実践開発レシピ

TensorFlow.jsで画像の特徴量を抽出し類似画像を検索する方法:Pythonの知識を応用する

Tags: TensorFlow.js, 画像認識, 特徴抽出, 類似度検索, JavaScript

画像認識AI開発の応用として、特定の画像と類似した画像を大量のデータの中から見つけ出す「画像類似度検索」は非常に有用な技術です。この技術は、ECサイトでの類似商品検索、コンテンツプラットフォームでの関連画像提案、著作権侵害チェックなど、様々な場面で活用されています。

Pythonと機械学習ライブラリ(TensorFlow/Keras等)を使った開発に慣れている方であれば、画像の特徴量を抽出する手法や、抽出した特徴量を使って類似度を計算する概念には馴染みがあるかもしれません。例えば、事前学習済みのCNNモデルの中間層の出力を特徴量として利用したり、抽出した特徴量ベクトル間の距離(ユークリッド距離やコサイン類似度など)を計算したりといった手法です。

本記事では、これらのPythonでの経験を活かしながら、TensorFlow.jsを用いてブラウザ上で画像の特徴量を抽出し、類似画像を検索する具体的な方法を解説します。Pythonで実現していた処理を、どのようにJavaScriptとTensorFlow.jsで実装するのか、具体的なコード例を交えてご紹介します。

画像の特徴量抽出とは

画像認識タスクにおいて、生のピクセル値をそのまま扱うことは困難かつ効率的ではありません。より抽象的で、画像の持つ本質的な情報を捉えた表現を用いることが一般的です。この「画像の持つ本質的な情報を捉えた表現」が、画像の特徴量や埋め込みベクトル(Embedding)と呼ばれるものです。

畳み込みニューラルネットワーク(CNN)は、この特徴量抽出において非常に強力な手法です。CNNは層を重ねるごとに、エッジやコーナーといった低レベルな特徴から、物体の部分やテクスチャといった高レベルな特徴へと、段階的に情報を抽象化していきます。

通常、画像分類タスクでは、CNNの最終層で特定のクラスへの確率を出力しますが、特徴量抽出のためには、この最終分類層の手前の層、例えばプーリング層や全結合層の直前の層の出力を利用します。この出力は、元の画像が持つ多様な情報を凝縮したベクトル表現となっており、異なる画像間での類似度を計算するために使用できます。

Python/Kerasでは、以下のようなコードで特定の層の出力を取得するモデルを構築することが可能です。

import tensorflow as tf
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.models import Model

# 事前学習済みMobileNetV2モデルをロード(分類層は含まない)
base_model = MobileNetV2(weights='imagenet', include_top=False, pooling='avg')

# 特徴量として使用する層を指定(ここではGlobalAveragePooling2Dの出力)
# KerasのFunctional APIを使って新しいモデルを定義
feature_extractor_model = Model(inputs=base_model.input, outputs=base_model.output)

# 特徴量を抽出したい画像(例: image_tensor)に対して予測を実行
# features = feature_extractor_model.predict(image_tensor)

TensorFlow.jsにおいても、これと同様のアプローチで事前学習済みモデルの特定層の出力を取得し、特徴量として利用することができます。

TensorFlow.jsでの特徴量抽出の実装

TensorFlow.jsで画像の特徴量を抽出するには、以下のステップを踏みます。

  1. 事前学習済みモデルのロード: tf.loadGraphModel または tf.loadLayersModel を使用して、特徴量抽出に適した事前学習済みモデル(例: MobileNet, EfficientNetなど)をロードします。通常、分類タスクのために学習されたモデルを利用しますが、最終の分類層は特徴量として不要なため、これを含まないモデルを使用するか、あるいは特定の層までを切り出して利用します。
  2. 特徴量抽出用モデルの構築: ロードしたモデルから、特徴量として使用したい層までの部分を切り出して、新しいモデルを構築します。これはPython/Kerasの tf.keras.models.Model をFunctional API的に使用するのと概念的に類似しています。
  3. 画像データの前処理: 特徴量抽出を行う前に、入力画像をモデルが要求する形式に変換する必要があります。これには、画像のリサイズ、正規化(ピクセル値のスケール調整など)、そしてTensorFlow.jsのTensor形式への変換が含まれます。
  4. 特徴量の抽出: 前処理した画像を構築した特徴量抽出用モデルに入力し、その出力を取得します。これが画像の特徴量ベクトルとなります。

具体的なコード例を見てみましょう。ここでは、TensorFlow.jsで利用可能な事前学習済みモデルである mobilenet_v2 を例に、特徴量抽出モデルを構築し、画像から特徴量を抽出する処理を記述します。

import * as tf from '@tensorflow/tfjs';
// 必要に応じて mobilenet をインポート(例:@tensorflow-models/mobilenet がインストールされている場合)
// import * as mobilenet from '@tensorflow-models/mobilenet';

// 事前学習済みモデル(MobileNetV2の特定バージョン)のURLを指定
// 通常、feature extraction用途では include_top: false のモデルを使用
const MOBILENET_V2_URL = 'https://tfhub.dev/google/tfjs-model/imagenet/mobilenet_v2_100_224/feature_vector/2/default/1';

let featureExtractorModel = null;

/**
 * 特徴量抽出モデルをロードする関数
 * @returns {Promise<tf.LayersModel>} ロードされたモデル
 */
async function loadFeatureExtractorModel() {
  if (featureExtractorModel) {
    return featureExtractorModel;
  }
  console.log('Loading feature extraction model...');
  // TF Hubからfeature_vectorモデルをロード
  // このモデルは、GlobalAveragePoolingD2の直前のConv層の出力を特徴量として返します
  featureExtractorModel = await tf.loadGraphModel(MOBILENET_V2_URL, { fromTFHub: true });
  console.log('Model loaded.');
  return featureExtractorModel;
}

/**
 * HTMLImageElementから特徴量を抽出する関数
 * @param {HTMLImageElement} imgElement 処理対象の画像要素
 * @returns {Promise<tf.Tensor>} 抽出された特徴量ベクトル
 */
async function extractFeaturesFromImage(imgElement) {
  await loadFeatureExtractorModel();

  // 画像データの前処理
  // 1. Tensorに変換
  // 2. モデルの入力サイズにリサイズ(例: MobileNetV2は224x224)
  // 3. チャンネルをRGB順にする(必要な場合)
  // 4. 0-1スケールまたは -1-1スケールに正規化(モデルによる)
  // 5. バッチ次元を追加 (shape: [1, height, width, channels])
  const imgTensor = tf.browser.fromPixels(imgElement).toFloat();
  const resizedTensor = tf.image.resizeBilinear(imgTensor, [224, 224]);

  // MobileNetV2 feature_vector モデルは入力が -1 から 1 の範囲を期待します
  // fromPixelsは0-255なので、まず0-1に、次に-1-1にスケール変換
  const offset = tf.scalar(127.5);
  const normalizedTensor = resizedTensor.sub(offset).div(offset);

  // バッチ次元を追加
  const batchedTensor = normalizedTensor.expandDims(0);

  // 特徴量の抽出
  // tf.tidy でメモリリークを防ぐ
  const features = tf.tidy(() => {
    const output = featureExtractorModel.predict(batchedTensor);
    // 出力が複数のテンソルになるモデルもあるので、必要に応じて適切なテンソルを選択
    // feature_vectorモデルの場合、出力は単一のテンソル (shape: [1, feature_size])
    return output;
  });

  // 元のTensorや中間Tensorを解放
  imgTensor.dispose();
  resizedTensor.dispose();
  normalizedTensor.dispose();
  batchedTensor.dispose();

  // 抽出された特徴量テンソルを返す
  return features;
}

// 使用例:
// <img id="myImage" src="path/to/image.jpg"> という要素があるとして
/*
const imgElement = document.getElementById('myImage');
extractFeaturesFromImage(imgElement).then(features => {
  console.log('Extracted features shape:', features.shape); // 例: [1, 1280]
  // 抽出された特徴量テンソル (features) を使って類似度計算などを行います
  // features.dispose(); // 使い終わったら必ず解放
});
*/

上記のコードでは、tf.browser.fromPixels でHTML <img> 要素からTensorを作成し、tf.image.resizeBilinear でモデルの入力サイズにリサイズしています。正規化ステップでは、subdiv を使ってピクセル値を -1 から 1 の範囲に変換しています。最後に expandDims(0) でバッチ次元を追加し、featureExtractorModel.predict に入力して特徴量を得ています。tf.tidy を使用することで、そのスコープ内で生成された中間テンソルが自動的に解放され、メモリリークを防ぐことができます。抽出された最終的な特徴量テンソルは手動で解放する必要があります。

Python/Kerasでの model.predict(image_tensor) と同様に、TensorFlow.jsでも model.predict(batchedTensor) でモデルの出力を取得できます。feature_vector モデルは最終分類層を持たないため、出力がそのまま特徴量ベクトルとなります。

類似度計算の実装

画像から特徴量ベクトルが得られたら、次はそれらを使って画像間の類似度を計算します。類似度を測る尺度としては、コサイン類似度(Cosine Similarity)やユークリッド距離(Euclidean Distance)などが一般的です。画像の特徴量ベクトルに対しては、多くの場合コサイン類似度が良い結果をもたらす傾向があります。

コサイン類似度は、2つのベクトルがなす角のコサイン値で定義されます。2つのベクトル AB のコサイン類似度は、内積をそれぞれのベクトルのノルム(長さ)の積で割ることで計算されます。

Cosine Similarity(A, B) = (A・B) / (||A|| * ||B||)

値の範囲は -1 から 1 で、1 に近いほどベクトルが同じ方向を向いており(類似度が高い)、-1 に近いほど反対方向を向いています。画像特徴量の場合、通常は 0 から 1 の範囲の値になります。

TensorFlow.jsでは、Tensor演算としてこれらの計算を行うことができます。

import * as tf from '@tensorflow/tfjs';

/**
 * 2つの特徴量テンソル間のコサイン類似度を計算する関数
 * @param {tf.Tensor} featuresA 1つ目の特徴量ベクトル (shape: [1, feature_size])
 * @param {tf.Tensor} featuresB 2つ目の特徴量ベクトル (shape: [1, feature_size])
 * @returns {Promise<number>} コサイン類似度 (0〜1の値)
 */
async function calculateCosineSimilarity(featuresA, featuresB) {
  // バッチ次元 ([1, feature_size]) を削除し、[feature_size] のベクトルにする
  const vectorA = featuresA.squeeze();
  const vectorB = featuresB.squeeze();

  // コサイン類似度 = (A・B) / (||A|| * ||B||)
  const similarity = tf.tidy(() => {
    // 内積 (Dot product)
    const dotProduct = vectorA.dot(vectorB);

    // ベクトルのノルム (L2-norm)
    const normA = vectorA.norm();
    const normB = vectorB.norm();

    // ノルムの積
    const normProduct = normA.mul(normB);

    // ゼロ除算を避ける
    if (normProduct.dataSync()[0] === 0) {
        return tf.scalar(0); // あるいはNaN、文脈による
    }

    // コサイン類似度
    return dotProduct.div(normProduct);
  });

  // 計算結果をJavaScriptの数値として取得
  const similarityValue = (await similarity.data())[0];

  // 中間テンソルを解放
  vectorA.dispose();
  vectorB.dispose();
  similarity.dispose();

  return similarityValue;
}

// 使用例:
/*
// featuresA と featuresB は extractFeaturesFromImage で得られたテンソル
// calculateCosineSimilarity(featuresA, featuresB).then(similarity => {
//   console.log('Cosine Similarity:', similarity); // 例: 0.85
//   featuresA.dispose(); // 使い終わったら解放
//   featuresB.dispose(); // 使い終わったら解放
// });
*/

このコードでは、featuresA.squeeze() でバッチ次元を削除し、1次元の特徴量ベクトルに変換しています。そして、tf.dot で内積、tf.norm でL2ノルムを計算し、それらを使ってコサイン類似度を計算しています。data() メソッドはTensorの内容を非同期で取得し、JavaScriptのTypedArrayとして返します。そこから最初の要素(スカラー値)を取り出して数値としています。ここでも tf.tidy を活用して、計算過程で生成される中間テンソルを自動的に解放しています。

画像類似度検索の全体フロー

画像類似度検索システムを構築するための基本的なフローは以下のようになります。

  1. 特徴量のデータベース化: 検索対象となる全ての画像について、事前に特徴量を抽出しておきます。これらの特徴量ベクトルと、対応する画像への参照(URLなど)をペアにして、検索可能な形式で保存します。ブラウザ上で行う場合、IndexedDBなどが考えられますが、大量の画像を扱う場合はサーバーサイドでの処理や、専用のベクトルデータベース(例: Faiss, Annoy)の利用が現実的です。
  2. 検索クエリの特徴量抽出: ユーザーが類似画像を検索したい対象画像(クエリ画像)から、同様の手法で特徴量ベクトルを抽出します。
  3. 類似度計算とランキング: 抽出したクエリ画像の特徴量と、データベースに保存されている各画像の特徴量との間で類似度(コサイン類似度など)を計算します。計算された類似度の高い順に画像をランキングします。
  4. 結果の表示: ランキング上位の画像をユーザーに表示します。

本記事では、ステップ2と3(特徴量抽出と類似度計算)に焦点を当てたコード例を示しました。ステップ1のデータベース化は、アプリケーションの規模や要件によって様々な方法が考えられます。小規模なデモやオフラインでの利用であればIndexedDBでも可能ですが、パフォーマンスやスケーラビリティを考慮すると、サーバーサイドや専用ライブラリ/サービスとの連携が重要になります。

まとめと次のステップ

本記事では、TensorFlow.jsを使って画像の特徴量を抽出し、コサイン類似度を用いて画像間の類似度を計算する基本的な方法を解説しました。Python/Kerasでの特徴抽出のアプローチ(事前学習済みモデルの中間層利用)が、TensorFlow.jsでも同様に実現できることを示し、具体的なJavaScriptコード例をご紹介しました。

画像類似度検索は、抽出した特徴量と類似度計算が核となりますが、実践的なシステム構築には、大量の特徴量データの効率的な管理・検索、そしてパフォーマンスの最適化が不可欠です。

次のステップとして、以下の点について検討してみることをお勧めします。

TensorFlow.jsを活用することで、ブラウザ上やNode.js環境で高度な画像特徴抽出・類似度計算処理を実装することが可能です。Pythonでの機械学習開発経験で培った知識や概念は、TF.jsの世界でも大いに役立ちます。ぜひ、これらの技術を応用して、様々な画像認識アプリケーション開発に挑戦してみてください。