Image segmentation background subtraction

Question:

I need to tweak this code from real-time to still image.

This code can already remove the background in real-time but I want to change it into still image. I need the code for my project.

# Data Flair background removal 

# import necessary packages
import os
import cv2
import numpy as np
import mediapipe as mp


# initialize mediapipe 
mp_selfie_segmentation = mp.solutions.selfie_segmentation
selfie_segmentation = mp_selfie_segmentation.SelfieSegmentation(model_selection=1)

# store background images in a list
image_path = 'images'
images = os.listdir(image_path)

image_index= 0
bg_image = cv2.imread(image_path+'/'+images[image_index])

# create videocapture object to access the webcam
cap = cv2.VideoCapture(0)

while cap.isOpened():
    _, frame = cap.read()

    # flip the frame to horizontal direction
    frame = cv2.flip(frame, 1)
    height , width, channel = frame.shape

    RGB = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    # get the result 
    results = selfie_segmentation.process(RGB)

    # extract segmented mask
    mask = results.segmentation_mask

    # it returns true or false where the condition applies in the mask
    condition = np.stack(
      (results.segmentation_mask,) * 3, axis=-1) > 0.6

    # resize the background image to the same size of the original frame
    bg_image = cv2.resize(bg_image, (width, height))

    # combine frame and background image using the condition
    output_image = np.where(condition, frame, bg_image)

    # show outputs
    #cv2.imshow("mask", mask)
    cv2.imshow("Output", output_image)
    cv2.imshow("Frame", frame)

    key = cv2.waitKey(1)
    if key == ord('q'):
        break

    # if 'd' key is pressed then change the background image
    elif key == ord('d'):

        if image_index != len(images)-1:
            image_index += 1
        else:
            image_index = 0
        bg_image = cv2.imread(image_path+'/'+images[image_index])


# release the capture object and close all active windows 
cap.release()
cv2.destroyAllWindows()
Asked By: rabbit b

||

Answers:

You have still images already there:

cv2.imshow("Output", output_image)

output_image is one of them, so add in your code:

key = cv2.waitKey(1) # your code line
if key == ord('q'):  # your code line
    break            # your code line
elif key == ord('c'): # <<< line to add
    cv2.imwrite('images/result.png', output_image) # <<< line to add

in order to capture the still image and write it to a file by pressing c on the keyboard.


If you want to change background of an image stored on the disk as file:

Below a function doing the job. Save the code below in a new Python script file in same directory as your Python script file with the code you put in the question is saved and run it to get the result image saved. Adjust the file names if necessary.

def rm_bckgrd(src_image_file_name, tgt_image_file_name, bgd_image_file_name):
    import cv2
    import numpy as np
    import mediapipe as mp

    frame    = cv2.imread(src_image_file_name)
    bg_image = cv2.imread(bgd_image_file_name)

    frame = cv2.flip(frame, 1)
    height, width, RGBAsize = frame.shape
    RGB = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    # initialize mediapipe 
    mp_selfie_segmentation = mp.solutions.selfie_segmentation
    selfie_segmentation = mp_selfie_segmentation.SelfieSegmentation(model_selection=1)
    # get the result 
    results = selfie_segmentation.process(RGB)
    # extract segmented mask
    mask = results.segmentation_mask
    # it returns true or false where the condition applies in the mask
    condition = np.stack(
      (results.segmentation_mask,) * 3, axis=-1) > 0.6
    # resize the background image to the same size of the original frame
    bg_image = cv2.resize(bg_image, (width, height))
    # combine frame and background image using the condition
    output_image = np.where(condition, frame, bg_image)
    # save the result to disk
    cv2.imwrite(tgt_image_file_name, output_image)

# Run the function to save the result in 'images/result.png' and
# Make sure that 'selfie.png' and 'background.png' exist in the directory [images]
rm_bckgrd('images/selfie.png','images/result.png','images/background.png')

Usage of the function (example):

rm_bckgrd('images/selfie.png','images/result.png','images/background.png')
Answered By: Claudio