使用Flask在B/S端部署YOLOv3（Pytorch）

文前白话

• Flask是一种用python实现轻量级的web微服务，其灵活性较强而且效率高，在深度学习方面，可以用来部署B/S检测模型。下面以yolov3模型为例，介绍基于Flask的简单部署过程。
• 有需要的话可以自行添加丰富的web展示界面。

相关链接传送：
【YOLO-v3 源码详细解读】https://blog.csdn.net/Kefenggewu_/article/details/122317868/
【基于YOLO-v3训练自己的数据与检测任务】
https://blog.csdn.net/Kefenggewu_/article/details/122336631

后端搭建启动服务

#  flask_server.py 
import flask
import io
import threading
import json
import base64
import numpy as np
import cv2
from models import *
from utils.utils import *
from utils.datasets import *
import os
import sys
import time
import datetime
import argparse
import torch
from torch.utils.data import DataLoader
from torchvision import datasets
from torch.autograd import Variable
 
 # 初始化Flask app
app = flask.Flask(__name__)

use_cuda = torch.cuda.is_available()

def load_model():
    """Load the pre-trained model, you can use your model just as easily.

    """
    ####### 重点配置点
    ###### 配置文件和权重文件路径（同一次训练得到的产物）############
    # yolov3.cfg 网络结构配置文件
    config_path = 'H:\\PyTorch-YOLOv3\\config\\yolov3.cfg'
    # yolov3.weights 网络的训练权重
    weights_path = 'H:\\PyTorch-YOLOv3\\weights\\yolov3.weights'
    # coco.names 网络训练时候的类别名称
    class_path = 'H:\\PyTorch-YOLOv3\\data\\coco.names'
    global classes
    classes = load_classes(class_path)
    
    global model
    #默认网络输入大小为416
    model = Darknet(config_path)
    #载入模型
    model.load_darknet_weights(weights_path)
    if use_cuda:
        model.cuda() 
    model.eval()
 
# class_path = 'H:\\唐宇迪\最新唐宇迪\\47-深度学习模型部署与剪枝优化实例-加密\\000资料\\深度学习模型部署与剪枝优化实例\\YOLO部署实例\\PyTorch-YOLOv3\\data\\coco.names'
# classes = load_classes(class_path)
 
Tensor = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor
 
# 数据预处理
def np2tensor(np_array, img_size):
    # 对图像进行padding
    h, w, _ = np.array(np_array).shape
    dim_diff = np.abs(h - w)
    pad1, pad2 = dim_diff // 2, dim_diff - dim_diff // 2
    pad = ((pad1, pad2), (0, 0), (0, 0)) if h <= w else ((0, 0), (pad1, pad2), (0, 0))
    # 进行resize 操作
    img_shape = (img_size, img_size) 
    input_img = np.pad(np_array, pad, 'constant', constant_values=127.5) / 255.
    input_img = cv2.resize(input_img, (img_size,img_size))
    input_img = np.transpose(input_img, (2, 0, 1))
    input_img = torch.from_numpy(input_img).float()
 
    return input_img
 
 
def yolo_detection(img_array, img_size = 416):
    img_array = np.array(img_array)
    img_tensor = np2tensor(img_array,img_size)
    img_tensor = Variable(img_tensor.type(Tensor))
    img_tensor = img_tensor.unsqueeze(0)
 
    # 线上不需要进行权重参数的更新
    with torch.no_grad():
        detections = model(img_tensor) # 输入图
        detections = non_max_suppression(detections) # 经过NMS，得到结果框
        
    ########### 以下是对检测结果进行坐标还原到在原图对应的绝对位置
    pad_x = max(img_array.shape[0] - img_array.shape[1], 0) * (img_size / max(img_array.shape))
    pad_y = max(img_array.shape[1] - img_array.shape[0], 0) * (img_size / max(img_array.shape))    
    unpad_h = img_size - pad_y
    unpad_w = img_size - pad_x
 
    results=[]
    if detections is not None:
        detection = detections[0]
        unique_labels = detection[:, -1].cpu().unique()
        n_cls_preds = len(unique_labels)
        for x1, y1, x2, y2, conf, cls_conf, cls_pred in detection:
            box_h = ((y2 - y1) / unpad_h) * img_array.shape[0]
            box_w = ((x2 - x1) / unpad_w) * img_array.shape[1]
            y1 = ((y1 - pad_y // 2) / unpad_h) * img_array.shape[0]
            x1 = ((x1 - pad_x // 2) / unpad_w) * img_array.shape[1]
 
            class_name = classes[int(cls_pred)]
            detect_result ={'class':class_name, 'x':x1.item(), 'y':y1.item(), 'h':box_h.item(), 'w':box_w.item()}
            results.append(detect_result)
        
    
    data_json = json.dumps(results,sort_keys=True, indent=4, separators=(',', ': '))
 
    return data_json


@app.route("/predict", methods=["POST"])
def predict():
    data = {"success": False}
    if flask.request.method == 'POST':
        if flask.request.files.get("image"):
            image = flask.request.files["image"].read()
            image = Image.open(io.BytesIO(image)).convert('RGB')
            res = yolo_detection(image)
            data['predictions'] = res
            data["success"] = True
    return flask.jsonify(data)

if __name__ == "__main__":
    print("Loading PyTorch model and Flask starting server ...")
    print("Please wait until server has fully started")
    load_model() # 加载模型
    app.run()
    

服务请求&调用模型

# flask_detect.py 

import requests
import argparse

# Initialize the PyTorch REST API endpoint URL.
PyTorch_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.
    r = requests.post(PyTorch_REST_API_URL, files=payload).json()
    # Ensure the request was successful.
    if r['success']:
        print(r['predictions'])



if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Classification demo')
    parser.add_argument('--file', type=str, help='test image file')

    args = parser.parse_args()
    predict_result(args.file)
    

• 终端环境下启动服务：
  
• 命令：

Python  flask_server.py

• 在编译器中，指定测试的图片路径，运行flask_detect.py：
  

• 检测结果：

• 完善前端页面，可以在Web 浏览器页面显示检测结果。

补充：一般的Pytorch模型部署，同样是两个步骤

开启服务

# run_pytorch_server.py

import io
import json
import flask
import torch
import torch
import torch.nn.functional as F
from PIL import Image
from torch import nn
from torchvision import transforms as T
from torchvision.models import resnet50   # 导入训练好的模型
from torch.autograd import Variable 

# 初始化Flask app
app = flask.Flask(__name__)
model = None
use_gpu = False

# 返回结果用的 
with open('imagenet_class.txt', 'r') as f:
    idx2label = eval(f.read())

# 加载模型进来
def load_model():
    """Load the pre-trained model, you can use your model just as easily.
    """
    global model
    #这里直接加载官方工具包里提供的训练好的模型（代码会自动下载）括号内参数为是否下载模型对应的配置信息
    model = resnet50(pretrained=True)   
    #将模型指定为测试格式
    model.eval()
    #是否使用gpu
    if use_gpu:
        model.cuda()

# 数据预处理（这里是对图像进行检测）
def prepare_image(image, target_size):
    """Do image preprocessing before prediction on any data.

    :param image:       original image
    :param target_size: target image size
    :return:
                        preprocessed image
    """
    #针对不同模型，image的格式不同，但需要统一至RGB格式
    if image.mode != 'RGB':
        image = image.convert("RGB")

    # Resize the input image and preprocess it.(按照所使用的模型将输入图片的尺寸修改，并转为tensor)
    image = T.Resize(target_size)(image)
    image = T.ToTensor()(image)

    # Convert to Torch.Tensor and normalize. mean与std   （RGB三通道）这里的参数和数据集中是对应的，训练过程中一致
    image = T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(image)  # 数据标准化

    # Add batch_size axis.增加一个维度，用于按batch测试   本次这里一次测试一张
    image = image[None]
    if use_gpu:
        image = image.cuda()
    return Variable(image, volatile=True) #不需要求导


# 开启服务   这里的predict只是一个名字，可自定义
@app.route("/predict", methods=["POST"])
def predict():
    # Initialize the data dictionary that will be returned from the view.
    #做一个标志，刚开始无图像传入时为false，传入图像时为true
    data = {"success": False}

    # Ensure an image was properly uploaded to our endpoint.
    # 如果收到请求
    if flask.request.method == 'POST':
        #判断是否为图像
        if flask.request.files.get("image"):
            # Read the image in PIL format
            # 将收到的图像进行读取
            image = flask.request.files["image"].read()
            image = Image.open(io.BytesIO(image)) #二进制数据

            # Preprocess the image and prepare it for classification.
            # 利用上面的预处理函数将读入的图像进行预处理
            image = prepare_image(image, target_size=(224, 224))

            # Classify the input image and then initialize the list of predictions to return to the client.
            preds = F.softmax(model(image), dim=1)
            results = torch.topk(preds.cpu().data, k=3, dim=1)
            results = (results[0].cpu().numpy(), results[1].cpu().numpy())

            #将data字典增加一个key,value,其中value为list格式
            data['predictions'] = list()

            # Loop over the results and add them to the list of returned predictions
            for prob, label in zip(results[0][0], results[1][0]):
                label_name = idx2label[label]
                r = {"label": label_name, "probability": float(prob)}
                #将预测结果添加至data字典
                data['predictions'].append(r)

            # Indicate that the request was a success.
            data["success"] = True

    # Return the data dictionary as a JSON response.
    # 将最终结果以json格式文件传出
    return flask.jsonify(data)


if __name__ == '__main__':
    print("Loading PyTorch model and Flask starting server ...")
    print("Please wait until server has fully started")
    #先加载模型
    load_model()
    #再开启服务
    app.run()

发送请求&响应

# simple_request.py
import requests
import argparse

# Initialize the PyTorch REST API endpoint URL.
PyTorch_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.
    r = requests.post(PyTorch_REST_API_URL, files=payload).json()

    # Ensure the request was successful.
    if r['success']:
        # Loop over the predictions and display them.
        for (i, result) in enumerate(r['predictions']):
            print('{}. {}: {:.4f}'.format(i + 1, result['label'],
                                          result['probability']))
    # Otherwise, the request failed.
    else:
        print('Request failed')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Classification demo')
    parser.add_argument('--file',default= './dog.jpg', type=str, help='test image file')

    args = parser.parse_args()
    predict_result(args.file)

• imagenet_class.txt