Android使用MediaPipe + FFmpeg生成绿幕视频

前言

事情的起因是最近在研究利用FFmpeg替换绿幕视频的背景，当然这个并不是本文要探讨的问题，主要是在这个过程中笔者在思考假设我现在拥有的素材背景不是纯色的，有没有什么方案可以将人抠出来呢？ps：类似的场景比较常见的应该是视频聊天时的背景替换吧。

于是就想到的机器学习，利用模型抠人像。通过GPT和百度等途径了解到Google的MediaPipe提供了一套人像分割的能力。所以本文主要记录的是笔者通过MediaPipe进行人像分割，再对每一帧进行背景替换，使用FFmpeg合成视频，最终生成一个绿幕视频这么一个过程。

MediaPipe

细节的东西就不介绍了，因为咱不熟悉，大家可以通过MediaPipe官网了解。

对于人像分割，MediaPipe也给出了开发引导，也可以通过MediaPipe Studio在线尝试效果。

Python实现

可以先PC上通过Python运行看看效果，通过询问GPT再稍作修改的python脚本是长这样的：

import cv2
import numpy as np
import mediapipe as mp
import os

from mediapipe.tasks import python
from mediapipe.tasks.python import vision

input_path = "input.mp4"
output_video_path = "output.mp4"

base_options = python.BaseOptions(model_asset_path='selfie_segmenter.tflite')
options = vision.ImageSegmenterOptions(base_options=base_options,
                                       running_mode=vision.RunningMode.VIDEO,
                                       output_category_mask=True)

# 创建保存分割后帧的目录
os.makedirs("temp_frames", exist_ok=True)
os.makedirs("segmented_frames", exist_ok=True)

cap = cv2.VideoCapture(input_path)
# 获取视频的帧率和大小
fps = cap.get(cv2.CAP_PROP_FPS)
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

print("fps: {}, width: {}, height: {}".format(fps, width, height))

frame_number = 0

segmenter = vision.ImageSegmenter.create_from_options(options)

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break

    # 保存帧为临时图像文件
    frame_path = f"temp_frames/temp_frame_{frame_number}.jpg"
    # print("frame: {}".format(frame))
    # 视频文件的当前位置，以ms为单位
    ts = cap.get(cv2.CAP_PROP_POS_MSEC)
    # print("ts: {}".format(ts))
    # cv2.imwrite(frame_path, frame)

    mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=frame)

    segmentation_result = segmenter.segment_for_video(mp_image, int(ts))
    category_mask = segmentation_result.category_mask

    image_data = mp_image.numpy_view()
    fg_image = np.zeros(image_data.shape, dtype=np.uint8)
    fg_image[:] = (255, 255, 255)

    # image_data = cv2.cvtColor(mp_image.numpy_view(), cv2.COLOR_BGR2RGB)
    condition = np.stack((category_mask.numpy_view(),) * 3, axis=-1) > 0.2
    output_image = np.where(condition, fg_image, image_data)

    # 保存分割后的帧
    segmented_frame_path = f"segmented_frames/segmented_frame_{frame_number}.jpg"
    cv2.imwrite(segmented_frame_path, output_image)

    frame_number += 1

# 使用 FFmpeg 将帧合成为视频
os.system(f"ffmpeg -framerate {fps} -i segmented_frames/segmented_frame_%d.jpg -c:v libx264 -pix_fmt yuv420p {output_video_path}")
# 清理临时文件
os.system("rm -r temp_frames")
os.system("rm -r segmented_frames")

cap.release()

大体思路是：

1. 通过OpenCV逐帧解析
2. 利用MediaPipe进行人像切割，这里使用的输出类型是Category Mask，它会输出一个数组可以理解为图片的像素点，一般情况下会存在数值0和255，0代表有人像部分，255代表其他。
3. 通过对输出结果结合原图处理，将不是0的部分处理成一个颜色，这里是黑色。
4. 最后利用ffmpeg将全部帧合成视频输出。

由于MediaPipe对Python版本和Mac版本有要求，笔者目前是MacOS 12.0.1，Python版本是3.11.5，所以使用的是MediaPipe 0.10.0版本。

从上述脚本中也能看出，使用的是TensorFlow Lite的模型，模型可以通过开发引导中下载。

参考文章

基于Mediapipe人像实时语义分割------抠图黑科技

Android集成实现

逐帧读取

那么怎样在Android上实现呢？参考上述脚本的思路，假设我们现在有一个视频的Uri，我们可以通过MediaMetadataRetriever来读取视频的信息。

比如可以读取到视频的播放时长、帧数、分辨率等。

val retriever = MediaMetadataRetriever()
retriever.setDataSource(this@MainActivity, uri)
// 视频时长
val videoLengthMs =
    retriever.extractMetadata(MediaMetadataRetriever.METADATA_KEY_DURATION)?.toLong()
// 帧数
val frameCount =
  retriever.extractMetadata(MediaMetadataRetriever.METADATA_KEY_VIDEO_FRAME_COUNT)?.toInt()
// 分辨率
val firstFrame = retriever.getFrameAtTime(0)
val width = firstFrame?.width
val height = firstFrame?.height

可以通过MediaMetadataRetriever#getFrameAtTime按照时长读取每一帧。

// @param timestampMs 视频文件的播放位置
val frame = retriever.getFrameAtTime(
    timestampMs * 1000, // convert from ms to micro-s
    MediaMetadataRetriever.OPTION_CLOSEST
)

然后把它们串起来

val retriever = MediaMetadataRetriever()
retriever.setDataSource(this@MainActivity, uri)
val videoLengthMs =
    retriever.extractMetadata(MediaMetadataRetriever.METADATA_KEY_DURATION)?.toLong()
val frameCount =
    retriever.extractMetadata(MediaMetadataRetriever.METADATA_KEY_VIDEO_FRAME_COUNT)?.toInt()

val firstFrame = retriever.getFrameAtTime(0)
val width = firstFrame?.width
val height = firstFrame?.height

if ((frameCount == null) || (videoLengthMs == null) || (width == null) || (height == null)) return

// 总共需要读取的帧数
val numberOfFrameToRead = frameCount
// 平均每帧的时长
val preFrameLengthMs = videoLengthMs.div(frameCount)

for (i in 0..numberOfFrameToRead) {
    val timestampMs = i * preFrameLengthMs // ms
    val frame = retriever.getFrameAtTime(
        timestampMs * 1000, // convert from ms to micro-s
        MediaMetadataRetriever.OPTION_CLOSEST
    )
}    

1. 通过读取的视频帧数，可以得知一共需要读取多少帧
2. 通过视频时长和帧数计算出平均每帧的时长
3. 循环读取每一帧

MediaPipe人像切割

在Android使用前，需要通过Gradle集成MediaPipe

//  MediaPipe
implementation 'com.google.mediapipe:tasks-vision:0.10.0'

在输入前，由于MediaPipe默认需要ARGB_8888的Bitmap所以这里需要进行一次转化，然后生成MediaPipe的MPImage对象作为输入

val argb8888Frame =
    if (frame.config == Bitmap.Config.ARGB_8888) frame
    else frame.copy(Bitmap.Config.ARGB_8888, false)

// Convert the input Bitmap object to an MPImage object to run inference
val mpImage = BitmapImageBuilder(argb8888Frame).build()

接下来就是进行人像切割，在此之前需要先初始化它的API，

• 使用到的模型和上述脚本一致
• 处理的是视频文件所以使用RunningMode.VIDEO
• 使用的输出类型是Category Mask，所以这里为true

private val imagesegmenter: ImageSegmenter by lazy {
    val baseOptions = BaseOptions.builder()
        .setDelegate(Delegate.CPU)
        .setModelAssetPath("selfie_segmenter.tflite")
        .build()

    val options = ImageSegmenter.ImageSegmenterOptions.builder()
        .setRunningMode(RunningMode.VIDEO)
        .setBaseOptions(baseOptions)
        .setOutputCategoryMask(true)
        .setOutputConfidenceMasks(false)
        .build()

    return@lazy ImageSegmenter.createFromOptions(this, options)
}

由于我们是对视频文件进行处理，所以使用的是segmentForVideo，timestampMs就是上述提到的视频文件的播放位置。

val result = imagesegmenter.segmentForVideo(mpImage, timestampMs)

val newImage = result.categoryMask().get()
val resultByteBuffer = ByteBufferExtractor.extract(newImage)

val pixels = IntArray(resultByteBuffer.capacity())

这里我们可以得到一个int数组，这个数组和上述脚本一致，存在数值0和255，0代表有人像部分，255代表其他。

最后通过这个输出结果pixels再转换成二维坐标，如果值是0则维持原帧在改点的颜色，如果是255则置为绿色，再生成一个最终的Bitmap进行保存。

val frameW = argb8888Frame.width    //  列数
for (index in pixels.indices) {
    // Using unsigned int here because selfie segmentation returns 0 or 255U (-1 signed)
    // with 0 being the found person, 255U for no label.
    pixels[index] =
        if (resultByteBuffer.get(index).toUInt() > 0U) Color.GREEN else {
            val x = index % frameW
            val y = index / frameW
            argb8888Frame.getPixel(x, y)
        }
}
val resultFrame = Bitmap.createBitmap(
    pixels,
    newImage.width,
    newImage.height,
    Bitmap.Config.RGB_565
)

合成结果视频

最后的最后就是合成结果视频了，这里笔者找了一个现成的FFmpeg应用层封装

implementation 'com.github.microshow:RxFFmpeg:4.9.0-lite'

视频合成

val commond = "ffmpeg -framerate 25 -i ${imageInputPath} -c:v libx264 -pix_fmt yuv420p ${videoOutputPath}"
RxFFmpegInvoke.getInstance().runCommand(commond.split(" ").toTypedArray(), object : RxFFmpegInvoke.IFFmpegListener {
        override fun onFinish() {
        }

        override fun onProgress(progress: Int, progressTime: Long) {
        }

        override fun onCancel() {
        }

        override fun onError(message: String?) {
        }
    })

这里将视频输出的fps写死成25，是因为笔者使用的测试视频无法通过MediaMetadataRetriever获取到视频原有的fps，这个由于没深入了解暂时不知道原因。

上述的代码有参考MediaPipe的官方示例

mediapipe-image_segmentation

输出效果

输出结果如图，左边是原视频，右边是结果视频

总结

以上就是基于MediaPipe进行人像分割，再通过FFmpeg合成视频的流程了。这个方案还是有明显的缺点的，譬如需要将每一帧处理的图片保存起来频繁进行IO，再譬如使用MediaMetadataRetriever按帧读取速度过慢，这些都会影响处理速度的，毕竟一个视频的总帧数非常大。所以说这个思路已经打开，但细节实现仍需完善，譬如通过自定义FFmpeg滤镜？