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