FGT_codes/RAFT/utils/flow_viz.py

# Flow visualization code used from https://github.com/tomrunia/OpticalFlow_Visualization


# 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

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

    Code follows the original C++ source code of Daniel Scharstein.
    Code follows the the Matlab source code of Deqing Sun.

    Returns:
        np.ndarray: 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.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_uv_to_colors(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

    Args:
        u (np.ndarray): Input horizontal flow of shape [H,W]
        v (np.ndarray): Input vertical flow of shape [H,W]
        convert_to_bgr (bool, optional): Convert output image to BGR. Defaults to False.

    Returns:
        np.ndarray: Flow visualization image of shape [H,W,3]
    """
    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)
    k0 = np.floor(fk).astype(np.int32)
    k1 = k0 + 1
    k1[k1 == ncols] = 0
    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_image(flow_uv, clip_flow=None, convert_to_bgr=False):
    """
    Expects a two dimensional flow image of shape.

    Args:
        flow_uv (np.ndarray): Flow UV image of shape [H,W,2]
        clip_flow (float, optional): Clip maximum of flow values. Defaults to None.
        convert_to_bgr (bool, optional): Convert output image to BGR. Defaults to False.

    Returns:
        np.ndarray: Flow visualization image of shape [H,W,3]
    """
    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_uv_to_colors(u, v, convert_to_bgr)


def flow_to_image_noNorm(flow_uvs, clip_flow=None, convert_to_bgr=False):
    """
    flow_uvs is a list that accomodates lots of flows
    All the flows in flow_uvs are normalized by the maximum value of in all the flows of flow_uvs
    """
    maximum_value = 0
    for flow_uv in flow_uvs:
        u = flow_uv[:, :, 0]
        v = flow_uv[:, :, 1]
        rad = np.sqrt(np.square(u) + np.square(v))
        rad_max = np.max(rad)
        maximum_value = max(maximum_value, rad_max)
    assert maximum_value > 0, 'Maximum value must be greater than 0'
    flow_colors = []
    for flow_uv in flow_uvs:
        u = flow_uv[:, :, 0]
        v = flow_uv[:, :, 1]
        epsilon = 1e-5
        u = u / (maximum_value + epsilon)
        v = v / (maximum_value + epsilon)
        flow_color = flow_uv_to_colors(u, v, convert_to_bgr)
        flow_colors.append(flow_color)
    return flow_colors


if __name__ == '__main__':
    # cvbase.read_flow test
    import cvbase
    import cv2
    import os
    import glob
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--flow_path', type=str, default='.')
    parser.add_argument('--out_path', type=str, default='.')
    args = parser.parse_args()
    flow_path = args.flow_path
    out_path = args.out_path
    videos = os.listdir(flow_path)
    video_nums = len(videos)
    v = 0
    for video in videos:
        v += 1
        # forward flow visualization
        forward_flows = sorted(glob.glob(os.path.join(flow_path, video, 'forward_flo', '*.flo')))
        backward_flows = sorted(glob.glob(os.path.join(flow_path, video, 'backward_flo', '*.flo')))
        assert len(forward_flows) == len(backward_flows), 'Unmatched number of flows, forward flow is {}, backward flow is {}'.format(len(forward_flows), len(backward_flows))
        forward_out_path = os.path.join(out_path, video, 'forward_flo')
        backward_out_path = os.path.join(out_path, video, 'backward_flo')
        if not os.path.exists(forward_out_path):
            os.makedirs(forward_out_path)
        if not os.path.exists(backward_out_path):
            os.makedirs(backward_out_path)
        for i in range(len(forward_flows)):
            forward_flow_data = cvbase.read_flow(forward_flows[i])
            backward_flow_data = cvbase.read_flow(backward_flows[i])
            forward_flow_image = flow_to_image(forward_flow_data, convert_to_bgr=True)
            backward_flow_image = flow_to_image(backward_flow_data, convert_to_bgr=True)
            cv2.imwrite(os.path.join(forward_out_path, '{:05d}.png'.format(i)), forward_flow_image)
            cv2.imwrite(os.path.join(backward_out_path, '{:05d}.png'.format(i)), backward_flow_image)
        print('[{}]/[{}] video {} has been processed'.format(v, video_nums, video))