cuda加速稀疏光流法optic-flow

import cv2
import numpy as np
import time

def initOpicalFlowPoints(p_start, p_end, num):

    if num > (p_end[0] - p_start[0]) or num > (p_end[1] - p_start[1]):
        print('err: num too large')
        return 0
    p0 = np.ones((num, num, 2))
    for i in range(num):
        p0[:, :, 0] = np.linspace(p_start[0], p_end[0], num).astype(int).reshape((1, num))
        p0[:, :, 1] = np.linspace(p_start[1], p_end[1], num).astype(int).reshape((num, 1))
    #p0 = p0.reshape((int(num * num), 1, 2))
    p0 = np.float32(p0).reshape((int(num * num), 1, 2))

    return p0

cap = cv2.VideoCapture(0)
feature_params = dict(maxCorners=100, qualityLevel=0.3, minDistance=7, blockSize=7)
lk_params = dict(winSize=(15, 15), maxLevel=2, criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.02))
color = np.random.randint(0,255,(100,3))
ret, old_frame = cap.read()
old_gray = cv2.cvtColor(old_frame, cv2.COLOR_BGR2GRAY)
#p0 = cv2.goodFeaturesToTrack(old_gray, mask=None, **feature_params)#选取好的特征点，返回特征点列表
p0 = initOpicalFlowPoints((200, 260), (400, 420), 2)
p00 = p0
#mask = np.zeros_like(old_frame)
brlist = []
brlist1 = []
optFlow_gpu = cv2.cuda_SparsePyrLKOpticalFlow.create()
while True:
    anow = time.time()
    ret, frame = cap.read()
    if frame is None:
        break
    print(time.time() - anow)
    frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    print(time.time() - anow)
    start_of = time.time()
    '''
    # calculate optical flow
    gpu_frame = cv2.cuda_GpuMat()
    gpu_frame.upload(frame_gray)
    # upload pre-processed frame to GPU
    gpu_previous = cv2.cuda_GpuMat()
    gpu_previous.upload(old_gray)
    gpu_p0 = cv2.cuda_GpuMat()
    gpu_p0.upload(p0.reshape(1, p0.shape[0], 2))
    gpu_p1, gpu_st, gpu_err = optFlow_gpu.calc(gpu_previous, gpu_frame, gpu_p0, None, )

    p1 = gpu_p1.download()
    '''
    p1, st, err = cv2.calcOpticalFlowPyrLK(old_gray, frame_gray, p0, None, )
    end_of = time.time()
    print(end_of - start_of)

    good_new = p1.reshape(-1, 2)
    good_old = p0
    good_old1 = p00
    for i,(new,old) in enumerate(zip(good_new,good_old)):
        a,b = new.ravel() #ravel()函数用于降维并且不产生副本
        c,d = old.ravel()
        #mask = cv2.line(mask, (int(a),int(b)),(int(c),int(d)), color[i].tolist(), 2)
        frame = cv2.circle(frame,(int(a),int(b)),5,(0,255,0),-1)#color[i].tolist(),-1)
    #img = cv2.add(frame,mask)

    cv2.imshow('frame', frame)
    k = cv2.waitKey(1) & 0xff
    if k == 27:
        break
    old_gray = frame_gray.copy()
    p0 = good_new.reshape(-1,1,2)
    end_of = time.time()
    print(end_of - start_of)
    print('fps:', time.time() - anow)
cap.release()