• TOP
• 機械学習

[Android][TensorFlow Lite]セグメンテーションを使い写真の背景を編集する

2022-11-23T21:27:00+09:00 2022-11-23T12:27:36Z Android 機械学習

Zoomやteamsなどのビデオ会議アプリでは、ビデオ会議時に背景をぼかしたらり変更したりすできます。
どのように実施しているのか気になり画像から人物領域を抽出する方法を見つけました。

BodyPix: Real-time Person Segmentation in the Browser with TensorFlow.js

驚くことに2019年の時点でtensorflowのBodyPixと言うライブラリで人物の領域・姿勢を認識できていました。
このライブラリはWebブラウザ上で動作できます。ほとんど会議アプリがWebブラウザベースですので、背景の変更が当然のようにできるわけですね。

Androidでも同じことができないかと思い調べてみました。
tensorflowの公式ページでtensorflow lite用にセグメンテーションAPIが公開されていました。

セグメンテーション

これを使うと画像から以下のカテゴリの領域を抽出できます。

'background', 'aeroplane', 'bicycle', 'bird', 'boat', 
'bottle', 'bus', 'car', 'cat', * 'chair', 'cow', 
'diningtable', 'dog', 'horse', 'motorbike', 
'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tv'

サンプルコードimage_segmentationのImageSegmentationHelperを使うと容易に利用することができます。
処理部分だけのコードを抜き出すと以下ようになります。

fun imageSegmentation(bitmap:Bitmap){
    //検出処理の実行
    imageSegmentationHelper.segment(bitmap!!, 0)
}

//検出結果の受け取りを行う
override fun onResults(
    results: List<Segmentation>,
    inferenceTime: Long,
    imageHeight: Int,
    imageWidth: Int
) {
        // 色と検出されたラベルのセットを使用してマスク ビットマップを作成します。
        // 1つのマスクのみを提供する OutputType CATEGORY_MASK を使用しているため、
        // このサンプルでは最初のマスクのみが必要です。
        val maskTensor = segmentResult[0].masks[0]
        val maskArray = maskTensor.buffer.array()
        val pixels = IntArray(maskArray.size)

        //マスクされたビットマップを作る
        for (i in maskArray.indices) {
            val colorLabels = segmentResult[0].coloredLabels.mapIndexed { index, coloredLabel ->
                ColorLabel(
                    index,
                    coloredLabel.getlabel(),
                    coloredLabel.argb
                )
            }

            // 人物のみのマスク画像を得る
            val colorLabel = colorLabels[maskArray[i].toInt()]
            //if(colorLabel.label != "person")continue
            val color = colorLabel.getColor()
            pixels[i] = color
        }

        //画像を作成する
        val image = Bitmap.createBitmap(
            pixels,
            maskTensor.width,
            maskTensor.height,
            Bitmap.Config.ARGB_8888
        )

        //元画像のサイズにあわせたマスク画像を作成する
        val maskedBitmap = Bitmap.createScaledBitmap(image, imageWidth, imageHeight, false)

}

この処理を組み合わせることで、以下のような画像を作れます。 上が元の画像で、下が背景を編集した画像です。