用Tensorflow和FastAPI構建圖像分類API

做者|Aniket Maurya
編譯|VK
來源|Towards Datas Science

這個博客的源代碼能夠從https://github.com/aniketmaurya/tensorflow-web-app-starter-pack得到

讓咱們從一個簡單的helloworld示例開始

首先，咱們導入FastAPI類並建立一個對象應用程序。這個類有一些有用的參數，好比咱們能夠傳遞swaggerui的標題和描述。

from fastapi import FastAPI
app = FastAPI(title='Hello world')

咱們定義一個函數並用@app.get. 這意味着咱們的API/index支持GET方法。這裏定義的函數是異步的，FastAPI經過爲普通的def函數建立線程池來自動處理異步和不使用異步方法，而且它爲異步函數使用異步事件循環。

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

圖像識別API

咱們將建立一個API來對圖像進行分類，咱們將其命名爲predict/image。咱們將使用Tensorflow來建立圖像分類模型。

Tensorflow圖像分類教程:https://aniketmaurya.ml/blog/tensorflow/deep learning/2019/05/12/image-classification-with-tf2.html

咱們建立了一個函數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

如今咱們將建立一個支持文件上傳的API/predict/image。咱們將過濾文件擴展名以僅支持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/

但願你喜歡這篇文章。

原文連接：https://towardsdatascience.com/image-classification-api-with-tensorflow-and-fastapi-fc85dc6d39e8

歡迎關注磐創AI博客站：
http://panchuang.net/

sklearn機器學習中文官方文檔：
http://sklearn123.com/

歡迎關注磐創博客資源彙總站：
http://docs.panchuang.net/