model-service/model_service/views/root_view.py
marcsello e819dcb26e
All checks were successful
continuous-integration/drone/push Build is passing
Added files field to model info
2020-10-02 03:44:23 +02:00

93 lines
2.8 KiB
Python

#!/usr/bin/env python3
from flask import jsonify, abort, request, url_for
from flask_classful import FlaskView, route
from marshmallow import ValidationError
from utils import json_required
from model import db, AIModel, AIModelType, Default
from schemas import AIModelSchema, DefaultSchema
class RootView(FlaskView):
route_base = '/'
aimodels_schema = AIModelSchema(many=True, exclude=['timestamp', 'details', 'target_class_name'])
aimodel_schema = AIModelSchema(many=False)
default_schema = DefaultSchema(many=False)
## Shared stuff goes here
def index(self):
models = AIModel.query.all()
return jsonify(self.aimodels_schema.dump(models))
@route('/<type_>')
def get_models(self, type_: str):
try:
aimodel_type = AIModelType[type_]
except KeyError:
abort(404, "Unknown type")
models = AIModel.query.filter_by(type=aimodel_type).all()
return jsonify(self.aimodels_schema.dump(models)), 200
@route('/<type_>/<id_>')
def get_model(self, type_: str, id_: str):
try:
aimodel_type = AIModelType[type_]
except KeyError:
abort(404, "Unknown type")
if id_ == "$default":
default = Default.query.filter_by(type=aimodel_type).first_or_404()
m = default.aimodel
else:
m = AIModel.query.filter_by(type=aimodel_type, id=id_).first_or_404()
# Append download links
details = self.aimodel_schema.dump(m)
# Vagy ez, vagy visszateszem a saját view-jébe és duplikálva lesz az egész
if aimodel_type == AIModelType.cnn:
details.update({
"files": {
"model": url_for("CNNView:get_file", id_=m.id),
"weights": url_for("CNNView:get_file", id_=m.id, weights=''),
}
})
elif aimodel_type == AIModelType.svm:
details.update({
"files": {
"model": url_for("SVMView:get_file", id_=m.id),
"means": url_for("SVMView:get_file", id_=m.id, means=''),
}
})
return jsonify(details)
@json_required
@route('/<type_>/$default', methods=['PUT'])
def put_default(self, type_: str):
try:
aimodel_type = AIModelType[type_]
except KeyError:
abort(404, "Unknown type")
try:
req = self.default_schema.load(request.json)
except ValidationError as e:
abort(400, str(e))
m = AIModel.query.filter_by(type=aimodel_type, id=req['id']).first_or_404()
Default.query.filter_by(type=aimodel_type).delete()
new_default = Default(type=aimodel_type, aimodel=m)
db.session.add(new_default)
db.session.commit()
return '', 204