AI 视觉实战 2：实时头发染色

2023-04-21
北京
本文字数：3289 字
阅读完需：约 11 分钟

1. 背景介绍

在实时视频编辑领域，头发变色、修改发型是很流行和受欢迎的场景。这种功能除了音视频相关的技术，还离不开 AI 能力的支持。而且这种场景本身对实时性要求高，很适合在端侧应用落地。上一篇文章我们基于谷歌的 MediaPipe 项目实现了本地实时人脸检测功能，本文我们再来一步一步跑通端侧实时染色功能。下面是效果：

2. 需求分析

上一篇中，人脸检测输入是一帧帧图片，输出是识别到的人脸数量，坐标及对应得分列表，我们可以通过得分与设置的阈值比较判断是否有人脸，还可以根据返回的坐标，给人脸标一个方框。

实时头发染色功能输入的仍然是一帧帧图片，因为头发本身是不规则的，如果输出坐标的话很难再去绘制，所以这次模型为我们返回了一个完整的图片内容，图片上是变色的头发的内容，并且和原图片头发位置坐标保持一直，这样我们可以先绘制原图像，再绘制变色的染色头发图片。

3. 代码实现

和上一篇类似，运行模型一般我们有以下几个步骤：

加载模型；
摄像头预览纹理转换为 RGBA
将图像数据 feed 到模型引擎进行推理
解析渲染结果

3.1 加载模型

hair_segmentation 模型加载时tflite::ops::builtin::BuiltinOpResolver新增了三个自定义 operations：

tflite::ops::builtin::BuiltinOpResolver  resolver;resolver.AddCustom("MaxPoolingWithArgmax2D",            mediapipe::tflite_operations::RegisterMaxPoolingWithArgmax2D());
resolver.AddCustom("MaxUnpooling2D",            mediapipe::tflite_operations::RegisterMaxUnpooling2D());
resolver.AddCustom("Convolution2DTransposeBias",            mediapipe::tflite_operations::RegisterConvolution2DTransposeBias());

复制代码

对应实现函数：

TfLiteRegistration* RegisterMaxPoolingWithArgmax2D() {  static TfLiteRegistration reg = {      [](TfLiteContext*, const char*, size_t) -> void* {        return new TfLitePaddingValues();      },      [](TfLiteContext*, void* buffer) -> void {        delete reinterpret_cast<TfLitePaddingValues*>(buffer);      },      Prepare, Eval};  return ®}


TfLiteRegistration* RegisterMaxUnpooling2D() {  static TfLiteRegistration reg = {      [](TfLiteContext*, const char*, size_t) -> void* {        return new TfLitePaddingValues();      },      [](TfLiteContext*, void* buffer) -> void {        delete reinterpret_cast<TfLitePaddingValues*>(buffer);      },      Prepare, Eval};  return ®}

TfLiteRegistration* RegisterConvolution2DTransposeBias() {  static TfLiteRegistration reg = {nullptr, nullptr, Prepare, Eval};  return ®}

复制代码

通过InterpreterBuilder创建执行器std::unique_ptr<tflite::Interpreter>后，获取模型输入输出函数：

static tflite_tensor_t      s_tensor_input;static tflite_tensor_t      s_tensor_segment;
tflite_get_tensor_by_name (&s_interpreter, 0, "input_1",  &s_tensor_input);tflite_get_tensor_by_name (&s_interpreter, 1, "conv2d_transpose_4",  &s_tensor_segment);

复制代码

tflite_tensor_t 结构有 ptr 指针成员，输入时存放图像信息，输出时存放被渲染过的头发的图像信息。

3.2 摄像头预览纹理转换为 RGBA

纹理转 RGBA 跟上一篇人脸检测一样，不在赘述。

3.3 将图像数据 feed 到模型引擎进行推理

feed 数据到模型跟上一篇人脸检测一样，不在赘述。feed 完后开始执行推理：

typedef struct _segmentation_result_t{    float *segmentmap;    int   segmentmap_dims[3];} segmentation_result_t;
int invoke_segmentation (segmentation_result_t *segment_result){    if (interpreter->Invoke() != kTfLiteOk)    {        DBG_LOGE ("ERR: %s(%d)\n", __FILE__, __LINE__);        return -1;    }
    segment_result->segmentmap         = (float *)s_tensor_segment.ptr;    segment_result->segmentmap_dims[0] = s_tensor_segment.dims[2];    segment_result->segmentmap_dims[1] = s_tensor_segment.dims[1];    segment_result->segmentmap_dims[2] = s_tensor_segment.dims[3];
    return 0;}

复制代码

结果主要包含被染发的图像数据。

3.4 解析渲染结果

绘制时先绘制原始图像纹理，然后绘制模型返回的修改后的数据：

void render_segment_result (int ofstx, int ofsty, int draw_w, int draw_h,                        texture_2d_t *srctex, segmentation_result_t *segment_ret){    float *segmap = segment_ret->segmentmap;    int segmap_w  = segment_ret->segmentmap_dims[0];    int segmap_h  = segment_ret->segmentmap_dims[1];    int segmap_c  = segment_ret->segmentmap_dims[2];    int x, y, c;    static unsigned int *imgbuf = NULL;    float hair_color[4] = {0};    float back_color[4] = {0};    static float s_hsv_h = 0.0f;
    if (imgbuf == NULL)    {        imgbuf = (unsigned int *)malloc (segmap_w * segmap_h * sizeof(unsigned int));    }
    s_hsv_h += 5.0f;    if (s_hsv_h >= 360.0f)        s_hsv_h = 0.0f;
    colormap_hsv (s_hsv_h / 360.0f, hair_color);
#if defined (RENDER_BY_BLEND)    float lumi = (hair_color[0] * 0.299f + hair_color[1] * 0.587f + hair_color[2] * 0.114f);    hair_color[3] = lumi;#endif
    /* find the most confident class for each pixel. */    for (y = 0; y < segmap_h; y ++)    {        for (x = 0; x < segmap_w; x ++)        {            int max_id;            float conf_max = 0;            for (c = 0; c < MAX_SEGMENT_CLASS; c ++)            {                float confidence = segmap[(y * segmap_w * segmap_c)+ (x * segmap_c) + c];                if (c == 0 || confidence > conf_max)                {                    conf_max = confidence;                    max_id = c;                }            }
            float *col = (max_id > 0) ? hair_color : back_color;            unsigned char r = ((int)(col[0] * 255)) & 0xff;            unsigned char g = ((int)(col[1] * 255)) & 0xff;            unsigned char b = ((int)(col[2] * 255)) & 0xff;            unsigned char a = ((int)(col[3] * 255)) & 0xff;
            imgbuf[y * segmap_w + x] = (a << 24) | (b << 16) | (g << 8) | (r);        }    }
    GLuint texid;    glGenTextures (1, &texid );    glBindTexture (GL_TEXTURE_2D, texid);
    glTexParameterf (GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);    glTexParameterf (GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);    glTexParameterf (GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);    glTexParameterf (GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
    glPixelStorei (GL_UNPACK_ALIGNMENT, 4);
    glTexImage2D (GL_TEXTURE_2D, 0, GL_RGBA,        segmap_w, segmap_h, 0, GL_RGBA,        GL_UNSIGNED_BYTE, imgbuf);
#if !defined (RENDER_BY_BLEND)    draw_colored_hair (srctex, texid, ofstx, ofsty, draw_w, draw_h, 0, hair_color);#else    draw_2d_texture_ex (srctex, ofstx, ofsty, draw_w, draw_h, 0);
    unsigned int blend_add  [] = {GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA, GL_ZERO, GL_ONE};    draw_2d_texture_blendfunc (texid, ofstx, ofsty, draw_w, draw_h, 0, blend_add);#endif
    glDeleteTextures (1, &texid);
    render_hsv_circle (ofstx + draw_w - 100, ofsty + 100, s_hsv_h);}

复制代码

4. 总结

本文介绍了 AI 技术的实时头发染色模型使用，主要应用于视频特效编辑等场景。该模型用到了 BuiltinOpResolver 的 AddCustom 新方法。

发布于: 47 分钟前阅读数: 7

原文链接:【http://xie.infoq.cn/article/9140e1d5bde04c6b4399abd1d】。文章转载请联系作者。

轻口味

关注

🏆2021年InfoQ写作平台-签约作者 🏆 2017-10-17 加入

Android、音视频、AI相关领域从业者。欢迎加我微信wodekouwei拉您进InfoQ音视频沟通群邮箱：qingkouwei@gmail.com

发布

暂无评论

创作场景

AI 视觉实战 2：实时头发染色

1. 背景介绍

2. 需求分析

3. 代码实现

3.1 加载模型

3.2 摄像头预览纹理转换为 RGBA

3.3 将图像数据 feed 到模型引擎进行推理

3.4 解析渲染结果

4. 总结

轻口味

评论