PyCS/pycs/project/project.py

#!/usr/bin/env python

"""project:  Encapsulates a project with all its properties."""

import json
import os

from pycs.pipeline.tf1 import Pipeline
from ..utils import Errorable
from ..utils import Video


class Project(Errorable):
    def __init__(self, project_file, config=None):
        Errorable.__init__(self)
        self._update_ui_fn = lambda: None

        new = config is not None

        # (1) Check if directory exists
        self.__project_file = project_file
        self.__project_path, self.__project_filename = os.path.split(project_file)

        if not os.path.exists(self.__project_path) and not new:
            self._report_error("Cannot open directory: %s" % self.__project_path)
            return
        elif not os.path.exists(self.__project_path) and new:
            try:
                os.mkdir(self.__project_path)
            except:
                self._report_error("Cannot create directory: %s" % self.__project_path)
                return

        # (2) Load project json
        if new:
            self._config = config
        else:
            try:
                with open(self.__project_file, 'r') as c:
                    self._config = json.load(c)
            except:
                self._report_error("Cannot open configuration for project: %s" % self.__project_path)
                return

        # (3) Load detection/extraction pipeline
        pipeline_config = self._config["pipeline"]

        self._pipeline = Pipeline(pipeline_config)
        self._err_children += [self._pipeline]

    def save(self):
        try:
            with open(self.__project_file, 'w') as c:
                json.dump(self._config, c, indent=2)
        except:
            self._report_error("Cannot save configuration for project: %s" % self.__project_path)

    def execute(self, jobs, callback=lambda progress: True):
        all_subjobs = []
        for job in jobs:
            job['prediction'] = {}
            job['jobinfo'] = {}
            for subjob in job['jobs']:
                if job['filetype'] == 'image':
                    pipeline_subjob = {'filetype': 'image',
                                       'filename': job['filename'],
                                       'prediction': job['prediction'],
                                       'jobinfo': job['jobinfo'],
                                       'subjob': subjob}
                    all_subjobs += [pipeline_subjob]
                # TODO split video into frames
                elif job['filetype'] == 'video':
                    cap = Video(job['filename'])
                    if cap.framecount > 0:
                        job['prediction-by-frame'] = [{} for frame in range(cap.framecount)]
                        job['jobinfo-by-frame'] = [{} for frame in range(cap.framecount)]
                        for frame in range(cap.framecount):
                            pipeline_subjob = {'filetype': 'video',
                                               'cap': cap,
                                               'frame': frame,
                                               'filename': job['filename'],
                                               'prediction': job['prediction-by-frame'][frame],
                                               'jobinfo': job['jobinfo-by-frame'][frame],
                                               'subjob': subjob}
                            all_subjobs += [pipeline_subjob]

        self._pipeline.execute(all_subjobs, callback)
        callback(1)

    def close(self):
        self.detector.close()
        self.features.close()

    def close(self):
        self._pipeline.close()