SIGGRAPH Asia2021の論文Layered Neural Atlases for Consistent Video Editingを動かす

今回はタイトルの論文のソースコードを動かしてみます。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などに動画をフレームごとに連番にした画像ファイルが入っています。
![00000.jpg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/935724/09a0b701-d659-bfcc-9c31-548aae34774a.jpeg)


# マスクの抽出

以下のコマンドで動画からマスク画像を抽出します。こちらのコマンドではdata/blackswanに入っている連番ファイルからbirdに分類される領域のマスクをdata/blackswan_maskrcnnへ出力します。

```
python preprocess_mask_rcnn.py --vid-path data/blackswan --class_name bird
```

![00000.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/935724/8fc62de0-3cc1-92be-684f-2f3a33a76522.png)


# オプティカルフローの抽出

サブモジュールを更新して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
```

![texture_edit1.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/935724/227b186f-b8f7-0230-0f5b-e395a1bde43d.png)

![texture_edit2.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/935724/d8683d5b-cff8-6253-5720-2afeced9a33a.png)


ここでの説明はblackswanでしたが、boat, libbyについても同様にmp4ファイルを出力できることを確認しました。

![edited_blackswan.PNG](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/935724/7a918167-b69f-3836-7ca5-ffa92d3e4a58.png)

![edited_boat.PNG](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/935724/62e4b4c2-fdc3-584a-ce02-2cecefba9240.png)

![edited_libby.PNG](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/935724/38fb6323-6ed3-ec4c-03bc-ff75dad18a08.png)