Webcam background replacer (can be used in MS Teams)

msteams-background-replacer.py

#
# Author: Rob Audenaerde - License: Apache 2.0
#
# Heavily inspired by the code by: ajaichemmanam@gmail.com, found here: https://github.com/ajaichemmanam/simple_bodypix_python
#

import os
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import math
import time
import cv2
import pyfakewebcam

from PIL import Image
from PIL import ImageFilter
from utils import load_graph_model, get_input_tensors, get_output_tensors
import tensorflow as tf


# make tensorflow stop spamming messages
os.environ['TF_CPP_MIN_LOG_LEVEL'] = "3"

# SETUP video devices. 
# Make sure the numbers are correct (/dev/video2 is my loopback device)

camera = pyfakewebcam.FakeWebcam('/dev/video2', 640, 480)
cap = cv2.VideoCapture(0)

# Setup the backgroud 

bg = cv2.imread('matrix.jpg')


# PATHS
modelPath = 'bodypix_mobilenet_float_050_model-stride16'

# CONSTANTS
OutputStride = 16


print("Loading model...", end="")
graph = load_graph_model(modelPath)  # downloaded from the link above
print("done.\nLoading sample image...", end="")


# Get input and output tensors
input_tensor_names = get_input_tensors(graph)
output_tensor_names = get_output_tensors(graph)
input_tensor = graph.get_tensor_by_name(input_tensor_names[0])


cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

# evaluate the loaded model directly
sess = tf.compat.v1.Session(graph=graph)

# Resize the background, and preprocess it to 32f.
bg_resized = cv2.resize(bg, (640,480), interpolation = cv2.INTER_AREA)
bg_resized_32f = np.float32(cv2.cvtColor(bg_resized, cv2.COLOR_BGR2RGB))

# Forever looping :)
while True:
  
	ret, frame = cap.read()
	frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
	img = Image.fromarray(frame)

	imgWidth, imgHeight = img.size

	targetWidth = (int(imgWidth) // OutputStride) * OutputStride 
	targetHeight = (int(imgHeight) // OutputStride) * OutputStride 

	img = img.resize((targetWidth, targetHeight))
	x = tf.keras.preprocessing.image.img_to_array(img, dtype=np.float32)
	InputImageShape = x.shape

	widthResolution = int((InputImageShape[1] - 1) / OutputStride) 
	heightResolution = int((InputImageShape[0] - 1) / OutputStride) 

        #mobile net preprocessing
	x = (x/127.5)-1

	sample_image = x[tf.newaxis, ...]

	output_tensor_names = ['float_segments:0']

	results = sess.run(output_tensor_names, feed_dict={input_tensor: sample_image})

	segments = np.squeeze(results[0], 0)

	# Segmentation MASk
	segmentation_threshold = 0.7
	segmentScores = tf.sigmoid(segments)
	mask = tf.math.greater(segmentScores, tf.constant(segmentation_threshold))
	segmentationMask = tf.dtypes.cast(mask, tf.int32)
	segmentationMask = np.reshape(
	    segmentationMask, (segmentationMask.shape[0], segmentationMask.shape[1]))

	# Create the mask image
	mask_img = Image.fromarray(segmentationMask * 255)
	mask_img = mask_img.resize(
	    (targetWidth, targetHeight), Image.LANCZOS).convert("RGB")
	    

	# Convert the segmentation mask to GRAY
	proc_out = cv2.cvtColor(np.asarray(mask_img), cv2.COLOR_RGB2GRAY)
	
	#Blur to smoothen the blocky input
	proc_out = cv2.GaussianBlur(proc_out,(151,151),0)

	#Threshold and blur to reduce the part that gets blended
	a, proc_out = cv2.threshold(proc_out,127,255,cv2.THRESH_BINARY)
	proc_out = cv2.GaussianBlur(proc_out,(11,11),0)	

        #Convert back to RGB and float32 for blending	
	proc_out = cv2.cvtColor(proc_out, cv2.COLOR_GRAY2RGB)
	mask_32f = np.float32(proc_out) / 255.0
	mask_32f_inv = 1.0 - mask_32f

        #blend
	img_in  = np.array(img)
	img_in_32f  = np.float32(img_in)
	img_fg = cv2.multiply(mask_32f    , img_in_32f    , 1.0/255.0)
	img_bg = cv2.multiply(mask_32f_inv, bg_resized_32f, 1.0/255.0)
	
	#add
	img_out = cv2.add(img_fg, img_bg);
	
	#convert back to 3channel 8 bit
	img_out = np.uint8(img_out)

	#put it in the fake webcam
	camera.schedule_frame(img_out)
	
	#Sleep a bit to keep the CPU usages acceptable for other stuff :)
	time.sleep(1/30.0)