前言

工作的原因是最近在研究运用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

输出效果

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

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

总结

以上便是根据MediaPipe进行人像切开，再经过FFmpeg组成视频的流程了。这个计划仍是有明显的缺点的，比如需求将每一帧处理的图片保存起来频繁进行IO，再比如运用MediaMetadataRetriever按帧读取速度过慢，这些都会影响处理速度的，究竟一个视频的总帧数非常大。所以说这个思路现已打开，但细节完成仍需完善，比如经过自定义FFmpeg滤镜？