Viết ứng dụng Android chuyển ảnh thành Anime sử dụng AnimeGANv2 + NCNN + OpenCV

Đặt vấn đề

Đợt này đang có trend chuyển ảnh chụp thành anime trên khắp các mạng xã hội. Chất lượng những bức vẽ được generate bởi AI thật sự quá đẹp và lôi cuốn khiến mình phải tò mò tìm hiểu xem công nghệ phía sau nó là gì và liệu có cách gì mình viết được một ứng dụng trên Android (chuyên môn hiện tại của mình) giống như vậy được hay không. Nghĩ là làm, mình bắt tay vào đào khắp các ngóc ngách trên Internet về cái này.

Kết quả là mình tìm thấy một mô hình tên AnimeGANv2 có thể làm được gần gần như ứng dụng kia. Tuy nhiên giữa một mô hình được nghiên cứu và một mô hình đã được thương mại vẫn là một khoảng cách rất xa. Dù sao đây cũng là một nguyên liệu ngon để bắt tay implement nó lên Android app rồi.

Ok, giờ làm sao để implement mô hình này lên Android đây nhỉ ? Có rất nhiều thư viện làm được việc này như TensorFlow Lite , Pytorch Lite... Sau khi tham khảo một lượt các thư viện khác nhau, mình tìm thấy một thư viện của Tencent tên là NCNN tối ưu performance cho các thiết bị di động. Không biết có ngon không nhưng cứ thử tìm cách implement bằng thư viện này để tăng độ khó cho game cái đã.

Model weights Face Potrait v2

Đây là một pretrained model của AnimeGANv2 dành cho việc chuyển đổi chân dung của một người thành hình vẽ anime. Chất lượng hình cho ra thực sự rất ấn tượng. Bạn có thể thử với model này tại đây.

Có một repository khác đã implement Face Potrait v2 trên mobile bằng PyTorch Lite được viết bằng React Native. Nhưng bạn cũng có thể clone và chuyển nó sang project Android Native nếu muốn. Mình đã thử và thấy tốc độ không được nhanh lắm sau khi đã so sánh với NCNN.

Do vậy mình sẽ tìm cách apply model weights này sang thư viện NCNN trên Android.

Implement với NCNN

1. NCNN là gì ?

NCNN là một framework tính toán suy luận neural network hiệu suất cao được tối ưu hóa cho các nền tảng di động. NCNN hoạt động đa nền tảng và chạy nhanh hơn tất cả các framework mã nguồn mở hiện có trên CPU di động. NCNN hiện đang được sử dụng trong nhiều ứng dụng của hãng Tencent, chẳng hạn như QQ, Qzone, WeChat, Pitu...

Với Android platform thì đây là một thư viện JNI/C++, bạn cần phải tải xuống và đưa nó vào app/src/main/jni để sử dụng nó trong mã C++. Về cơ bản các file model weights của thư viện này có đuôi là .bin, trong khi các file tham số có đuôi là .param. Bạn cũng có thể nạp parameters cho neural network theo kiểu programmatically.

Thư viện này chạy được cả trên CPU và GPU. Với GPU, do Android đã tích hợp sẵn Vulkan (thay cho OpenGL ES trước đây) và thư viện này có một bản hỗ trợ sẵn Vulkan nên chúng ta có thể dùng bản này để chạy infer trên GPU.

Thông tin thêm các bạn có thể tham khảo tại đây:

2. Chuyển model weights từ Pytorch (.pt) sang NCNN (.bin)

PNNX là một công cụ cho phép chuyển đổi model PyTorch sang các framework khác bao gồm NCNN.

Mình tìm được một repo đã có sẵn các model weights cho một số model neural network phổ biến cho NCNN và hướng dẫn convert từ model PyTorch sang NCNN bằng PNNX luôn. Dưới đây là README của Face Potrait v2.

Có thể repository bên trên tổng hợp từ repository dưới đây

Như vậy chúng ta đã có file pretrained của Face Potrait v2 dành cho NCNN, một code sample C++ chưa chạy được ngay trên Android mà cần modify một chút.

3. Bắt tay vào code

Bước 1: Thiết lập môi trường và các dependencies cần thiết

• Tạo project Android mới.
• Thêm framework NCNN

Các bạn tải về file ncnn-20230816-android-vulkan.zip nhé. Sau đó giải nén vào đưa vào project Android đã tạo.

Sau đó bạn cần tải tiếp OpenCV Mobile (một bản rút gọn của OpenCV SDK).

Các bạn tải về file opencv-mobile-4.8.0-android.zip sau đó giải nén và cũng đưa vào thư mục jni nhé.

Bây giờ trong thư mục jni chúng ta tạo file CMakeLists.txt. Nội dung file này như sau:

project(facepotraitv2)

cmake_minimum_required(VERSION 3.4.1)

set(OpenCV_DIR ${CMAKE_SOURCE_DIR}/opencv-mobile-4.8.0-android/sdk/native/jni)
find_package(OpenCV REQUIRED core imgproc highgui)

set(ncnn_DIR ${CMAKE_SOURCE_DIR}/ncnn-20230816-android-vulkan/${ANDROID_ABI}/lib/cmake/ncnn)
find_package(ncnn REQUIRED)

add_library(facepotraitv2 SHARED facepotraitv2_jni.cpp)

target_link_libraries(facepotraitv2 ncnn ${OpenCV_LIBS})

Trong file build.gradle.kts của module app, ta thêm đoạn khai báo sau:

android {
	...
    defaultConfig {
        ...   

        ndk {
            moduleName = "ncnn"
            abiFilters += listOf("armeabi-v7a", "arm64-v8a", "x86_64", "x86")
        }
    }
    ...
    externalNativeBuild {
        cmake {
            version = "3.22.1"
            path = file("src/main/jni/CMakeLists.txt")
        }
    }
}

• Thêm file model, param

Các bạn tại ở đây ha, tải file animeganv2.zip nhé.

Sau khi giải nén thì đưa vào thư mục assets. Vị trí của các file và thư mục như dưới đây:

Bước 2: Code

Tạo một interface IFacePortraitV2Converter khai báo các phương thức cần thiết.

interface IFacePortraitV2Converter {

    fun init(mgr: AssetManager): Boolean
    fun convert(bitmap: Bitmap): Boolean
}

Trong đó:

• init: Nơi khởi tạo network, nạp model weights, parameters...
• convert: Hàm thực hiện việc convert Bitmap, sau khi hoàn tất sẽ update nội dung của chính Bitmap đó.

Tạo một class implement interface trên:

class FacePortraitV2Converter : IFacePortraitV2Converter {

    init {
        System.loadLibrary("faceportraitv2")
    }

    external override fun init(mgr: AssetManager): Boolean
    external override fun convert(bitmap: Bitmap): Boolean
}

Trong đó faceportraitv2 là tên thư viện native được khai báo trong CMakeList.txt. 2 phương thức implement sử dụng từ khoá external tương đương với native method của Java. Như vậy mỗi khi 2 phương thức này được gọi, Android sẽ gọi đến phương thức JNI với tên tương ứng.

Trong file faceportraitv2ncnn_jni.cpp chúng ta khai báo 2 phương thức JNI tương ứng với 2 phương thức external ở trên.

JNIEXPORT jboolean JNICALL
Java_com_phucynwa_faceportraitv2ncnn_FacePortraitV2Converter_init(JNIEnv *env, jobject thiz,
                                                                  jobject assetManager) {
    ncnn::Option opt;
    opt.lightmode = true;
    opt.num_threads = 4;
    opt.blob_allocator = &g_blob_pool_allocator;
    opt.workspace_allocator = &g_workspace_pool_allocator;

    // use vulkan compute
    if (ncnn::get_gpu_count() != 0)
        opt.use_vulkan_compute = true;

    AAssetManager *mgr = AAssetManager_fromJava(env, assetManager);

    face_portrait_v2_net.opt = opt;
    int ret0 = face_portrait_v2_net.load_param(mgr, "face_paint_512_v2.param");
    int ret1 = face_portrait_v2_net.load_model(mgr, "face_paint_512_v2.bin");

    __android_log_print(ANDROID_LOG_DEBUG, "FacePortraitV2", "load %d %d", ret0, ret1);

    return JNI_TRUE;
}


JNIEXPORT jboolean JNICALL
Java_com_phucynwa_faceportraitv2ncnn_FacePortraitV2Converter_convert(JNIEnv *env, jobject thiz,
                                                                     jobject bitmap) {
    AndroidBitmapInfo info;
    AndroidBitmap_getInfo(env, bitmap, &info);
    const int width = info.width;
    const int height = info.height;

    int target_w = 512;
    int target_h = 512;

    const float mean_vals[3] = {127.5f, 127.5f, 127.5f};
    const float norm_vals[3] = {1 / 127.5f, 1 / 127.5f, 1 / 127.5f};
    ncnn::Mat in = ncnn::Mat::from_android_bitmap_resize(env, bitmap, ncnn::Mat::PIXEL_RGB, target_w, target_h);
    in.substract_mean_normalize(mean_vals, norm_vals);
    ncnn::Mat out;
    {
        ncnn::Extractor ex = face_portrait_v2_net.create_extractor();
        ex.set_vulkan_compute(true);
        ex.input("in0", in);
        ex.extract("out0", out);
    }

    __android_log_print(ANDROID_LOG_DEBUG, "FacePortraitV2","w%d x h%d", out.w, out.h);

    cv::Mat result(out.h, out.w, CV_32FC3);
    for (int c = 0; c < out.c; c++) {
        float *out_data = out.channel(c);
        for (int row = 0; row < out.h; row++) {
            for (int col = 0; col < out.w; col++) {
                result.at<cv::Vec3f>(row, col)[c] = out_data[row * out.h + col];
            }
        }
    }
    cv::Mat result8U(out.h, out.w, CV_8UC3);
    result.convertTo(result8U, CV_8UC3, 127.5, 127.5);
    cv::Mat dst(height, width, result8U.type());
    resize(result8U, dst, dst.size(), 0, 0, cv::INTER_CUBIC);
    MatToBitmap(env, dst, bitmap, false);
    return JNI_TRUE;
}

}

Bước 3: Thêm glue code và chạy thử

Phần này hơi rườm rà và không liên quan đến nội dung bài vì vậy mình sẽ để link sample code ở đây để các bạn clone về chạy thử. Bạn cũng có thể tải về file APK tại đây.

Ảnh chụp trong ứng dụng:

Kết luận

• Sau khi đã nghịch ngợm framework NCNN này mình thấy rằng nó khá dễ dùng và performance cũng rất tốt. Tuy nhiên các pretrained model sẽ không sẵn như PyTorch hay TFLite.
• Hiện nay có rất nhiều mô hình GAN được pretrained có chất lượng tốt có thể áp dụng cho các product đơn giản.
• Lĩnh vực mobile app không chỉ có call API, show data và bo tròn các Button, tuy không hàn lâm được như các lĩnh vực khác nhưng việc làm chủ và tối ưu hoá performance cho các implementation của các công nghệ khác là hoàn toàn có thể và vẫn luôn đòi hỏi chuyên môn và công sức nghiên cứu.