summaryrefslogtreecommitdiff
path: root/megapixels/commands/processor/videos_to_frames.py
blob: 0b56c46a48a134bab5f418462d48d9d7575eb6ce (plain)
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
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