support extract frames by seconds

Author: Cathy0908

data_juicer/ops/mapper/video_extract_frames_mapper.py

@@ -6,10 +6,11 @@
 
 from data_juicer.utils.constant import Fields
 from data_juicer.utils.file_utils import create_directory_if_not_exists
-from data_juicer.utils.mm_utils import (SpecialTokens, close_video,
-                                        extract_key_frames,
-                                        extract_video_frames_uniformly,
-                                        load_data_with_context, load_video)
+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, Mapper
 from ..op_fusion import LOADED_VIDEOS
@@ -44,6 +45,7 @@ def __init__(
         self,
         frame_sampling_method: str = 'all_keyframes',
         frame_num: PositiveInt = 3,
+        duration: float = 0,
         frame_dir: str = None,
         frame_key=Fields.video_frames,
         *args,
@@ -57,13 +59,19 @@ def __init__(
             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.
@@ -82,6 +90,7 @@ def __init__(
         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'
 
@@ -109,7 +118,7 @@ def process_single(self, sample, context=False):
         loaded_video_keys = sample[self.video_key]
         sample, videos = load_data_with_context(sample, context,
                                                 loaded_video_keys, load_video)
-        video_to_frames = {}
+        video_to_frame_dir = {}
         text = sample[self.text_key]
         offset = 0
 
@@ -124,10 +133,18 @@ def process_single(self, sample, context=False):
                     video = videos[video_key]
                     # extract frame videos
                     if self.frame_sampling_method == 'all_keyframes':
-                        frames = extract_key_frames(video)
+                        if self.duration:
+                            frames = extract_key_frames_by_seconds(
+                                video, self.duration)
+                        else:
+                            frames = extract_key_frames(video)
                     elif self.frame_sampling_method == 'uniform':
-                        frames = extract_video_frames_uniformly(
-                            video, self.frame_num)
+                        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}`.')
@@ -141,23 +158,20 @@ def process_single(self, sample, context=False):
                         # 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
 
-                    video_to_frames[video_key] = []
                     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)
 
-                        video_to_frames[video_key].append(frame_path)
-
                 offset += video_count
 
         if not context:
             for vid_key in videos:
                 close_video(videos[vid_key])
 
-        sample[self.frame_key] = json.dumps(video_to_frames)
-        # sample[self.frame_key] = video_to_frames
+        sample[self.frame_key] = json.dumps(video_to_frame_dir)
 
         return sample

data_juicer/utils/mm_utils.py

@@ -1,5 +1,6 @@
 import base64
 import datetime
+import io
 import os
 import re
 import shutil
@@ -321,7 +322,11 @@ def cut_video_by_seconds(
         container = input_video
 
     # create the output video
-    output_container = load_video(output_video, 'w')
+    if output_video:
+        output_container = load_video(output_video, 'w')
+    else:
+        output_buffer = io.BytesIO()
+        output_container = av.open(output_buffer, mode='w', format='mp4')
 
     # add the video stream into the output video according to input video
     input_video_stream = container.streams.video[0]
@@ -390,6 +395,11 @@ def cut_video_by_seconds(
     if isinstance(input_video, str):
         close_video(container)
     close_video(output_container)
+
+    if not output_video:
+        output_buffer.seek(0)
+        return output_buffer
+
     if not os.path.exists(output_video):
         logger.warning(f'This video could not be successfully cut in '
                        f'[{start_seconds}, {end_seconds}] seconds. '
@@ -463,6 +473,39 @@ def process_each_frame(input_video: Union[str, av.container.InputContainer],
                 if isinstance(input_video, str) else input_video.name)
 
 
+def extract_key_frames_by_seconds(
+        input_video: Union[str, av.container.InputContainer],
+        duration: float = 1):
+    """Extract key frames by seconds.
+        :param input_video: input video path or av.container.InputContainer.
+        :param duration: duration of each video split in seconds.
+    """
+    # load the input video
+    if isinstance(input_video, str):
+        container = load_video(input_video)
+    elif isinstance(input_video, av.container.InputContainer):
+        container = input_video
+    else:
+        raise ValueError(f'Unsupported type of input_video. Should be one of '
+                         f'[str, av.container.InputContainer], but given '
+                         f'[{type(input_video)}].')
+
+    video_duration = get_video_duration(container)
+    timestamps = np.arange(0, video_duration, duration).tolist()
+
+    all_key_frames = []
+    for i in range(1, len(timestamps)):
+        output_buffer = cut_video_by_seconds(container, None,
+                                             timestamps[i - 1], timestamps[i])
+        if output_buffer:
+            cut_inp_container = av.open(output_buffer, format='mp4', mode='r')
+            key_frames = extract_key_frames(cut_inp_container)
+            all_key_frames.extend(key_frames)
+            close_video(cut_inp_container)
+
+    return all_key_frames
+
+
 def extract_key_frames(input_video: Union[str, av.container.InputContainer]):
     """
     Extract key frames from the input video. If there is no keyframes in the
@@ -516,6 +559,43 @@ def get_key_frame_seconds(input_video: Union[str,
     return ts
 
 
+def extract_video_frames_uniformly_by_seconds(
+        input_video: Union[str, av.container.InputContainer],
+        frame_num: PositiveInt,
+        duration: float = 1):
+    """Extract video frames uniformly by seconds.
+        :param input_video: input video path or av.container.InputContainer.
+        :param frame_num: the number of frames to be extracted uniformly from
+            each video split by duration.
+        :param duration: duration of each video split in seconds.
+    """
+    # load the input video
+    if isinstance(input_video, str):
+        container = load_video(input_video)
+    elif isinstance(input_video, av.container.InputContainer):
+        container = input_video
+    else:
+        raise ValueError(f'Unsupported type of input_video. Should be one of '
+                         f'[str, av.container.InputContainer], but given '
+                         f'[{type(input_video)}].')
+
+    video_duration = get_video_duration(container)
+    timestamps = np.arange(0, video_duration, duration).tolist()
+
+    all_frames = []
+    for i in range(1, len(timestamps)):
+        output_buffer = cut_video_by_seconds(container, None,
+                                             timestamps[i - 1], timestamps[i])
+        if output_buffer:
+            cut_inp_container = av.open(output_buffer, format='mp4', mode='r')
+            key_frames = extract_video_frames_uniformly(cut_inp_container,
+                                                        frame_num=frame_num)
+            all_frames.extend(key_frames)
+            close_video(cut_inp_container)
+
+    return all_frames
+
+
 def extract_video_frames_uniformly(
     input_video: Union[str, av.container.InputContainer],
     frame_num: PositiveInt,

tests/ops/mapper/test_video_extract_frames_mapper.py

@@ -1,5 +1,6 @@
 import os
 import os.path as osp
+import re
 import copy
 import unittest
 import json
@@ -25,21 +26,63 @@ def tearDown(self):
         super().tearDown()
         shutil.rmtree(self.tmp_dir)
 
-    def _run_video_extract_frames_mapper(self,
-                                            op,
-                                            source_list,
-                                            target_list,
-                                            num_proc=1):
-        dataset = Dataset.from_list(source_list)
-        dataset = dataset.map(op.process, batch_size=2, num_proc=num_proc)
-        res_list = dataset.to_list()
-        self.assertEqual(res_list, target_list)
-
     def _get_frames_list(self, filepath, frame_dir, frame_num):
         frames_dir = osp.join(frame_dir, osp.splitext(osp.basename(filepath))[0])
         frames_list = [osp.join(frames_dir, f'frame_{i}.jpg') for i in range(frame_num)]
         return frames_list
 
+    def _get_frames_dir(self, filepath, frame_dir):
+        frames_dir = osp.join(frame_dir, osp.splitext(osp.basename(filepath))[0])
+        return frames_dir
+
+    def _sort_files(self, file_list):
+        return sorted(file_list, key=lambda x: int(re.search(r'(\d+)', x).group()))
+
+    def test_duration(self):
+        ds_list = [{
+            'text': f'{SpecialTokens.video} 白色的小羊站在一旁讲话。旁边还有两只灰色猫咪和一只拉着灰狼的猫咪。',
+            'videos': [self.vid1_path]
+        }, {
+            'text':
+            f'{SpecialTokens.video} 身穿白色上衣的男子，拿着一个东西，拍打自己的胃部。{SpecialTokens.eoc}',
+            'videos': [self.vid2_path]
+        }, {
+            'text':
+            f'{SpecialTokens.video} 两个长头发的女子正坐在一张圆桌前讲话互动。 {SpecialTokens.eoc}',
+            'videos': [self.vid3_path]
+        }]
+
+        frame_num = 2
+        frame_dir=os.path.join(self.tmp_dir, 'test1')
+        vid1_frame_dir =  self._get_frames_dir(self.vid1_path, frame_dir)
+        vid2_frame_dir =  self._get_frames_dir(self.vid2_path, frame_dir)
+        vid3_frame_dir =  self._get_frames_dir(self.vid3_path, frame_dir)
+
+        tgt_list = copy.deepcopy(ds_list)
+        tgt_list[0].update({Fields.video_frames: json.dumps({self.vid1_path: vid1_frame_dir})})
+        tgt_list[1].update({Fields.video_frames: json.dumps({self.vid2_path: vid2_frame_dir})})
+        tgt_list[2].update({Fields.video_frames: json.dumps({self.vid3_path: vid3_frame_dir})})
+
+        op = VideoExtractFramesMapper(
+            frame_sampling_method='uniform',
+            frame_num=frame_num,
+            duration=0,
+            frame_dir=frame_dir)
+
+        dataset = Dataset.from_list(ds_list)
+        dataset = dataset.map(op.process, batch_size=2, num_proc=1)
+        res_list = dataset.to_list()
+        self.assertEqual(res_list, tgt_list)
+        self.assertListEqual(
+            self._sort_files(os.listdir(vid1_frame_dir)),
+            [f'frame_{i}.jpg' for i in range(frame_num)])
+        self.assertListEqual(
+            self._sort_files(os.listdir(vid2_frame_dir)),
+            [f'frame_{i}.jpg' for i in range(frame_num)])
+        self.assertListEqual(
+            self._sort_files(os.listdir(vid3_frame_dir)),
+            [f'frame_{i}.jpg' for i in range(frame_num)])
+
     def test_uniform_sampling(self):
         ds_list = [{
             'text': f'{SpecialTokens.video} 白色的小羊站在一旁讲话。旁边还有两只灰色猫咪和一只拉着灰狼的猫咪。',
@@ -55,22 +98,35 @@ def test_uniform_sampling(self):
         }]
         frame_num = 3
         frame_dir=os.path.join(self.tmp_dir, 'test1')
+        vid1_frame_dir =  self._get_frames_dir(self.vid1_path, frame_dir)
+        vid2_frame_dir =  self._get_frames_dir(self.vid2_path, frame_dir)
+        vid3_frame_dir =  self._get_frames_dir(self.vid3_path, frame_dir)
 
         tgt_list = copy.deepcopy(ds_list)
-        tgt_list[0].update({Fields.video_frames: 
-            json.dumps({self.vid1_path: self._get_frames_list(self.vid1_path, frame_dir, frame_num)})})
-        tgt_list[1].update({Fields.video_frames: 
-            json.dumps({self.vid2_path: self._get_frames_list(self.vid2_path, frame_dir, frame_num)})})
-        tgt_list[2].update({Fields.video_frames: 
-            json.dumps({self.vid3_path: self._get_frames_list(self.vid3_path, frame_dir, frame_num)})})
-        
+        tgt_list[0].update({Fields.video_frames: json.dumps({self.vid1_path: vid1_frame_dir})})
+        tgt_list[1].update({Fields.video_frames: json.dumps({self.vid2_path: vid2_frame_dir})})
+        tgt_list[2].update({Fields.video_frames: json.dumps({self.vid3_path: vid3_frame_dir})})
+
         op = VideoExtractFramesMapper(
             frame_sampling_method='uniform',
             frame_num=frame_num,
+            duration=10,
             frame_dir=frame_dir)
-        self._run_video_extract_frames_mapper(op, ds_list, tgt_list)
 
-    
+        dataset = Dataset.from_list(ds_list)
+        dataset = dataset.map(op.process, batch_size=2, num_proc=1)
+        res_list = dataset.to_list()
+        self.assertEqual(res_list, tgt_list)
+        self.assertListEqual(
+            self._sort_files(os.listdir(vid1_frame_dir)),
+            [f'frame_{i}.jpg' for i in range(3)])
+        self.assertListEqual(
+            self._sort_files(os.listdir(vid2_frame_dir)),
+            [f'frame_{i}.jpg' for i in range(6)])
+        self.assertListEqual(
+            self._sort_files(os.listdir(vid3_frame_dir)),
+            [f'frame_{i}.jpg' for i in range(12)])
+
     def test_all_keyframes_sampling(self):
         ds_list = [{
             'text': f'{SpecialTokens.video} 白色的小羊站在一旁讲话。旁边还有两只灰色猫咪和一只拉着灰狼的猫咪。',
@@ -86,22 +142,38 @@ def test_all_keyframes_sampling(self):
             'videos': [self.vid3_path]
         }]
         frame_dir=os.path.join(self.tmp_dir, 'test2')
+        vid1_frame_dir =  self._get_frames_dir(self.vid1_path, frame_dir)
+        vid2_frame_dir =  self._get_frames_dir(self.vid2_path, frame_dir)
+        vid3_frame_dir =  self._get_frames_dir(self.vid3_path, frame_dir)
 
         tgt_list = copy.deepcopy(ds_list)
         tgt_list[0].update({Fields.video_frames: 
-            json.dumps({self.vid1_path: self._get_frames_list(self.vid1_path, frame_dir, 3)})})
+            json.dumps({self.vid1_path: vid1_frame_dir})})
         tgt_list[1].update({Fields.video_frames: json.dumps({
-            self.vid2_path: self._get_frames_list(self.vid2_path, frame_dir, 3),
-            self.vid3_path: self._get_frames_list(self.vid3_path, frame_dir, 6)
+            self.vid2_path: vid2_frame_dir,
+            self.vid3_path: vid3_frame_dir
             })})
         tgt_list[2].update({Fields.video_frames: 
-            json.dumps({self.vid3_path: self._get_frames_list(self.vid3_path, frame_dir, 6)})})
+            json.dumps({self.vid3_path: vid3_frame_dir})})
 
         op = VideoExtractFramesMapper(
             frame_sampling_method='all_keyframes',
-            frame_dir=frame_dir)
-        self._run_video_extract_frames_mapper(op, ds_list, tgt_list)
+            frame_dir=frame_dir,
+            duration=5)
 
+        dataset = Dataset.from_list(ds_list)
+        dataset = dataset.map(op.process, batch_size=2, num_proc=2)
+        res_list = dataset.to_list()
+        self.assertEqual(res_list, tgt_list)
+        self.assertListEqual(
+            self._sort_files(os.listdir(vid1_frame_dir)),
+            [f'frame_{i}.jpg' for i in range(4)])
+        self.assertListEqual(
+            self._sort_files(os.listdir(vid2_frame_dir)),
+            [f'frame_{i}.jpg' for i in range(5)])
+        self.assertListEqual(
+            self._sort_files(os.listdir(vid3_frame_dir)),
+            [f'frame_{i}.jpg' for i in range(13)])
 
 
 if __name__ == '__main__':