使用falsk部署pytorch版YOLOv3

序言

训练好的模型如何部署在云端，如何上传图片进行侦测并返回预测结果，这是一个很令人困惑的问题，之前有写过一篇PyTorch从零开始编写YOLOv3网络和侦测代码并加载官方权重 的文章，介绍了YOLOv3使用pytorch的侦测过程，今天介绍一下如何使用python的flask框架部署YOLOv3并返回预测框信息的json文件，文件结构做个对比，可以看到多出了三个文件：

···························

环境配置

• windows10
• Pychram
• pytorch1.2
• flask
• python3.7

其他的配置就不用再多说，主要是flask，先安装：

pip install flask

一、server部分

我们知道每次启动模型，load参数是一件非常费时间的事情，而每次做前向传播的时候模型其实都是一样的，所以我们最好的办法就是load一次模型，然后做完前向传播之后仍然保留这个load好的模型，下一次有新的数据进来，我们就可以不用重新load模型，可以直接做前向传播得到结果，这样无疑节约了很多load模型的时间。

所以我们需要建立一个类似于服务器的机制，将模型在服务器上load好，方便我们不断去调用模型做前向传播，那么怎么能够达到这个目的呢？在server.py文件中，利用函数model.load_weights(）来读取模型，并加载为全局模型，这样一来就不用每次上传一张图片就要加载一次模型。这里的model.load_weights(）函数是定义在网络里的，用于darknet权重加载到pytorch模型上，如果是pytorch的pt权重的话直接正常load就可以。

然后将model传到Detector() 中进行相关初始化，并打开端口监听模式，这里监听端口是5000，主机采用0.0.0.0，以便外网可以直接访问，如果采用127.0.0.1就只能在本机上运行，如果是在自己的pc机上运行，采用0.0.0.0所提到的外网，就指的是在同一个路由下的其他的主机，如果想要使用广域网（路由器局域网之外的外网）进行访问，就必须使用内网穿透进行代理如果是租用云端服务器进行部署，需要进行安全组的配置，开放所需的端口。

if __name__ == "__main__":
    print(("* Loading Pytorch model and Flask starting server..."
        "please wait until server has fully started"))

    path = r"E:\YOLO\model\yolov3.weights"
    model = Darknet(80)
    model.load_weights(path)
    detector = Detector(model)
    print('..... Finished loading model! ......')
    app.run(host='0.0.0.0', port =5000,debug=True )

利用函数prediict() 作为模型主要的预测部分：

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

    # ensure an image was properly uploaded to our endpoint
    if request.method == "POST":
        # print("Hello")
        if request.files.get("image"):
            # print("world")
            now = time.strftime("%Y-%m-%d-%H_%M_%S",time.localtime(time.time()))
            # read the image in PIL format

            image = request.files["image"].read()
            image = Image.open(io.BytesIO(image)).convert('RGB')
            image.save(now + '.jpg')
            
            # preprocess the image and prepare it for classification
            W, H, scale, img = narrow_image(image)  # 传入缩放函数中，得到W,H,缩放比例，缩放后的416*416的图片
            img_data = transforms(img).unsqueeze(0)

            box = detector(img_data, 0.25, anchors_cfg.ANCHORS_GROUP)       # 将图片传入侦测
            box = enlarge_box(W, H, scale, box)    # 得到的box是基于416图片，需反算回到原图
            res = json_text(box,W,H)     # 将box整理成json格式
            
            return res

具体过程在注释中已经写明，涉及的函数比较多，不做过多解释，server完整代码：

import io
from PIL import Image
from flask import Flask, request, jsonify

import torch
import json
import time
from darknet import Darknet
from detect_yolo import *

mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
input_size = 300

app = Flask(__name__)

transforms = torchvision.transforms.Compose([  # 归一化，Tensor处理
    torchvision.transforms.ToTensor()
])


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

    # ensure an image was properly uploaded to our endpoint
    if request.method == "POST":
        # print("Hello")
        if request.files.get("image"):
            # print("world")
            now = time.strftime("%Y-%m-%d-%H_%M_%S",time.localtime(time.time()))
            # read the image in PIL format

            image = request.files["image"].read()
            image = Image.open(io.BytesIO(image)).convert('RGB')
            image.save(now + '.jpg')
            # preprocess the image and prepare it for classification
            W, H, scale, img = narrow_image(image)  # 传入缩放函数中，得到W,H,缩放比例，缩放后的416*416的图片
            img_data = transforms(img).unsqueeze(0)

            # classify the input image and then initialize the list
            # of predictions to return to the client

            box = detector(img_data, 0.25, anchors_cfg.ANCHORS_GROUP)
            box = enlarge_box(W, H, scale, box)
            res = json_text(box,W,H)

            return res
    # return jsonify(data)

# if this is the main thread of execution first load the model and
# then start the server
if __name__ == "__main__":
    print(("* Loading Pytorch model and Flask starting server..."
        "please wait until server has fully started"))

    path = r"E:\YOLO\model\yolov3.weights"
    model = Darknet(80)
    model.load_weights(path)
    detector = Detector(model)
    print('..... Finished loading model! ......')
    app.run(host='0.0.0.0', port =5000,debug=True )

二、client部分

服务端搭建好后，我们需要一个客户端来发送数据，这里命名为client.py。

发送数据请求并不难，首先我们需要知道上面定义好的 flask server 的地址，因为这就是我们在本地定义的，所以地址是：

PyTorch_REST_API_URL = 'http://127.0.0.1:5000/predict'

上面的 /predict 是因为我们前面使用了 @app.route("/predict", methods=[“POST”])。

接着我们定义一个函数来发送数据请求：

def predict_result(image_path):
    # Initialize image path
    image = open(image_path, 'rb').read()
    payload = {'image': image}

    # Submit the request.
    try:
        r = requests.post(REST_API_URL, files=payload).json()   # 得到返回结果
        pprint.pprint(r)  #将结果输出

    except:
        pass

传入的参数 image_path 是图片路径，然后使用 requests.post(PyTorch_REST_API_URL, files=payload).json() 向服务器传入数据，同时得到服务器计算的结果，最后将结果 print 出来就可以了，这里使用了pprint，这个函数作用了print是一样的，不过可以在输出的时候会自动换行，看着会舒服一些。

client完整代码：

import requests
import time
import pprint

# Initialize the keras REST API endpoint URL.
REST_API_URL = 'http://127.0.0.1:5000/predict'


def predict_result(image_path):
    # Initialize image path
    image = open(image_path, 'rb').read()
    payload = {'image': image}

    # Submit the request.
    try:
        r = requests.post(REST_API_URL, files=payload).json()

        pprint.pprint(r)
    except:
        pass


if __name__ == '__main__':

    t1 = time.time()
    img_path = r'E:\YOLO\test_image\timg1.jpg'
    predict_result(img_path)
    t2 = time.time()
    print(t2-t1)

三、启动server，开始测试

首先运行server，看到如下界面：

然后运行client，图片路径可以自己更改，这里加载的是yolo官方的weights文件：

运行结束之后可以看到输出了server返回的结果已经预测的时间，这里主要返回的json参数为：分类类别和左上角、右下角坐标。

文章写道这里就差不多结束了，当然这只是一个简单的flask部署演示过程，要实际部署的话还要做很多后续处理，就需要根据自己的环境进行修改调整了。

项目完整代码：github

—————————————————————————————————————————

参考文章：

• 如何用flask部署pytorch模型
• flask部署深度学习模型