import numpy as np
import cv2
from sklearn.cluster import KMeans
CLUSTER=10
colorlist = [(0,0,0),(0,0,255),(0,255,0),(0,255,255),(255,0,0),(255,0,255),(255,255,0),(255,255,255)]
ESC_KEY = 27
def get_bounding_box(X, Y, result, cluster):
NO = np.where(result==cluster)
x_max = int(np.max(X[NO]))
x_min = int(np.min(X[NO]))
y_max = int(np.max(Y[NO]))
y_min = int(np.min(Y[NO]))
return (x_min, y_min, x_max, y_max)
def draw_bounding_box(img, bbox):
img = cv2.rectangle(img, (bbox[0], bbox[1]), (bbox[2], bbox[3],), (0, 255, 0), 3)
return img
def draw_circle(img, X, Y, result, cluster):
NO = np.where(result==cluster)
for p in NO[0] :
img = cv2.circle(img, (X[p], Y[p]), 5, colorlist[cluster%8],-1)
return img
font = cv2.FONT_HERSHEY_SIMPLEX
def draw_number(img, X, Y, result, cluster):
NO = np.where(result==cluster)
for p in NO[0] :
img = cv2.putText(img, str(cluster),(X[p], Y[p]), font, 1,colorlist[cluster%8],1,cv2.LINE_AA)
return img
fgbg = cv2.bgsegm.createBackgroundSubtractorMOG()
cap = cv2.VideoCapture("video/vtest.avi")
w = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
h = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
fps = cap.get(cv2.CAP_PROP_FPS)
while True:
ret, frame = cap.read()
if frame is None:
break
frame = cv2.GaussianBlur(frame, (5,5), 0)
fgmask = fgbg.apply(frame)
img1=frame
img2=fgmask
Y, X = np.where(fgmask > 200)
if X.size < w :
continue
means = KMeans(n_clusters=CLUSTER, init='k-means++', n_jobs=1).fit_predict(np.array([X,Y]))
for i in range(CLUSTER) :
bbox = get_bounding_box(X, Y, means, i)
img2 = draw_bounding_box(img1, bbox)
img2 = draw_circle(img1, X, Y, means, i)
<em># img2 = draw_number(img1, X, Y, means, i)</em>
cv2.imshow('img1', img2)
cv2.imshow('fgmask', fgmask)
keyboard = cv2.waitKey(30)
if keyboard == ESC_KEY:
break
cap.release()
cv2.destroyAllWindows()