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()

96
main.py
View file

@ -1,42 +1,78 @@
import zmq import zmq
import cv2 import cv2
from pyfakewebcam import FakeWebcam
import numpy as np 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() ctx = zmq.Context()
sock = ctx.socket(zmq.REQ) sock = ctx.socket(zmq.REQ)
sock.connect('ipc:///tmp/bodypix') sock.connect('ipc:///tmp/bodypix')
cap = cv2.VideoCapture(0)
while True: background, (width, height) = get_background()
_, frame = cap.read() fake = FakeWebcam('/dev/video20', height, width)
running = True
print("Running...")
while running:
frame = camera.read()
_, image = cv2.imencode('.jpg', frame) _, image = cv2.imencode('.jpg', frame)
frame = frame.astype(np.uint8)
print("LENGTH", len(image.tostring()), flush=True) # Process image to find body masks
sock.send(image.tostring()) sock.send(image.tostring())
convereted_img = sock.recv() mask = process_mask(sock.recv(), (width, height))
mask = np.frombuffer(convereted_img, dtype=np.uint8) frame = composite(frame, background, mask)
mask = mask.reshape((frame.shape[0], frame.shape[1], 4))
mask = mask[:, :, 0]
# post-process mask and frame # Send to the fake camera device
mask = cv2.UMat(mask) frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
mask = cv2.dilate(mask, np.ones((10, 10), np.uint8), iterations=1) fake.schedule_frame(frame)
mask = cv2.blur(cv2.UMat(mask.get().astype(np.float32)), (30, 30))
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
cv2.imshow('frame', frame)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()

View file

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