model-service/model_service/views/model_view.py
2020-04-14 19:42:43 +02:00

170 lines
5.8 KiB
Python

#!/usr/bin/env python3
import tempfile
import os
from flask import request, jsonify, current_app, abort, Response
from flask_classful import FlaskView, route
from model import db, Default, AIModel
from minio.error import BucketAlreadyExists, BucketAlreadyOwnedByYou, ResponseError, NoSuchKey
from schemas import AIModelSchema, DefaultSchema, InfoSchema
from marshmallow.exceptions import ValidationError
from utils import json_required, storage
from pyAudioAnalysis.audioTrainTest import load_model, load_model_knn
class ModelView(FlaskView):
aimodel_schema = AIModelSchema(many=False)
aimodels_schema = AIModelSchema(many=True,
exclude=['timestamp', 'mid_window', 'mid_step', 'short_window', 'short_step',
'compute_beat'])
default_schema = DefaultSchema(many=False)
info_schema = InfoSchema(many=False)
def _ensure_buckets(self):
for bucket_name in [current_app.config['MINIO_MEANS_BUCKET_NAME'],
current_app.config['MINIO_MODEL_BUCKET_NAME']]:
try:
storage.connection.make_bucket(bucket_name)
except BucketAlreadyOwnedByYou as err:
pass
except BucketAlreadyExists as err:
pass
# Everything else should be raised
def index(self):
models = AIModel.query.all()
return jsonify(self.aimodels_schema.dump(models)), 200
def post(self):
# get important data from the request
try:
info = self.info_schema.loads(request.form.get('info'))
except ValidationError as e:
abort(400, str(e))
# check for conflict
m = AIModel.query.filter_by(id=info['id']).first()
if m:
abort(409)
# get and validate file
model_file = request.files['modelFile']
if model_file.content_length <= 0:
abort(411, f"Content length for modelFile is not a positive integer or missing.")
means_file = request.files['meansFile']
if means_file.content_length <= 0:
abort(411, f"Content length for meansFile is not a positive integer or missing.")
# create bucket if necessary
self._ensure_buckets()
# Temporarily save the file, because pyAudioAnalysis can only read files
_, temp_model_filename = tempfile.mkstemp()
temp_means_filename = temp_model_filename + "MEANS"
model_file.save(temp_model_filename)
means_file.save(temp_means_filename)
try:
if info['type'] == 'knn':
_, _, _, _, mid_window, mid_step, short_window, short_step, compute_beat \
= load_model_knn(temp_model_filename)
else:
_, _, _, _, mid_window, mid_step, short_window, short_step, compute_beat \
= load_model(temp_model_filename)
# Because of pyAudiomeme the files already saved, so we just use the file uploader functions
storage.connection.fput_object(
current_app.config['MINIO_MODEL_BUCKET_NAME'],
str(info['id']),
temp_model_filename
)
storage.connection.fput_object(
current_app.config['MINIO_MEANS_BUCKET_NAME'],
str(info['id']),
temp_means_filename
)
finally:
os.remove(temp_model_filename)
os.remove(temp_means_filename)
m = AIModel(id=info['id'], mid_window=mid_window, mid_step=mid_step, short_window=short_window,
short_step=short_step, compute_beat=compute_beat, type=info['type'])
db.session.add(m)
db.session.commit()
return jsonify(self.aimodel_schema.dump(m)), 200
def get(self, _id: str):
if _id == "$default":
default = Default.query.first_or_404() # TODO: Kitalálni, hogy inkább a latestestest-el térjen-e vissza
m = default.aimodel
else:
m = AIModel.query.filter_by(id=_id).first_or_404()
if "means" in request.args:
bucket = current_app.config['MINIO_MEANS_BUCKET_NAME']
else:
bucket = current_app.config['MINIO_MODEL_BUCKET_NAME']
try:
data = storage.connection.get_object(bucket, str(m.id))
except NoSuchKey:
abort(500, "The ID is stored in the database but not int the Object Store")
return Response(data.stream(), mimetype=data.headers['Content-type'])
@route('<_id>/details')
def get_details(self, _id: str):
if _id == "$default":
default = Default.query.first_or_404() # TODO: Kitalálni, hogy inkább a latestestest-el térjen-e vissza
m = default.aimodel
else:
m = AIModel.query.filter_by(id=_id).first_or_404()
return jsonify(self.aimodel_schema.dump(m))
def delete(self, _id: str):
if _id == '$default':
default = Default.query.first_or_404()
m = default.aimodel
else:
m = AIModel.query.filter_by(id=_id).first_or_404()
storage.connection.remove_object(current_app.config['MINIO_MODEL_BUCKET_NAME'], str(m.id))
storage.connection.remove_object(current_app.config['MINIO_MEANS_BUCKET_NAME'], str(m.id))
db.session.delete(m)
db.session.commit()
return '', 204
@json_required
@route('$default', methods=['PUT'])
def put_default(self):
try:
req = self.default_schema.load(request.json)
except ValidationError as e:
abort(404, str(e))
m = AIModel.query.filter_by(id=req['id']).first_or_404()
Default.query.delete()
new_default = Default(aimodel=m)
db.session.add(new_default)
db.session.commit()
return '', 204