PyCS/pycs/frontend/endpoints/pipelines/PredictModel.py

from typing import Any

from flask import make_response, request, abort
from flask.views import View

from pycs.database.Database import Database
from pycs.frontend.notifications.NotificationList import NotificationList
from pycs.frontend.notifications.NotificationManager import NotificationManager
from pycs.interfaces.MediaFile import MediaFile
from pycs.interfaces.MediaStorage import MediaStorage
from pycs.jobs.JobGroupBusyException import JobGroupBusyException
from pycs.jobs.JobRunner import JobRunner
from pycs.util.PipelineUtil import load_from_root_folder as load_pipeline


class PredictModel(View):
    """
    load a model and create predictions
    """
    # pylint: disable=arguments-differ
    methods = ['POST']

    def __init__(self, db: Database, nm: NotificationManager, jobs: JobRunner):
        # pylint: disable=invalid-name
        self.db = db
        self.nm = nm
        self.jobs = jobs

    def dispatch_request(self, project_id):
        # extract request data
        data = request.get_json(force=True)

        if 'predict' not in data or data['predict'] not in ['all', 'new']:
            return abort(400)

        # find project
        project = self.db.project(project_id)
        if project is None:
            return abort(404)

        # create job
        try:
            notifications = NotificationList(self.nm)

            self.jobs.run(project,
                          'Model Interaction',
                          f'{project.name} (create predictions)',
                          f'{project.name}/model-interaction',
                          self.load_and_predict,
                          self.db, project.identifier, data['predict'], notifications,
                          progress=self.progress)
        except JobGroupBusyException:
            return abort(400)

        return make_response()

    @staticmethod
    def load_and_predict(database: Database, project_id: int, file_filter: Any,
                         notifications: NotificationList):
        db = None
        pipeline = None

        # create new database instance
        try:
            db = database.copy()
            project = db.project(project_id)
            model = project.model()
            storage = MediaStorage(db, project_id, notifications)

            # create a list of MediaFile
            if isinstance(file_filter, str):
                if file_filter == 'new':
                    length = project.count_files_without_results()
                    files = map(lambda f: MediaFile(f, notifications),
                                project.files_without_results())
                else:
                    length = project.count_files()
                    files = map(lambda f: MediaFile(f, notifications),
                                project.files())
            else:
                files = map(lambda f: MediaFile(project.file(f.identifier), notifications),
                            file_filter)
                length = len(file_filter)

            # load pipeline
            try:
                pipeline = load_pipeline(model.root_folder)

                # iterate over files
                index = 0
                for file in files:
                    # remove old predictions
                    file.remove_predictions()

                    # create new predictions
                    pipeline.execute(storage, file)

                    # commit changes and yield progress
                    db.commit()
                    yield index / length, notifications

                    index += 1
            finally:
                if pipeline is not None:
                    pipeline.close()
        finally:
            if db is not None:
                db.close()

    @staticmethod
    def progress(progress: float, notifications: NotificationList):
        notifications.fire()
        return progress