Switch to side-by-side view

--- a
+++ b/tools/misc/flow_extraction.py
@@ -0,0 +1,187 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import argparse
+import os
+import os.path as osp
+
+import cv2
+import numpy as np
+
+
+def flow_to_img(raw_flow, bound=20.):
+    """Convert flow to gray image.
+
+    Args:
+        raw_flow (np.ndarray[float]): Estimated flow with the shape (w, h).
+        bound (float): Bound for the flow-to-image normalization. Default: 20.
+
+    Returns:
+        np.ndarray[uint8]: The result list of np.ndarray[uint8], with shape
+                        (w, h).
+    """
+    flow = np.clip(raw_flow, -bound, bound)
+    flow += bound
+    flow *= (255 / float(2 * bound))
+    flow = flow.astype(np.uint8)
+    return flow
+
+
+def generate_flow(frames, method='tvl1'):
+    """Estimate flow with given frames.
+
+    Args:
+        frames (list[np.ndarray[uint8]]): List of rgb frames, with shape
+                                        (w, h, 3).
+        method (str): Use which method to generate flow. Options are 'tvl1'
+                    and 'farneback'. Default: 'tvl1'.
+
+    Returns:
+        list[np.ndarray[float]]: The result list of np.ndarray[float], with
+                                shape (w, h, 2).
+    """
+    assert method in ['tvl1', 'farneback']
+    gray_frames = [cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) for frame in frames]
+
+    if method == 'tvl1':
+        tvl1 = cv2.optflow.DualTVL1OpticalFlow_create()
+
+        def op(x, y):
+            return tvl1.calc(x, y, None)
+    elif method == 'farneback':
+
+        def op(x, y):
+            return cv2.calcOpticalFlowFarneback(x, y, None, 0.5, 3, 15, 3, 5,
+                                                1.2, 0)
+
+    gray_st = gray_frames[:-1]
+    gray_ed = gray_frames[1:]
+
+    flow = [op(x, y) for x, y in zip(gray_st, gray_ed)]
+    return flow
+
+
+def extract_dense_flow(path,
+                       dest,
+                       bound=20.,
+                       save_rgb=False,
+                       start_idx=0,
+                       rgb_tmpl='img_{:05d}.jpg',
+                       flow_tmpl='{}_{:05d}.jpg',
+                       method='tvl1'):
+    """Extract dense flow given video or frames, save them as gray-scale
+    images.
+
+    Args:
+        path (str): Location of the input video.
+        dest (str): The directory to store the extracted flow images.
+        bound (float): Bound for the flow-to-image normalization. Default: 20.
+        save_rgb (bool): Save extracted RGB frames. Default: False.
+        start_idx (int): The starting frame index if use frames as input, the
+            first image is path.format(start_idx). Default: 0.
+        rgb_tmpl (str): The template of RGB frame names, Default:
+            'img_{:05d}.jpg'.
+        flow_tmpl (str): The template of Flow frame names, Default:
+            '{}_{:05d}.jpg'.
+        method (str): Use which method to generate flow. Options are 'tvl1'
+            and 'farneback'. Default: 'tvl1'.
+    """
+
+    frames = []
+    assert osp.exists(path)
+    video = cv2.VideoCapture(path)
+    flag, f = video.read()
+    while flag:
+        frames.append(f)
+        flag, f = video.read()
+
+    flow = generate_flow(frames, method=method)
+
+    flow_x = [flow_to_img(x[:, :, 0], bound) for x in flow]
+    flow_y = [flow_to_img(x[:, :, 1], bound) for x in flow]
+
+    if not osp.exists(dest):
+        os.system('mkdir -p ' + dest)
+    flow_x_names = [
+        osp.join(dest, flow_tmpl.format('x', ind + start_idx))
+        for ind in range(len(flow_x))
+    ]
+    flow_y_names = [
+        osp.join(dest, flow_tmpl.format('y', ind + start_idx))
+        for ind in range(len(flow_y))
+    ]
+
+    num_frames = len(flow)
+    for i in range(num_frames):
+        cv2.imwrite(flow_x_names[i], flow_x[i])
+        cv2.imwrite(flow_y_names[i], flow_y[i])
+
+    if save_rgb:
+        img_names = [
+            osp.join(dest, rgb_tmpl.format(ind + start_idx))
+            for ind in range(len(frames))
+        ]
+        for frame, name in zip(frames, img_names):
+            cv2.imwrite(name, frame)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='Extract flow and RGB images')
+    parser.add_argument(
+        '--input',
+        help='videos for frame extraction, can be'
+        'single video or a video list, the video list should be a txt file '
+        'and just consists of filenames without directories')
+    parser.add_argument(
+        '--prefix',
+        default='',
+        help='the prefix of input '
+        'videos, used when input is a video list')
+    parser.add_argument(
+        '--dest',
+        default='',
+        help='the destination to save '
+        'extracted frames')
+    parser.add_argument(
+        '--save-rgb', action='store_true', help='also save '
+        'rgb frames')
+    parser.add_argument(
+        '--rgb-tmpl',
+        default='img_{:05d}.jpg',
+        help='template filename of rgb frames')
+    parser.add_argument(
+        '--flow-tmpl',
+        default='{}_{:05d}.jpg',
+        help='template filename of flow frames')
+    parser.add_argument(
+        '--start-idx',
+        type=int,
+        default=1,
+        help='the start '
+        'index of extracted frames')
+    parser.add_argument(
+        '--method',
+        default='tvl1',
+        help='use which method to '
+        'generate flow')
+    parser.add_argument(
+        '--bound', type=float, default=20, help='maximum of '
+        'optical flow')
+
+    args = parser.parse_args()
+    return args
+
+
+if __name__ == '__main__':
+    args = parse_args()
+    if args.input.endswith('.txt'):
+        lines = open(args.input).readlines()
+        lines = [x.strip() for x in lines]
+        videos = [osp.join(args.prefix, x) for x in lines]
+        dests = [osp.join(args.dest, x.split('.')[0]) for x in lines]
+        for video, dest in zip(videos, dests):
+            extract_dense_flow(video, dest, args.bound, args.save_rgb,
+                               args.start_idx, args.rgb_tmpl, args.flow_tmpl,
+                               args.method)
+    else:
+        extract_dense_flow(args.input, args.dest, args.bound, args.save_rgb,
+                           args.start_idx, args.rgb_tmpl, args.flow_tmpl,
+                           args.method)