Splitting up and fake webcam

This commit is contained in:
Brandon Rozek 2020-05-03 16:03:48 -04:00
parent c3f2738e78
commit 97a758bdb3
3 changed files with 110 additions and 32 deletions

41
camera.py Normal file
View file

@ -0,0 +1,41 @@
from threading import Thread, Event, Lock
import cv2
class Camera(Thread):
def __init__(self):
Thread.__init__(self, daemon=True)
self.finished = Event()
self.ready = Event()
self._cap = cv2.VideoCapture(0)
self._frame = None
self.frame_lock = Lock()
def stop(self):
"""Stop the timer if it hasn't finished yet."""
self.finished.set()
def read(self):
f = None
self.ready.wait()
f = self.frame
self.ready.clear()
return f
@property
def frame(self):
f = None
with self.frame_lock:
f = self._frame
return f
@frame.setter
def frame(self, f):
with self.frame_lock:
self._frame = f
def run(self):
while not self.finished.is_set():
_, self.frame = self._cap.read()
self.ready.set()
self._cap.release()

98
main.py
View file

@ -1,42 +1,78 @@
import zmq
import cv2
from pyfakewebcam import FakeWebcam
import numpy as np
from camera import Camera
from signal import signal, SIGINT
ctx = zmq.Context()
sock = ctx.socket(zmq.REQ)
sock.connect('ipc:///tmp/bodypix')
cap = cv2.VideoCapture(0)
while True:
_, frame = cap.read()
_, image = cv2.imencode('.jpg', frame)
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()
print("LENGTH", len(image.tostring()), flush=True)
sock.send(image.tostring())
convereted_img = sock.recv()
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
mask = np.frombuffer(convereted_img, dtype=np.uint8)
mask = mask.reshape((frame.shape[0], frame.shape[1], 4))
mask = mask[:, :, 0]
# post-process mask and frame
mask = cv2.UMat(mask)
mask = cv2.dilate(mask, np.ones((10, 10), np.uint8), iterations=1)
mask = cv2.blur(cv2.UMat(mask.get().astype(np.float32)), (30, 30))
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)
frame = cv2.UMat(frame)
background = cv2.GaussianBlur(frame, (221, 221), sigmaX=20, sigmaY=20)
# composite the foreground and background
frame = frame.get().astype(np.uint8)
mask = mask.get().astype(np.float32)
background = background.get().astype(np.uint8)
inv_mask = 1 - mask
for c in range(frame.shape[2]):
frame[:, :, c] = frame[:, :, c] * mask + background[:, :, c] * inv_mask
def stop(signal_received, stack_frame):
"""Gracefully stops the application."""
global running
running = False
print("Stopping...")
cv2.imshow('frame', frame)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
# 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)

View file

@ -25,6 +25,7 @@ async function run() {
const image = decodeJpeg(msg)
const segmentation = await net.segmentPerson(image);
await sock.send(segmentation.data);
tf.dispose(image)
}
}