78 lines
2.1 KiB
Python
78 lines
2.1 KiB
Python
import zmq
|
|
import cv2
|
|
from pyfakewebcam import FakeWebcam
|
|
import numpy as np
|
|
from camera import Camera
|
|
from signal import signal, SIGINT
|
|
|
|
def process_mask(m, shape):
|
|
w, h = shape
|
|
m = np.frombuffer(m, dtype=np.uint8).reshape((w, h, -1))
|
|
m = m[:, :, 0].astype(np.float32)
|
|
m = cv2.blur(cv2.UMat(m), (30, 30))
|
|
return m.get()
|
|
|
|
def composite(foreground, backdrop, fore_mask):
|
|
inv_mask = 1 - fore_mask
|
|
for c in range(foreground.shape[2]):
|
|
foreground[:, :, c] = foreground[:, :, c] * fore_mask + backdrop[:, :, c] * inv_mask
|
|
return foreground
|
|
|
|
|
|
def get_background(uri=None):
|
|
"""
|
|
Either grabs the image specified or
|
|
blurs the current background.
|
|
"""
|
|
initial_frame = camera.read()
|
|
w, h = initial_frame.shape[0], initial_frame.shape[1]
|
|
backdrop = None
|
|
if uri is not None:
|
|
backdrop = cv2.resize(
|
|
cv2.UMat(cv2.imread('/home/rozek/Pictures/IMG_20191013_181848.jpg')),
|
|
(h, w)
|
|
).get().astype(np.uint8)
|
|
else:
|
|
backdrop = cv2.GaussianBlur(initial_frame, (221, 221), sigmaX=20, sigmaY=20).astype(np.uint8)
|
|
return backdrop, (w, h)
|
|
|
|
|
|
def stop(signal_received, stack_frame):
|
|
"""Gracefully stops the application."""
|
|
global running
|
|
running = False
|
|
print("Stopping...")
|
|
|
|
|
|
# Start Camera
|
|
camera = Camera()
|
|
camera.start()
|
|
|
|
if __name__ == "__main__":
|
|
signal(SIGINT, stop)
|
|
|
|
# Setup ZeroMQ
|
|
ctx = zmq.Context()
|
|
sock = ctx.socket(zmq.REQ)
|
|
sock.connect('ipc:///tmp/bodypix')
|
|
|
|
background, (width, height) = get_background()
|
|
fake = FakeWebcam('/dev/video20', height, width)
|
|
|
|
running = True
|
|
print("Running...")
|
|
|
|
while running:
|
|
frame = camera.read()
|
|
_, image = cv2.imencode('.jpg', frame)
|
|
frame = frame.astype(np.uint8)
|
|
|
|
# Process image to find body masks
|
|
sock.send(image.tostring())
|
|
mask = process_mask(sock.recv(), (width, height))
|
|
|
|
frame = composite(frame, background, mask)
|
|
|
|
# Send to the fake camera device
|
|
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
|
fake.schedule_frame(frame)
|