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