viz.py 1.3 KB
Newer Older
1
import numpy as np
2
import cv2 as cv
3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19
import argparse
import os


def main():
    argparser = argparse.ArgumentParser(description='Vizualization of the LSBP/GSOC background subtraction algorithm.')

    argparser.add_argument('-g', '--gt', help='Directory with ground-truth frames', required=True)
    argparser.add_argument('-f', '--frames', help='Directory with input frames', required=True)
    argparser.add_argument('-l', '--lsbp', help='Display LSBP instead of GSOC', default=False)
    args = argparser.parse_args()

    gt = map(lambda x: os.path.join(args.gt, x), os.listdir(args.gt))
    gt.sort()
    f = map(lambda x: os.path.join(args.frames, x), os.listdir(args.frames))
    f.sort()

20 21
    gt = np.uint8(map(lambda x: cv.imread(x, cv.IMREAD_GRAYSCALE), gt))
    f = np.uint8(map(lambda x: cv.imread(x, cv.IMREAD_COLOR), f))
22 23

    if not args.lsbp:
24
        bgs = cv.bgsegm.createBackgroundSubtractorGSOC()
25
    else:
26
        bgs = cv.bgsegm.createBackgroundSubtractorLSBP()
27 28

    for i in xrange(f.shape[0]):
29 30
        cv.imshow('Frame', f[i])
        cv.imshow('Ground-truth', gt[i])
31 32
        mask = bgs.apply(f[i])
        bg = bgs.getBackgroundImage()
33 34 35
        cv.imshow('BG', bg)
        cv.imshow('Output mask', mask)
        k = cv.waitKey(0)
36 37 38 39 40 41
        if k == 27:
            break


if __name__ == '__main__':
    main()