New Optical Flow OP & Allow to save the computed optical flows

data_juicer/ops/filter/video_motion_score_filter.py

@@ -5,7 +5,7 @@
 import numpy as np
 from pydantic import PositiveFloat, PositiveInt
 
-from data_juicer.utils.constant import Fields, StatsKeys
+from data_juicer.utils.constant import Fields, MetaKeys, StatsKeys
 from data_juicer.utils.lazy_loader import LazyLoader
 from data_juicer.utils.mm_utils import calculate_resized_dimensions
 
@@ -57,6 +57,8 @@ def __init__(
         divisible: PositiveInt = 1,
         relative: bool = False,
         any_or_all: str = "any",
+        if_output_optical_flow: bool = False,
+        optical_flow_key: str = MetaKeys.video_optical_flow,
         *args,
         **kwargs,
     ):
@@ -84,6 +86,11 @@ def __init__(
             all videos. 'any': keep this sample if any videos meet the
             condition. 'all': keep this sample only if all videos meet the
             condition.
+        :param if_output_camera_parameters: Determines whether to output
+            the computed optical flows into the metas. The optical flows for each
+            video will be stored in the shape of (num_frame, H, W, 2)
+        :param optical_flow_key: The field name to store the optical flows. It's
+            "video_optical_flow" in default.
         :param args: extra args
         :param kwargs: extra args
         """
@@ -113,6 +120,9 @@ def __init__(
             raise ValueError(f"Keep strategy [{any_or_all}] is not supported. " f'Can only be one of ["any", "all"].')
         self.any = any_or_all == "any"
 
+        self.if_output_optical_flow = if_output_optical_flow
+        self.optical_flow_key = optical_flow_key
+
     def setup_model(self, rank=None):
         self.model = cv2.calcOpticalFlowFarneback
 
@@ -141,12 +151,14 @@ def compute_stats_single(self, sample, rank=None, context=False):
         # load videos
         loaded_video_keys = sample[self.video_key]
         unique_motion_scores = {}
+        video_optical_flows = {}
         for video_key in loaded_video_keys:
             # skip duplicate videos
             if video_key in unique_motion_scores:
                 continue
 
             video_motion_scores = []
+            optical_flows = []
             with VideoCapture(video_key) as cap:
                 if cap.isOpened():
                     fps = cap.get(cv2.CAP_PROP_FPS)
@@ -176,6 +188,7 @@ def compute_stats_single(self, sample, rank=None, context=False):
                     flow, prev_frame = self.compute_flow(prev_frame, frame)
                     if flow is None:
                         continue
+                    optical_flows.append(flow)
                     mag, _ = cv2.cartToPolar(flow[..., 0], flow[..., 1])
                     frame_motion_score = np.mean(mag)
                     if self.relative:
@@ -192,7 +205,11 @@ def compute_stats_single(self, sample, rank=None, context=False):
             else:
                 unique_motion_scores[video_key] = np.mean(video_motion_scores or [-1])
 
+            video_optical_flows[video_key] = np.stack(optical_flows)
+
         sample[Fields.stats][StatsKeys.video_motion_score] = [unique_motion_scores[key] for key in loaded_video_keys]
+        if self.if_output_optical_flow:
+            sample[Fields.meta][self.optical_flow_key] = [video_optical_flows[key] for key in loaded_video_keys]
         return sample
 
     def process_single(self, sample):

data_juicer/ops/filter/video_motion_score_raft_filter.py

@@ -4,6 +4,7 @@
 from pydantic import PositiveFloat, PositiveInt
 
 from data_juicer.ops.filter.video_motion_score_filter import VideoMotionScoreFilter
+from data_juicer.utils.constant import MetaKeys
 from data_juicer.utils.lazy_loader import LazyLoader
 from data_juicer.utils.resource_utils import cuda_device_count
 
@@ -55,11 +56,24 @@ def __init__(
         divisible: PositiveInt = 8,
         relative: bool = False,
         any_or_all: str = "any",
+        if_output_optical_flow: bool = False,
+        optical_flow_key: str = MetaKeys.video_optical_flow,
         *args,
         **kwargs,
     ):
         super().__init__(
-            min_score, max_score, sampling_fps, size, max_size, divisible, relative, any_or_all, *args, **kwargs
+            min_score,
+            max_score,
+            sampling_fps,
+            size,
+            max_size,
+            divisible,
+            relative,
+            any_or_all,
+            if_output_optical_flow,
+            optical_flow_key,
+            *args,
+            **kwargs,
         )
 
     def setup_model(self, rank=None):

data_juicer/utils/constant.py

@@ -61,6 +61,8 @@ class MetaKeys(object):
     video_audio_tags = "video_audio_tags"
     # # video frames
     video_frames = "video_frames"
+    # # video optical flow
+    video_optical_flow = "video_optical_flow"
     # # info extracted by VGGT
     vggt_tags = "vggt_tags"
     # # image tags

tests/ops/filter/test_video_motion_score_filter.py

@@ -5,7 +5,7 @@
 
 from data_juicer.ops.filter.video_motion_score_filter import \
     VideoMotionScoreFilter
-from data_juicer.utils.constant import Fields
+from data_juicer.utils.constant import Fields, MetaKeys
 from data_juicer.utils.unittest_utils import DataJuicerTestCaseBase
 
 
@@ -177,6 +177,38 @@ def test_parallel(self):
         op = VideoMotionScoreFilter(min_score=1.5, max_score=3.0)
         self._run_helper(op, ds_list, tgt_list, np=2)
 
+    def test_output_optical_flow(self):
+        ds_list = [{
+            'videos': [self.vid1_path]
+        }, {
+            'videos': [self.vid2_path]
+        }, {
+            'videos': [self.vid3_path]
+        }]
+        tgt_list = [{
+            'videos': [self.vid1_path]
+        }, {
+            'videos': [self.vid2_path]
+        }, {
+            'videos': [self.vid3_path]
+        }]
+        op = VideoMotionScoreFilter(if_output_optical_flow=True)
+        dataset = Dataset.from_list(ds_list)
+        if Fields.stats not in dataset.features:
+            dataset = dataset.add_column(name=Fields.stats,
+                                         column=[{}] * dataset.num_rows)
+        if Fields.meta not in dataset.features:
+            dataset = dataset.add_column(name=Fields.meta,
+                                         column=[{}] * dataset.num_rows)
+        dataset = dataset.map(op.compute_stats, num_proc=1)
+        dataset = dataset.filter(op.process, num_proc=1)
+        metas = dataset.select_columns(column_names=[Fields.meta])
+        self.assertIn(MetaKeys.video_optical_flow, metas.features)
+
+        dataset = dataset.select_columns(column_names=[op.video_key])
+        res_list = dataset.to_list()
+        self.assertEqual(res_list, tgt_list)
+
 
 if __name__ == '__main__':
     unittest.main()

tests/ops/filter/test_video_motion_score_raft_filter.py

@@ -5,7 +5,7 @@
 
 from data_juicer.ops.filter.video_motion_score_raft_filter import \
     VideoMotionScoreRaftFilter
-from data_juicer.utils.constant import Fields
+from data_juicer.utils.constant import Fields, MetaKeys
 from data_juicer.utils.unittest_utils import DataJuicerTestCaseBase
 class VideoMotionScoreRaftFilterTest(DataJuicerTestCaseBase):
 
@@ -179,6 +179,38 @@ def test_parallel(self):
         op = VideoMotionScoreRaftFilter(min_score=3, max_score=10.2)
         self._run_helper(op, ds_list, tgt_list, np=2)
 
+    def test_output_optical_flow(self):
+        ds_list = [{
+            'videos': [self.vid1_path]
+        }, {
+            'videos': [self.vid2_path]
+        }, {
+            'videos': [self.vid3_path]
+        }]
+        tgt_list = [{
+            'videos': [self.vid1_path]
+        }, {
+            'videos': [self.vid2_path]
+        }, {
+            'videos': [self.vid3_path]
+        }]
+        op = VideoMotionScoreRaftFilter(if_output_optical_flow=True)
+        dataset = Dataset.from_list(ds_list)
+        if Fields.stats not in dataset.features:
+            dataset = dataset.add_column(name=Fields.stats,
+                                         column=[{}] * dataset.num_rows)
+        if Fields.meta not in dataset.features:
+            dataset = dataset.add_column(name=Fields.meta,
+                                         column=[{}] * dataset.num_rows)
+        dataset = dataset.map(op.compute_stats, num_proc=1)
+        dataset = dataset.filter(op.process, num_proc=1)
+        metas = dataset.select_columns(column_names=[Fields.meta])
+        self.assertIn(MetaKeys.video_optical_flow, metas.features)
+
+        dataset = dataset.select_columns(column_names=[op.video_key])
+        res_list = dataset.to_list()
+        self.assertEqual(res_list, tgt_list)
+
 
 if __name__ == '__main__':
     unittest.main()