【mediapipe】人体分割抠图

前言

由于mediapipe的分割模块，可以实现人物分割，对人体图片进行抠图

代码实现

导入库

import cv2
import mediapipe as mp
from tqdm import tqdm
import time
import matplotlib.pyplot as plt
%matplotlib inline

定义查看图片函数

def look_img(img):
    img_RGB=cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
    plt.imshow(img_RGB)
    plt.show()

读入图像

img=cv2.imread('C:\\Users\\123\\Desktop\\2022-6-28.jpg')
img_RGB=cv2.cvtColor(img,cv2.COLOR_BGR2RGB)#BGR转RGB

导入分割模块

seg = mp.solutions.selfie_segmentation.SelfieSegmentation(model_selection=0)
results = seg.process(img_RGB)#输入模型，获取预测结果
mask = results.segmentation_mask

plt.imshow(mask)
plt.show()

初次分割后如下图

以0.5为阈值过滤像素

mask = mask > 0.5
plt.imshow(mask)
plt.show()

import numpy as np
mask_3=np.stack((mask,mask,mask),axis=-1)#单通道转三通道
MASK_COLOR=[0,200,0]
fg_image=np.zeros(img.shape,dtype=np.uint8)
fg_image[:]=MASK_COLOR
FG_img=np.where(mask_3,img,fg_image)
BG_img=np.where(~mask_3,img,fg_image)

查看人体前景

look_img(FG_img)

如下图所示

查看人体背景

look_img(BG_img)

如下图所示