PyCS/pycs/frontend/WebServer.py

import typing as T

from glob import glob
from pathlib import Path

import eventlet
import munch
import socketio

from flask import send_from_directory

from pycs.database.Model import Model
from pycs.database.LabelProvider import LabelProvider
from pycs.frontend.endpoints.ListJobs import ListJobs
from pycs.frontend.endpoints.ListLabelProviders import ListLabelProviders
from pycs.frontend.endpoints.ListModels import ListModels
from pycs.frontend.endpoints.ListProjects import ListProjects
from pycs.frontend.endpoints.additional.FolderInformation import FolderInformation
from pycs.frontend.endpoints.data.GetCroppedFile import GetCroppedFile
from pycs.frontend.endpoints.data.GetFile import GetFile
from pycs.frontend.endpoints.data.GetPreviousAndNextFile import GetPreviousAndNextFile
from pycs.frontend.endpoints.data.GetResizedFile import GetResizedFile
from pycs.frontend.endpoints.data.RemoveFile import RemoveFile
from pycs.frontend.endpoints.data.UploadFile import UploadFile
from pycs.frontend.endpoints.jobs.RemoveJob import RemoveJob
from pycs.frontend.endpoints.labels.CreateLabel import CreateLabel
from pycs.frontend.endpoints.labels.EditLabelName import EditLabelName
from pycs.frontend.endpoints.labels.EditLabelParent import EditLabelParent
from pycs.frontend.endpoints.labels.ListLabelTree import ListLabelTree
from pycs.frontend.endpoints.labels.ListLabels import ListLabels
from pycs.frontend.endpoints.labels.RemoveLabel import RemoveLabel
from pycs.frontend.endpoints.pipelines.FitModel import FitModel
from pycs.frontend.endpoints.pipelines.PredictBoundingBox import PredictBoundingBox
from pycs.frontend.endpoints.pipelines.PredictFile import PredictFile
from pycs.frontend.endpoints.pipelines.PredictModel import PredictModel
from pycs.frontend.endpoints.projects.CreateProject import CreateProject
from pycs.frontend.endpoints.projects.EditProjectDescription import EditProjectDescription
from pycs.frontend.endpoints.projects.EditProjectName import EditProjectName
from pycs.frontend.endpoints.projects.ExecuteExternalStorage import ExecuteExternalStorage
from pycs.frontend.endpoints.projects.ExecuteLabelProvider import ExecuteLabelProvider
from pycs.frontend.endpoints.projects.GetProjectModel import GetProjectModel
from pycs.frontend.endpoints.projects.ListProjectCollections import ListProjectCollections
from pycs.frontend.endpoints.projects.ListProjectFiles import ListProjectFiles
from pycs.frontend.endpoints.projects.RemoveProject import RemoveProject
from pycs.frontend.endpoints.results.ConfirmResult import ConfirmResult
from pycs.frontend.endpoints.results.CreateResult import CreateResult
from pycs.frontend.endpoints.results.EditResultData import EditResultData
from pycs.frontend.endpoints.results.EditResultLabel import EditResultLabel
from pycs.frontend.endpoints.results.GetProjectResults import GetProjectResults
from pycs.frontend.endpoints.results.GetResults import GetResults
from pycs.frontend.endpoints.results.RemoveResult import RemoveResult
from pycs.frontend.endpoints.results.ResetResults import ResetResults
from pycs.frontend.notifications.NotificationManager import NotificationManager
from pycs.frontend.util.JSONEncoder import JSONEncoder
from pycs.jobs.JobRunner import JobRunner
from pycs.util.PipelineCache import PipelineCache


class WebServer:
    """
    wrapper class for flask and socket.io which initializes most networking
    """

    index: Path = Path.cwd() / 'webui' / 'index.html'

    def __init__(self, app, settings: munch.Munch, discovery: bool = True):

        self.app = app
        # set json encoder so database objects are serialized correctly
        self.app.json_encoder = JSONEncoder

        # initialize web server
        if self.is_production:
            app.logger.info('production build')

            # overwrite root path to serve index.html
            @self.app.route('/', methods=['GET'])
            def index():
                # pylint: disable=unused-variable
                return send_from_directory(str(self.index.parent), self.index.name)

        else:
            app.logger.info('development build')

            # set access control header to allow requests from Vue.js development server
            @self.app.after_request
            def after_request(response):
                # pylint: disable=unused-variable
                response.headers['Access-Control-Allow-Origin'] = '*'
                return response

        # create service objects
        self.sio = socketio.Server(**self.sio_kwargs(settings.allowedOrigins))
        self.wsgi_app = socketio.WSGIApp(self.sio, app, static_files=self.static_files)

        self.host = settings.host
        self.port = settings.port

        # create notification manager
        self.jobs = JobRunner()
        self.pipelines = PipelineCache(self.jobs, settings.get("pipeline_cache_time"))
        self.notifications = NotificationManager(self.sio)

        self.jobs.on_create(self.notifications.create_job)
        self.jobs.on_start(self.notifications.edit_job)
        self.jobs.on_progress(self.notifications.edit_job)
        self.jobs.on_finish(self.notifications.edit_job)
        self.jobs.on_remove(self.notifications.remove_job)

        self.define_routes()

        if discovery:
            Model.discover("models/")
            LabelProvider.discover("labels/")

    def sio_kwargs(self, allowed_origins) -> T.Dict[str, T.Union[str, list]]:
        """keyword arguments for the socketio.Server depending on the mode"""
        kwargs: T.Dict[str, T.Union[str, list]] = dict(async_mode="eventlet")

        if self.is_production:
            if isinstance(allowed_origins, list) and len(allowed_origins) > 0:
                kwargs["cors_allowed_origins"] = allowed_origins
        else:
            kwargs["cors_allowed_origins"] = "*"

        return kwargs

    @property
    def is_production(self) -> bool:
        """property checking, whether the UI is built (production mode)
           or served by npm serve (development mode)"""
        return self.index.exists()

    @property
    def static_files(self) -> T.Optional[T.Dict[str, T.Union[str, dict]]]:
        """returns a dictionary of static files (production mode)
           or None (development mode)"""

        if not self.is_production:
            return None

        # find static files and folders
        static_files: T.Dict[str, T.Union[str, dict]] = {}

        for file_path in glob('webui/*'):
            file_path = file_path.replace('\\', '/')
            static_files[file_path[5:]] = file_path

        # separately add svg files and set their correct mime type
        for svg_path in glob('webui/img/*.svg'):
            svg_path = svg_path.replace('\\', '/')
            static_files[svg_path[5:]] = {'content_type': 'image/svg+xml', 'filename': svg_path}

        return static_files

    def define_routes(self):
        """ defines app routes """

        # additional
        self.app.add_url_rule(
            '/folder',
            view_func=FolderInformation.as_view('folder_information')
        )

        # jobs
        self.app.add_url_rule(
            '/jobs',
            view_func=ListJobs.as_view('list_jobs', self.jobs)
        )
        self.app.add_url_rule(
            '/jobs/<job_id>/remove',
            view_func=RemoveJob.as_view('remove_job', self.jobs)
        )

        # models
        self.app.add_url_rule(
            '/models',
            view_func=ListModels.as_view('list_models')
        )
        self.app.add_url_rule(
            '/projects/<int:project_id>/model',
            view_func=GetProjectModel.as_view('get_project_model')
        )

        # labels
        self.app.add_url_rule(
            '/label_providers',
            view_func=ListLabelProviders.as_view('label_providers')
        )
        self.app.add_url_rule(
            '/projects/<int:project_id>/labels',
            view_func=ListLabels.as_view('list_labels')
        )
        self.app.add_url_rule(
            '/projects/<int:project_id>/labels/tree',
            view_func=ListLabelTree.as_view('list_label_tree')
        )
        self.app.add_url_rule(
            '/projects/<int:project_id>/labels',
            view_func=CreateLabel.as_view('create_label', self.notifications)
        )
        self.app.add_url_rule(
            '/projects/<int:project_id>/labels/<int:label_id>/remove',
            view_func=RemoveLabel.as_view('remove_label', self.notifications)
        )
        self.app.add_url_rule(
            '/projects/<int:project_id>/labels/<int:label_id>/name',
            view_func=EditLabelName.as_view('edit_label_name', self.notifications)
        )
        self.app.add_url_rule(
            '/projects/<int:project_id>/labels/<int:label_id>/parent',
            view_func=EditLabelParent.as_view('edit_label_parent', self.notifications)
        )

        # collections
        self.app.add_url_rule(
            '/projects/<int:project_id>/collections',
            view_func=ListProjectCollections.as_view('list_collections')
        )
        self.app.add_url_rule(
            '/projects/<int:project_id>/data/<int:collection_id>/<int:start>/<int:length>',
            view_func=ListProjectFiles.as_view('list_collection_files')
        )

        # data
        self.app.add_url_rule(
            '/projects/<int:project_id>/data',
            view_func=UploadFile.as_view('upload_file', self.notifications)
        )
        self.app.add_url_rule(
            '/projects/<int:project_id>/data',
            view_func=ListProjectFiles.as_view('list_all_files')
        )
        self.app.add_url_rule(
            '/projects/<int:project_id>/data/<int:start>/<int:length>',
            view_func=ListProjectFiles.as_view('list_files')
        )
        self.app.add_url_rule(
            '/data/<int:file_id>/remove',
            view_func=RemoveFile.as_view('remove_file', self.notifications)
        )
        self.app.add_url_rule(
            '/data/<int:file_id>',
            view_func=GetFile.as_view('get_file')
        )
        self.app.add_url_rule(
            '/data/<int:file_id>/<resolution>',
            view_func=GetResizedFile.as_view('get_resized_file')
        )
        self.app.add_url_rule(
            '/data/<int:file_id>/<resolution>/<crop_box>',
            view_func=GetCroppedFile.as_view('get_cropped_file')
        )
        self.app.add_url_rule(
            '/data/<int:file_id>/previous_next',
            view_func=GetPreviousAndNextFile.as_view('get_previous_and_next_file')
        )

        # results
        self.app.add_url_rule(
            '/projects/<int:project_id>/results',
            view_func=GetProjectResults.as_view('get_project_results')
        )
        self.app.add_url_rule(
            '/data/<int:file_id>/results',
            view_func=GetResults.as_view('get_results')
        )
        self.app.add_url_rule(
            '/data/<int:file_id>/results',
            view_func=CreateResult.as_view('create_result', self.notifications)
        )
        self.app.add_url_rule(
            '/data/<int:file_id>/reset',
            view_func=ResetResults.as_view('reset_results', self.notifications)
        )
        self.app.add_url_rule(
            '/results/<int:result_id>/remove',
            view_func=RemoveResult.as_view('remove_result', self.notifications)
        )
        self.app.add_url_rule(
            '/results/<int:result_id>/confirm',
            view_func=ConfirmResult.as_view('confirm_result', self.notifications)
        )

        self.app.add_url_rule(
            '/results/<int:result_id>/label',
            view_func=EditResultLabel.as_view('edit_result_label', self.notifications)
        )
        self.app.add_url_rule(
            '/results/<int:result_id>/data',
            view_func=EditResultData.as_view('edit_result_data', self.notifications)
        )

        # projects
        self.app.add_url_rule(
            '/projects',
            view_func=ListProjects.as_view('list_projects')
        )
        self.app.add_url_rule(
            '/projects',
            view_func=CreateProject.as_view('create_project', self.notifications, self.jobs)
        )
        self.app.add_url_rule(
            '/projects/<int:project_id>/label_provider',
            view_func=ExecuteLabelProvider.as_view('execute_label_provider',
                                                   self.notifications, self.jobs)
        )
        self.app.add_url_rule(
            '/projects/<int:project_id>/external_storage',
            view_func=ExecuteExternalStorage.as_view('execute_external_storage',
                                                     self.notifications, self.jobs)
        )
        self.app.add_url_rule(
            '/projects/<int:project_id>/remove',
            view_func=RemoveProject.as_view('remove_project', self.notifications)
        )
        self.app.add_url_rule(
            '/projects/<int:project_id>/name',
            view_func=EditProjectName.as_view('edit_project_name', self.notifications)
        )
        self.app.add_url_rule(
            '/projects/<int:project_id>/description',
            view_func=EditProjectDescription.as_view('edit_project_description', self.notifications)
        )

        # pipelines
        self.app.add_url_rule(
            '/projects/<int:project_id>/pipelines/fit',
            view_func=FitModel.as_view('fit_model', self.jobs, self.pipelines)
        )
        self.app.add_url_rule(
            '/projects/<int:project_id>/pipelines/predict',
            view_func=PredictModel.as_view('predict_model', self.notifications, self.jobs,
                                           self.pipelines)
        )
        self.app.add_url_rule(
            '/data/<int:file_id>/predict',
            view_func=PredictFile.as_view('predict_file', self.notifications,
                                          self.jobs, self.pipelines)
        )
        self.app.add_url_rule(
            '/data/<int:file_id>/<int:bbox_id>/predict_bounding_box',
            view_func=PredictBoundingBox.as_view('predict_bounding_box', self.notifications,
                                          self.jobs, self.pipelines)
        )

    def run(self):
        """ start web server """
        self.pipelines.start()
        eventlet.wsgi.server(eventlet.listen((self.host, self.port)), self.wsgi_app)