Upload selfiseg.py

selfiseg.py

@@ -0,0 +1,42 @@
+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