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

from contextlib import closing

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

from pycs.database.Database import Database
from pycs.frontend.notifications.NotificationManager import NotificationManager
from pycs.interfaces.MediaFile import MediaFile
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)

        # get model
        model = project.model()

        # get data and results
        if data['predict'] == 'new':
            files = project.files_without_results()
        else:
            files = project.files()

        objects = list(map(MediaFile, files))

        # create job
        def store(index, length, result):
            with self.db:
                for remove in files[index].results():
                    if remove.origin == 'pipeline':
                        remove.remove()
                        self.nm.remove_result(remove)

                for entry in result:
                    file_type = entry['type']
                    del entry['type']

                    if 'label' in entry:
                        label = entry['label']
                        del entry['label']
                    else:
                        label = None

                    if file_type == 'labeled-image':
                        for remove in files[index].results():
                            remove.remove()
                            self.nm.remove_result(remove)

                    created = files[index].create_result('pipeline', file_type, label, entry)
                    self.nm.create_result(created)

            return (index + 1) / length

        try:
            self.jobs.run(project,
                          'Model Interaction',
                          f'{project.name} (create predictions)',
                          f'{project.name}/model-interaction',
                          self.load_and_predict, model, objects,
                          progress=store)
        except JobGroupBusyException:
            return abort(400)

        return make_response()

    @staticmethod
    def load_and_predict(model, files):
        with closing(load_pipeline(model.root_folder)) as pipeline:
            length = len(files)
            for index in range(length):
                result = pipeline.execute(files[index])
                yield index, length, result