python

FastAPIとTensorFlowで画像認識APIを作成してみる

 

機械学習系のアプリを作成する場合は、APIを作成してフロントから呼ぶ仕組みが基本です。

Flaskを使ってもいいのですが、Flaskの場合はviewを作る必要があります。

そこで、FastAPIの自動ドキュメントを作成してくれる機能を使って簡単に試していこうと思います。

今回作るのは本当に簡単で、クラインアントから画像をアップロードしてサーバーから予測結果を返すというものです。

サーバー機能としてこの程度です。

サーバー機能
  • アップロードされた画像のロード
  • 画像の前処理
  • モデルによる予測
  • 予測結果を返す

セットアップ

必要なパッケージをインストールしておきます。

※現在、tensorflowのバージョンが2.6.0でインストールされ、numpyが最新の1.21.2に対応していないので状況に合わせてバージョン指定してください

pip install fastapi
pip install uvicorn[standard]
pip install pillow
pip install numpy==1.19.2
pip install tensorflow

まずは、FastAPIの方でエンドポイントを作成しておきます。

動作確認用と実際に使用するものと２つ作っています。動作確認用は適宜削除してください。

main.py

from fastapi import FastAPI
from fastapi import UploadFile, File
import uvicorn

app = FastAPI()

@app.get('/index')
def hello_world(name: str):
    return f"Hello {name}!"

@app.post('/api/predict')
def predict_image(file: UploadFile = File(...)):
    pass


if __name__ == "__main__":
    uvicorn.run(app, port=8080, host='0.0.0.0')

続いて推論するためのモジュールを作成していきます。

上記で説明したサーバーの機能を実装します。

モデルとパラメーターは既に提供されているもの使用します。

難しいことは特にやっていないので、公式サイトのチュートリアルを終えた方は問題なく理解できると思います。

prediction.py

from PIL import Image
from io import BytesIO
import numpy as np
import tensorflow as tf
from tensorflow.keras.applications import imagenet_utils


model = None

def load_model():
    model = tf.keras.applications.MobileNetV2(weights="imagenet")

    return model

def read_image(image_encoded: Image.Image):
    pil_image = Image.open(BytesIO(image_encoded))

    return pil_image

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

    return image

def predict(image: np.ndarray):
    global model
    if model is None:
        model = load_model()
    pred = imagenet_utils.decode_predictions(model.predict(image), 2)[0]

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

        response.append(resp)

    return response

後は、main.pyの方で作成したモジュールを読み込んで、組み込んでいきます。

画像以外をアップロードされても困るので、拡張子を指定してエラーが起きないようにしておきます。

main.py

from fastapi import FastAPI
from fastapi import UploadFile, File
import uvicorn
import prediction as pd

app = FastAPI()

@app.get('/index')
def hello_world(name: str):
    return f"Hello {name}!"

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

    pred = pd.predict(image)
    print(pred)
    return pred


if __name__ == "__main__":
    uvicorn.run(app, port=8080, host='0.0.0.0')

サーバーを起動して、機能するか確認していきます。

python main.py
#　公式サイトに準拠するなら下記コマンド
# uvicorn main:app --reload

エンドポイントがエラーなく表示されているか確認しましょう。

無事確認できたら、Try it outをクリックし画像を選択してExecuteしてみます。

今回は犬の画像を選択してどのような結果が返ってくるのか見てみましょう。

ちゃんとゴールデンレトリバーであると予測が86%程度と出ましたね。

他にもラブラドールレトリバーの可能性も1%程と出てますが、気にしなくていいでしょう。

まとめ

FastAPIと画像認識を組み合わせたAPI作成を実装してきました。

大分シンプルですが、もう少し規模が大きいアプリケーションでも一連の流れは変わりません。

やらなきゃいけないことは増えますが。。。

WebアプリケーションとしてUIを整えたい場合にはフロントエンドから今回作成したバックエンドにリクエストを送ればいいので、簡単な話です。

以上簡単な機械学習API作成でした。

fastAPI tensorflow
RELATED POST