Python+OpenCV 实现Farneback光流法从RGB图像序列中提取光流

写在前面：

  前段时间处理行为识别的数据集要做对比实验，久违的又用到了光流图像，这篇把生成光流图像的代码存一下。
  首先，我们用到的是Farneback光流法提取图像序列间的光流特征，这个算法在OpenCV库中是有对应的实现函数的，叫做cv2.calcOpticalFlowFarneback()。在完成光流特征的提取之后，要做的就是光流图像的生成。关于光流图像的生成，本文代码是直接照搬了参考文章中的实现方式，不过肯定还有很多不同的实现，可以多查阅一下其他资料。

参考文章：

  这一篇文章里作者对于cv2.calcOpticalFlowFarneback()函数的参数介绍很详细。
  https://blog.csdn.net/qq_33757398/article/details/124834092
  本文中光流图像生成部分的代码直接用的这一篇。
  https://blog.csdn.net/qq_34535410/article/details/89976801

代码：

import cv2
import numpy as np
from PIL import Image

def make_color_wheel():
    """``
    Generate color wheel according Middlebury color code
    :return: Color wheel
    """
    RY = 15
    YG = 6
    GC = 4
    CB = 11
    BM = 13
    MR = 6

    ncols = RY + YG + GC + CB + BM + MR

    colorwheel = np.zeros([ncols, 3])

    col = 0

    # RY
    colorwheel[0:RY, 0] = 255
    colorwheel[0:RY, 1] = np.transpose(np.floor(255*np.arange(0, RY) / RY))
    col += RY

    # YG
    colorwheel[col:col+YG, 0] = 255 - np.transpose(np.floor(255*np.arange(0, YG) / YG))
    colorwheel[col:col+YG, 1] = 255
    col += YG

    # GC
    colorwheel[col:col+GC, 1] = 255
    colorwheel[col:col+GC, 2] = np.transpose(np.floor(255*np.arange(0, GC) / GC))
    col += GC

    # CB
    colorwheel[col:col+CB, 1] = 255 - np.transpose(np.floor(255*np.arange(0, CB) / CB))
    colorwheel[col:col+CB, 2] = 255
    col += CB

    # BM
    colorwheel[col:col+BM, 2] = 255
    colorwheel[col:col+BM, 0] = np.transpose(np.floor(255*np.arange(0, BM) / BM))
    col += + BM

    # MR
    colorwheel[col:col+MR, 2] = 255 - np.transpose(np.floor(255 * np.arange(0, MR) / MR))
    colorwheel[col:col+MR, 0] = 255

    return colorwheel

def compute_color(u, v):
    """
    compute optical flow color map
    :param u: optical flow horizontal map
    :param v: optical flow vertical map
    :return: optical flow in color code
    """
    [h, w] = u.shape
    img = np.zeros([h, w, 3])
    nanIdx = np.isnan(u) | np.isnan(v)
    u[nanIdx] = 0
    v[nanIdx] = 0

    colorwheel = make_color_wheel()
    ncols = np.size(colorwheel, 0)

    rad = np.sqrt(u**2+v**2)

    a = np.arctan2(-v, -u) / np.pi

    fk = (a+1) / 2 * (ncols - 1) + 1

    k0 = np.floor(fk).astype(int)

    k1 = k0 + 1
    k1[k1 == ncols+1] = 1
    f = fk - k0

    for i in range(0, np.size(colorwheel,1)):
        tmp = colorwheel[:, i]
        col0 = tmp[k0-1] / 255
        col1 = tmp[k1-1] / 255
        col = (1-f) * col0 + f * col1

        idx = rad <= 1
        col[idx] = 1-rad[idx]*(1-col[idx])
        notidx = np.logical_not(idx)

        col[notidx] *= 0.75
        img[:, :, i] = np.uint8(np.floor(255 * col*(1-nanIdx)))

    return img

def flow_to_image(flow):
    """
    Convert flow into middlebury color code image
    :param flow: optical flow map
    :return: optical flow image in middlebury color
    """
    u = flow[:, :, 0]
    v = flow[:, :, 1]

    maxu = -999.
    maxv = -999.
    minu = 999.
    minv = 999.
    UNKNOWN_FLOW_THRESH = 1e7
    SMALLFLOW = 0.0
    LARGEFLOW = 1e8

    idxUnknow = (abs(u) > UNKNOWN_FLOW_THRESH) | (abs(v) > UNKNOWN_FLOW_THRESH)
    u[idxUnknow] = 0
    v[idxUnknow] = 0

    maxu = max(maxu, np.max(u))
    minu = min(minu, np.min(u))

    maxv = max(maxv, np.max(v))
    minv = min(minv, np.min(v))

    rad = np.sqrt(u ** 2 + v ** 2)
    maxrad = max(-1, np.max(rad))

    u = u/(maxrad + np.finfo(float).eps)
    v = v/(maxrad + np.finfo(float).eps)

    img = compute_color(u, v)

    idx = np.repeat(idxUnknow[:, :, np.newaxis], 3, axis=2)
    img[idx] = 0

    return np.uint8(img)

# 读取相邻两张RGB图像
prev = cv2.imread(r"./1/00000.jpg",cv2.IMREAD_UNCHANGED)
next = cv2.imread(r"./1/00001.jpg",cv2.IMREAD_UNCHANGED)

prev_img =  cv2.cvtColor(prev,cv2.COLOR_RGB2GRAY)
next_img =  cv2.cvtColor(next,cv2.COLOR_RGB2GRAY)

# poly_n= 7 or 5
# poly_sigma= 1.5 or 1.1
# 光流提取
flows = cv2.calcOpticalFlowFarneback(prev_img, next_img, None, pyr_scale=0.5, levels=3, winsize=55, iterations=3, poly_n=7, poly_sigma=1.5, flags=cv2.OPTFLOW_FARNEBACK_GAUSSIAN)

# show flows
flow_img = flow_to_image(flows)
cv2.imwrite('./flows.jpg', flow_img)  # save flow_img

效果展示：

  图片太小了显得这个水印多少有点。。。先这样吧hh