[mediapipe] human body segmentation cutout

foreword

Due to the segmentation module of mediapipe, it is possible to achieve character segmentation and cut out human body pictures

Code

import library

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

Define the view image function

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

read image

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

import split module

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

plt.imshow(mask)
plt.show()

After the initial split, as shown below

Filter pixels with a threshold of 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)

View Human Foreground

look_img(FG_img)

As shown below

View human body background

look_img(BG_img)

As shown below