text: # AUTOGENERATED! DO NOT EDIT! File to edit: nbs/01a_face_extraction.ipynb (unless otherwise specified).
  
  __all__ = ['video2frames', 'detect_facenet_pytorch', 'DetectionPipeline']
  
  # Cell
  from fastai.vision import *
  from .nb_00 import *
  from IPython.display import HTML
  import cv2
  from tqdm import tqdm
  from facenet_pytorch import MTCNN
  from nbdev.export import *
  
  # Cell
  def video2frames(fname, *fs):
      '''
      fname - path to mp4 file
      fs - fractional lengths to resize original image to. e.g. (.5, .25)
      '''
      capture = cv2.VideoCapture(fname)
      imgs = []
      for i in tqdm(range(int(capture.get(cv2.CAP_PROP_FRAME_COUNT)))):
          _, img0 = capture.read()
          img0 = cv2.cvtColor(img0, cv2.COLOR_BGR2RGB)
          imgs.append([img0] + [cv2.resize(img0, (0, 0), fx=f, fy=f) for f in fs])
      capture.release()
      return [np.stack(imgsz) for imgsz in zip(*imgs)]
  
  # Cell
  def detect_facenet_pytorch(detector, images, batch_size):
      '''
      detector - facenet_pytorch.MTCNN
      images:  numpy.array
        array of images
      batch_size: int
        Number of images to be processed by `detector` in one go.
      '''
      faces = []
      for lb in np.arange(0, len(images), batch_size):
          imgs_pil = [PIL.Image.fromarray(image) for image in images[lb:lb+batch_size]]
          faces.extend(detector(imgs_pil))
      return torch.stack(faces)
  
  # Cell
  
  #comes from https://www.kaggle.com/unkownhihi/starter-kernel-with-cnn-model-ll-lb-0-69235
  class DetectionPipeline:
      """Pipeline class for detecting faces in the frames of a video file."""
  
      def __init__(self, detector, n_frames=None, batch_size=60, resize=None):
          """Constructor for DetectionPipeline class.
  
          Keyword Arguments:
              n_frames {int} -- Total number of frames to load. These will be evenly spaced
                  throughout the video. If not specified (i.e., None), all frames will be loaded.
                  (default: {None})
              batch_size {int} -- Batch size to use with MTCNN face detector. (default: {32})
              resize {float} -- Fraction by which to resize frames from original prior to face
                  detection. A value less than 1 results in downsampling and a value greater than
                  1 result in upsampling. (default: {None})
          """
          self.detector = detector
          self.n_frames, self.batch_size, self.resize = n_frames, batch_size, resize
  
      def __call__(self, filename, label=None, save_dir=None):
          """Load frames from an MP4 video and detect faces.
          Arguments:
              filename {str} -- Path to video.
          """
          v_cap = cv2.VideoCapture(filename)
          v_len = int(v_cap.get(cv2.CAP_PROP_FRAME_COUNT))
  
          if self.n_frames is None: sample = np.arange(0, v_len)
          else: sample = np.linspace(0, v_len - 1, self.n_frames).astype(int)
  
          faces = []
          idxs, frames = [], []
          for j in range(v_len):
              success = v_cap.grab()
              if j in sample:
                  success, frame = v_cap.retrieve()
                  if not success: continue
                  frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                  frame = PIL.Image.fromarray(frame)
  
                  if self.resize is not None:
                      frame = frame.resize([int(d * self.resize) for d in frame.size])
                  idxs.append(j); frames.append(frame)
  
                  if len(frames) % self.batch_size == 0 or j == sample[-1]:
                      if save_dir is not None:
                          save_paths = self.get_savepaths(filename, idxs, label, save_dir)
                          faces.extend(self.detector(frames, save_path=save_paths))
                      else: faces.extend(self.detector(frames))
                      idxs, frames = [], []
  
          v_cap.release()
          return faces
  
      def get_savepaths(self, filename, idxs, label=None, save_dir=None):
          if isinstance(filename, str): filename = Path(filename)
          if save_dir is None: save_dir = Path('./')
          if label is None: save_paths = [save_dir/f'{filename.stem}_{i:03d}.png' for i in idxs]
          else: save_paths = [save_dir/f'{filename.stem}_{i:03d}_{label}.png' for i in idxs]
          return [str(o) for o in save_paths]