opencv python background subtraction with mask

BackGroundSegWithMask.py

import cv2
import numpy as np


class BackGroundSegmentor(object):
    def __init__(self):
        self.fgbg = cv2.createBackgroundSubtractorMOG2(
            history=450, varThreshold=50, detectShadows=True)
        self.fgbg.setNMixtures(3)
        self.vbg = None

    def set_varThreshold(self, varThreshold):
        self.fgbg.setVarThreshold(varThreshold)

    def get_varThreshold(self):
        return self.fgbg.getVarThreshold()

    def set_history(self, history):
        self.fgbg.setHistory(history)

    def get_history(self):
        return self.fgbg.getHistory()

    def set_mixtures(self, mixtures):
        self.fgbg.setNMixtures(mixtures)

    def get_mixtures(self):
        return self.fgbg.getNMixtures()

    def mask(self, image, mask):
        if np.sum(mask) == 0:
            self.vbg = image.copy()
        else:
            image[mask, :] = self.vbg[mask, :]

        fgmask = self.fgbg.apply(image, learningRate=-1)
        fgmask[fgmask < 255] = 0
        return fgmask


class MaskGenerator(object):
    def __init__(self, size, period, loc=None):
        self.period = period
        self.size = size
        self.counter = 0
        self.x = None
        self.y = None
        self.loc = loc

    def gen(self, height, width):
        mask = np.zeros((height, width), np.uint8)
        self.counter += 1
        if self.counter > 2 * self.period:
            self.counter = 0
            self.x = None
            self.y = None

        if self.counter > self.period and (self.x is None or self.y is None):
            if self.loc is None:
                self.x = np.random.randint(0, width - self.size, 1)[0]
                self.y = np.random.randint(0, height - self.size, 1)[0]
            else:
                self.x = self.loc[0]
                self.y = self.loc[1]

        if self.x is not None or self.y is not None:
            mask[self.y:self.y + self.size, self.x:self.x + self.size] = 1

        return mask > 0


def app():
    cap = cv2.VideoCapture(1)

    size = 100
    p = (200, 200)
    q = (p[0] + size, p[1] + size)
    fgbg = BackGroundSegmentor()
    mask = MaskGenerator(size=size, period=100, loc=p)

    cv2.namedWindow('frame', cv2.WINDOW_NORMAL)
    cv2.namedWindow('mask', cv2.WINDOW_NORMAL)

    while True:
        # Capture frame-by-frame
        ret, frame = cap.read()

        hmask = mask.gen(height=frame.shape[0], width=frame.shape[1])
        fgmask = fgbg.mask(frame, hmask)

        # Our operations on the frame come here
        # gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        # Display the resulting frame
        cv2.rectangle(frame, p, q, (0, 255, 0), 2)
        frame[:, :, 2] = hmask * 255
        cv2.imshow('frame', frame)

        fgmask = cv2.cvtColor(fgmask, cv2.COLOR_GRAY2BGR)
        cv2.rectangle(fgmask, p, q, (0, 0, 255), 2)
        cv2.imshow('mask', fgmask)

        key = cv2.waitKey(1)
        if key == 27:
            break
        elif key == ord('.'):
            hist = fgbg.get_history()
            print("hist = {}".format(hist))
            hist += 50
            fgbg.set_history(hist)
        elif key == ord(','):
            hist = fgbg.get_history()
            print("hist = {}".format(hist))
            hist -= 50
            if hist < 10:
                hist = 10
            fgbg.set_history(hist)

    # When everything done, release the capture
    cap.release()
    cv2.destroyAllWindows()


if __name__ == "__main__":
    app()