iOSで動く姿勢推定モデルを色々と試す #Swift

iOSのCore MLモデルは、Appleが公式で色々と配布していますが、6種類のうちなんと5種類が物体認識（imagenetの1000クラス分類）。しかし、GitHub等をみると、多種多様なモデルが公開されてきています。

今回は姿勢推定（Pose Estimation）のCore MLモデルで公開されているものをいくつかピックアップし、それらがどれぐらい使えるのか、試してみました。

PoseNet-CoreML

まずはこちら。

サンプルをそのまま試す

こちらはサンプルが同梱されていて、すぐにビルドして試せました。

結果はこんな感じです。

（実際にiPhone XS上で実行してスクリーンショットを撮った結果を正方形にクロップして載せています）

後ろの観客の姿勢推定はできているようですが、肝心の前方のプレイヤーの姿勢が取れていません。。

他にも画像が同梱されていたので、変えてみました。

あれ・・・片足や腕など、諸々検出できておらず、正直精度はイマイチです。

モデルを入れ替えてみる

プロジェクトを見てみると、mlmodelが複数入っていました。

ソースコードを見てみると、使っているのはposenet513_v1_075.mlmodelのようです。

let model = posenet513_v1_075()

"_075"というのは何か精度よりも軽量化あるいは高速化に振ったような感じを受ける命名なので、他のモデルを試してみます。実際、モデルサイズは他のものより小さいようでした。

posenet513.mlmodelというモデルを見ると、

入力形式は同じなので、そのまま実行できそうです。こちらに入れ替えて試してみましょう。

let model = posenet513()

なんと！プレイヤーの姿勢もばっちりとれるようになりました。

重量挙げの方も見てみましょう。

こちらは腕や足はしっかりとれるようになったものの、胴がとれなくなりました。まだまだ満足の行く精度とはいえないようです。

実装

入力としては画像をCVPixelBuffer型で受け取ります。

func runCoreML(_ img: CVPixelBuffer) -> [Pose]{        
    let result = try? model.prediction(image__0: img)

ここは普通ですが、出力側の扱いがかなり従来の物体認識等のモデルと違っていました。

let tensors = result?.featureNames.reduce(into: [String: Tensor]()) {
    $0[$1] = getTensor(
        result?.featureValue(for: $1)?.multiArrayValue)
}
let sum = tensors!["heatmap__0"]!.reduce(0, +) / (17 * 33 * 33)
print(sum)

let poses = posenet.decodeMultiplePoses(
    scores: tensors!["heatmap__0"]!,
    offsets: tensors!["offset_2__0"]!,
    displacementsFwd: tensors!["displacement_fwd_2__0"]!,
    displacementsBwd: tensors!["displacement_bwd_2__0"]!,
    outputStride: 16, maxPoseDetections: 15,
    scoreThreshold: 0.5,nmsRadius: 20)

最終的にPoseという構造体を結果として生成しているのですが、

struct Pose {
    var keypoints: [Keypoint]
    var score: Float
}

上に載せたコードの通り、Core MLモデルからの出力をアプリ側でゴニョゴニョとした上でこのPose構造体にしています。これらの意味は理解できてないので使いこなすにはもっとモデル自体への理解を深める必要がありそうです。精度もこのへんの扱いが影響してそうです。

・・・とここまで書いてREADMEを見て気付いたのですが、どうやら作者の方は日本人で、御本人が書かれた記事を見つけてしまいました。

PoseNetをマルチプラットフォームで実装してみた - Qiita

後ほど拝読します。

PoseEstimationForMobile

こちらはAndroidしかデモが同梱されておらず、iOSのデモは別の方が別リポジトリで公開しているようです。

こちらのリポジトリにはmlmodelが入っていません。「PoseEstimationForMobile」リポジトリのreleaseフォルダに生成済みのmlmodelファイルが入っているので、ここからサンプル配下のフォルダにコピーしてきます。サンプルプロジェクトには

model_hourglass.mlmodel
model_cpm.mlmodel

の２つのモデルへの参照が追加されていますが、コードを見るとmodel_hourglassの方はテストでしか使われていないので、いったんプロジェクトから削除してしまってもサンプルは動作します。

model_cpm.mlmodelの入力・出力形式は次のようになっていました。

サンプルを動かしてみると、こちらはリアルタイムカメラ式で、比較のため「PoseNet-CoreML」のサンプルに入っていたテスト用画像をモニタに移して姿勢推定させてみました。

かなり良さそうです。

せっかくなのでリアルタイムで動いている様子も録りました。まわりに人がいないので対象がGoogle画像検索結果ですが。。

実装

入力は普通にカメラから得たCVPixelBufferを渡しています。

重要なのは出力のMultiArrayオブジェクトをどう処理するかですが、

JointViewController.swift

if let observations = request.results as? [VNCoreMLFeatureValueObservation],
    let heatmap = observations.first?.featureValue.multiArrayValue {

    // convert heatmap to [keypoint]
    let n_kpoints = heatmap.convertHeatmapToBodyPoint()

    DispatchQueue.main.sync {
        // draw line
        self.jointView.bodyPoints = n_kpoints

        // show key points description
        self.showKeypointsDescription(with: n_kpoints)

        // end of measure
        self.👨‍🔧.🎬🤚()
    }
}

こんな感じになっていて、MultiArrayオブジェクトから取得した諸々をゴニョゴニョして人間の姿勢の各点を表すBodyPointの配列を返しているようです。MultiArrayからBodyPointの配列を得る処理convertHeatmapToBodyPointはMLMultiArray+PoseEstimation.swiftにあります。