Spaces:
Running
Running
Mohammed Innat
commited on
Commit
•
3510c12
1
Parent(s):
890c650
Upload selfiseg.py
Browse files- selfiseg.py +42 -0
selfiseg.py
ADDED
@@ -0,0 +1,42 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import mediapipe as mp
|
2 |
+
from utils import read_n_resize
|
3 |
+
from random import sample
|
4 |
+
import cv2, numpy as np
|
5 |
+
|
6 |
+
BG_IMG = [
|
7 |
+
'examples/back1.jpg',
|
8 |
+
'examples/back2.jpg',
|
9 |
+
'examples/back3.jpg',
|
10 |
+
'examples/back4.jpg',
|
11 |
+
'examples/back5.jpg',
|
12 |
+
'examples/back6.jpg'
|
13 |
+
]
|
14 |
+
|
15 |
+
def mp_selfi_segment_fn(image):
|
16 |
+
mp_selfie_segmentation = mp.solutions.selfie_segmentation
|
17 |
+
|
18 |
+
with mp_selfie_segmentation.SelfieSegmentation(
|
19 |
+
model_selection=0) as selfie_segmentation:
|
20 |
+
image = read_n_resize(image, read=False)
|
21 |
+
image_height, image_width, _ = image.shape
|
22 |
+
|
23 |
+
# get a random background picture to fill original background
|
24 |
+
backs = cv2.imread(sample(BG_IMG, 1)[0])
|
25 |
+
backs = cv2.resize(backs, (image_width, image_height))
|
26 |
+
backs = cv2.cvtColor(backs, cv2.COLOR_BGR2RGB)
|
27 |
+
|
28 |
+
# pass to model
|
29 |
+
results = selfie_segmentation.process(image)
|
30 |
+
|
31 |
+
# Draw selfie segmentation on the background image.
|
32 |
+
# To improve segmentation around boundaries, consider applying a joint
|
33 |
+
# bilateral filter to "results.segmentation_mask" with "image".
|
34 |
+
condition = np.stack((results.segmentation_mask,) * 3, axis=-1) > 0.1
|
35 |
+
|
36 |
+
# Generate solid color images for showing the output selfie segmentation mask.
|
37 |
+
fg_image = np.zeros(image.shape, dtype=np.uint8)
|
38 |
+
fg_image[:] = image
|
39 |
+
bg_image = np.zeros(image.shape, dtype=np.uint8)
|
40 |
+
bg_image[:] = backs
|
41 |
+
output_image = np.where(condition, fg_image, bg_image)
|
42 |
+
return output_image
|