今回はタイトルの論文のソースコードを動かしてみます。githubはこちら。
https://github.com/ykasten/layered-neural-atlases
PyTorchとdetectron2のバージョンさえ合わせれば素直に動くようです。
動作環境
RTX3080Tiを搭載したUbuntuマシンで試します。RTX3000系で動かしたいのでCUDA11系を使っています。
ホストマシン
- OS: Ubuntu20.04
- GPU: RTX3080Ti
- nvidia driver: 510.47.03
docker環境
- docker image: nvidia/cuda:11.3.0-devel-ubuntu20.04
- torch: 1.10.2+cu113(stable)
- CUDA: 11.3
PyTorchのバージョン
detectron2とtorchのバージョンを合わせる必要がありました。detectron2が対応している最新のtorchは1.10なので、1.10.2+cu113を使いました。以下のコマンドでインストールします。torchが新しすぎるとdetectron2のimportのところで止まってしまうようです。
pip install torch==1.10.2+cu113 torchvision==0.11.3+cu113 torchaudio==0.10.2+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html
detectron2のインストール
今回はここがちょっと難しかったです。元のソースコードではCUDA10系を使っていますが、CUDA11系を使いたい場合はこちらのページのInstall Pre-Built Detectron2 (Linux only)のテーブルからcuda11.3とtorch 1.10の組み合わせのところのインストールコマンドを利用します。
https://detectron2.readthedocs.io/en/latest/tutorials/install.html
python -m pip install detectron2 -f \
https://dl.fbaipublicfiles.com/detectron2/wheels/cu113/torch1.10/index.html
その他のモジュール
Open-CVはopencv-pythonでインストールしています。
requirements.txt
numpy
tqdm
opencv-python
scikit-image
tensorboard
tensorflow
scipy
matplotlib
gdown
imageio-ffmpeg
データの準備
gdownコマンドではダウンロードできなかったので、ブラウザからGoogleドライブを開いてダウンロードします。
https://drive.google.com/uc?id=1WipZR9LaANTNJh764ukznXXAANJ5TChe
ダウンロード後に展開します。
unzip data.zip
こちらのdata/blackswanなどに動画をフレームごとに連番にした画像ファイルが入っています。
マスクの抽出
以下のコマンドで動画からマスク画像を抽出します。こちらのコマンドではdata/blackswanに入っている連番ファイルからbirdに分類される領域のマスクをdata/blackswan_maskrcnnへ出力します。
python preprocess_mask_rcnn.py --vid-path data/blackswan --class_name bird
オプティカルフローの抽出
サブモジュールを更新してRAFTのダウンロードを実行します。
git submodule update --init
cd thirdparty/RAFT/
./download_models.sh
cd ../..
以下のコマンドを実行するとdata/blackswan_flowへフレームごとのオプティカルフローが出力されます。*.npy形式のファイルなので画像のようには表示できません。
python preprocess_optical_flow.py --vid-path data/blackswan --max_long_edge 768
訓練済みモデルのダウンロード
Googleドライブからダウンロードします。
https://drive.google.com/uc?id=10voSCdMGM5HTIYfT0bPW029W9y6Xij4D
unzipコマンドで展開。
unzip pretrained_models.zip
学習
ここは時間がかかりそうなので飛ばします。
評価
そのままで動きました。最後に出るのが精度でしょうか？
python only_evaluate.py --trained_model_folder=pretrained_models/checkpoints/blackswan --video_name=blackswan --data_folder=data --output_folder=evaluation_outputs
1
2
3
...
47
48
49
29.98105892457371
編集
以下のコマンドを実行するとediting_outputsへ結果のmp4ファイルが出力されます。差し替え用の画像はpretrained_models/edit_inputs/blackswanのものを使います。
python only_edit.py --trained_model_folder=pretrained_models/checkpoints/blackswan --video_name=blackswan --data_folder=data --output_folder=editing_outputs --edit_foreground_path=pretrained_models/edit_inputs/blackswan/edit_blackswan_foreground.png --edit_background_path=pretrained_models/edit_inputs/blackswan/edit_blackswan_background.png
ここでの説明はblackswanでしたが、boat, libbyについても同様にmp4ファイルを出力できることを確認しました。
