Posenet (TensorFlow Lite) のヒートマップをOpenGLESで可視化

概要

　Posenetが出力する各関節ごとの確信度(ヒートマップ)を、OpenGLESのシェーダを用いて高速に可視化します。
　組み込み向けを意識しているので、Posenet は TensorFlowLite C++ 環境で実行します。

Posenetについて

　下記のように、人体の関節位置を抽出し、姿勢推定を行うことができるDNNモデルです。

　Posenetは、学習済みモデルが公開されており、また、様々な環境向けのソースコードも公開されているので、誰でも簡単に試すことができます。

参考

• [GoogleのPosenet紹介ページ] (https://www.tensorflow.org/lite/models/pose_estimation/overview)
• [ブラウザ内でPosenet実行(TFJS)] (https://github.com/tensorflow/tfjs-models/tree/master/posenet)
• [Android端末でPosenet実行] (https://github.com/tensorflow/examples/tree/master/lite/examples/posenet/android)
• [EdgeTPUを用いてPosenet実行] (https://github.com/google-coral/project-posenet)

TensorFlow Lite C++ 版 Posenet の実装

　組み込み環境でPosenet実行することを想定した場合、上記サンプルのような JavaScript や kotolin ではなく、TensorFlow Lite C++ で実装しておいたほうが環境依存が少なくて済みそうです。
　が、少し時間をかけてWeb検索しても、TensorFlow Lite C++で実装されているコードが見つけられなかったので、Android用ソース(kotolin実装) を参考に、自前でC++へ移植しました。

　開発したものは [GitHub] (https://github.com/terryky/tflite_gles_app/tree/master/gl2posenet) で公開しています。
　※TensorFlow Lite 依存部は300行程度のすごくシンプルなコードですので、移植に関する説明はスキップとさせてください^^;

実行画面

下図は今回開発した TensorFlow Lite C++ 版の画面キャプチャです。

• 処理実行時間を画面左に表示するようにしました。
• 処理時間が変動することを想定し、処理時間を時系列表示するグラフも描画しています。
• 全ての描画はOpenGLESを使っています。

Posenetが出力するヒートマップについて

　本稿の本題であるヒートマップの説明に入る前に、Posenetネットワークモデルの出力層を再確認します。

　このモデルは４つのテンソルを出力しますが、Android用ソース(kotolin実装)によれば、１番左のものがヒートマップのようです。
　ヒートマップ出力用テンソルの次元は (1x9x9x17) です。これは下図のように、(9x9)解像度のヒートマップが各関節ごとに並んでるイメージだと考えるとわかりやすいと思います。１枚のヒートマップは、ある関節（例えば「右の足首」）に注目した時に、その関節が画像中のどこに存在すると推定されるかを表現したものであり、より具体的には、入力画像を(9x9)の領域にタイル分割して、個々のタイル毎にその関節が存在するか否かを確信度として算出したものです。以下同様に全ての関節（17個）についてヒートマップを作成したものが、出力テンソルに格納されています。

OpenGLES シェーダによる Jet カラーマップの生成

Web上の情報では、DNNの推論結果を可視化するのに、OpenCVや各種pythonモジュールを使った例が多くみられます。が、下記の理由（個人的な思い込み含む）からOpenGLESを使って可視化を行います。

• 餅は餅屋。やっぱり描画するなら OpenGLES だよね。
• 組み込み環境でも OpenGLESはサポートされている環境が多いし安心
• 描画する目的でOpenCVを使うのは道具が違う気がする。
• 組み込み環境でpython動かすのはまだ時期尚早な気がする。

MATLAB Jet カラーマップ

　やりたいことは、レンジ(0.0～1.0)の入力値を、下記のカラーマップ (MATLAB Jet)に従って色変換することです。一番簡単な実装方法は、下記グラデーション画像を (256x1) サイズのテクスチャとして保持しておいて、フラグメントシェーダでテクスチャ引きすることで直接色値を取得する方法だと思います。が、今回はグラデーションを試行錯誤で変更したかったこともあり、フラグメントシェーダで動的にカラーマップを生成することにしました。

というわけで、MATLAB Jet カラーマップ描画をOpenGLESで描画するための、GLSLシェーダコードと、ホスト側コードの紹介です。

GLSL フラグメントシェーダ

レンジ(0.0～1.0)の入力値は、LUMINANCEテクスチャとして設定されることを想定しています。
グラデーションの計算式自体はWebに情報があったものを適用しました。元画像に重畳するためのブレンド係数(A値)はホストプログラムから uniform 変数として設定されることを想定しています。

precision mediump float;
varying     vec2      v_TexCoord;
uniform     sampler2D u_sampler;
uniform     float     u_alpha;

float cmap_jet_red(float x) {
    if (x < 0.7) {
        return 4.0 * x - 1.5;
    } else {
        return -4.0 * x + 4.5;
    }
}

float cmap_jet_green(float x) {
    if (x < 0.5) {
        return 4.0 * x - 0.5;
    } else {
        return -4.0 * x + 3.5;
    }
}

float cmap_jet_blue(float x) {
    if (x < 0.3) {
       return 4.0 * x + 0.5;
    } else {
       return -4.0 * x + 2.5;
    }
}

vec4 colormap_jet(float x) {
    float r = clamp(cmap_jet_red(x),   0.0, 1.0);
    float g = clamp(cmap_jet_green(x), 0.0, 1.0);
    float b = clamp(cmap_jet_blue(x),  0.0, 1.0);
    return vec4(r, g, b, 1.0);
}

/* 入力は (0～1)の値を持つLUMINANCEテクスチャ。これをカラーマップ化する */
void main (void)
{
    vec4 src_col = texture2D (u_sampler, v_TexCoord);
    gl_FragColor = colormap_jet (src_col.r);
    gl_FragColor.a *= u_alpha;
}

ホスト側コード

　TensorFlow Lite で Invoke() した後、出力テンソル領域にヒートマップ情報が格納されるので、そのポインタからデータを読み出して、ヒートマップの解像度と同じの大きさ (ここでは 9x9) のLUMINANCEテクスチャを生成します。
　フラグメントシェーダに入力するヒートマップデータは [0.0～1.0] のレンジに収める必要がありますが、TensorFlow Lite が出力する値のレンジはそれよりも広いです。本来なら、グラデーション描画するレンジ（最小値と最大値)を慎重に決める必要がありますが、試行錯誤の結果、TensorFlow Lite の出力値のうち [-5.0～1.0] のレンジの値をグラデーション描画するといい感じになったのでそれを採用しています。

    /* Posenet推論実行結果が格納される、ヒートマップテンソルのポインタ */
    float *heatmap = pose_ret->pose[0].heatmap;         /* TensorFlow Lite出力ポインタ */
    int heatmap_w  = pose_ret->pose[0].heatmap_dims[0]; /* ヒートマップ解像度(W) */
    int heatmap_h  = pose_ret->pose[0].heatmap_dims[1]; /* ヒートマップ解像度(H) */
    unsigned char imgbuf[heatmap_w * heatmap_h];
    int key_id = kRightShoulder; /* 可視化したい関節(ここでは右肩) */

    /* グラデーションで表現可能なレンジを（エイヤで固定的に）設定 */
    float conf_min = -5.0f;
    float conf_max =  1.0f;

    /* ヒートマップの値を (0～1) に正規化 */
    for (int y = 0; y < heatmap_h; y ++)
    {
        for (int x = 0; x < heatmap_w; x ++)
        {
            float confidence = heatmap[(y * heatmap_w * kPoseKeyNum) +
                                       (x             * kPoseKeyNum) + key_id];
            confidence = (confidence - conf_min) / (conf_max - conf_min);
            if (confidence < 0.0f) confidence = 0.0f;
            if (confidence > 1.0f) confidence = 1.0f;
            imgbuf[y * heatmap_w + x] = confidence * 255;
        }
    }

    /* カラーマップ描画用テクスチャ生成 */
    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, 1);
    glTexImage2D (GL_TEXTURE_2D, 0, GL_LUMINANCE,
        heatmap_w, heatmap_h, 0, GL_LUMINANCE,
        GL_UNSIGNED_BYTE, imgbuf);

    /* グラデーションテクスチャを元画像の上にブレンド描画 */
    draw_2d_colormap (texid, ofstx, ofsty, draw_w, draw_h, 0.8f, 0);

ヒートマップ描画結果

　ヒートマップ描画結果が下記です。
　17枚のヒートマップ画像をアニメーションGIFとして順に表示しています。
　このテスト画像の場合、関節位置はほぼ正しい位置に推定できていますが、ヒートマップをみると、ところどころ、左半身なのか右半身なのかを迷っているような気配が見られます。（例えば左足首のヒートマップに、本来の左足首だけでなく、右足首にもホットスポットが生じているなど）。

別の例

　別のテスト画像では、ふり上げた右足に追随できませんでした。

　ヒートマップを見てみると、右足首用のヒートマップにおいて、確かに正解位置である右足首の位置に反応している様子が見て取れますが、誤った位置である左足首の位置に強く反応しているため、誤推定しているようです。

まとめ

　DNN全体に言えることですが、誤認識した原因を探るのは困難な作業となります。
　Posenetを利用する際には、左半身/右半身の判断があやうい可能性がある、といった特性を理解したうえで、注意深くチューニングしていく必要があると感じました。

今回作ったソースコード

GitHubで公開しています。
https://github.com/terryky/tflite_gles_app/tree/master/gl2posenet