128 lines
3.3 KiB
Python
128 lines
3.3 KiB
Python
|
# coding: utf-8
|
||
|
|
"""Common tools to optical flow algorithms.
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
from skimage.transform import pyramid_reduce
|
||
|
|
from skimage.util.dtype import _convert
|
||
|
|
from scipy import ndimage as ndi
|
||
|
|
|
||
|
|
|
||
|
|
def resize_flow(flow, shape):
|
||
|
|
"""Rescale the values of the vector field (u, v) to the desired shape.
|
||
|
|
|
||
|
|
The values of the output vector field are scaled to the new
|
||
|
|
resolution.
|
||
|
|
|
||
|
|
Parameters
|
||
|
|
----------
|
||
|
|
flow : ndarray
|
||
|
|
The motion field to be processed.
|
||
|
|
shape : iterable
|
||
|
|
Couple of integers representing the output shape.
|
||
|
|
|
||
|
|
Returns
|
||
|
|
-------
|
||
|
|
rflow : ndarray
|
||
|
|
The resized and rescaled motion field.
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
scale = [n / o for n, o in zip(shape, flow.shape[1:])]
|
||
|
|
scale_factor = np.array(scale, dtype=flow.dtype)
|
||
|
|
|
||
|
|
for _ in shape:
|
||
|
|
scale_factor = scale_factor[..., np.newaxis]
|
||
|
|
|
||
|
|
rflow = scale_factor*ndi.zoom(flow, [1] + scale, order=0,
|
||
|
|
mode='nearest', prefilter=False)
|
||
|
|
|
||
|
|
return rflow
|
||
|
|
|
||
|
|
|
||
|
|
def get_pyramid(I, downscale=2.0, nlevel=10, min_size=16):
|
||
|
|
"""Construct image pyramid.
|
||
|
|
|
||
|
|
Parameters
|
||
|
|
----------
|
||
|
|
I : ndarray
|
||
|
|
The image to be preprocessed (Gray scale or RGB).
|
||
|
|
downscale : float
|
||
|
|
The pyramid downscale factor.
|
||
|
|
nlevel : int
|
||
|
|
The maximum number of pyramid levels.
|
||
|
|
min_size : int
|
||
|
|
The minimum size for any dimension of the pyramid levels.
|
||
|
|
|
||
|
|
Returns
|
||
|
|
-------
|
||
|
|
pyramid : list[ndarray]
|
||
|
|
The coarse to fine images pyramid.
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
pyramid = [I]
|
||
|
|
size = min(I.shape)
|
||
|
|
count = 1
|
||
|
|
|
||
|
|
while (count < nlevel) and (size > downscale * min_size):
|
||
|
|
J = pyramid_reduce(pyramid[-1], downscale, multichannel=False)
|
||
|
|
pyramid.append(J)
|
||
|
|
size = min(J.shape)
|
||
|
|
count += 1
|
||
|
|
|
||
|
|
return pyramid[::-1]
|
||
|
|
|
||
|
|
|
||
|
|
def coarse_to_fine(I0, I1, solver, downscale=2, nlevel=10, min_size=16,
|
||
|
|
dtype=np.float32):
|
||
|
|
"""Generic coarse to fine solver.
|
||
|
|
|
||
|
|
Parameters
|
||
|
|
----------
|
||
|
|
I0 : ndarray
|
||
|
|
The first gray scale image of the sequence.
|
||
|
|
I1 : ndarray
|
||
|
|
The second gray scale image of the sequence.
|
||
|
|
solver : callable
|
||
|
|
The solver applyed at each pyramid level.
|
||
|
|
downscale : float
|
||
|
|
The pyramid downscale factor.
|
||
|
|
nlevel : int
|
||
|
|
The maximum number of pyramid levels.
|
||
|
|
min_size : int
|
||
|
|
The minimum size for any dimension of the pyramid levels.
|
||
|
|
dtype : dtype
|
||
|
|
Output data type.
|
||
|
|
|
||
|
|
Returns
|
||
|
|
-------
|
||
|
|
flow : ndarray
|
||
|
|
The estimated optical flow components for each axis.
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
if I0.shape != I1.shape:
|
||
|
|
raise ValueError("Input images should have the same shape")
|
||
|
|
|
||
|
|
if np.dtype(dtype).char not in 'efdg':
|
||
|
|
raise ValueError("Only floating point data type are valid"
|
||
|
|
" for optical flow")
|
||
|
|
|
||
|
|
pyramid = list(zip(get_pyramid(_convert(I0, dtype),
|
||
|
|
downscale, nlevel, min_size),
|
||
|
|
get_pyramid(_convert(I1, dtype),
|
||
|
|
downscale, nlevel, min_size)))
|
||
|
|
|
||
|
|
# Initialization to 0 at coarsest level.
|
||
|
|
flow = np.zeros((pyramid[0][0].ndim, ) + pyramid[0][0].shape,
|
||
|
|
dtype=dtype)
|
||
|
|
|
||
|
|
flow = solver(pyramid[0][0], pyramid[0][1], flow)
|
||
|
|
|
||
|
|
for J0, J1 in pyramid[1:]:
|
||
|
|
flow = solver(J0, J1, resize_flow(flow, J0.shape))
|
||
|
|
|
||
|
|
return flow
|