PyCS/pycs/pipeline/tf1/pipeline.py

#!/usr/bin/env python

"""pipeline:  Detection and other model pipeline."""
from PIL import Image

from pycs.pipeline.Pipeline import Pipeline as PipelineInterface
from pycs.utils import Video
from .detection import Detector
from .features import Features
from ..Job import Job
from ..Result import Result


class Pipeline(PipelineInterface):
    def __init__(self):
        self.__detector = None
        self.__features = None

    def load(self, root: str, distribution: dict):
        print('tf1 load')
        detector_config = distribution['detection']
        features_config = distribution['features']

        self.__detector = Detector(config={
            **detector_config,
            'distribution-root': root
        })
        self.__features = Features(config={
            **features_config,
            'distribution-root': root
        })

    def close(self):
        print('tf1 close')
        if self.__detector is not None:
            self.__detector.close()
        if self.__features is not None:
            self.__features.close()

    def execute(self, job: Job) -> Result:
        subjobs = [{
            'subjob': job.type,
            'prediction': {},
            'jobinfo': {},
            'filetype': 'image',
            'filename': job.object_full_path
        }]
        self.__execute(subjobs, lambda x: 0)

        return Result(job, subjobs[0]['prediction']['faces'])

    def __execute(self, subjobs, callback):
        callback(0)
        subjob_count = float(len(subjobs))
        for index, subjob in enumerate(subjobs):
            subjob_name = subjob['subjob']
            prediction = subjob['prediction']
            jobinfo = subjob['jobinfo']
            jobinfo[subjob_name] = {'done-by': 'pipeline'}

            if subjob_name == 'detect-faces':
                # Run face detection
                if self.__detector.last_error is not None:
                    jobinfo[subjob_name]['error'] = self.__detector.last_error
                    jobinfo[subjob_name]['result'] = False

                else:
                    filename = subjob['filename']
                    print(filename)

                    # Acquire image
                    if subjob['filetype'] == 'image':
                        img = Image.open(filename)
                    elif subjob['filetype'] == 'video':
                        if 'cap' in subjob.keys():
                            cap = subjob['cap']
                        else:
                            cap = Video(filename)
                        if cap.last_error is None:
                            jobinfo[subjob_name]['error'] = cap.last_error
                        else:
                            jobinfo[subjob_name]['result'] = False
                            continue

                        img = cap.get_frame(subjob['frame'])
                    else:
                        jobinfo[subjob_name]['error'] = 'File format not supported!'
                        jobinfo[subjob_name]['result'] = False
                        continue

                    faces = self.__detector.detect_faces(img)

                    if self.__detector.last_error is not None:
                        jobinfo[subjob_name]['error'] = self.__detector.last_error
                        jobinfo[subjob_name]['result'] = False
                    else:
                        prediction['faces'] = faces
                        jobinfo[subjob_name]['result'] = True
            else:
                jobinfo[subjob_name]['result'] = False

            callback(float(index) / subjob_count)
        callback(1)