PyCS/models/haarcascade_frontalface_default/Pipeline.py

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

import cv2

from pycs.interfaces.MediaFile import MediaFile
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):
        print('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):
        print('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, file: MediaFile) -> List[dict]:
        print('hcffdv1 execute')

        # load image and convert to grayscale
        image = cv2.imread(file.path)
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

        height, width = gray.shape
        min_size = int(min(width, height) / 10)

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

        # convert faces to result list
        result = []
        for x, y, w, h in faces:
            result.append(self.create_bounding_box_result(
                x / width,
                y / height,
                w / width,
                h / height
            ))

        # set file collection
        if len(result) > 0:
            result.append(self.create_collection_result('face'))
        else:
            result.append(self.create_collection_result('none'))

        return result