PyCS/models/haarcascade_frontalface_default/Pipeline.py

from os import path
from typing import List
from urllib.request import urlretrieve

import cv2

from pycs import app
from pycs.interfaces.MediaFile import MediaFile
from pycs.interfaces.MediaStorage import MediaStorage
from pycs.interfaces.Pipeline import Pipeline as Interface


class Pipeline(Interface):
    URL = 'https://raw.githubusercontent.com/opencv/opencv/master/data/haarcascades/haarcascade_frontalface_default.xml'

    def __init__(self, root_folder, distribution):
        app.logger.debug('hcffdv1 init')

        # get path to xml file
        xml_file = path.join(root_folder, 'haarcascade_frontalface_default.xml')

        # download
        if not path.exists(xml_file):
            urlretrieve(self.URL, xml_file)

        # load
        self.face_cascade = cv2.CascadeClassifier(xml_file)

    def close(self):
        app.logger.debug('hcffdv1 close')

    def collections(self) -> List[dict]:
        return [
            self.create_collection('face', 'face detected', autoselect=True),
            self.create_collection('none', 'no face detected')
        ]

    def execute(self, storage: MediaStorage, file: MediaFile):
        app.logger.debug('hcffdv1 execute')

        # load file and analyze frames
        found = False

        if file.type == 'image':
            found = self.__find(file, cv2.imread(file.path))
        else:
            video = cv2.VideoCapture(file.path)
            index = 0

            ret, image = video.read()
            while ret:
                if self.__find(file, image, index):
                    found = True

                ret, image = video.read()
                index += 1

            video.release()

        # set file collection
        if found:
            file.set_collection('face')
        else:
            file.set_collection('none')

    def __find(self, file: MediaFile, image, frame=None):
        # convert to grayscale, scale down
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        height, width = gray.shape

        scale_factor = min(2048 / width, 2048 / height, 1.0)
        scale_height, scale_width = int(height * scale_factor), int(width * scale_factor)
        scaled = cv2.resize(gray, (scale_width, scale_height))

        # detect faces
        faces = self.face_cascade.detectMultiScale(
            scaled,
            scaleFactor=1.1,
            minNeighbors=5,
            minSize=(192, 192)
        )

        # add faces to results
        for x, y, w, h in faces:
            file.add_bounding_box(x / scale_width,
                                  y / scale_height,
                                  w / scale_width,
                                  h / scale_height,
                                  frame=frame)

        return len(faces) > 0