vibe python實現

參考VIBE背景建模算法之C語言實現

寫了一個python版本，vibe在靜態背景檢測效果還是不錯的

# -*- coding:utf-8  -*-
import cv2 as cv
import numpy as np

#64, 20, 3, 8
class vibe():
	def __init__(self, gVibeRandom = 29, gRadius = 20, gMin = 3, gVibeSimpleSize = 5):
		self.gVibeRandom = gVibeRandom
		self.gRadius = gRadius
		self.gMin = gMin
		self.gVibeSimpleSize = gVibeSimpleSize

	def initModelData(self, img):
		self.rows, self.cols = img.shape
		self.gVibeSize = self.rows * self.cols
		self.model = []
		for i in range(self.gVibeSimpleSize):
			self.model.append(img.copy())
		self.model.append(np.zeros(img.shape[:2], dtype=np.uint8))
		self.CountNum = 0

	def update(self, img):

		self.CountNum += 1
		if(self.CountNum > self.gVibeRandom):
			self.CountNum = 0

		noffset = self.CountNum // 5

		dst = np.zeros(img.shape[:2], dtype=np.uint8)

		for r in range(self.rows):
			for c in range(self.cols):
				nCount = 0
				dst[r,c] = 255
				for j in range(self.gVibeSimpleSize):
					model_img = self.model[j]
					nDiff = int(model_img[r,c]) - int(img[r,c])
					if abs(nDiff) <= self.gRadius:
						nCount += 1
						if nCount >= self.gMin:
							dst[r,c] = 0
							break
				if dst[r,c]:
					self.model[-1][r,c] += 1
					if(self.model[-1][r,c] > 100):
						self.model[noffset][r,c] = img[r,c]
				else:
					self.model[noffset][r,c] = img[r,c]
					self.model[-1][r,c] = 0
		return dst

def get_rects_on_bitmap(img):
	mat, contours, hierarchy = cv.findContours(img, cv.RETR_TREE, cv.CHAIN_APPROX_NONE)
	cv.drawContours(img, contours, -1, (0, 255, 255), 2)
	rects = []
	for i in range(len(contours)):
		rects.append(cv.boundingRect(contours[i]))
	return rects
def main(file):
	detector = vibe()
	cap = cv.VideoCapture(file)
	initflag = True
	while 1:
		ret, frame = cap.read()
		if ret == False:
			break
		gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
		rows, cols = gray.shape
		gray = cv.resize(gray, (int(cols/2), int(rows/2)))

		if initflag:
			detector.initModelData(gray)
			initflag = False
			continue
		else:
			dst = detector.update(gray)
		rects = get_rects_on_bitmap(dst)

		for i in range(len(rects)):
			x, y, w, h = rects[i]
			x1, y1 = 2*x, 2*y
			x2, y2 = 2*(x+w), 2*(y+h)
			cv.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 255), 2)
		cv.imshow("detect.jpg", frame)
		cv.waitKey(20)
	cap.release()

if __name__ == '__main__':
	main('vtest.avi')

vtest_vibe.avi