135 lines
4.7 KiB
Python
135 lines
4.7 KiB
Python
#!/usr/bin/env python3
|
|
from flask import request, jsonify, current_app, abort, Response
|
|
from flask_classful import FlaskView, route
|
|
from model import db, Default, AIModel, AIModelType
|
|
from minio.error import NoSuchKey
|
|
from schemas import AIModelSchema, DefaultSchema, InfoSchema
|
|
from marshmallow.exceptions import ValidationError
|
|
from utils import json_required, storage, ensure_buckets
|
|
|
|
|
|
class CNNView(FlaskView):
|
|
route_base = 'cnn'
|
|
|
|
aimodel_schema = AIModelSchema(many=False)
|
|
aimodels_schema = AIModelSchema(many=True, exclude=['timestamp', 'details'])
|
|
default_schema = DefaultSchema(many=False)
|
|
info_schema = InfoSchema(many=False)
|
|
|
|
def index(self):
|
|
models = AIModel.query.filter_by(type=AIModelType.CNN).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.")
|
|
|
|
weights_file = request.files['weightsFile']
|
|
|
|
if weights_file.content_length <= 0:
|
|
abort(411, f"Content length for weightsFile is not a positive integer or missing.")
|
|
|
|
# create bucket if necessary
|
|
ensure_buckets()
|
|
|
|
# Create the entry in the db
|
|
if info['id']:
|
|
m = AIModel(id=info['id'], type=AIModelType.CNN)
|
|
else:
|
|
m = AIModel(type=AIModelType.CNN)
|
|
|
|
# Put files into MinIO
|
|
storage.connection.put_object(current_app.config['MINIO_CNN_BUCKET_NAME'], "model/" + str(m.id), model_file,
|
|
model_file.content_length, content_type=model_file.content_type)
|
|
|
|
storage.connection.put_object(current_app.config['MINIO_CNN_BUCKET_NAME'], "weights/" + str(m.id), weights_file,
|
|
weights_file.content_length, content_type=weights_file.content_type)
|
|
|
|
db.session.add(m)
|
|
db.session.commit()
|
|
|
|
return jsonify(self.aimodel_schema.dump(m)), 200
|
|
|
|
def get(self, _id: str):
|
|
|
|
if _id == "$default":
|
|
# TODO: Kitalálni, hogy inkább a latestestest-el térjen-e vissza
|
|
default = Default.query.filter_by(type=AIModelType.CNN).first_or_404()
|
|
m = default.aimodel
|
|
else:
|
|
m = AIModel.query.filter_by(type=AIModelType.CNN, id=_id).first_or_404()
|
|
|
|
if "weights" in request.args:
|
|
path = "weights/" + str(m.id)
|
|
else:
|
|
path = "model/" + str(m.id)
|
|
|
|
try:
|
|
data = storage.connection.get_object(current_app.config['MINIO_CNN_BUCKET_NAME'], path)
|
|
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":
|
|
# TODO: Kitalálni, hogy inkább a latestestest-el térjen-e vissza
|
|
default = Default.query.filter_by(type=AIModelType.CNN).first_or_404()
|
|
m = default.aimodel
|
|
else:
|
|
m = AIModel.query.filter_by(type=AIModelType.CNN, id=_id).first_or_404()
|
|
|
|
return jsonify(self.aimodel_schema.dump(m))
|
|
|
|
def delete(self, _id: str):
|
|
|
|
if _id == "$default":
|
|
# TODO: Kitalálni, hogy inkább a latestestest-el térjen-e vissza
|
|
default = Default.query.filter_by(type=AIModelType.CNN).first_or_404()
|
|
m = default.aimodel
|
|
else:
|
|
m = AIModel.query.filter_by(type=AIModelType.CNN, id=_id).first_or_404()
|
|
|
|
storage.connection.remove_object(current_app.config['MINIO_CNN_BUCKET_NAME'], "weights/" + str(m.id))
|
|
storage.connection.remove_object(current_app.config['MINIO_CNN_BUCKET_NAME'], "model/" + 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(400, str(e))
|
|
|
|
m = AIModel.query.filter_by(type=AIModelType.CNN, id=req['id']).first_or_404()
|
|
|
|
Default.query.filter_by(type=AIModelType.CNN).delete()
|
|
new_default = Default(type=AIModelType.CNN, aimodel=m)
|
|
db.session.add(new_default)
|
|
db.session.commit()
|
|
|
|
return '', 204
|