initial segmentation code

README.md

@@ -0,0 +1,12 @@
+---
+title: SelfieSegmentation
+emoji: 📐
+colorFrom: gray
+colorTo: red
+sdk: gradio
+sdk_version: 4.17.0
+app_file: app.py
+pinned: false
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -7,16 +7,33 @@ import numpy as np
 mp_drawing = mp.solutions.drawing_utils
 mp_selfie_segmentation = mp.solutions.selfie_segmentation
 
+seg = mp_selfie_segmentation.SelfieSegmentation()
+
+BG_COLOR = (0, 0, 0)
+
+
+def close_segmentation_model():
+    print("Closing segmentation model")
+    seg.close()
+
 
 def snap(video):
-    print(video)
-    return video
+    image = cv2.cvtColor(video, cv2.COLOR_RGB2BGR)
+    results = seg.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
+    condition = np.stack((results.segmentation_mask,) * 3, axis=-1) > 0.5
+
+    bg_image = np.zeros(image.shape, dtype=np.uint8)
+    bg_image[:] = BG_COLOR
+    output_image = np.where(condition, video, bg_image)
+
+    return output_image
 
 
 demo = gr.Interface(
     fn=snap,
-    inputs=gr.Video(sources=["webcam"]),
-    outputs=["video"],
+    inputs=gr.Image(sources=["webcam"], streaming=True),
+    outputs="image",
+    live=True,
 )
 
 if __name__ == "__main__":

requirements.txt

@@ -0,0 +1,3 @@
+gradio
+mediapipe
+opencv-python