Web部署AI模型-Flask1（测试resnet模型）

需求：异地启动一个服务，后端工程师将影像设备（摄像机）获取的图像（输入）通过post方式推送过来，服务通过后台模型预测，把结果存入dict，然后由后端工程师通过get方式获取AI模型预测结果。

# 预处理
import torch
import numpy as np
from skimage.transform import resize


def preprocess(inp_arr):
    x = np.array(inp_arr)
    x = resize(x, (224, 224))
    x = np.expand_dims(x, axis=0)
    x = np.transpose(x, (0, 3, 1, 2))
    inp_arr_pre = torch.tensor(x, dtype=torch.float)
    return inp_arr_pre

# 加载模型
import torch
from torchvision import models


def load_model_resnet(mode='cuda'):
    with torch.no_grad():
        if mode == 'cuda':
            model = models.resnet18().cuda().eval()
        else:
            model = models.resnet18().eval()
    return model

# 主代码
from skimage import io
from flask import Flask, request, jsonify
from gevent.pywsgi import WSGIServer
from models.load_model import load_model_resnet
from transforms.process import preprocess

app = Flask(__name__)


@app.route('/')
def hello():
    return "Deploy"


@app.route('/predict', methods=["GET", "POST"])
def predict():
    # Initialize the data dictionary that will be returned from the view.
    data = {"success": False}

    if request.method == 'POST':
        if request.files.get('image'):
            # 传入url图像需要decode("utf-8")解码，传入本地图像不需要
            input_dir = request.files['image'].read().decode("utf-8")

            input_arr = io.imread(input_dir)
            input_arr_pre = preprocess(input_arr).cuda()
            output = load_model_resnet()(input_arr_pre)
            _, output_cls = output.max(dim=1)

            # data['predictions'] = 0
            data['predictions'] = output_cls.cpu().numpy().tolist()

            data['success'] = True
    return jsonify(data)


if __name__ == '__main__':
     # app.run(debug=True)
     http_server = WSGIServer(('0.0.0.0', 5002), app)
     http_server.serve_forever()