DIVX テックブログ

catch-img

ResNetとOpenSearchでの類似間取り画像検索システムについて

はじめに

こんにちは、DIVX R&Dエンジニア兼広報室室長のyasunaです。

私の所属するR&D部では、AIを活用して新しい価値を生み出すための研究開発を日々行っています。その一環で、大量の画像データの中から似ているものを瞬時に探し出す、いわゆる「類似画像検索」の機能を実装する機会がありました。

今回は、畳み込みニューラルネットワークの代表的なモデルであるResNetと、全文検索エンジンとして知られるOpenSearchを組み合わせて、間取り画像の類似検索システムを構築した際の過程と考え方を整理してみようと思います。

この記事を読んで分かること

この記事では、特定の画像(今回は間取り図)に似た画像をデータベースから検索するシステムの、企画から実装、評価までの一連の流れを解説します。画像の特徴量抽出から、それを高速に検索するためのインデックス作成、そして実際に検索を実行するまでの具体的な手順を追うことができます。

背景

今回の目的は、「ある間取り画像をアップロードすると、それに似た間取り画像の候補をいくつか表示する」というシンプルな機能を実現することでした。これを実現するためには、まずコンピューターが画像の「類似性」を判断できる形に変換する必要があります。

その手法として、画像の特徴を数値のベクトルに変換するResNetという深層学習モデルを採用しました。そして、生成された大量のベクトルデータの中から、特定のベクトルに最も近いものを高速に見つけ出すために、OpenSearchのk-NN(k-nearest neighbor)プラグインを利用することにしました。

試したこと

開発は、大きく分けて「ベクトル生成」「OpenSearch連携」「UI連携」の3つのステップで進めました。

まず初めに行ったのは、手元にある間取り画像をすべてベクトルに変換する処理です。ResNetモデルを使い、各画像から512次元の特徴量ベクトルを抽出するPythonスクリプトを作成しました。このスクリプトは、指定されたディレクトリ内の画像を一枚ずつ処理し、物件IDなどのメタデータとベクトル情報をセットにして、後続の処理で扱いやすいJSON Lines形式で出力するように実装しました。

Python

# src/ai_ocr/resnet_processing.py の概念コード
import torch
from torchvision import models, transforms
from PIL import Image

class ResNetEncoder:
    def __init__(self):
        # 事前学習済みのResNet18モデルをロード
        self.model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
        self.model.eval()
        # 最終層を入れ替えて特徴量抽出器として利用
        self.model.fc = torch.nn.Identity()
        self.transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])

    def encode(self, image_path):
        image = Image.open(image_path).convert("RGB")
        image_tensor = self.transform(image).unsqueeze(0)
        with torch.no_grad():
            embedding = self.model(image_tensor)
        # ベクトルをL2正規化
        normalized_embedding = torch.nn.functional.normalize(embedding, p=2, dim=1)
        return normalized_embedding.squeeze().numpy()

# scripts/generate_layout_vectors.py で上記クラスを利用
# encoder = ResNetEncoder()
# vector = encoder.encode("houselayout/layout0.jpg")

なぜL2正規化を行うかというと、ベクトルの長さを1に揃えることで、ベクトル間の距離計算(類似度判定)がより安定して行えるようになるためです。

次に、生成したベクトルデータをOpenSearchに登録しました。これにはk-NN検索に対応した特別なインデックスマッピングが必要です。embeddingフィールドをknn_vector型として定義し、ベクトルの次元数や類似度計算のアルゴリズム(今回はhnsw)を指定しました。

JSON

{
  "settings": {
    "index": {
      "knn": true
    }
  },
  "mappings": {
    "properties": {
      "embedding": {
        "type": "knn_vector",
        "dimension": 512,
        "method": {
          "name": "hnsw",
          "engine": "nmslib",
          "space_type": "l2"
        }
      },
      "property_id": {"type": "keyword"},
      "image_path": {"type": "keyword"},
      "thumbnail_path": {"type": "keyword"}
    }
  }
}

このマッピングを定義した後、先ほど作成したJSON LinesファイルをOpenSearchの_bulk APIを使って一括で投入するスクリプトを実行し、データの登録を完了させました。

データ登録が完了すれば、あとは検索処理を実装するだけです。ユーザーがアップロードした画像を同様にResNetでベクトル化し、そのベクトルを使ってOpenSearchにk-NNクエリを投げます。クエリは非常にシンプルで、検索対象のベクトルと、取得したい件数(k)を指定するだけです。

JSON

{
  "size": 5,
  "query": {
    "knn": {
      "embedding": {
        "vector": [0.12, 0.45, ...],
        "k": 5
      }
    }
  }
}

このクエリの実行結果として、類似度の高い順に物件IDやサムネイルのパスが返ってくるので、それをGradioなどのUIフレームワークを使って画面に表示させることで、類似画像検索システムが完成しました。

考察

実際にシステムを構築してみて、ResNetによる特徴量抽出とOpenSearchのk-NN検索の組み合わせは、類似画像検索を実装する上で非常に効率的で強力な手段だと感じました。特に、あらかじめタスクを細分化し、ベクトル生成パイプラインとデータ投入のスクリプトを整備したことで、再現性高く開発を進めることができました。

一方で、ただ動くものを作るだけでなく、その「検索精度」をどう評価し、改善していくかが重要であることにも気づきました。手動での評価リストを作成し、検索結果が妥当かどうかを定期的にチェックするワークフローを定義しました。この評価結果をもとに、画像の前処理を工夫したり、将来的には間取りデータでResNetを追加学習(Fine-tuning)させたりといった改善策を検討しています。

まとめ

今回は、ResNetとOpenSearchを用いて、間取りの類似画像検索システムを構築するプロセスを紹介しました。深層学習モデルと検索エンジンという異なる技術を組み合わせることで、実用的なAIアプリケーションを迅速に開発できることを実感しました。

この一連のプロセスは、画像だけでなく、テキストや音声など様々なデータに応用が可能です。この記事が、皆さんの開発のヒントになれば嬉しく思います。

お気軽にご相談ください


ご不明な点はお気軽に
お問い合わせください

サービス資料や
お役立ち資料はこちら

DIVXブログ

テックブログ タグ一覧

人気記事ランキング

関連記事