本記事では、End-to-End Referring Video Object Segmentation with Multimodal Transformersを使用して、自然言語表現で指定した物体をビデオから検出する方法をご紹介します。

Multimodal Transformers(MMTR)

概要

2021年11月に論文発表された参照ビデオオブジェクトセグメンテーションタスク(RVOS:The referring video object segmentation task)を実現する技術です。

このタスクでは、自然言語表現で入力されたテキストに応じたオブジェクトをビデオから物体検出、及び物体追跡します。

このタスクは、テキストやビデオの理解、インスタンスセグメンテーション、オブジェクトトラッキングなど様々なサブタスクから構成されています。

従来技術は、複雑なパイプラインでこのタスクを実現していましたが、MMTRでは、Transformerを用いたシンプルな構成で実現しています。このためtext-related inductive biasコンポーネントがなく、mask-refinement処理も必要ありません。

出典: End-to-End Referring Video Object Segmentation with Multimodal Transformers

結果として、毎秒76フレームを処理しながら良好な結果を示していると論文では述べられています。

詳細はこちらの論文をご確認ください。

本記事では、MMTRを用いてビデオから入力テキストに応じた物体の検出方法をご紹介します。

デモ(Colaboratory)

それでは、実際に動かしながらMMTRを用いたRVOSを行っていきます。
ソースコードは本記事にも記載していますが、下記のGitHubでも取得可能です。
GitHub - Colaboratory demo

また、下記から直接Google Colaboratoryで開くこともできます。

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

環境セットアップ

それではセットアップしていきます。 Colaboratoryを開いたら下記を設定しGPUを使用するようにしてください。

まずライブラリのインストールします。

!pip install av ruamel.yaml einops timm transformers
!pip install moviepy==0.2.3.5 imageio==2.4.1
!pip install yt-dlp
!pip install --upgrade gdown

次に、ライブラリをインポートします。

%cd /content

import torch
import torchvision
import torchvision.transforms.functional as F
from einops import rearrange
import numpy as np
from PIL import Image, ImageDraw, ImageOps, ImageFont
from yt_dlp import YoutubeDL
from moviepy.editor import VideoFileClip, AudioFileClip, ImageSequenceClip
from moviepy.video.io.ffmpeg_tools import ffmpeg_extract_subclip
from IPython.display import HTML
from base64 import b64encode
from tqdm.notebook import trange, tqdm
from transformers import logging
logging.set_verbosity_error()

以上で、環境セットアップは完了です。

テスト動画のセットアップ

次にRVOSを行う動画のセットアップを行います。
対象の動画をYoutubeより取得するため、任意のYoutube動画のURL及び、切り抜き位置を指定します。

#@markdown テスト動画に使用するYoutubeのリンクを設定してください。
video_url = 'https://www.youtube.com/watch?v=cREP_TiI13A' #@param {type:"string"}

#@markdown 動画の切り抜き範囲(秒)を指定してください。
start_sec = 23 #@param {type:"integer"}
end_sec =  33 #@param {type:"integer"}

#@markdown テキストクエリをカンマ区切りで入力してください。
text = "a person in white clothes with a racket,tennis racket" #@param {type:"string"}

(start_pt, end_pt) = (start_sec, end_sec)

text_queries = text.split(",")

assert  0 < end_pt - start_pt <= 10, 'error - 動画の長さは10秒以内にしてください。'
assert  1 <= len(text_queries) <= 2, 'error - テキストクエリは2個までにしてください。'

設定したパラメータに応じて動画を取得します。

download_resolution = 360
full_video_path = 'full_video.mp4'
input_clip_path = 'input_clip.mp4'

# 動画ダウンロード
ydl_opts = {'format': f'best[height<={download_resolution}]', 'overwrites': True, 'outtmpl': full_video_path}
with YoutubeDL(ydl_opts) as ydl:
    ydl.download([video_url])

設定したパラメータに応じて動画を切り抜き、表示します。

# extract the relevant subclip:
with VideoFileClip(full_video_path) as video:
    subclip = video.subclip(start_pt, end_pt)
    subclip.write_videofile(input_clip_path)
    
# visualize the input clip:
input_clip = open(input_clip_path,'rb').read()
data_url = "data:video/mp4;base64," + b64encode(input_clip).decode()
HTML("""<video width=720 controls><source src="%s" type="video/mp4"></video>""" % data_url)

ここまでで、ご自身の設定に従って切り抜き部分動画が表示されていることを確認ください。

The referring video object segmentation task

それではRVOSを行っていきます。
始めにモデルをロードします。

model, postprocessor = torch.hub.load('mttr2021/MTTR:main','mttr_refer_youtube_vos', force_reload=True)
model = model.cuda()

マスク処理や、Tensor定義用関数を定義します。

class NestedTensor(object):
    def __init__(self, tensors, mask):
        self.tensors = tensors
        self.mask = mask

def nested_tensor_from_videos_list(videos_list):
    def _max_by_axis(the_list):
      maxes = the_list[0]
      for sublist in the_list[1:]:
          for index, item in enumerate(sublist):
              maxes[index] = max(maxes[index], item)
      return maxes

    max_size = _max_by_axis([list(img.shape) for img in videos_list])
    padded_batch_shape = [len(videos_list)] + max_size
    b, t, c, h, w = padded_batch_shape
    dtype = videos_list[0].dtype
    device = videos_list[0].device
    padded_videos = torch.zeros(padded_batch_shape, dtype=dtype, device=device)
    videos_pad_masks = torch.ones((b, t, h, w), dtype=torch.bool, device=device)
    for vid_frames, pad_vid_frames, vid_pad_m in zip(videos_list, padded_videos, videos_pad_masks):
        pad_vid_frames[:vid_frames.shape[0], :, :vid_frames.shape[2], :vid_frames.shape[3]].copy_(vid_frames)
        vid_pad_m[:vid_frames.shape[0], :vid_frames.shape[2], :vid_frames.shape[3]] = False
    return NestedTensor(padded_videos.transpose(0, 1), videos_pad_masks.transpose(0, 1))

def apply_mask(image, mask, color, transparency=0.7):
    mask = mask[..., np.newaxis].repeat(repeats=3, axis=2)
    mask = mask * transparency
    color_matrix = np.ones(image.shape, dtype=np.float) * color
    out_image = color_matrix * mask + image * (1.0 - mask)
    return out_image

RVOSを実行します。

window_length = 24  # length of window during inference
window_overlap = 6  # overlap (in frames) between consecutive windows

with torch.inference_mode():
  # read and preprocess the video clip:
  video, audio, meta = torchvision.io.read_video(filename=input_clip_path)
  video = rearrange(video, 't h w c -> t c h w')
  input_video = F.resize(video, size=360, max_size=640).cuda()
  input_video = input_video.to(torch.float).div_(255)
  input_video = F.normalize(input_video, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
  video_metadata = {'resized_frame_size': input_video.shape[-2:], 'original_frame_size': video.shape[-2:]}
  
  # partition the clip into overlapping windows of frames:
  windows = [input_video[i:i+window_length] for i in range(0, len(input_video), window_length - window_overlap)]
  # clean up the text queries:
  text_queries = [" ".join(q.lower().split()) for q in text_queries]

  pred_masks_per_query = []
  t, _, h, w = video.shape
  for text_query in tqdm(text_queries, desc='text queries'):
    pred_masks = torch.zeros(size=(t, 1, h, w))
    for i, window in enumerate(tqdm(windows, desc='windows')):
      window = nested_tensor_from_videos_list([window])
      valid_indices = torch.arange(len(window.tensors)).cuda()
      outputs = model(window, valid_indices, [text_query])
      window_masks = postprocessor(outputs, [video_metadata], window.tensors.shape[-2:])[0]['pred_masks']
      win_start_idx = i*(window_length-window_overlap)
      pred_masks[win_start_idx:win_start_idx + window_length] = window_masks
    pred_masks_per_query.append(pred_masks)

推論結果のセグメンテーションマスクをビデオにマッピングします。

# RGB colors for instance masks:
light_blue = (41, 171, 226)
purple = (237, 30, 121)
dark_green = (35, 161, 90)
orange = (255, 148, 59)
colors = np.array([light_blue, purple, dark_green, orange])

# width (in pixels) of the black strip above the video on which the text queries will be displayed:
text_border_height_per_query = 36

video_np = rearrange(video, 't c h w -> t h w c').numpy() / 255.0
# del video
pred_masks_per_frame = rearrange(torch.stack(pred_masks_per_query), 'q t 1 h w -> t q h w').numpy()
masked_video = []
for vid_frame, frame_masks in tqdm(zip(video_np, pred_masks_per_frame), total=len(video_np), desc='applying masks...'):
  # apply the masks:
  for inst_mask, color in zip(frame_masks, colors):
    vid_frame = apply_mask(vid_frame, inst_mask, color / 255.0)
  vid_frame = Image.fromarray((vid_frame * 255).astype(np.uint8))
  # visualize the text queries:
  vid_frame = ImageOps.expand(vid_frame, border=(0, len(text_queries)*text_border_height_per_query, 0, 0))
  W, H = vid_frame.size
  draw = ImageDraw.Draw(vid_frame)
  font = ImageFont.truetype(font='LiberationSans-Regular.ttf', size=30)
  for i, (text_query, color) in enumerate(zip(text_queries, colors), start=1):
      w, h = draw.textsize(text_query, font=font)
      draw.text(((W - w) / 2, (text_border_height_per_query * i) - h - 3),
                text_query, fill=tuple(color) + (255,), font=font)
  masked_video.append(np.array(vid_frame))

# generate and save the output clip:
output_clip_path = 'output_clip.mp4'
clip = ImageSequenceClip(sequence=masked_video, fps=meta['video_fps'])
clip = clip.set_audio(AudioFileClip(input_clip_path))
clip.write_videofile(output_clip_path, fps=meta['video_fps'], audio=True)
del masked_video

推論結果を表示します。

# visualize the output clip:
output_clip = open(output_clip_path,'rb').read()
data_url = "data:video/mp4;base64," + b64encode(output_clip).decode()
HTML("""<video width=720 controls><source src="%s" type="video/mp4"></video>""" % data_url)

出力結果は以下の通りです。

手前の人物とラケットは正確に検出、追跡できているようです。
奥の人物は検出がやや不明瞭ですね。ラケットがかなり小さいので、ラケットを持っていない人として判定されている可能性がありそうです。

まとめ

本記事では、MMTRを用いたRVOSを実施する方法をご紹介しました。

自然言語処理に端を発したTransformerですが、画像分類、物体検出、物体追跡など様々な分野に応用され パフォーマンスを改善していますね。

これを機に機械学習に興味を持つ方が一人でもいらっしゃいましたら幸いです。

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

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

参考文献

1.  論文 - End-to-End Referring Video Object Segmentation with Multimodal Transformers

2. GitHub - mttr2021/MTTR