commons/selfie_seg.py

import os
import cv2
import mediapipe as mp
import numpy as np
from moviepy.editor import (
  VideoFileClip, AudioFileClip)

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

imgs = os.listdir("bg_imgs/")
rnd_img = "bg_imgs/" + np.random.choice(imgs)
#IMAGE_FILES = ["/home/samuel/Documents/Computer Vision Codes/istockphoto-1193994027-170667a.jpg"]



# For webcam input:
def load_from_webcam(bg_type: str = "blur"):
  cap = cv2.VideoCapture(0)
  with mp_selfie_segmentation.SelfieSegmentation(
      model_selection=1) as selfie_segmentation:
    
    while cap.isOpened():
      success, image = cap.read()
      if not success:
        print("Ignoring empty camera frame.")
        # If loading a video, use 'break' instead of 'continue'.
        continue

      # Flip the image horizontally for a later selfie-view display, and convert
      # the BGR image to RGB.
      image = cv2.cvtColor(cv2.flip(image, 1), cv2.COLOR_BGR2RGB)
      # To improve performance, optionally mark the image as not writeable to
      # pass by reference.
      image.flags.writeable = False
      results = selfie_segmentation.process(image)

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

      # Draw selfie segmentation on the background image.
      # To improve segmentation around boundaries, consider applying a joint
      # bilateral filter to "results.segmentation_mask" with "image".
      condition = np.stack(
        (results.segmentation_mask,) * 3, axis=-1) > 0.7
      # 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)
      image_height, image_width, _ = image.shape
      
      if bg_type == "blur":
        bg_image = cv2.GaussianBlur(image,(55,55),0)

      if bg_type == "random_image":
        bg_image = cv2.resize(cv2.imread(rnd_img), (image_width, image_height))
      
      if (bg_image is None) or (bg_image == "solid_colors"):
        bg_image = np.zeros(image.shape, dtype=np.uint8)
        bg_image[:] = np.random.randint(0, high=256, size=(3,)).tolist()

      output_image = np.where(condition, image, bg_image)

      cv2.imshow('MediaPipe Selfie Segmentation', output_image)
      if cv2.waitKey(5) & 0xFF == ord("q"):
        break
  cap.release()



# For static images
def load_from_static_image(file: cv2.Mat, bg_type: str = "solid_colors"):
  with mp_selfie_segmentation.SelfieSegmentation(
      model_selection=0) as selfie_segmentation:
    #for idx, file in enumerate(IMAGE_FILES):
    #print (file)
    image = file #cv2.imread(file)
    image_height, image_width, _ = image.shape
    # Convert the BGR image to RGB before processing.
    results = selfie_segmentation.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))

    # Draw selfie segmentation on the background image.
    # To improve segmentation around boundaries, consider applying a joint
    # bilateral filter to "results.segmentation_mask" with "image".
    # increase threshold to 0.8 or reduce
    condition = np.stack((results.segmentation_mask,) * 3, axis=-1) > 0.8

    if bg_type == "blur":
      bg_image = cv2.GaussianBlur(image, (55,55), 0)
    if bg_type == "random_image":
      bg_image = cv2.resize(cv2.imread(rnd_img), (image_width, image_height))
    if (bg_type is None) or (bg_type == "solid_colors"):
      bg_image = np.zeros(image.shape, dtype=np.uint8)
      bg_image[:] = np.random.randint(0, high=256, size=(3,)).tolist()

    output_image = np.where(condition, image, bg_image)
    return output_image


# For Videos
def load_from_video(file: str, bg_type: str = "solid_colors"):
  vcap = cv2.VideoCapture(file)
  # Get video properties
  frame_width = int(vcap.get(3))
  frame_height = int(vcap.get(4))
  vid_fps = int(vcap.get(5))

  vid_size = (frame_width, frame_height)

  audio_path = "audio.mp3"
  video_path = "output_video_from_file.mp4"
  # *'h264'
  output = cv2.VideoWriter(video_path, cv2.VideoWriter_fourcc(*'avc1'), vid_fps, vid_size)

  selfie_segmentation = mp_selfie_segmentation.SelfieSegmentation(model_selection=1)
  solid_bg = np.random.randint(0, high=256, size=(3,)).tolist()

  while True:
    success, image = vcap.read()
    if success == True:

      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.7

      image_height, image_width = image.shape[:2]

      if bg_type == "blur":
        bg_image = cv2.GaussianBlur(image, (55,55),0)

      if bg_type == "random_image":
        bg_image = cv2.resize(cv2.imread(rnd_img), (image_width, image_height))

      if (bg_type == None) | (bg_type == "solid_colors"):
        bg_image = np.zeros(image.shape, dtype=np.uint8)
        bg_image[:] = solid_bg

      output_image = np.where(condition, image, bg_image)
      output.write(output_image)

    else:
      print ("Video stream disconnected")
      break
  vcap.release()
  output.release()

  try:
    clip = VideoFileClip(file)
    clip.audio.write_audiofile(audio_path)

    video_clip = VideoFileClip(video_path)
    audio_clip = AudioFileClip(audio_path)

    if video_clip.end > audio_clip.end:
      final_clip = video_clip.set_audio(audio_clip)
      final_clip.write_videofile("final.mp4")
    else:
      audio_clip = audio_clip.subclip(0, video_clip.end)
      final_clip = video_clip.set_audio(audio_clip)
      final_clip.write_videofile("final.mp4")
  
    os.remove(video_path)
    os.remove(audio_path)  
  except AttributeError:    #i.e there's no audio in the video
    return "/home/samuel/Documents/Computer Vision Codes/selfie_seg/output_video_from_file.mp4"


  return "final.mp4"


if __name__ == "__main__":

  vp = "/home/samuel/Documents/Computer Vision Codes/Course Overview_5.mp4"
  load_from_video(vp, bg_type="solid_colors")

  vp = "/home/samuel/Documents/Computer Vision Codes/Course Overview_5.mp4"


  """ vcap = cv2.VideoCapture(vp)
  frame_width = int(vcap.get(3))
  frame_height = int(vcap.get(4))
  frame_size = (frame_width,frame_height)
  fps = int(vcap.get(5))

  audio_path = "audio.mp3"
  video_path = "output_video_from_file.mp4"

  output = cv2.VideoWriter(video_path, cv2.VideoWriter_fourcc('M','J','P','G'), fps, frame_size)

  clip = VideoFileClip(vp)
  clip.audio.write_audiofile(audio_path)

  while True:
    ret, frame = vcap.read()
    if ret == True:
      output.write(frame)
    else:
      print ("Video stram disconnected")
      break

  vcap.release()
  output.release()


  video_clip = VideoFileClip(video_path)
  audio_clip = AudioFileClip(audio_path)


  if video_clip.end > audio_clip.end:
    final_clip = video_clip.set_audio(audio_clip)
    final_clip.write_videofile("final.mp4")
  else:
    audio_clip = audio_clip.subclip(0, video_clip.end)
    final_clip = video_clip.set_audio(audio_clip)
    final_clip.write_videofile("final.mp4")


  os.remove(video_path)
  os.remove(audio_path) """


  #final_output = os.system("ffmpeg -i " + video_path+" -i "+audio_path+" -c:v copy -c:a aac "+output_path)