File size: 6,294 Bytes
d3dbf03 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 |
# 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)
|