update vide_face_blur

data_juicer/ops/mapper/video_face_blur_mapper.py

@@ -1,17 +1,20 @@
+import os
+
 import av
 
 from data_juicer.utils.availability_utils import AvailabilityChecking
 from data_juicer.utils.file_utils import transfer_filename
-from data_juicer.utils.mm_utils import (load_data_with_context, load_video,
-                                        pil_to_opencv, process_each_frame)
+from data_juicer.utils.mm_utils import (detect_faces, load_data_with_context,
+                                        load_video, process_each_frame)
+from data_juicer.utils.model_utils import get_model, prepare_model
 
 from ..base_op import OPERATORS, Mapper
 from ..op_fusion import LOADED_VIDEOS
 
 OP_NAME = 'video_face_blur_mapper'
 
-with AvailabilityChecking(['dlib', 'Pillow'], OP_NAME):
-    import dlib
+with AvailabilityChecking(['opencv-python', 'Pillow'], OP_NAME):
+    import cv2
     from PIL import ImageFilter
 
 
@@ -21,16 +24,24 @@ class VideoFaceBlurMapper(Mapper):
     """Mapper to blur faces detected in videos.
     """
 
-    _default_kwargs = {'upsample_num_times': 0}
+    _default_kwargs = {
+        'scaleFactor': 1.1,
+        'minNeighbors': 3,
+        'minSize': None,
+        'maxSize': None,
+    }
 
     def __init__(self,
+                 cv_classifier='',
                  blur_type: str = 'gaussian',
                  radius: float = 2,
                  *args,
                  **kwargs):
         """
         Initialization method.
 
+        :param cv_classifier: OpenCV classifier path for face detection.
+            By default, we will use 'haarcascade_frontalface_alt.xml'.
         :param blur_type: Type of blur kernel, including
             ['mean', 'box', 'gaussian'].
         :param radius: Radius of blur kernel.
@@ -40,6 +51,9 @@ def __init__(self,
         super().__init__(*args, **kwargs)
         self._init_parameters = self.remove_extra_parameters(locals())
 
+        if cv_classifier == '':
+            cv_classifier = os.path.join(cv2.data.haarcascades,
+                                         'haarcascade_frontalface_alt.xml')
         if blur_type not in ['mean', 'box', 'gaussian']:
             raise ValueError(
                 f'Blur_type [{blur_type}] is not supported. '
@@ -62,8 +76,8 @@ def __init__(self,
             if key in self.extra_kwargs:
                 self.extra_kwargs[key] = kwargs[key]
 
-        # Initialize face detector
-        self.detector = dlib.get_frontal_face_detector()
+        self.model_key = prepare_model(model_type='opencv_classifier',
+                                       model_path=cv_classifier)
 
     def process(self, sample, context=False):
         # there is no video in this sample
@@ -74,6 +88,19 @@ def process(self, sample, context=False):
         sample, videos = load_data_with_context(sample, context,
                                                 loaded_video_keys, load_video)
 
+        model = get_model(self.model_key)
+
+        def _blur_func(frame):
+            image = frame.to_image()
+            dets = detect_faces(image, model, **self.extra_kwargs)
+            if len(dets) > 0:
+                for (x, y, w, h) in dets:
+                    box = (x, y, x + w, y + h)
+                    blured_roi = image.crop(box).filter(self.blur)
+                    image.paste(blured_roi, box)
+                frame = av.VideoFrame.from_image(image)
+            return frame
+
         processed_video_keys = {}
         for video_key in loaded_video_keys:
             # skip duplicate
@@ -84,7 +111,7 @@ def process(self, sample, context=False):
             blured_video_key = transfer_filename(video_key, OP_NAME,
                                                  **self._init_parameters)
             output_video_key = process_each_frame(video, blured_video_key,
-                                                  self._blur_face)
+                                                  _blur_func)
             processed_video_keys[video_key] = output_video_key
 
             if not context:
@@ -94,18 +121,3 @@ def process(self, sample, context=False):
             processed_video_keys[key] for key in loaded_video_keys
         ]
         return sample
-
-    def _blur_face(self, frame):
-        image = frame.to_image()
-        img = pil_to_opencv(image)
-        dets = self.detector(img, **self.extra_kwargs)
-        if len(dets) > 0:
-            for det in dets:
-                x1 = max(det.left(), 0)
-                y1 = max(det.top(), 0)
-                x2 = min(det.right(), image.width)
-                y2 = min(det.bottom(), image.height)
-                blured_roi = image.crop((x1, y1, x2, y2)).filter(self.blur)
-                image.paste(blured_roi, (x1, y1, x2, y2))
-            frame = av.VideoFrame.from_image(image)
-        return frame

tests/ops/mapper/test_video_face_blur_mapper.py

@@ -71,6 +71,9 @@ def test_mean(self):
         self._run_helper(op, ds_list)
 
     def test_gaussian_radius_parallel(self):
+        import multiprocess as mp
+        mp.set_start_method('forkserver', force=True)
+
         ds_list = [{
             'videos': [self.vid1_path]
         }, {