この記事は
ViViTモデルを活用して、動画に映っている動作を認識するコードを説明します。ビールを飲んでいるお兄さんの動画で推論させてみたところ、drinkingという出力を得ました。
ViViTとは
Googleの研究者が2021年に発表した動画認識モデル Video Vision Transformerのことです。動画（2〜3秒程度）に映っている人物の動作を識別するための機械学習モデルです。
ViViTモデルはTransformerモデルを基盤としています。Transformerモデルとは、自然言語処理の分野で革新をもたらしたモデルで、BERTやGPTなど最新のLLMの基盤技術になっています。ChatGPTを使っているなら、確実にTransformerモデルのお世話になっています。
Transformerは単語の予測に使用されていますが、このTransformerを画像認識に拡張したのが Vision Transfomer(ViT)です。ViTはCNNによる画像認識と同じように、画像に写っているものが何なのかを識別することができます。
ViTをさらに動画認識に拡張したのが、この投稿で取り上げているViViTモデルです。ViTとViViTはどちらもTransformerのアーキテクチャを踏襲しています。ViViTでは動画をトークンと呼ばれる単位に分割し、Transfomerアーキテクチャ内の自己注意機構と呼ばれる仕組みで相互参照させています。さらに、トークンの分割方法を工夫することで、動画内の時系列情報を認識させることを可能としています。
ViViTの学習済みモデル
HuggingFace ViViT
ViViTはその基盤となるTransformerの仕組み上、CNNが持つ帰納的バイアス（例えば、物体が画像上を平行移動していても正しく認識できる性能）を持っていないため、CNNに比べて大量のデータセットでの学習が必要となるようです。そのため、学習済みViViTモデルを利用したり転移学習するのが良いです。この投稿では、Hugging Faceの学習済みViViTモデルを利用します。
データセット
Hugging Faceの学習済みViViTモデルは、Kinetics-400というデータセットで学習されています。Kinetics-400は、GoogleがYouTubeから収集した400種類の動作、306,245本の動画クリップで構成された大規模なデータセットです。収録されている動作とラベルの一覧は下記のデータセットの論文にまとめられています。
ここでやってみること
数秒の動画クリップを自前で用意し、その動画の人物が何をしているのかを識別します。1動画に1動作のみが収められている動画です。
実行環境
Google Colabを利用します。
用意したビデオ
この投稿用に、若いお兄さんがビールを飲む動画を用意して検証しました（TikTokから適当に拾ってきたものです）。みなさんがお試しの際は、ご自身で動画を用意するか、以下のコードでHugginFaceが提供している動画クリップを利用できます（スパゲッティを食べているようです）。
from huggingface_hub import hf_hub_download
file_path = hf_hub_download(
repo_id="nielsr/video-demo", filename="eating_spaghetti.mp4", repo_type="dataset"
)
container = av.open(file_path)
コード
ここで説明するコードは、https://huggingface.co/docs/transformers/model_doc/vivit から引用したものです。sample_frame_indices関数と read_video_pyav関数は、ユーザ側で notebook のセルにコピペしておく必要があります。その後、メインコードを同じ notebook のセルにコピペして実行することで動画認識が可能です。
sample_frame_indices関数
動画の総フレーム数 seg_len から clip_len で指定した数だけのフレームをサンプリングし、そのインデックスのリストを返します。frame_sample_rate はフレームのサンプリングレートで、この数ごとにフレームをサンプリングします。
def sample_frame_indices(clip_len, frame_sample_rate, seg_len):
converted_len = int(clip_len * frame_sample_rate)
end_idx = np.random.randint(converted_len, seg_len)
start_idx = end_idx - converted_len
indices = np.linspace(start_idx, end_idx, num=clip_len)
indices = np.clip(indices, start_idx, end_idx - 1).astype(np.int64)
return indices
clip_len と frame_sample_rate でサンプリングするフレームの区間が決まると、動画のどの部分をサンプリングするかはランダムに決定されます。
read_video_pyav関数
sample_frame_indices関数で取得したインデックスのリストを元に、ndarrayでデコードされた動画フレームのリストを返します。
def read_video_pyav(container, indices):
frames = []
container.seek(0)
start_index = indices[0]
end_index = indices[-1]
for i, frame in enumerate(container.decode(video=0)):
if i > end_index:
break
if i >= start_index and i in indices:
frames.append(frame)
return np.stack([x.to_ndarray(format="rgb24") for x in frames])
メインコード
# 学習済みモデルをロード
model = VivitForVideoClassification.from_pretrained("google/vivit-b-16x2-kinetics400")
# 動画をロード
file_path = './drinking_beer.mp4'
container = av.open(file_path)
# インデックスの取得
indices = sample_frame_indices(clip_len=32, frame_sample_rate=10, seg_len=container.streams.video[0].frames)
# インデックスからフレームのリストを取得
video = read_video_pyav(container=container, indices=indices)
# イメージプロセッサを初期化
image_processor = VivitImageProcessor.from_pretrained("google/vivit-b-16x2-kinetics400")
# イメージプロセッサで動画を変換
inputs = image_processor(list(video), return_tensors="pt")
# 推論
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
# 最大のラベルを表示
predicted_label = logits.argmax(-1).item()
print(model.config.id2label[predicted_label])
メインコードの各行を説明します。まず、以下のコードで学習済みViViTモデルを取得します。
model = VivitForVideoClassification.from_pretrained("google/vivit-b-16x2-kinetics400")
次に動画を取得します。このコードでは、自前で用意した動画を適当な場所に保存してから取得しています。
file_path = './drinking_beer.mp4'
container = av.open(file_path)
取得した動画から、必要なフレームを取得します。上のセクションで定義した sample_frame_indices 関数と read_video_pyav 関数を利用します。
indices = sample_frame_indices(clip_len=32, frame_sample_rate=10, seg_len=container.streams.video[0].frames)
video = read_video_pyav(container=container, indices=indices)
続いて、取得したフレームのリストをViViTモデルの入力に合わせてリサイズし、正規化を行います。以下のイメージプロセッサは、モデルに合わせてその処理を行ってくれているようです。 return_tensors="pt" は、戻り値の各要素を torch.Tensor に指定しています。
image_processor = VivitImageProcessor.from_pretrained("google/vivit-b-16x2-kinetics400")
inputs = image_processor(list(video), return_tensors="pt")
準備できた動画フレームで推論を行います。
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
最後に一番確率値の大きい要素を取得し、そのラベルを取得します。
predicted_label = logits.argmax(-1).item()
print(model.config.id2label[predicted_label])
結果
今回、私が用意した「若いお兄さんがビールを飲む動画」で推論したところ、
LABEL_101
という出力値を得ました。データセットの論文によると、ラベル101はdrinkingとのことです。
まあ「飲んでいる」動画なので正解と言えば正解ですが、ラベル102にdrinking beerってあるのが惜しいです。推論にかかる時間は、CPU(Intel Core i7)で30〜50秒、Google ColabのA100で数秒です。けっこう重いですね。