# coding=utf-8

"""

4. 使用光流估计方法，在前述测试视频上计算特征点，进一步进行特征点光流估计。
结果分析：
1) 测试视频（vtest.avi）共有 795 帧
2) 每 50 帧截取一次前景(含特征点光流估计: 利用圆圈和短直线显示移动向量)
https://gitee.com/skyrookieyu/csdn_ai25_week3/blob/master/OF-50.jpg
https://gitee.com/skyrookieyu/csdn_ai25_week3/blob/master/OF-100.jpg
https://gitee.com/skyrookieyu/csdn_ai25_week3/blob/master/OF-150.jpg
https://gitee.com/skyrookieyu/csdn_ai25_week3/blob/master/OF-200.jpg
https://gitee.com/skyrookieyu/csdn_ai25_week3/blob/master/OF-250.jpg
https://gitee.com/skyrookieyu/csdn_ai25_week3/blob/master/OF-300.jpg
https://gitee.com/skyrookieyu/csdn_ai25_week3/blob/master/OF-350.jpg
https://gitee.com/skyrookieyu/csdn_ai25_week3/blob/master/OF-400.jpg
https://gitee.com/skyrookieyu/csdn_ai25_week3/blob/master/OF-450.jpg
https://gitee.com/skyrookieyu/csdn_ai25_week3/blob/master/OF-500.jpg
https://gitee.com/skyrookieyu/csdn_ai25_week3/blob/master/OF-550.jpg
https://gitee.com/skyrookieyu/csdn_ai25_week3/blob/master/OF-600.jpg
https://gitee.com/skyrookieyu/csdn_ai25_week3/blob/master/OF-650.jpg
https://gitee.com/skyrookieyu/csdn_ai25_week3/blob/master/OF-700.jpg
https://gitee.com/skyrookieyu/csdn_ai25_week3/blob/master/OF-750.jpg
3) 遗漏许多肉眼明显可见的追踪，符合老师在课堂强调 L-K 方法在实际上需要合并其他的方法。

"""

import cv2 as cv
import copy
import numpy as np

videoFilename = r"vtest.avi"  # 测试视频

# 角点检测所需参数,在进行光流估计计算之前，需要先进行角点检测，然后再把检测出的特征点进行光流估计。
# maxCorners : 设置最多返回的关键点数量。
# qualityLevel : 角点的品质因子，决定能不能成为角点。
# minDistance : 关键点之间的最少像素点，此邻域范围内如果存在更强角点，则删除此角点。
# blockSize : 计算一个角点是否为关键点时所取的区域大小。
feature_params = dict(maxCorners=100,
                      qualityLevel=0.3,
                      minDistance=7,
                      blockSize=7)
# lucas Kanade参数
# winSize：每个金字塔级别的搜索窗口大小。
# maxLevel：从0开始的最大金字塔等级数字；如果设置为0，则不使用金字塔（单个级别）；如果设置为1，则使用两个级别。
# criteria：参数，指定迭代搜索算法的终止条件（在指定的最大迭代次数之后 - criteria.maxCount 或当搜索窗口移动小于时 - criteria.epsilon）。
lk_params = dict(winSize=(15, 15),
                 maxLevel=2,
                 criteria=(cv.TERM_CRITERIA_EPS | cv.TERM_CRITERIA_COUNT, 10, 0.03))
COLORS = ((0xff, 0, 0),  # BLUE 0
          (0, 0xff,  0),  # GREEN 1
          (0, 0, 0xff),  # RED 2
          (0xff, 0x7f, 0),  # ORANGE 3
          (0xff, 0xff, 0),  # YELLOW 4
          (0x82, 0, 0x4b),  # CYAN 5
          (0xd3, 0, 0x94))  # PURPLE 6

cap = cv.VideoCapture(videoFilename)

# 拿到第一帧图像
ret, prev = cap.read()
prevGray = cv.cvtColor(prev, cv.COLOR_BGR2GRAY)

# 返回所有检测特征点，第一个参数输入图像，maxCorners：角点最大数量（效率），qualityLevel：品质因子（特征值越大的越好，来筛选）
# minDistance：距离，相当于这区间有比这个角点强的，就不要这个弱的了。
p0 = cv.goodFeaturesToTrack(prevGray, mask=None, **feature_params)  # 获取图像中最好的角点特征, None表示在整幅图上寻找角点。

mask = np.zeros_like(prev)  # 创建一个mask
frameCount = 0
while True:
    ret, frame = cap.read()
    if not ret:
        break

    frameCount += 1
    gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
    # 需要传入前一帧和当前图像以及前一帧检测到的角点
    p1, st, err = cv.calcOpticalFlowPyrLK(prevGray, gray, p0, None, **lk_params)

    goodPoints = p1[st == 1]  # st=1表示当前帧图像检测到了上一帧的角点特征
    goodPrevPoints = p0[st == 1]

    res = frame.copy()
    drawColor = (0, 0, 255)
    for i, (cur, prev) in enumerate(zip(goodPoints, goodPrevPoints)):
        x0, y0 = cur.ravel()  # 新坐标
        x1, y1 = prev.ravel()  # 旧坐标
        cv.line(res, (x0, y0), (x1, y1), COLORS[4], 2)  # 直线的起点（x0，y0），直线的终点坐标（x1，y1）
        cv.circle(res, (x0, y0), 3, COLORS[6])
        #mask = cv.line(mask, (x0, y0), (x1, y1), drawColor, 2)  # 直线的起点（a，b），直线的终点坐标（c，d）
        #res = cv.circle(res, (x0, y0), 5, drawColor, -1)
    #img = cv.add(res, mask)

    prevGray = gray.copy()
    p0 = goodPoints.reshape(-1, 1, 2)

    cv.imshow('Result', res)
    if frameCount % 50 == 0:
        cv.imwrite('{}-{}.jpg'.format("OF", frameCount), res)
    key = cv.waitKey(30)
    if key == 27:
        break

cap.release()
cv.destroyAllWindows()