图像编辑器 Monica 之使用 ONNX Runtime 在端侧部署生成素描画的模型

一.  图像编辑器 Monica

Monica 是一款跨平台的桌面图像编辑软件，使用 Kotlin Compose Desktop 作为 UI 框架。应用层使用 Kotlin 编写，基于 mvvm 架构，使用 koin 作为依赖注入框架。

大部分算法也是用 Kotlin 编写，少部分通过 jni 调用 OpenCV C++ 实现的图像处理或调用深度学习的模型。

Monica 目前还处于开发阶段，当前版本的可以参见 github 地址：https://github.com/fengzhizi715/Monica

二. ONNX Runtime 部署模型

ONNX Runtime 是一个高性能的推理引擎，专门用于加速和优化基于 ONNX（Open Neural Network Exchange）模型格式的机器学习模型的推理。ONNX 是一个开源的中间表示格式，它可以使不同深度学习框架（如 PyTorch、TensorFlow 等）的模型互相转换并在不同的硬件平台上运行。

ONNX Runtime 常用于部署和推理阶段，能够快速高效地运行机器学习模型。

ONNX Runtime 在 mac 下可以通过 brew 或者直接下载源码的方式进行安装。我采用直接下载源码，这样可以安装我想要的版本。

在 cmake 做如下配置即可使用 ONNX Runtime

# 指定ONNX Runtime的路径
set(ONNXRUNTIME_ROOT "/Users/Tony/onnxruntime/onnxruntime-osx-x86_64-1.10.0")
# 包含ONNX Runtime头文件
include_directories(${ONNXRUNTIME_ROOT}/include/)
# 查找ONNX Runtime库文件
find_library(ONNXRUNTIME_LIBS onnx-runtime PATHS ${ONNXRUNTIME_ROOT}/lib)

target_link_libraries(MonicaImageProcess ${ONNXRUNTIME_ROOT}/lib/libonnxruntime.1.10.0.dylib)

安装完成后，先封装一个比较通用的 OnnxRuntimeBase 类，下面是 OnnxRuntimeBase.h

#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>
#include <vector>
#include <onnxruntime_cxx_api.h>
#include <onnxruntime_c_api.h>

class OnnxRuntimeBase
{

public:
    OnnxRuntimeBase(std::string modelPath, const char* logId, const char* provider);
    ~OnnxRuntimeBase();

    virtual std::vector<Ort::Value> forward(Ort::Value& inputTensors);

protected:
    Ort::Env env;
    Ort::Session ort_session{ nullptr };
    Ort::SessionOptions sessionOptions = Ort::SessionOptions();
    std::vector<char*> input_names;
    std::vector<char*> output_names;
    std::vector<std::vector<int64_t>> input_node_dims;  // >=1 outputs
    std::vector<std::vector<int64_t>> output_node_dims; // >=1 outputs

    Ort::MemoryInfo memory_info_handler = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
};

OnnxRuntimeBase 类的实现：

#include <iostream>
#include "../../include/common/OnnxRuntimeBase.h"
#include "../../include/Constants.h"
#include "../../include/utils/Utils.h"

using namespace cv;
using namespace std;
using namespace Ort;


OnnxRuntimeBase::OnnxRuntimeBase(string modelPath, const char* logId, const char* provider)
{
    sessionOptions.SetGraphOptimizationLevel(ORT_ENABLE_BASIC);

    env = Ort::Env(ORT_LOGGING_LEVEL_WARNING, logId);

    std::vector<std::string> availableProviders = Ort::GetAvailableProviders();
    auto cudaAvailable = std::find(availableProviders.begin(), availableProviders.end(), "CUDAExecutionProvider");
    OrtCUDAProviderOptions cudaOption;
    if (provider == OnnxProviders::CUDA.c_str()) {
        if (cudaAvailable == availableProviders.end()) {
            std::cout << "CUDA is not supported by your ONNXRuntime build. Fallback to CPU." << std::endl;
            std::cout << "Inference device: CPU" << std::endl;
        }
        else {
            std::cout << "Inference device: GPU" << std::endl;
            sessionOptions.AppendExecutionProvider_CUDA(cudaOption);
        }
    }
    else if (provider == OnnxProviders::CPU.c_str()) {
        // "cpu" by default
        std::cout << "Inference device: CPU" << std::endl;
    }
    else
    {
        throw std::runtime_error("NotImplemented provider=" + std::string(provider));
    }

    const char* model_path = "";
    #ifdef _WIN32
        auto modelPathW = get_win_path(modelPath);  // For Windows (wstring)
        model_path = modelPathW.c_str();
    #else
        model_path = modelPath.c_str();             // For Linux、MacOS (string)
    #endif

    ort_session = Ort::Session(env, model_path, sessionOptions);

    size_t numInputNodes = ort_session.GetInputCount();
    size_t numOutputNodes = ort_session.GetOutputCount();
    AllocatorWithDefaultOptions allocator;

    for (int i = 0; i < numInputNodes; i++)
    {
        input_names.push_back(ort_session.GetInputName(i, allocator));
        Ort::TypeInfo input_type_info = ort_session.GetInputTypeInfo(i);
        auto input_tensor_info = input_type_info.GetTensorTypeAndShapeInfo();
        auto input_dims = input_tensor_info.GetShape();
        input_node_dims.push_back(input_dims);
    }
    for (int i = 0; i < numOutputNodes; i++)
    {
        output_names.push_back(ort_session.GetOutputName(i, allocator)); 
        Ort::TypeInfo output_type_info = ort_session.GetOutputTypeInfo(i);
        auto output_tensor_info = output_type_info.GetTensorTypeAndShapeInfo();
        auto output_dims = output_tensor_info.GetShape();
        output_node_dims.push_back(output_dims);
    }
}

OnnxRuntimeBase::~OnnxRuntimeBase() {
    sessionOptions.release();
    ort_session.release();
}

std::vector<Ort::Value> OnnxRuntimeBase::forward(Ort::Value& inputTensors)
{
    return ort_session.Run(RunOptions{ nullptr }, input_names.data(), &inputTensors, 1, output_names.data(), output_names.size());
}

我使用的 ONNX Runtime 版本是 1.10.0，如果使用更高的版本 API 会略微有不同。

接下来是定义 SketchDrawing 类，继承自 OnnxRuntimeBase 类，用于部署模型和完成推理。

#include <iostream>
#include <fstream>
#include <string>
#include <math.h>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>
#include <onnxruntime_cxx_api.h>
#include "../common/OnnxRuntimeBase.h"

using namespace cv;
using namespace std;
using namespace Ort;

class SketchDrawing: public OnnxRuntimeBase {

public:
    SketchDrawing(std::string modelPath, const char* logId, const char* provider);
    Mat detect(Mat& image);

private:
    vector<float> input_image_;
    int inpWidth;
    int inpHeight;
    int outWidth;
    int outHeight;
};

SketchDrawing 类的实现可以参考项目的源码。

三. 应用层调用

定义和实现完 SketchDrawing 类后，需要通过 jni 暴露给应用层使用。下面是 jni 初始化相关的模型，生成素描画的方法。

SketchDrawing *sketchDrawing = nullptr;

JNIEXPORT void JNICALL Java_cn_netdiscovery_monica_opencv_ImageProcess_initSketchDrawing
        (JNIEnv* env, jobject,jstring jModelPath) {

     const char* modelPath = env->GetStringUTFChars(jModelPath, JNI_FALSE);
     const std::string& onnx_logid = "Sketch Drawing";
     const std::string& onnx_provider = OnnxProviders::CPU;
     sketchDrawing = new SketchDrawing(modelPath, onnx_logid.c_str(), onnx_provider.c_str());

     env->ReleaseStringUTFChars(jModelPath, modelPath);
}

JNIEXPORT jintArray JNICALL Java_cn_netdiscovery_monica_opencv_ImageProcess_sketchDrawing
        (JNIEnv* env, jobject,jbyteArray array) {

    Mat image = byteArrayToMat(env,array);
    Mat dst;

    try {
        dst = sketchDrawing->detect(image);
    } catch(...) {
    }

    jthrowable mException = NULL;
    mException = env->ExceptionOccurred();

    if (mException != NULL) {
       env->ExceptionClear();
       jclass exceptionClazz = env->FindClass("java/lang/Exception");
       env->ThrowNew(exceptionClazz, "jni exception");
       env->DeleteLocalRef(exceptionClazz);

       return env->NewIntArray(0);
    }

    if (dst.channels() == 1) {
        cvtColor(dst,dst,COLOR_GRAY2BGR);
    }

    return matToIntArray(env,dst);
}

对于应用层，需要编写好调用 jni 层的代码：

object ImageProcess {

    val loadPath = System.getProperty("compose.application.resources.dir") + File.separator

    init {
        // 需要先加载图像处理库，否则无法通过 jni 调用算法
        LoadManager.load()
    }

    ......

    /**
     * 初始化生成素描画模块
     */
    external fun initSketchDrawing(modelPath:String)

    /**
     * 生成素描画
     */
    external fun sketchDrawing(src: ByteArray):IntArray
}

在 Monica 启动时，先加载必须的模型

runInBackground { // 初始化生成素描画的模块
            DLManager.initSketchDrawingModule()
        }

最后，终于可以在应用层调用了

val (width,height,byteArray) = state.currentImage!!.getImageInfo()

                    try {
                        val outPixels = ImageProcess.sketchDrawing(byteArray)
                        state.addQueue(state.currentImage!!)
                        state.currentImage = BufferedImages.toBufferedImage(outPixels,width,height)
                    } catch (e:Exception) {
                        logger.error("faceDetect is failed", e)
                    }

我们来看看在 Monica 中使用的效果

image1.png

image1的效果.png

image2.png

image2的效果.png

通过 ONNX Runtime 加速，推理的时间大概花费了0.5s

image: 4596x3064 538.9ms
speed: 5.8ms preprocess, 529.1ms inference, 4.0ms postprocess per image

四. 总结

Monica 快要到 1.0.0 版本，后续的重点依然是优化软件的架构，其次才会考虑引入一些比较有趣的深度学习的模型。如果模型过大的话，后续也可能会考虑放在服务端部署。

Monica github 地址：https://github.com/fengzhizi715/Monica

【Java与Android技术栈】公众号

关注 Java/Kotlin 服务端、桌面端 、Android 、机器学习、端侧智能

更多精彩内容请关注：