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