動画から動作認識 TimeSformer【Python】

はじめに
前提条件
インストール手順
example.py で確認
カスタムデータで学習させてみる
推論コード
おわりに

はじめに

前回は、三次元物体認識の omni3D について説明しました。

今回は、動画から動作認識を行う TimeSformer について、実際に動かすまでの手順を説明していきます。

前提条件

前提条件は以下の通りです。

Python3.9
Windows11
venv 環境

インストール手順

python -m venv venv
python -m pip install -U pip
pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu117
pip install 'git+https://github.com/facebookresearch/fvcore'
pip install simplejson
pip install einops
pip install timm
pip install psutil
pip install av
pip install scikit-learn
pip install opencv-python
pip install tensorboard

git clone https://github.com/facebookresearch/TimeSformer.git
cd TimeSformer
python setup.py build develop

example.py で確認

example.py を以下のように作成します。

import torch
from timesformer.models.vit import TimeSformer
import time
import numpy as np

model_file = './TimeSformer_divST_8x32_224_K600.pyth'

model = TimeSformer(img_size=224, num_classes=600, num_frames=8, attention_type='divided_space_time',  pretrained_model=str(model_file))
model.cuda()
device = torch.device("cuda:0")
dummy_video = torch.randn(1, 3, 8, 224, 224, device=device) # (batch x channels x frames x height x width)

start = time.time()
pred = model(dummy_video,) # (2, 600)
pred = pred.detach().cpu().numpy()
max_index = np.argmax(pred[0])
print(max_index, pred[0][max_index])
print("prediction time: ", time.time()-start)

上記を実行してエラーが出なければ、TimeSformer の準備は完了です。

カスタムデータで学習させてみる

学習のために、demo フォルダに動画と train.csv, test.csv, val.csv を配置します。
各 csv ファイルは以下のようにします。

./demo/demo.mp4 0
./demo/demo-1.mp4 0
./demo/demo2.mp4 1

左側はファイルパス、右側はクラス番号です。

各動画は、こちらから拝借しました。

また、configs/Kinetics/TimeSformer_divST_8x32_224.yaml を以下のようにします。

TRAIN:
  ENABLE: True
  DATASET: kinetics
  BATCH_SIZE: 1
  EVAL_PERIOD: 100
  CHECKPOINT_PERIOD: 100
  AUTO_RESUME: True
DATA:
  PATH_TO_DATA_DIR: ./demo
  NUM_FRAMES: 8
  SAMPLING_RATE: 32
  TRAIN_JITTER_SCALES: [256, 320]
  TRAIN_CROP_SIZE: 224
  TEST_CROP_SIZE: 224
  INPUT_CHANNEL_NUM: [3]
TIMESFORMER:
  ATTENTION_TYPE: 'divided_space_time'
SOLVER:
  BASE_LR: 0.01
  LR_POLICY: steps_with_relative_lrs
  STEPS: [0, 101, 201]
  LRS: [1, 0.1, 0.01]
  MAX_EPOCH: 500
  MOMENTUM: 0.9
  WEIGHT_DECAY: 1e-4
  OPTIMIZING_METHOD: sgd
MODEL:
  MODEL_NAME: vit_base_patch16_224
  NUM_CLASSES: 2
  ARCH: vit
  LOSS_FUNC: cross_entropy
  DROPOUT_RATE: 0.05
TEST:
  ENABLE: True
  DATASET: kinetics
  BATCH_SIZE: 1
  NUM_ENSEMBLE_VIEWS: 1
  NUM_SPATIAL_CROPS: 3
DATA_LOADER:
  NUM_WORKERS: 1
  PIN_MEMORY: True
NUM_GPUS: 1
NUM_SHARDS: 1
RNG_SEED: 0
OUTPUT_DIR: .

以下のコマンドで学習可能です。

python tools/run_net.py --cfg configs/Kinetics/TimeSformer_divST_8x32_224.yaml DATA.PATH_TO_DATA_DIR ./demo NUM_GPUS 1 TRAIN.BATCH_SIZE 1

これで loss が減少していけばいいのですが、loss に変化はなく、学習もできていませんでした…。

推論コード

学習がうまくいきませんでしたが、備忘録として推論コードを残しておきます。

import torch
from timesformer.models.vit import TimeSformer
import time
import numpy as np
import cv2

model_file = './checkpoints/checkpoint_epoch_00050.pyth'

model = TimeSformer(img_size=224, num_classes=2, num_frames=8, attention_type='divided_space_time',  pretrained_model=str(model_file))
# model.cuda()

movie_path = "./demo/demo.mp4"
cap = cv2.VideoCapture(movie_path)
FRAMES = 8
BATCH = 1
SCALE = 1

while cap.isOpened():
    frame_list = []
    for i in range(FRAMES):
        for _ in range(1):
            ret = cap.grab()
        ret, frame = cap.retrieve()
        height, width, channel = frame.shape
        height = int(height/SCALE)
        width = int(width/SCALE)
        frame = cv2.resize(frame, (width,height))

        if not ret:
            break
        frame_list.append(frame)

    input_movie = np.array(frame_list).reshape(BATCH,channel,FRAMES,height,width)/255
    input_movie_tensor = torch.from_numpy(input_movie).float()

    pred = model(input_movie_tensor,)
    pred = pred.detach().cpu().numpy()
    max_index = np.argmax(pred[0])
    print(max_index, pred[0][max_index], pred.shape, pred)

    cv2.imshow("frame", frame)
    lastkey = cv2.waitKey(1)
    if lastkey == ord("q"):
        cv2.destroyAllWindows()
        break