.

1 files changed, 73 insertions, 0 deletions

diff --git a/megapixels/commands/cv/videos_to_frames.py b/megapixels/commands/cv/videos_to_frames.py
new file mode 100644
index 00000000..0b56c46a
--- /dev/null
+++ b/megapixels/commands/cv/videos_to_frames.py
@@ -0,0 +1,73 @@
+from glob import glob
+import os
+from os.path import join
+from pathlib import Path
+
+import click
+
+from app.settings import types
+from app.utils import click_utils
+from app.settings import app_cfg as cfg
+from app.utils import logger_utils
+
+import dlib
+import pandas as pd
+from PIL import Image, ImageOps, ImageFilter
+from app.utils import file_utils, im_utils
+
+
+log = logger_utils.Logger.getLogger()
+
+@click.command()
+@click.option('-i', '--input', 'opt_fp_in', required=True,
+  help='Input directory')
+@click.option('-o', '--output', 'opt_fp_out', required=True,
+  help='Output directory')
+@click.option('--size', 'opt_size', default=(320, 240),
+  help='Inference size for face detection' )
+@click.option('--interval', 'opt_frame_interval', default=20,
+  help='Number of frames before saving next face')
+@click.pass_context
+def cli(ctx, opt_fp_in, opt_fp_out, opt_size, opt_frame_interval):
+  """Converts videos to frames with faces"""
+
+  # -------------------------------------------------
+  # process  
+  
+  from tqdm import tqdm
+  import cv2 as cv
+  from tqdm import tqdm
+  from app.processors import face_detector
+  
+  detector = face_detector.DetectorDLIBCNN()
+
+  # get file list
+  fp_videos = glob(join(opt_fp_in, '*.mp4'))
+  fp_videos += glob(join(opt_fp_in, '*.webm'))
+  fp_videos += glob(join(opt_fp_in, '*.mkv'))
+  
+  frame_interval_count = 0
+  frame_count = 0
+
+  file_utils.mkdirs(opt_fp_out)
+
+  for fp_video in tqdm(fp_videos):
+
+    video = cv.VideoCapture(fp_video)
+    
+    while video.isOpened():
+      res, frame = video.read()
+      if not res:
+        break
+
+      frame_count += 1  # for naming
+      frame_interval_count += 1  # for interval
+      
+      bboxes = detector.detect(frame, opt_size=opt_size, opt_pyramids=0)
+      if len(bboxes) > 0 and frame_interval_count >= opt_frame_interval:
+        # save frame
+        fname = file_utils.zpad(frame_count)
+        fp_frame = join(opt_fp_out, '{}_{}.jpg'.format(Path(fp_video).stem, fname))
+        cv.imwrite(fp_frame, frame)
+        frame_interval_count = 0
+