Google-MediaPipe / selfiseg.py
Mohammed Innat
Upload selfiseg.py
3510c12
import mediapipe as mp
from utils import read_n_resize
from random import sample
import cv2, numpy as np
BG_IMG = [
'examples/back1.jpg',
'examples/back2.jpg',
'examples/back3.jpg',
'examples/back4.jpg',
'examples/back5.jpg',
'examples/back6.jpg'
]
def mp_selfi_segment_fn(image):
mp_selfie_segmentation = mp.solutions.selfie_segmentation
with mp_selfie_segmentation.SelfieSegmentation(
model_selection=0) as selfie_segmentation:
image = read_n_resize(image, read=False)
image_height, image_width, _ = image.shape
# get a random background picture to fill original background
backs = cv2.imread(sample(BG_IMG, 1)[0])
backs = cv2.resize(backs, (image_width, image_height))
backs = cv2.cvtColor(backs, cv2.COLOR_BGR2RGB)
# pass to model
results = selfie_segmentation.process(image)
# Draw selfie segmentation on the background image.
# To improve segmentation around boundaries, consider applying a joint
# bilateral filter to "results.segmentation_mask" with "image".
condition = np.stack((results.segmentation_mask,) * 3, axis=-1) > 0.1
# Generate solid color images for showing the output selfie segmentation mask.
fg_image = np.zeros(image.shape, dtype=np.uint8)
fg_image[:] = image
bg_image = np.zeros(image.shape, dtype=np.uint8)
bg_image[:] = backs
output_image = np.where(condition, fg_image, bg_image)
return output_image