stack_foreground_segmentation.py

import numpy as np
import cv2

videoPath='C:\\Users\\lav singh\\workplace\\yolov4\\stack_segmentation_video.mp4'
output_videoPath='C:\\Users\\lav singh\\workplace\\yolov4\\output_video.mp4'
create=None
cap=cv2.VideoCapture(videoPath)
while True:
    _, frame=cap.read()
    frame=cv2.resize(frame, (720, 720), fx=None, fy=None)
    mask = np.zeros(frame.shape[:2], np.uint8)
    '''Specify the background and foreground model using numpy the array is constructed of 1 row and 65 columns, 
       and all array elements are 0. Data type for the array is np.float64 (default)'''
    backgroundModel = np.zeros((1, 65), np.float64)
    foregroundModel = np.zeros((1, 65), np.float64)
    '''Define the Region of Interest (ROI) as the coordinates of the rectangle where the values are entered as 
       (startingPoint_x, startingPoint_y, width, height).These coordinates are according to the input image, it 
       may vary for different images.'''
    rectangle = (210, 1, 270, 640)
    # Apply the Grabcut algorithm with appropriate values as parameters, number of iterations = 3
    # cv2.GC_INIT_WITH_RECT is used because of the rectangle mode is used.
    cv2.grabCut(frame, mask, rectangle,backgroundModel, foregroundModel, 3, cv2.GC_INIT_WITH_RECT)
    mask2 = np.where((mask == 2)|(mask == 0), 0, 1).astype('uint8')
    # The final mask is multiplied with the input frame to give the segmented frames.
    frame = frame * mask2[:, :, np.newaxis]
    cv2.imshow("output", frame)
    if create is None:
        fourcc = cv2.VideoWriter_fourcc(*'MP4V')
        create = cv2.VideoWriter(output_videoPath, fourcc, 15, (720, 720), True)
    create.write(frame)
    key=cv2.waitKey(1)
    if key==ord('q'):
        break
cap.release()
create.release()
cv2.distroyAllWindows()