概要
ncnnで、写真からアニメ絵への推定を行う、AnimeGANを動かしてみます。
環境構築
まず、M5UnitV2のクロスコンパイル環境を構築する手順で、ncnnやOpenCVをUbuntuにインストールしました。
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での変換を行います。
% gedit 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
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という異なる学習モデルで生成された3枚のアニメ絵が作成されます。
$ 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
プログラムの格納先
このプログラムは以下に格納しています。
注記
AnimeGANは大容量のメモリを消費するため、M5UnitV2で動作させることは困難です。
参考資料
この記事を作成するにあたり、以下のウェブサイトを参考にしました。