Flask + Pytorch (以yolov3为例)构建算法微服务

之前做服务用的是golang调用C,因为go语言和C有比较好的兼容性。但是自从pytorch1.0问世及tensorflow发布了2.0,python框架的简洁和易用性，以及python的可读性，我觉得有必要去尝试一下用python做成算法API服务。Django和Flask都是python写成的服务框架，但据说Flask相较于Django更加轻量级，因此我首先尝试用Flask构建微服务的可能性。

另外这里用的深度学习框架是Pytorch，并以yolov3目标检测作为例子。

准备工作

安装pytorch：

这里建议用conda安装，版本0.4.1以上即可（建议用1.0）。按照官网建议的安装方式即可，如果遇到下载过慢情况，可以参考这篇博客。

安装flask:

我是用conda 安装的

conda install flask

下载pytorch-yolov3：

git clone https://github.com/eriklindernoren/PyTorch-YOLOv3

我用的是这个版本，测试过检测效果是接近于原版的。

服务代码

from flask import Flask, request
from gevent.pywsgi import WSGIServer
from geventwebsocket.handler import WebSocketHandler
import threading
import json
import base64
import numpy as np
import cv2

from yolo import *

app=Flask(__name__)

@app.route('/test/',methods=['POST'])
def detection():
    if request.method == 'POST':
        data = request.get_data()
        json_data = json.loads(data)
        base = json_data['image']
        image_str = base64.b64decode(base)
        img_array = np.fromstring(image_str,np.uint8)
        image_np = cv2.imdecode(img_array,cv2.COLOR_RGB2BGR)
        res = yolo_detection(image_np)      #算法检测接口
        return res
    else:
        return 'bad'
if __name__ == "__main__":
    http_server = WSGIServer(('0.0.0.0',15000), app, handler_class=WebSocketHandler)
    print('Listening on address: http//192.168.18.85:15000')
    http_server.serve_forever()

算法API代码

重构了一个检测API yolo.py，代码如下：

from __future__ import division

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

import json
import numpy as np

#判断gpu是否可用
use_cuda = torch.cuda.is_available()

#配置文件和权重文件路径
config_path = '/home/yzy/yzy/PyTorch-YOLOv3/PyTorch-YOLOv3-master/config/yolov3.cfg'
weights_path = '/home/yzy/yzy/PyTorch-YOLOv3/PyTorch-YOLOv3-master/weights/yolov3.weights'

#载入模型
model = Darknet(config_path)        #默认网络输入大小为416
model.load_weights(weights_path)

if use_cuda:
    model.cuda()

model.eval()

class_path = '/home/yzy/yzy/PyTorch-YOLOv3/PyTorch-YOLOv3-master/data/coco.names'
classes = load_classes(class_path)

Tensor = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor

#numpy array转换成pytorch tensor，同时注意yolo特有的resize方式
def np2tensor(np_array, img_size):
    h, w, _ = 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))
    img_shape = (img_size, img_size)
    input_img = np.pad(np_array, pad, 'constant', constant_values=127.5) / 255.
    input_img = resize(input_img, (*img_shape, 3), mode='reflect')
    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_tensor = np2tensor(img_array,img_size)
    img_tensor = Variable(img_tensor.type(Tensor))
    img_tensor = img_tensor.unsqueeze(0)

    #Get detections
    with torch.no_grad():
        #print("Type of img_tensor: ",type(img_tensor))
        #print("size of img_tensor: ",img_tensor.shape)
        detections = model(img_tensor)
        detections = non_max_suppression(detections, len(classes), 0.8, 0.5)
    
    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:
        #print("type of detections: ",type(detections))
        #print(detections)
        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)
        
        #print("results: ~~~~~~~~~~~~~~~~~~~~~",results)
    
    data_json = json.dumps(results,sort_keys=True, indent=4, separators=(',', ': '))

    return data_json

if __name__ == "__main__":
    img = np.array(Image.open("/home/yzy/yzy/PyTorch-YOLOv3/PyTorch-YOLOv3-master/data/samples/dog.jpg"))

    res = yolo_detection(img)

调用接口yolo_detection即可，注意传入参数为np array和网络输入大小（默认416）。