mean shift算法

import cv2
import numpy as np
import matplotlib.pyplot as plt
import time

xs,ys,ws,hs = 0,0,0,0  #selection.x selection.y
xo,yo=0,0 #origin.x origin.y
selectObject = False
trackObject = 0         #追踪目标，0代表没有，1代表有，-1代表需要更新目标

#鼠标移动时会进入这个函数
def onMouse(event, x, y, flags, prams): 
    global xs,ys,ws,hs,selectObject,xo,yo,trackObject
    #实时更新鼠标捕捉矩形大小，设置左上角为起点坐标
    if selectObject == True:
        xs = min(x, xo)
        ys = min(y, yo)
        ws = abs(x-xo)
        hs = abs(y-yo)
    if event == cv2.EVENT_LBUTTONDOWN:  #按下
        xo,yo = x, y                    #按下时的origin坐标                     
        xs,ys,ws,hs= x, y, 0, 0         #矩形大小初始化
        selectObject = True             #标记开始选中
    elif event == cv2.EVENT_LBUTTONUP:  #松开
        selectObject = False
        if((ws>0) and (hs>0)):
            trackObject = -1                #需要追踪目标


# 设置初始化的窗口位置
# r,h,c,w = 0,100,0,100 # 设置初试窗口位置和大小
# track_window = (c,r,w,h)

cap = cv2.VideoCapture(0)
cv2.namedWindow('origin')   #新建一个窗口
cv2.setMouseCallback('origin',onMouse)  
ret, frame= cap.read()

# 设置追踪的区域
# roi = frame[r:r+h, c:c+w]
# roi区域的hsv图像
hsv_roi = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
# 取值hsv值在(0,60,32)到(180,255,255)之间的部分
mask = cv2.inRange(hsv_roi, np.array((0., 60.,32.)), np.array((180.,255.,255.)))
# 计算直方图,参数为 图片(可多)，通道数，蒙板区域，直方图长度，范围
roi_hist = cv2.calcHist([hsv_roi],[0],mask,[180],[0,180])
# 归一化
cv2.normalize(roi_hist,roi_hist,0,255,cv2.NORM_MINMAX)

# 设置终止条件，迭代10次或者至少移动1次
term_crit = ( cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 1 )

while(1):
    ret, frame = cap.read()
    if trackObject != 0:
        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
        #cv2.imshow('hsv',hsv)
        mask = cv2.inRange(hsv, np.array((0., 30.,10.)), np.array((180.,256.,255.)))    #阈值分割
        #cv2.imshow('mask',mask)
        if trackObject == -1:
            hsv_roi = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
            track_window = (xs,ys,ws,hs)
            maskroi = mask[ys:ys+hs, xs:xs+ws]  #在mask图像中截取track_window
            hsv_roi = hsv[ys:ys+hs, xs:xs+ws]   #在hsv图像中截取track_window
            roi_hist = cv2.calcHist([hsv_roi],
                [0],    #计算直方图的通道，这里使用颜色计算直方图，所以就直接使用第一个通道；
                maskroi,
                [180],  #直方图分成180份
                [0.,180.])#表示直方图中需要统计的各个像素的值，[0.0, 180.0]表示直方图能表示RGB所有颜色。
            cv2.normalize(roi_hist,roi_hist,0,255,cv2.NORM_MINMAX) #线性归一化
            roi_hist_f=roi_hist.flatten()   #value转成list
            plt.bar(range(0,len(roi_hist_f)*2,2), roi_hist_f) #显示方便，这里x轴*2，间隔也*2
            plt.show()
            trackObject=1
        dst = cv2.calcBackProject([hsv],[0],roi_hist,[0,180],1)
        cv2.imshow('dst',dst)
        ret, track_window = cv2.meanShift(dst, track_window, term_crit)
        x,y,w,h = track_window
        frame = cv2.rectangle(frame, (x,y), (x+w,y+h), 255,2)
    cv2.imshow('origin',frame)
    if  cv2.waitKey(10)==27:
        break
cap.release()
cv2.destroyAllWindows()
猜你喜欢
