ncnnでAnimeGANを動かしてみる

## 概要
ncnnで、写真からアニメ絵への推定を行う、AnimeGANを動かしてみます。

![Screenshot from 2023-11-19 21-42-12.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/73020/c47a0c52-d6a7-3299-4ded-ffc24aeb3960.png)


## 環境構築
まず、M5UnitV2のクロスコンパイル環境を構築する手順で、ncnnやOpenCVをUbuntuにインストールしました。

[M5UnitV2のクロスコンパイル環境を構築する手順](https://qiita.com/nnn112358/pbulic/9e15450265b689300216)

```
Ubuntu 22.04.3 LTS(x64)
OpenCV ver 0.4.4
ncnn ver 231027
```
## animegan2-ncnn-vulkanのダウンロード

GitHubからanimegan2-ncnn-vulkanのコードをダウンロードします。

```
$ git clone https://github.com/FeiGeChuanShu/animegan2-ncnn-vulkan
$ cd animegan2-ncnn-vulkan
```


## animegavn2の学習モデルをPNNXで変換する

pythonのスクリプトで、PyTorchのHubからanimegavn2の学習モデルをダウンロードしてきます。そして、pnnxでの変換を行います。


```sh
% gedit convert.py
```

```convert.py
import os
import torch

model1 = torch.hub.load("bryandlee/animegan2-pytorch:main", "generator", pretrained="celeba_distill")
model2 = torch.hub.load("bryandlee/animegan2-pytorch:main", "generator", pretrained="face_paint_512_v1")
model3 = torch.hub.load("bryandlee/animegan2-pytorch:main", "generator", pretrained="face_paint_512_v2")
model4 = torch.hub.load("bryandlee/animegan2-pytorch:main", "generator", pretrained="paprika")

model1.cpu()
model2.cpu()
model3.cpu()
model4.cpu()
model1.eval()
model2.eval()
model3.eval()
model4.eval()

traced_script_module1 = torch.jit.trace(model1, torch.randn(1, 3, 512, 512), strict=False)
traced_script_module1.save("celeba_distill.pt")

traced_script_module2 = torch.jit.trace(model2, torch.randn(1, 3, 512, 512), strict=False)
traced_script_module2.save("face_paint_512_v1.pt")

traced_script_module3 = torch.jit.trace(model3, torch.randn(1, 3, 512, 512), strict=False)
traced_script_module3.save("face_paint_512_v2.pt")

traced_script_module4 = torch.jit.trace(model4, torch.randn(1, 3, 512, 512), strict=False)
traced_script_module4.save("paprika.pt")

os.system("pnnx celeba_distill.pt inputshape=[1,3,512,512] device=cpu") 
os.system("pnnx face_paint_512_v1.pt inputshape=[1,3,512,512] device=cpu") 
os.system("pnnx face_paint_512_v2.pt inputshape=[1,3,512,512] device=cpu") 
os.system("pnnx paprika.pt inputshape=[1,3,512,512] device=cpu") 
```

ncnnのRobustVideoMattingの学習モデルである、celeba_distill.ncnn.celeba_distill.ncnn.param、などが生成されていることを確認します。




## CMakeLists.txt ファイルの作成

cmakeでコンパイルするために、CMakeLists.txtを編集します。

```
gedit CMakeLists.txt
```

```CMakeLists.txt
SET(CMAKE_BUILD_TYPE "Debug")

SET(CMAKE_CXX_FLAGS_DEBUG "$ENV{CXXFLAGS} -O0 -Wall -g2 -ggdb")
SET(CMAKE_CXX_FLAGS_RELEASE "$ENV{CXXFLAGS} -O3 -Wall")

if(TARGET_COMPILER STREQUAL "arm")
    message(STATUS "TARGET_COMPILE STREQUAL arm")
    SET(CMAKE_CXX_COMPILER arm-none-linux-gnueabihf-g++)
    SET(CMAKE_C_COMPILER arm-none-linux-gnueabihf-gcc)
    SET(NCNN_INSTALL_DIR /opt/external/ncnn/install/arm/)
    SET(ZBAR_INSTALL_DIR /opt/external/zbar-0.10/build/arm)
    SET(OpenCV_DIR /opt/external/opencv/build/arm)
else()
    message(STATUS "TARGET_COMPILE STREQUAL X64")
    SET(NCNN_INSTALL_DIR /opt/external/ncnn/install/x64/)
    SET(ZBAR_INSTALL_DIR /opt/external/zbar-0.10/build/x64/)
    SET(OpenCV_DIR /opt/external/opencv/build/x64) 
endif()


PROJECT(animegan2-ncnn)
cmake_minimum_required(VERSION 3.5)

set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR}/bin)

if(OpenCV_FOUND)
 message(OPENCV_VERSION)
 message(STATUS "OpenCV_LIBS: ${OpenCV_LIBS}")
 message(STATUS "OpenCV_INCLUDE_DIRS: ${OpenCV_INCLUDE_DIRS}")
endif(OpenCV_FOUND)

find_package(OpenCV REQUIRED)

link_directories(${OpenCV_LIBRARY_DIRS})

include_directories(${OpenCV_INCLUDE_DIRS} )
include_directories(${PROJECT_SOURCE_DIR})
include_directories(${PROJECT_SOURCE_DIR}/include)
include_directories(${NCNN_INSTALL_DIR}/include/ncnn)

add_executable(animegan2-ncnn main.cpp)
target_link_libraries(animegan2-ncnn ${OpenCV_LIBRARIES})
target_link_libraries(animegan2-ncnn ${NCNN_INSTALL_DIR}/lib/libncnn.a -fopenmp)
```


## ビルド
Ubuntu用のバイナリをcmakeでコンパイルして、生成されたバイナリを実行します。
celeba/face_paint_512_v1/face_paint_512_v2という異なる学習モデルで生成された３枚のアニメ絵が作成されます。

```shell
$ cmake .
$ make
$ ./bin/animegan2-ncnn test.png
load 0 0
load 0 0
load 0 0
styletransfer
Extractor
result
styletransfer
Extractor
result
styletransfer
Extractor
result
```

![Screenshot from 2023-11-19 21-51-40.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/73020/ab02cbd9-730c-6698-9ad7-93d9904c59ae.png)


## プログラムの格納先
このプログラムは以下に格納しています。

https://github.com/nnn112358/animegan2-ncnn-vulkan

## 注記
AnimeGANは大容量のメモリを消費するため、M5UnitV2で動作させることは困難です。


## 参考資料
この記事を作成するにあたり、以下のウェブサイトを参考にしました。

[eiGeChuanShu/animegan2-ncnn-vulkan](https://github.com/FeiGeChuanShu/animegan2-ncnn-vulkan)

[bryandlee/animegan2-pytorch](https://github.com/bryandlee/animegan2-pytorch)