meanShift

meanShift-demo.py

import numpy as np
import imageio
import cv2

reader = imageio.get_reader(r"W:\Hudi\trypsin7[0_0_12 - 0_0_24].avi")
fps = reader.get_meta_data()['fps']

writer = imageio.get_writer('track1.mp4', fps=fps)

# take first frame of the video
frame = reader.get_next_data()

# setup initial location of window
#### this is the "hard" part
r,h,c,w = 85,18,165,18  # simply hardcoded the values


track_window = (c,r,w,h)
frame2 = frame.copy()
cv2.rectangle(frame2, (c,r), (c+w, r+h), (255, 0, 0), 2)


# set up the ROI for tracking
roi = frame[r:r+h, c:c+w, :]
hsv_roi =  cv2.cvtColor(roi, cv2.COLOR_BGR2HSV)
cv2.imshow('img1', frame2)
# also a bit of voodoo
mask = cv2.inRange(hsv_roi, np.array((0., 0.,65.)), np.array((0.,0.,100.)))
roi_hist = cv2.calcHist([hsv_roi],[2],mask,[180],[0,180])
cv2.normalize(roi_hist,roi_hist,0,255,cv2.NORM_MINMAX)

# Setup the termination criteria, either 10 iteration or move by atleast 1 pt
term_crit = ( cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 1 )

for frame in reader:
    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
    dst = cv2.calcBackProject([hsv],[2],roi_hist,[0,180],1)

    # ret, track_window = cv2.CamShift(dst, track_window, term_crit)
    # print ret
    # print track_window
    ret, track_window = cv2.meanShift(dst, track_window, term_crit)

    # Draw it on image
    cv2.rectangle(frame, (track_window[0], track_window[1]), (track_window[0]+track_window[2], track_window[1]+track_window[3]), (0, 255, 0), 2)
    cv2.imshow('img2',frame)
    writer.append_data(frame)

    k = cv2.waitKey(1) & 0xff
    if k == 27:
        break
    else:
        cv2.imwrite(chr(k)+".jpg",frame)

cv2.destroyAllWindows()
reader.close()
writer.close()