本記事では、Meta(旧 facebook)発表のOmnivoreを使って画像分類・動画分類する方法を紹介します。

Omnivore

概要

Omnivoreは、まったく同じモデルパラメータを使用して画像・動画・シングルビュー3Dを分類する技術です。

従来技術では、画像では画像専用の、動画は動画専用の、シングルビュー3Dはシングルビュー3D専用のそれぞれ個別のモデルを生成していました。
Meta(旧facebook)が発表したOmnivoreは、これらの入力データを単一のモデルで分類可能にしました。

まさしくOmnivore(雑食動物)ですね。

さらに、単一のモデルで分類が可能でありながら、従来のそれぞれ専用のモデルと同等かそれ以上のパフォーマンスを発揮します。

出典: OMNIVORE: A Single Model for Many Visual Modalities

出典: OMNIVORE: A Single Model for Many Visual Modalities

デモ(Colaboratory)

それでは早速OmnivoreをGoogle Colaboratoryで動かしていきます。

なお、これから紹介するソースコードは全てこちらのGitHubに掲載しております。以下のボタンをクリックするとColaboratoryを開くことも可能です。

なお、このデモはPythonで実装しています。
Pythonの実装に不安がある方、Pythonを使った機械学習について詳しく勉強したい方は、以下の書籍やオンライン講座などがおすすめです。

環境セットアップ

まず、Pytorchや、その他ライブラリをインストールします。

# Pytorch 1.9.0をインストール
#!pip uninstall torch -y
#!pip install torch==1.9.0+cu111 torchvision==0.10.0+cu111 torchaudio==0.9.0 -f https://download.pytorch.org/whl/torch_stable.html
!pip install einops timm moviepy

# インストールしたPytorchの動作確認(バージョン、CUDAの認識可否、GPU数確認)
import torch
!python -c "import torch; print(torch.__version__); print(torch.cuda.is_available()); print(torch.cuda.device_count());"

!pip install "git+https://github.com/facebookresearch/pytorchvideo.git"
!git clone https://github.com/facebookresearch/omnivore.git

ColaboratoryにはデフォルトでPytorch1.10.0がインストールされていますが、Omnivoreの推奨バージョンが1.9.0のためPytorchをインストールし直しています。

モデルのロード

以下でモデルをロードしています。
omnivoreのmodel zooに学習済みモデルの一覧があります。本記事ではOmnivore Swin Bを使用します。

# Device on which to run the model
# Set to cuda to load on GPU
device = "cuda" if torch.cuda.is_available() else "cpu"
print("set device->", device)

# Pick a pretrained model 
model_name = "omnivore_swinB"
model = torch.hub.load("facebookresearch/omnivore:main", model=model_name)

# Set to eval mode and move to desired device
model = model.to(device)
model = model.eval()

画像分類

それでは、画像分類を行っていきます。本記事ではこちらのフリー画像を使用します。

prediction = model(image, input_type="image")
prediction = F.softmax(prediction, dim=1)
values, pred_classes = prediction.topk(k=5)

print("上位5番目までの検出ラベル")
for value, pred_class in zip(values[0], pred_classes[0]):
  print('ラベル:', imagenet_id_to_classname[str(pred_class.item())], ', Score:', value.item())

# GPU Memoryがギリギリのため予測が終わったらGPUに配置したimageを削除
del image
torch.cuda.empty_cache()

画像分類結果は以下の通りです。

アナログ時計の最もスコアが高い結果となりました。画像を正確に分類できていますね。
一点注意点ですが、次の動画分類で10GBほどRAMを消費します。Colaboratory環境ですと結構RAMがカツカツになるため使用済みのデータは適宜削除しています。

動画分類

続いて動画分類を行います。
単一モデルで分類が可能であるためモデルを再びロードする必要はありません。

本記事では、こちらのフリー動画を使用します。

# Pass the input clip through the model 
prediction = model(video_input, input_type="video")

# Get the predicted classes 
prediction = F.softmax(prediction, dim=1)
values, pred_classes = prediction.topk(k=5)

print("上位5番目までの検出ラベル")
for value, pred_class in zip(values[0], pred_classes[0]):
  print('ラベル:', kinetics_id_to_classname[int(pred_class)], ', Score:', value.item())

鳥の餌やりが最も高いスコアとなりました。鳥であることは分類されていますね。

シングルビュー3D画像分類

最後にシングルビュー3D画像分類を行います。
本記事では、こちらの画像を使用します。

prediction = model(rgbd_input, input_type="rgbd")
prediction = F.softmax(prediction, dim=1)
pred_classes = prediction.topk(k=5).indices

print("上位5番目までの検出ラベル")
for value, pred_class in zip(values[0], pred_classes[0]):
  print('ラベル:', sunrgbd_id_to_classname[str(pred_class.item())], ', Score:', value.item())

シングルビュー3D画像分類結果は以下の通りです。

furniture_store(家具屋)が最も高いスコアとなりました。この画像はIKEAの店内の画像なので正しく分類できています。

まとめ

本記事では、Omnivoreと呼ばれる機械学習手法を利用し、単一モデルで画像分類・動画分類・シングルビュー3D画像分類を行う方法を紹介しました。
一つのモデルで様々な分類が可能になることにより、モデルを切り替える必要がなくなりシステム全体の実行速度が向上したり、保存するモデル数が減ることによりディスク容量削減や、モデルアップデート等の管理コストの削減等が見込めますね。

また本記事では、機械学習を動かすことにフォーカスしてご紹介しました。
機械学習について学術的に体系立てて学びたいという方には以下の書籍などがお勧めです。ぜひご一読下さい。

また動かせるだけから理解して応用できるエンジニアの足掛かりに下記のUdemyなどもお勧めです。

参考文献

1.  論文 - Omnivore: A Single Model for Many Visual Modalities

2. GitHub - facebookresearch/omnivore