276 lines
7.0 KiB
Python
276 lines
7.0 KiB
Python
# MIT License
|
|
#
|
|
# Copyright (c) 2018 Tom Runia
|
|
#
|
|
# Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
# of this software and associated documentation files (the "Software"), to deal
|
|
# in the Software without restriction, including without limitation the rights
|
|
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
# copies of the Software, and to permit persons to whom the Software is
|
|
# furnished to do so, subject to conditions.
|
|
#
|
|
# Author: Tom Runia
|
|
# Date Created: 2018-08-03
|
|
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
|
|
import numpy as np
|
|
|
|
|
|
def make_colorwheel():
|
|
'''
|
|
Generates a color wheel for optical flow visualization as presented in:
|
|
Baker et al. "A Database and Evaluation Methodology for Optical Flow" (ICCV, 2007)
|
|
URL: http://vision.middlebury.edu/flow/flowEval-iccv07.pdf
|
|
According to the C++ source code of Daniel Scharstein
|
|
According to the Matlab source code of Deqing Sun
|
|
'''
|
|
|
|
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(u, v, convert_to_bgr=False):
|
|
'''
|
|
Applies the flow color wheel to (possibly clipped) flow components u and v.
|
|
According to the C++ source code of Daniel Scharstein
|
|
According to the Matlab source code of Deqing Sun
|
|
:param u: np.ndarray, input horizontal flow
|
|
:param v: np.ndarray, input vertical flow
|
|
:param convert_to_bgr: bool, whether to change ordering and output BGR instead of RGB
|
|
:return:
|
|
'''
|
|
|
|
flow_image = np.zeros((u.shape[0], u.shape[1], 3), np.uint8)
|
|
|
|
colorwheel = make_colorwheel() # shape [55x3]
|
|
ncols = colorwheel.shape[0]
|
|
|
|
rad = np.sqrt(np.square(u) + np.square(v))
|
|
a = np.arctan2(-v, -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?
|
|
|
|
# Note the 2-i => BGR instead of RGB
|
|
ch_idx = 2-i if convert_to_bgr else i
|
|
flow_image[:,:,ch_idx] = np.floor(255 * col)
|
|
|
|
return flow_image
|
|
|
|
|
|
def flow_to_color(flow_uv, clip_flow=None, convert_to_bgr=False):
|
|
'''
|
|
Expects a two dimensional flow image of shape [H,W,2]
|
|
According to the C++ source code of Daniel Scharstein
|
|
According to the Matlab source code of Deqing Sun
|
|
:param flow_uv: np.ndarray of shape [H,W,2]
|
|
:param clip_flow: float, maximum clipping value for flow
|
|
:return:
|
|
'''
|
|
|
|
assert flow_uv.ndim == 3, 'input flow must have three dimensions'
|
|
assert flow_uv.shape[2] == 2, 'input flow must have shape [H,W,2]'
|
|
|
|
if clip_flow is not None:
|
|
flow_uv = np.clip(flow_uv, 0, clip_flow)
|
|
|
|
u = flow_uv[:,:,0]
|
|
v = flow_uv[:,:,1]
|
|
|
|
rad = np.sqrt(np.square(u) + np.square(v))
|
|
rad_max = np.max(rad)
|
|
|
|
epsilon = 1e-5
|
|
u = u / (rad_max + epsilon)
|
|
v = v / (rad_max + epsilon)
|
|
|
|
return flow_compute_color(u, v, convert_to_bgr)
|
|
|
|
|
|
|
|
UNKNOWN_FLOW_THRESH = 1e7
|
|
SMALLFLOW = 0.0
|
|
LARGEFLOW = 1e8
|
|
|
|
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
|
|
|
|
# from https://github.com/gengshan-y/VCN
|
|
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.
|
|
|
|
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)
|