この記事は「PY - PARTY TECH STUDIO Advent Calendar 2024」の2日目の記事です。
はじめに
私は普段、PARTY TECH STUDIOのPYで主にサービス系のテクニカルディレクターやバックエンドまわりを担当しています。しかし最近、ひょんなことからAIを使った画像解析のプロジェクトに参加することになりました。
そこでNVIDIA Jetson Orin NanoでNanoOWLを使ってリアルタイム画像解析を試してみて、なかなか面白かったので使い方と中身について簡単に紹介します。
NanoOWLとは
NVIDIA自身がAIモデルやライブラリをJetson Orinシリーズ向けに最適化したチュートリアルをNVIDIA Jetson AI Labというサイトに公開しています。
ここではLLMやVision Transformersのものが公開されており、そのなかのひとつがNanoOWLです。
NanoOWLはOWL-ViTをリアルタイムにうごかすための最適化が施されています。OWL-ViTは簡単にいうと画像内の対象物をテキストで検索・検出できるAIモデルです。
そのためこれらを使うことで検出対象物の学習なしに対象物を表すテキスト（例えば car など）を入力するだけで画像内の写るその対象物を検出できます。
百聞は一見にしかず。上記ページのNanoOWLのデモ動画を見ればその簡単さが実感できます。
NanoOWLを使ってみる
利用にはJetson Orinシリーズが必要です（Jetson Orin Nanoでも9万円前後します。ちょっとお高い…）。Jetson Orin Nanoが手に入ってしまえば、NanoOWLをとりあえず動かすのはめちゃくちゃ簡単です。
Jetson Orin Nanoの初期セットアップ後、上記ページの手順に従って jetson-containers をクローンして指定のコマンドを実行するだけでうごきます。簡単！
ただ簡単に使えるようにまとめられてる分、実際の処理の内容が見えにくくなっています。
そこで jetson-containers がなんなのか、NanoOWLはどういう処理をしているのか、を書いていきます。
jetson-containersとは
jetson-containers はNVIDIAが提供しているJetson向けのコンテナイメージを管理するためのリポジトリです。Jetsonで使えるAIモデルやライブラリがコンテナイメージとして提供されており、jetson-containers を通せばコマンド一発でDockerコンテナを起動してAIモデルを使うことができます。
NanoOWLは jetson-containers run --workdir /opt/nanoowl $(autotag nanoowl) というコマンドで実行できますが、このとき下記にあるDcokerfileが使われています。
そしてNanoOWLのソースコードは下記リポジトリで管理されています
NanoOWLの処理を読む
ここからはNanoOWLの処理をかいつまんで読んでいきます。
エンジンファイルの生成
jsetson-containers で起動したコンテナでNanoOWLの実行するとき引数に ../../data/owl_image_encoder_patch32.engine を指定します。この.engineファイルは推論モデルから生成されたTensorRTでの処理を高速化するためのランタイムです。これがあることでJetson上で高速な論が可能になります。
このファイル自体は jetson-containers を利用している場合は起動時に下記コマンドで生成されています。
プロンプトの分解
NanoOWLでは検出したい対象を下記のようなプロンプトで指定します。
[a face [a nose, an eye, a mouth]]
[a face (interested, yawning / bored)]
(indoors, outdoors)
たとえば [a face [a nose, an eye, a mouth]] のように入れ子にすると、画像内に顔が写っている場合にその顔の部分について鼻、目、口が検出できます。
また () を使うとクラスタライズされたものをできます。たとえば [a face (interested, yawning / bored)] のようにすると、顔が写っている場合にその顔が興味を持っているか、あくびをしているか、退屈しているかを検出できます。
NanoOWLのサンプルコードでは、起動したウェブページで入力したプロンプトはWebsocketで送信され、Treeオブジェクトに変換されます。
Treeオブジェクトの中では以下のような構造に分解されます。
# 入力
[a person, a car(black)]
# オブジェクト
{'nodes': [{'op': 'detect', 'input': 0, 'outputs': [1, 2]}, {'op': 'classify', 'input': 2, 'outputs': [3]}], 'labels': ['image', 'a person', 'a car', 'black']}
分解された各ラベルは predictor.encode_clip_text(tree)(クラスタライズラベル用) predictor.encode_owl_text(tree)（物体ラベル用）でそれぞれベクトル化されます。
この構造を知っておくと、各検出対象ごとに閾値を変更したり、検出結果のプレビューにスコアを表示するなどのカスタマイズがやりやすくなります。
物体検出
OpenCVで取得したフレーム画像は下処理されたあと、ROIを切り出してベクトル化し、物体検出をしています。このとき画像内に写るすべての物体（と思われる範囲）を出力しています。
このあと、ベクトル化したプロンプトのラベルと画像で関連性の評価を行い、スコアが閾値よりも高いものだけを抽出し、結果として出力しています。
サンプルコードのままでは閾値はすべてのラベルで共通になっているので、ラベルごとに閾値を変更したい場合はこのTreeオブジェクトの構造とこのdecodeの処理をカスタマイズすることで可能になります。
ちなみに、プロンプトで入れ子構造を使っている場合、親のラベルが検出された範囲で画像を切り抜き、再度画像の下処理をして、子のラベルで評価を行います。そのため、入れ子を多用すると処理がかなり重くなります。
参考（試したカスタマイズ）
かなりかいつまんで処理を紹介しましたが、（NVIDIAのサンプルにしては！）きれいなコードで読みやすく、カスタマイズがしやすくなっていました。色々検証していく中で、現場でも使えそうなものがあったので簡単に紹介します。
OWLv2への変更
NanoOWLはOWL-Vitがベースになっていますが、コードほぼそのままでOWLv2に変更することができました。
OWLv2のモデルでエンジンファイルを生成して、NanoOWLにそれを渡せば実行できるはずです。結果は非常によく、検出精度が驚くほど上がります。
ただ、実行時のメモリ容量が増えるため、Jetson Orin Nanoでは実行できませんでした。試すにはJetson AGX Orinなどのメモリ容量が大きいモデルが必要です（お高い…）。
リファレンス画像による検出
NanoOWLはテキストラベルによって物体を検出できますが、カスタマイズすると画像検索のようにリファレンス画像を指定してそれに近い物体をリアルタイムで検出することもできました。
実は本家OWL-ViTにはリファレンス画像の指定した範囲の特徴と近い部分を検出するサンプルがあります。
それを参考に、画像内の特定範囲をあらかじめベクトル化しておき、それを関連性の評価に利用することでテキストと同時に画像よる検出も可能になりました。テキストラベルで表現しにくいものや、検出してくれないものを対象としたい場合に使えそうです。
Discussion