使用FastAPI构建机器学习API

FastAPI是用于在Python中构建API的高性能异步框架。
它提供了对Swagger UI开箱即用的支持。

该博客的源代码可用https://github.com/aniketmaurya/tensorflow-web-app-starter-pack

让我们从一个简单的hello-world示例开始

首先，我们导入FastAPIclass并创建一个object app。此类具有有用的参数，例如我们可以传递Swagger UI的标题和描述。

from fastapi import FastAPI
app = FastAPI(title='Hello world', description='This is a hello world example', version='0.0.1')

我们定义一个函数并用修饰它@app.get。这意味着我们的API /index支持GET方法。此处定义的函数是async，FastAPI通过为常规def函数创建线程池来自动处理async，而无需使用async方法，并且async事件循环用于async函数。

@app.get('/index')
async def hello_world():
    return "hello world"

金字塔支持

FastAPI提供的我最喜欢的功能之一是Pydantic支持。我们可以定义Pydantic模型，并且FastAPI将为这些模型处理请求-响应。
让我们创建一个COVID-19症状检查器API来理解这一点。

Covid-19（新冠肺炎）症状检查器API

我们创建一个请求主体，它是客户端发送请求的格式。它将由Swagger UI使用。

from pydantic import BaseModel


class Symptom(BaseModel):
    fever: bool = False
    dry_cough: bool = False
    tiredness: bool = False
    breathing_problem: bool = False

让我们创建一个函数，根据输入分配风险级别。

这仅用于学习，不应在现实生活中使用，最好请教医生。

def get_risk_level(symptom: Symptom):
    if not (symptom.fever or symptom.dry_cough or symptom.tiredness or symptom.breathing_problem):
        return 'Low risk level. THIS IS A DEMO APP'


    if not (symptom.breathing_problem or symptom.dry_cough):
        if symptom.fever:
            return 'moderate risk level. THIS IS A DEMO APP'


    if symptom.breathing_problem:
        return 'High-risk level. THIS IS A DEMO APP'


    return 'THIS IS A DEMO APP'

让我们创建用于检查症状的API

@app.post('/api/covid-symptom-check')
def check_risk(symptom: Symptom):
    return get_risk_level(symptom)

图像识别API

我们将创建一个API来对图像进行分类，我们将其命名为predict/image。
我们将使用Tensorflow创建图像分类模型。

使用Tensorflow进行图像分类的教程

我们创建一个函数load_model，该函数将返回具有预先训练的权重的MobileNet CNN模型，即，它已经过训练，可以对1000种独特的图像类别进行分类。

import tensorflow as tf


def load_model():
    model = tf.keras.applications.MobileNetV2(weights="imagenet")
    print("Model loaded")
    return model


model = load_model()

我们定义了一个predict函数，该函数将接受图像并返回预测。
我们将图像调整为224x224并将像素值标准化为[-1，1]。

from tensorflow.keras.applications.imagenet_utils import decode_predictions

decode_predictions用于解码预测对象的类名称。
在这里，我们将返回前2个可能的类。

def predict(image: Image.Image):


    image = np.asarray(image.resize((224, 224)))[..., :3]
    image = np.expand_dims(image, 0)
    image = image / 127.5 - 1.0


    result = decode_predictions(model.predict(image), 2)[0]


    response = []
    for i, res in enumerate(result):
        resp = {}
        resp["class"] = res[1]
        resp["confidence"] = f"{res[2]*100:0.2f} %"


        response.append(resp)


    return response

现在，我们将创建一个/predict/image支持文件上传的API 。我们将过滤文件扩展名以仅支持jpg，jpeg和png格式的图像。

我们将使用枕头Pillow加载图像。

def read_imagefile(file) -> Image.Image:
    image = Image.open(BytesIO(file))
    return image

@app.post("/predict/image")
async def predict_api(file: UploadFile = File(...)):
    extension = file.filename.split(".")[-1] in ("jpg", "jpeg", "png")
    if not extension:
        return "Image must be jpg or png format!"
    image = read_imagefile(await file.read())
    prediction = predict(image)


    return prediction

最终代码

import uvicorn
from fastapi import FastAPI, File, UploadFile


from application.components import predict, read_imagefile


app = FastAPI()


@app.post("/predict/image")
async def predict_api(file: UploadFile = File(...)):
    extension = file.filename.split(".")[-1] in ("jpg", "jpeg", "png")
    if not extension:
        return "Image must be jpg or png format!"
    image = read_imagefile(await file.read())
    prediction = predict(image)


    return prediction




@app.post("/api/covid-symptom-check")
def check_risk(symptom: Symptom):
    return symptom_check.get_risk_level(symptom)




if __name__ == "__main__":
    uvicorn.run(app, debug=True)

FastAPI文档是了解更多有关框架核心概念的最佳场所。（https://fastapi.tiangolo.com/）