# Blurring the Webcam using OpenCV and Mediapipe
# Partialy from here: https://pythonrepo.com/repo/HxnDev-Live-Background-Blur
# See also https://google.github.io/mediapipe/solutions/selfie_segmentation

# Key Bindings
# f = show time
# g = gray
# space = play old frames (max 60)
# b = blur (none, background, all)
 
import cv2
import mediapipe as mp
import numpy as np
import time

mp_drawing = mp.solutions.drawing_utils
mp_selfie_segmentation = mp.solutions.selfie_segmentation

# Screen size
screenx = 1920
screeny = 1080

# create a black image as big as the screen
bg_img = np.zeros((screeny, screenx, 3), np.uint8)
bg_img[:, 0:screeny] = (0, 0, 0)

BG_COLOR = (192, 192, 192) # gray

cap = cv2.VideoCapture(0)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 320);
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 240);

# Some experiments:
#cap.set(cv2.CAP_PROP_FPS, 0.1); # set fps
#cap.set(cv2.CAP_PROP_FOURCC, cv2.FOURCC('M', 'J', 'P', 'G'));
#codec = 0x47504A4D # MJPG
#cap.set(cv2.CAP_PROP_FOURCC, codec)

# Experiment with these for checking lag:
cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
cap.set(cv2.CAP_PROP_BUFFERSIZE, 2)

window_name = "window"
cv2.namedWindow(window_name, cv2.WND_PROP_FULLSCREEN)
cv2.setWindowProperty(window_name, cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)

font = cv2.FONT_HERSHEY_SIMPLEX

# for calculation of time per frame
framecount = 0
starttime = 0
captime = 0
proctime = 0

# status
blur = 0 # start with no blur
showtime = False
hide = False # space pressed
gray = False

# for storing old frames
old_frames = []
old_frameindex = 0
old_framedir = 1
MAX_NUM_OLD_FRAMES = 60

with mp_selfie_segmentation.SelfieSegmentation(
    model_selection=1) as selfie_segmentation:
    
  bg_image = None
  
  while cap.isOpened():
      
    starttime = time.time()
    
    success, image = cap.read()
    
    captime = time.time()
    
    if not success:
        print("Ignoring empty camera frame.")
        continue

    if not hide: # no space pressed
        
        if blur == 1:
            # segment and blur background
            image = cv2.cvtColor(cv2.flip(image, 1), cv2.COLOR_BGR2RGB)
            image.flags.writeable = False
            results = selfie_segmentation.process(image)

            image.flags.writeable = True
            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

            condition = np.stack(
              (results.segmentation_mask,) * 3, axis=-1) > 0.1
            # The background can be customized.
            #   a) Load an image (with the same width and height of the input image) to
            #      be the background, e.g., bg_image = cv2.imread('/path/to/image/file')
            #   b) Blur the input image by applying image filtering, e.g.,
            #      bg_image = cv2.GaussianBlur(image,(55,55),0)
            if bg_image is None:
              # create new bg image
              bg_image = cv2.GaussianBlur(image,(55,55),0)
              
            if (framecount % 30) == 0:
                bg_image = cv2.GaussianBlur(image,(55,55),0) # update bg image
            framecount += 1
            
            output_image = np.where(condition, image, bg_image)
        elif blur == 0:
            # no blur, just flip
            output_image = cv2.flip(image, 1)
        else:
            # blur all
            output_image = cv2.flip(image, 1)
            output_image = cv2.GaussianBlur(output_image,(55,55),0)

    if hide:
        # space was pressed, play old frames
        output_image = old_frames[old_frameindex]
        old_frameindex += old_framedir
        if old_frameindex == 0 or old_frameindex == len(old_frames)-1: # bounds reached, change direction
            old_framedir *= -1;
    else:
        # remember old frames
        if len(old_frames) > MAX_NUM_OLD_FRAMES: 
            old_frames = old_frames[1:]
        old_frames.append(output_image)
        

    # calc resized size
    scale = 4.0
    nw = int(output_image.shape[1] * scale)
    nh = int(output_image.shape[0] * scale)
    
    # resize
    output_image = cv2.resize(output_image, (nw, nh), interpolation = cv2.INTER_NEAREST)
    
    proctime = time.time()
    
    if showtime:
        cv2.putText(output_image, "C "+str((captime-starttime)*100//1)+", P "+str((proctime-captime)*100//1), (7, 70), font, 3, (100, 255, 0), 3, cv2.LINE_AA)

    # position in black image:
    # calculate the offsets required to center the image
    x_offset = (screenx - nw)//2
    y_offset = (screeny - nh)//2
    
    # insert the smaller image into the larger image, centered
    bg_img[y_offset:y_offset+output_image.shape[0], x_offset:x_offset+output_image.shape[1]] = output_image
    
    disp_img = bg_img;
    
    if gray:
        disp_img = cv2.cvtColor(bg_img, cv2.COLOR_BGR2GRAY)
        
	# show result
    cv2.imshow(window_name, disp_img)
        
    # check keyboard input
    key = cv2.waitKey(5) & 0xFF
    if key == 27: # escape
        break
    elif key == 98: # b
        blur = (blur + 1)%3
    elif key == 32: # space
        hide = not hide
        old_frameindex = len(old_frames)-1
        old_framedir = -1
    elif key == 102: # f
        showtime = not showtime
    elif key == 103: # g
        gray = not gray
    #else:
    #    print(key)
    
cap.release()
cv2.destroyAllWindows()