Switch to side-by-side view

--- a
+++ b/tools/data/build_rawframes.py
@@ -0,0 +1,278 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import argparse
+import glob
+import os
+import os.path as osp
+import sys
+import warnings
+from multiprocessing import Lock, Pool
+
+import mmcv
+import numpy as np
+
+
+def extract_frame(vid_item):
+    """Generate optical flow using dense flow.
+
+    Args:
+        vid_item (list): Video item containing video full path,
+            video (short) path, video id.
+
+    Returns:
+        bool: Whether generate optical flow successfully.
+    """
+    full_path, vid_path, vid_id, method, task, report_file = vid_item
+    if '/' in vid_path:
+        act_name = osp.basename(osp.dirname(vid_path))
+        out_full_path = osp.join(args.out_dir, act_name)
+    else:
+        out_full_path = args.out_dir
+
+    run_success = -1
+
+    if task == 'rgb':
+        if args.use_opencv:
+            # Not like using denseflow,
+            # Use OpenCV will not make a sub directory with the video name
+            try:
+                video_name = osp.splitext(osp.basename(vid_path))[0]
+                out_full_path = osp.join(out_full_path, video_name)
+
+                vr = mmcv.VideoReader(full_path)
+                for i, vr_frame in enumerate(vr):
+                    if vr_frame is not None:
+                        w, h, _ = np.shape(vr_frame)
+                        if args.new_short == 0:
+                            if args.new_width == 0 or args.new_height == 0:
+                                # Keep original shape
+                                out_img = vr_frame
+                            else:
+                                out_img = mmcv.imresize(
+                                    vr_frame,
+                                    (args.new_width, args.new_height))
+                        else:
+                            if min(h, w) == h:
+                                new_h = args.new_short
+                                new_w = int((new_h / h) * w)
+                            else:
+                                new_w = args.new_short
+                                new_h = int((new_w / w) * h)
+                            out_img = mmcv.imresize(vr_frame, (new_h, new_w))
+                        mmcv.imwrite(out_img,
+                                     f'{out_full_path}/img_{i + 1:05d}.jpg')
+                    else:
+                        warnings.warn(
+                            'Length inconsistent!'
+                            f'Early stop with {i + 1} out of {len(vr)} frames.'
+                        )
+                        break
+                run_success = 0
+            except Exception:
+                run_success = -1
+        else:
+            if args.new_short == 0:
+                cmd = osp.join(
+                    f"denseflow '{full_path}' -b=20 -s=0 -o='{out_full_path}'"
+                    f' -nw={args.new_width} -nh={args.new_height} -v')
+            else:
+                cmd = osp.join(
+                    f"denseflow '{full_path}' -b=20 -s=0 -o='{out_full_path}'"
+                    f' -ns={args.new_short} -v')
+            run_success = os.system(cmd)
+    elif task == 'flow':
+        if args.input_frames:
+            if args.new_short == 0:
+                cmd = osp.join(
+                    f"denseflow '{full_path}' -a={method} -b=20 -s=1 -o='{out_full_path}'"  # noqa: E501
+                    f' -nw={args.new_width} --nh={args.new_height} -v --if')
+            else:
+                cmd = osp.join(
+                    f"denseflow '{full_path}' -a={method} -b=20 -s=1 -o='{out_full_path}'"  # noqa: E501
+                    f' -ns={args.new_short} -v --if')
+        else:
+            if args.new_short == 0:
+                cmd = osp.join(
+                    f"denseflow '{full_path}' -a={method} -b=20 -s=1 -o='{out_full_path}'"  # noqa: E501
+                    f' -nw={args.new_width} --nh={args.new_height} -v')
+            else:
+                cmd = osp.join(
+                    f"denseflow '{full_path}' -a={method} -b=20 -s=1 -o='{out_full_path}'"  # noqa: E501
+                    f' -ns={args.new_short} -v')
+        run_success = os.system(cmd)
+    else:
+        if args.new_short == 0:
+            cmd_rgb = osp.join(
+                f"denseflow '{full_path}' -b=20 -s=0 -o='{out_full_path}'"
+                f' -nw={args.new_width} -nh={args.new_height} -v')
+            cmd_flow = osp.join(
+                f"denseflow '{full_path}' -a={method} -b=20 -s=1 -o='{out_full_path}'"  # noqa: E501
+                f' -nw={args.new_width} -nh={args.new_height} -v')
+        else:
+            cmd_rgb = osp.join(
+                f"denseflow '{full_path}' -b=20 -s=0 -o='{out_full_path}'"
+                f' -ns={args.new_short} -v')
+            cmd_flow = osp.join(
+                f"denseflow '{full_path}' -a={method} -b=20 -s=1 -o='{out_full_path}'"  # noqa: E501
+                f' -ns={args.new_short} -v')
+        run_success_rgb = os.system(cmd_rgb)
+        run_success_flow = os.system(cmd_flow)
+        if run_success_flow == 0 and run_success_rgb == 0:
+            run_success = 0
+
+    if run_success == 0:
+        print(f'{task} {vid_id} {vid_path} {method} done')
+        sys.stdout.flush()
+
+        lock.acquire()
+        with open(report_file, 'a') as f:
+            line = full_path + '\n'
+            f.write(line)
+        lock.release()
+    else:
+        print(f'{task} {vid_id} {vid_path} {method} got something wrong')
+        sys.stdout.flush()
+
+    return True
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='extract optical flows')
+    parser.add_argument('src_dir', type=str, help='source video directory')
+    parser.add_argument('out_dir', type=str, help='output rawframe directory')
+    parser.add_argument(
+        '--task',
+        type=str,
+        default='flow',
+        choices=['rgb', 'flow', 'both'],
+        help='which type of frames to be extracted')
+    parser.add_argument(
+        '--level',
+        type=int,
+        choices=[1, 2],
+        default=2,
+        help='directory level of data')
+    parser.add_argument(
+        '--num-worker',
+        type=int,
+        default=8,
+        help='number of workers to build rawframes')
+    parser.add_argument(
+        '--flow-type',
+        type=str,
+        default=None,
+        choices=[None, 'tvl1', 'warp_tvl1', 'farn', 'brox'],
+        help='flow type to be generated')
+    parser.add_argument(
+        '--out-format',
+        type=str,
+        default='jpg',
+        choices=['jpg', 'h5', 'png'],
+        help='output format')
+    parser.add_argument(
+        '--ext',
+        type=str,
+        default='avi',
+        choices=['avi', 'mp4', 'webm'],
+        help='video file extensions')
+    parser.add_argument(
+        '--mixed-ext',
+        action='store_true',
+        help='process video files with mixed extensions')
+    parser.add_argument(
+        '--new-width', type=int, default=0, help='resize image width')
+    parser.add_argument(
+        '--new-height', type=int, default=0, help='resize image height')
+    parser.add_argument(
+        '--new-short',
+        type=int,
+        default=0,
+        help='resize image short side length keeping ratio')
+    parser.add_argument('--num-gpu', type=int, default=8, help='number of GPU')
+    parser.add_argument(
+        '--resume',
+        action='store_true',
+        default=False,
+        help='resume optical flow extraction instead of overwriting')
+    parser.add_argument(
+        '--use-opencv',
+        action='store_true',
+        help='Whether to use opencv to extract rgb frames')
+    parser.add_argument(
+        '--input-frames',
+        action='store_true',
+        help='Whether to extract flow frames based on rgb frames')
+    parser.add_argument(
+        '--report-file',
+        type=str,
+        default='build_report.txt',
+        help='report to record files which have been successfully processed')
+    args = parser.parse_args()
+
+    return args
+
+
+def init(lock_):
+    global lock
+    lock = lock_
+
+
+if __name__ == '__main__':
+    args = parse_args()
+
+    if not osp.isdir(args.out_dir):
+        print(f'Creating folder: {args.out_dir}')
+        os.makedirs(args.out_dir)
+
+    if args.level == 2:
+        classes = os.listdir(args.src_dir)
+        for classname in classes:
+            new_dir = osp.join(args.out_dir, classname)
+            if not osp.isdir(new_dir):
+                print(f'Creating folder: {new_dir}')
+                os.makedirs(new_dir)
+
+    if args.input_frames:
+        print('Reading rgb frames from folder: ', args.src_dir)
+        fullpath_list = glob.glob(args.src_dir + '/*' * args.level)
+        print('Total number of rgb frame folders found: ', len(fullpath_list))
+    else:
+        print('Reading videos from folder: ', args.src_dir)
+        if args.mixed_ext:
+            print('Extension of videos is mixed')
+            fullpath_list = glob.glob(args.src_dir + '/*' * args.level)
+        else:
+            print('Extension of videos: ', args.ext)
+            fullpath_list = glob.glob(args.src_dir + '/*' * args.level + '.' +
+                                      args.ext)
+        print('Total number of videos found: ', len(fullpath_list))
+
+    if args.resume:
+        done_fullpath_list = []
+        with open(args.report_file) as f:
+            for line in f:
+                if line == '\n':
+                    continue
+                done_full_path = line.strip().split()[0]
+                done_fullpath_list.append(done_full_path)
+        done_fullpath_list = set(done_fullpath_list)
+        fullpath_list = list(set(fullpath_list).difference(done_fullpath_list))
+
+    if args.level == 2:
+        vid_list = list(
+            map(
+                lambda p: osp.join(
+                    osp.basename(osp.dirname(p)), osp.basename(p)),
+                fullpath_list))
+    elif args.level == 1:
+        vid_list = list(map(osp.basename, fullpath_list))
+
+    lock = Lock()
+    pool = Pool(args.num_worker, initializer=init, initargs=(lock, ))
+    pool.map(
+        extract_frame,
+        zip(fullpath_list, vid_list, range(len(vid_list)),
+            len(vid_list) * [args.flow_type],
+            len(vid_list) * [args.task],
+            len(vid_list) * [args.report_file]))
+    pool.close()
+    pool.join()