PyCS/pycs/frontend/WebServer.py

from glob import glob
from json import dumps
from os import path, getcwd
from os.path import exists
from time import time

import eventlet
import socketio
from flask import Flask, make_response, send_from_directory, request
from werkzeug import formparser

from pycs.ApplicationStatus import ApplicationStatus
from pycs.util.GenericWrapper import GenericWrapper
from pycs.util.ProgressFileWriter import ProgressFileWriter
from pycs.util.RecursiveDictionary import set_recursive


class WebServer:
    def __init__(self, app_status: ApplicationStatus):
        # initialize web server
        if exists('webui/index.html'):
            print('production build')

            # find svg icons and add them as separate static files to
            # set their correct mime type / content_type
            static_files = {}

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

            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}

            self.__sio = socketio.Server()
            self.__flask = Flask(__name__)
            self.__app = socketio.WSGIApp(self.__sio, self.__flask, static_files=static_files)

            def response(data=None):
                if data is None:
                    return make_response()
                else:
                    return make_response(data)

            @self.__flask.route('/', methods=['GET'])
            def index():
                return send_from_directory(path.join(getcwd(), 'webui'), 'index.html')
        else:
            print('development build')

            self.__sio = socketio.Server(cors_allowed_origins='*')
            self.__flask = Flask(__name__)
            self.__app = socketio.WSGIApp(self.__sio, self.__flask)

            def response(data=None):
                if data is None:
                    rsp = make_response()
                else:
                    rsp = make_response(data)

                rsp.headers['Access-Control-Allow-Origin'] = '*'
                return rsp

        # save every change in application status and send it to the client
        app_status.subscribe(self.__update_application_status, immediate=True)

        # define events
        @self.__sio.event
        def connect(id, msg):
            self.__sio.emit('app_status', {
                'keys': [],
                'value': self.__status
            }, to=id)

        @self.__flask.route('/settings', methods=['POST'])
        def edit_settings():
            data = request.get_json(force=True)
            set_recursive(data, app_status['settings'])

            return response()

        @self.__flask.route('/projects', methods=['POST'])
        def create_project():
            data = request.get_json(force=True)
            app_status['projects'].create_project(data['name'], data['description'], data['model'], data['label'], data['unmanaged'])

            return response()

        @self.__flask.route('/projects/<identifier>', methods=['POST'])
        def edit_project(identifier):
            data = request.get_json(force=True)

            if 'delete' in data.keys():
                app_status['projects'].delete_project(identifier)
            elif 'fit' in data.keys():
                app_status['projects'].fit(identifier)
            elif 'predictAll' in data.keys():
                app_status['projects'].predict(identifier)
            elif 'predictUnlabeled' in data.keys():
                app_status['projects'].predict(identifier, unlabeled=True)
            elif 'predict' in data.keys():
                app_status['projects'].predict(identifier, data['predict'])
            else:
                app_status['projects'].update_project(identifier, data)

            return response()

        @self.__flask.route('/projects/<identifier>/data', methods=['POST'])
        def upload_file(identifier):
            # abort if project id is not valid
            if identifier not in app_status['projects'].keys():
                # TODO return 404
                return make_response('project does not exist', 500)

            # get project and upload path
            project = app_status['projects'][identifier]
            upload_path, file_uuid = project.new_media_file_path()

            # prepare wrapper objects
            job = GenericWrapper()
            file_name = GenericWrapper()
            file_extension = GenericWrapper()
            file_size = GenericWrapper(0)

            # save upload to file
            def custom_stream_factory(total_content_length, filename, content_type, content_length=None):
                file_name.value, file_extension.value = path.splitext(filename)
                file_path = path.join(upload_path, f'{file_uuid}{file_extension.value}')

                # add job to app status
                project['jobs'][file_uuid] = {
                    'id': file_uuid,
                    'type': 'upload',
                    'progress': 0,
                    'filename': filename,
                    'created': int(time()),
                    'finished': None
                }
                job.value = project['jobs'][file_uuid]

                # define progress callback
                length = content_length if content_length is not None and content_length != 0 else total_content_length

                def callback(progress):
                    file_size.value += progress
                    relative = progress / length

                    if relative - job.value['progress'] > 0.02:
                        job.value['progress'] = relative

                # open file handler
                return ProgressFileWriter(file_path, 'wb', callback)

            stream, form, files = formparser.parse_form_data(request.environ, stream_factory=custom_stream_factory)

            if 'file' not in files.keys():
                return make_response('no file uploaded', 500)

            # set progress to 1 after upload is done
            job = job.value

            job['progress'] = 1
            job['finished'] = int(time())

            # add to project files
            project.add_media_file(file_uuid, file_name.value, file_extension.value, file_size.value, job['created'])

            # return default success response
            return response()

        @self.__flask.route('/projects/<project_identifier>/unmanaged/<int:file_number>', methods=['GET'])
        def get_unmanaged_file(project_identifier, file_number):
            # abort if project id is not valid
            if project_identifier not in app_status['projects'].keys():
                return make_response('project does not exist', 500)

            project = app_status['projects'][project_identifier]

            # abort if file id is not valid
            if file_number >= len(project.unmanaged_files_keys):
                return make_response('file does not exist', 500)

            file_identifier = project.unmanaged_files_keys[file_number]
            if file_identifier not in project.unmanaged_files_keys:
                return make_response('file does not exist', 500)

            target_object = project.unmanaged_files[file_identifier]

            # return element data
            return response(target_object.get_data())

        @self.__flask.route('/projects/<project_identifier>/data/<file_identifier>', defaults={'size': None}, methods=['GET'])
        @self.__flask.route('/projects/<project_identifier>/data/<file_identifier>/<size>', methods=['GET'])
        def get_file(project_identifier, file_identifier, size):
            # abort if project id is not valid
            if project_identifier not in app_status['projects'].keys():
                return make_response('project does not exist', 500)

            project = app_status['projects'][project_identifier]

            # abort if file id is not valid
            target_object = project.get_media_file(file_identifier)
            if target_object is None:
                return make_response('file does not exist', 500)

            # resize image to requested size
            if size is not None:
                target_object = target_object.resize(size)

            # construct directory and filename
            file_directory = path.join(getcwd(), target_object.directory)
            file_name = target_object.full_name

            # return data
            return send_from_directory(file_directory, file_name)

        @self.__flask.route('/projects/<project_identifier>/data/<file_identifier>', methods=['POST'])
        def add_result(project_identifier, file_identifier):
            # abort if project id is not valid
            if project_identifier not in app_status['projects'].keys():
                return make_response('project does not exist', 500)

            project = app_status['projects'][project_identifier]

            # abort if file id is not valid
            target_object = project.get_media_file(file_identifier)
            if target_object is None:
                return make_response('file does not exist', 500)

            # add result
            result = request.get_json(force=True)
            if result:
                if 'delete' in result:
                    project.remove_media_file(file_identifier)
                elif 'reset' in result:
                    target_object.remove_results()
                elif 'x' not in result:
                    if result['label']:
                        result['type'] = 'labeled-image'
                        target_object.add_global_result(result)
                    else:
                        target_object.remove_global_result()
                else:
                    result['type'] = 'labeled-bounding-box' if 'label' in result else 'bounding-box'
                    target_object.add_result(result)

            # return default success response
            return response()

        @self.__flask.route('/projects/<project_identifier>/data/<file_identifier>/<result_identifier>', methods=['POST'])
        def edit_result(project_identifier, file_identifier, result_identifier):
            # abort if project id is not valid
            if project_identifier not in app_status['projects'].keys():
                return make_response('project does not exist', 500)

            project = app_status['projects'][project_identifier]

            # abort if file id is not valid
            target_object = project.get_media_file(file_identifier)
            if target_object is None:
                return make_response('file does not exist', 500)

            # parse post data
            result = request.get_json(force=True)
            if result:
                # remove result
                if 'delete' in result.keys():
                    target_object.remove_result(result_identifier)

                # update result
                else:
                    target_object.update_result(result_identifier, result)

            # return default success response
            return response()

        @self.__flask.route('/projects/<project_identifier>/labels', methods=['POST'])
        def create_label(project_identifier):
            # abort if project id is not valid
            if project_identifier not in app_status['projects'].keys():
                return make_response('project does not exist', 500)

            project = app_status['projects'][project_identifier]

            # add result
            result = request.get_json(force=True)
            if result:
                project.add_label(result['name'])

            # return default success response
            return response()

        @self.__flask.route('/projects/<project_identifier>/labels/<label_identifier>', methods=['POST'])
        def edit_label(project_identifier, label_identifier):
            # abort if project id is not valid
            if project_identifier not in app_status['projects'].keys():
                return make_response('project does not exist', 500)

            project = app_status['projects'][project_identifier]

            # parse post data
            result = request.get_json(force=True)
            if result:
                # remove label
                if 'delete' in result.keys():
                    project.remove_label(label_identifier)

                # update label
                else:
                    project.update_label(label_identifier, result['name'])

            # return default success response
            return response()

        @self.__flask.route('/projects/<project_identifier>/predictions', methods=['GET'])
        def download_predictions(project_identifier):
            # abort if project id is not valid
            if project_identifier not in app_status['projects'].keys():
                return make_response('project does not exist', 404)

            project = app_status['projects'][project_identifier]

            # create export
            result = []

            def mk_obj(type, name, extension, predictions):
                data_res = {
                    'type': type,
                    'filename': name + extension,
                    'predictions': []
                }

                for result_key in predictions:
                    result_obj = predictions[result_key]
                    data_res['predictions'].append(result_obj)

                return data_res

            for data_key in project['data']:
                data_obj = project['data'][data_key]
                result.append(mk_obj(
                    data_obj['type'],
                    data_obj['name'],
                    data_obj['extension'],
                    data_obj['predictionResults']
                ))

            for data_key in project.unmanaged_files:
                data_obj = project.unmanaged_files[data_key].get_data()
                result.append(mk_obj(
                    data_obj['type'],
                    data_obj['id'],
                    data_obj['extension'],
                    data_obj['predictionResults']
                ))

            # send to user
            rsp = make_response(dumps(result))
            rsp.headers['Content-Type'] = 'text/json;charset=UTF-8'
            rsp.headers['Content-Disposition'] = 'attachment;filename=predictions.json'
            return rsp

        # finally start web server
        host = app_status['settings']['frontend']['host']
        port = app_status['settings']['frontend']['port']

        eventlet.wsgi.server(eventlet.listen((host, port)), self.__app)

    def __update_application_status(self, status, keys):
        value = status
        for key in keys[:-1]:
            value = value[key]

        self.__status = status
        self.__sio.emit('app_status', {
            'keys': keys[:-1],
            'value': value
        })