-
Notifications
You must be signed in to change notification settings - Fork 197
/
Copy pathvideo_extract_frames_mapper.py
174 lines (152 loc) · 7.31 KB
/
video_extract_frames_mapper.py
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
import json
import os
import os.path as osp
from pydantic import PositiveInt
from data_juicer.utils.constant import Fields, MetaKeys
from data_juicer.utils.file_utils import dict_to_hash
from data_juicer.utils.mm_utils import (
SpecialTokens, close_video, extract_key_frames,
extract_key_frames_by_seconds, extract_video_frames_uniformly,
extract_video_frames_uniformly_by_seconds, load_data_with_context,
load_video)
from ..base_op import OPERATORS, TAGGING_OPS, Mapper
from ..op_fusion import LOADED_VIDEOS
OP_NAME = 'video_extract_frames_mapper'
@TAGGING_OPS.register_module(OP_NAME)
@OPERATORS.register_module(OP_NAME)
@LOADED_VIDEOS.register_module(OP_NAME)
class VideoExtractFramesMapper(Mapper):
"""Mapper to extract frames from video files according to specified methods.
Extracted Frames Data Format:
The data format for the extracted frames is a dictionary mapping
video key to extracted frames directory where the extracted
frames are saved. The dictionary follows the structure:
{
"video_key_1": "/${frame_dir}/video_key_1_filename/",
"video_key_2": "/${frame_dir}/video_key_2_filename/",
...
}
"""
_batched_op = True
def __init__(
self,
frame_sampling_method: str = 'all_keyframes',
frame_num: PositiveInt = 3,
duration: float = 0,
frame_dir: str = None,
frame_key=MetaKeys.video_frames,
*args,
**kwargs,
):
"""
Initialization method.
:param frame_sampling_method: sampling method of extracting frame
videos from the videos. Should be one of
["all_keyframes", "uniform"].
The former one extracts all key frames (the number
of which depends on the duration of the video) and the latter
one extract specified number of frames uniformly from the video.
If "duration" > 0, frame_sampling_method acts on every segment.
Default: "all_keyframes".
:param frame_num: the number of frames to be extracted uniformly from
the video. Only works when frame_sampling_method is "uniform". If
it's 1, only the middle frame will be extracted. If it's 2, only
the first and the last frames will be extracted. If it's larger
than 2, in addition to the first and the last frames, other frames
will be extracted uniformly within the video duration.
If "duration" > 0, frame_num is the number of frames per segment.
:param duration: The duration of each segment in seconds.
If 0, frames are extracted from the entire video.
If duration > 0, the video is segmented into multiple segments
based on duration, and frames are extracted from each segment.
:param frame_dir: Output directory to save extracted frames.
If None, a default directory based on the video file path is used.
:param frame_key: The name of field to save generated frames info.
:param args: extra args
:param kwargs: extra args
"""
super().__init__(*args, **kwargs)
self._init_parameters = self.remove_extra_parameters(locals())
if frame_sampling_method not in ['all_keyframes', 'uniform']:
raise ValueError(
f'Frame sampling method '
f'[{frame_sampling_method}] is not supported. '
f'Can only be one of ["all_keyframes", "uniform"].')
self.frame_dir = frame_dir
self.frame_sampling_method = frame_sampling_method
self.frame_num = frame_num
self.duration = duration
self.frame_key = frame_key
self.frame_fname_template = 'frame_{}.jpg'
def _get_default_frame_dir(self, original_filepath):
original_dir = os.path.dirname(original_filepath)
dir_token = f'/{Fields.multimodal_data_output_dir}/'
if dir_token in original_dir:
original_dir = original_dir.split(dir_token)[0]
saved_dir = os.path.join(
original_dir, f'{Fields.multimodal_data_output_dir}/{OP_NAME}')
original_filename = osp.splitext(osp.basename(original_filepath))[0]
hash_val = dict_to_hash(self._init_parameters)
return osp.join(saved_dir,
f'{original_filename}__dj_hash_#{hash_val}#')
def process_single(self, sample, context=False):
# check if it's generated already
if self.frame_key in sample[Fields.meta]:
return sample
# there is no videos in this sample
if self.video_key not in sample or not sample[self.video_key]:
return []
# load videos
loaded_video_keys = sample[self.video_key]
sample, videos = load_data_with_context(sample, context,
loaded_video_keys, load_video)
video_to_frame_dir = {}
text = sample[self.text_key]
offset = 0
for chunk in text.split(SpecialTokens.eoc):
video_count = chunk.count(SpecialTokens.video)
# no video or no text
if video_count == 0 or len(chunk) == 0:
continue
else:
for video_key in loaded_video_keys[offset:offset +
video_count]:
video = videos[video_key]
# extract frame videos
if self.frame_sampling_method == 'all_keyframes':
if self.duration:
frames = extract_key_frames_by_seconds(
video, self.duration)
else:
frames = extract_key_frames(video)
elif self.frame_sampling_method == 'uniform':
if self.duration:
frames = extract_video_frames_uniformly_by_seconds(
video, self.frame_num, duration=self.duration)
else:
frames = extract_video_frames_uniformly(
video, self.frame_num)
else:
raise ValueError(f'Not support sampling method \
`{self.frame_sampling_method}`.')
frames = [frame.to_image() for frame in frames]
if self.frame_dir:
frame_dir = osp.join(
self.frame_dir,
osp.splitext(osp.basename(video_key))[0])
else:
# video path as frames directory
frame_dir = self._get_default_frame_dir(video_key)
os.makedirs(frame_dir, exist_ok=True)
video_to_frame_dir[video_key] = frame_dir
for i, frame in enumerate(frames):
frame_path = osp.join(
frame_dir, self.frame_fname_template.format(i))
if not os.path.exists(frame_path):
frame.save(frame_path)
offset += video_count
if not context:
for vid_key in videos:
close_video(videos[vid_key])
sample[Fields.meta][self.frame_key] = json.dumps(video_to_frame_dir)
return sample