1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
|
# Visualize optical flow
# =======================
# Code modified from:
# https://github.com/tomrunia/OpticalFlow_Visualization/blob/master/flow_vis.py
import numpy as np
def make_colorwheel():
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.floor(255 * np.arange(0, RY) / RY)
col = col + RY
# YG
colorwheel[col : col + YG, 0] = 255 - np.floor(255 * np.arange(0, YG) / YG)
colorwheel[col : col + YG, 1] = 255
col = col + YG
# GC
colorwheel[col : col + GC, 1] = 255
colorwheel[col : col + GC, 2] = np.floor(255 * np.arange(0, GC) / GC)
col = col + GC
# CB
colorwheel[col : col + CB, 1] = 255 - np.floor(255 * np.arange(CB) / CB)
colorwheel[col : col + CB, 2] = 255
col = col + CB
# BM
colorwheel[col : col + BM, 2] = 255
colorwheel[col : col + BM, 0] = np.floor(255 * np.arange(0, BM) / BM)
col = col + BM
# MR
colorwheel[col : col + MR, 2] = 255 - np.floor(255 * np.arange(MR) / MR)
colorwheel[col : col + MR, 0] = 255
return colorwheel
def flow_compute_color(flow_u, flow_v):
H, W = flow_u.shape[:2]
flow_image = np.zeros((H, W, 3), np.uint8)
colorwheel = make_colorwheel() # shape [55x3]
ncols = colorwheel.shape[0]
rad = np.sqrt(np.square(flow_u) + np.square(flow_v))
a = np.arctan2(-flow_v, -flow_u) / np.pi
fk = (a + 1) / 2 * (ncols - 1) + 1
k0 = np.floor(fk).astype(np.int32)
k1 = k0 + 1
k1[k1 == ncols] = 1
f = fk - k0
for i in range(colorwheel.shape[1]):
tmp = colorwheel[:, i]
col0 = tmp[k0] / 255.0
col1 = tmp[k1] / 255.0
col = (1 - f) * col0 + f * col1
idx = rad <= 1
col[idx] = 1 - rad[idx] * (1 - col[idx])
col[~idx] = col[~idx] * 0.75 # out of range?
ch_idx = i
flow_image[:, :, ch_idx] = np.floor(255 * col)
return flow_image
def flow2rgb(flow_uv):
"""Visualize optical flow.
Parameters
----------
flow_uv: numpy.ndarray, (H, W, 2), float
Optical flow.
Returns
-------
dst: numpy.ndarray
RGB image.
"""
assert flow_uv.ndim == 3, "flow must be 3 dimensional"
assert flow_uv.shape[2] == 2, "flow must have shape (H, W, 2)"
assert np.issubdtype(flow_uv.dtype, np.floating), "float must be float type"
flow_u = flow_uv[:, :, 0]
flow_v = flow_uv[:, :, 1]
rad = np.sqrt(np.square(flow_u) + np.square(flow_v))
rad_max = np.max(rad)
epsilon = 1e-5
flow_u = flow_u / (rad_max + epsilon)
flow_v = flow_v / (rad_max + epsilon)
return flow_compute_color(flow_u, flow_v)
|
