FastAPI 와 Tensorflow의 mobileNet V2 모델을 활용해서

웹페이지에 이미지를 업로드하면 바로 모델을 사용해볼 수 있도록 데모페이지를 제작해 보았다.

main.py

기본적으로 FastAPI 웹페이지를 호스팅할 main.py가 필요하다.

# 필요한 라이브러리와 모듈을 임포트합니다.
from fastapi import FastAPI, File, UploadFile  # FastAPI와 파일 업로드를 위한 클래스를 임포트합니다.
from PIL import Image  # 이미지 처리를 위해 PIL 라이브러리의 Image 모듈을 임포트합니다.
from io import BytesIO  # 바이트 스트림을 다루기 위해 BytesIO를 임포트합니다.
from predict import predict  # 예측 함수를 포함하고 있는 사용자 정의 모듈을 임포트합니다.

# FastAPI 애플리케이션 인스턴스를 생성합니다.
app = FastAPI()

# "/predict/image" 경로에 POST 요청이 들어왔을 때 실행할 함수를 데코레이터로 등록합니다.
@app.post("/predict/image")
async def predict_api(file: UploadFile = File(...)):  # 비동기 함수로, 파일을 업로드 파일로 받습니다.
    # 업로드된 파일의 확장자를 확인하여 jpg, jpeg, png인지 검사합니다.
    extension = file.filename.split(".")[-1] in ("jpg", "jpeg", "png")
    if not extension:  # 확장자가 이들 중 하나가 아니라면 에러 메시지를 반환합니다.
        return "Image must be jpg or png format!"
    
    # 업로드된 파일을 PIL 이미지로 변환합니다. 이를 위해 파일을 바이트로 읽고, BytesIO로 이미지 데이터를 메모리에 로드합니다.
    image = Image.open(BytesIO(await file.read()))
    
    # 사용자 정의 `predict` 함수를 사용하여 이미지에 대한 예측을 수행합니다.
    prediction = predict(image)
    
    # 예측 결과를 반환합니다.
    return prediction

# 스크립트가 직접 실행될 때만 아래 코드가 실행됩니다.
if __name__ == "__main__":
    import uvicorn  # ASGI 서버인 uvicorn을 임포트합니다.
    # uvicorn을 사용하여 앱을 호스트합니다. 여기서 "main:app"는 이 스크립트(main.py)에 정의된 app 인스턴스를 가리킵니다.
    uvicorn.run("main:app", port=8000, log_level="info")

predict.py

# 필요한 라이브러리를 임포트합니다.
from PIL import Image  # 이미지 처리를 위한 PIL 라이브러리
import numpy as np  # 수치 연산을 위한 NumPy 라이브러리
from tensorflow.keras.applications.imagenet_utils import decode_predictions  # 예측 결과를 해석하기 위한 함수
from model_loader import model  # 사용자 정의 모듈에서 학습된 모델을 임포트합니다.

# 이미지를 입력받아 예측 결과를 반환하는 함수를 정의합니다.
def predict(image: Image):
    # 입력 이미지를 224x224 크기로 조정하고, RGB 채널만 유지하며 NumPy 배열로 변환합니다.
    image = np.asarray(image.resize((224, 224)))[..., :3]  # RGB
    
    # 변환된 이미지 배열에 배치 차원을 추가합니다. 이는 모델이 배치 처리를 기대하기 때문입니다.
    image = np.expand_dims(image, 0)
    
    # 이미지 데이터를 -1에서 1 사이의 값으로 정규화합니다. 이는 많은 사전 훈련된 모델의 입력 요구사항입니다.
    image = image / 127.5 - 1.0  # Scaler(정규화)
    
    # 로드된 모델을 사용하여 이미지에 대한 예측을 수행하고, 상위 3개의 예측 결과를 디코드합니다.
    result = decode_predictions(model.predict(image), 3)[0]  # 상위 3개의 결과 반환

    # 예측 결과를 저장할 리스트를 초기화합니다.
    result_list = []
    
    # 예측 결과를 순회하며 각 클래스의 이름과 신뢰도를 결과 리스트에 추가합니다.
    for res in result:
        print(res)  # 콘솔에 결과를 출력합니다.
        # 결과 리스트에 클래스 이름과 신뢰도(확률)를 사전 형태로 추가합니다. 신뢰도는 백분율로 변환됩니다.
        result_list.append({"class": res[1], "confidence": f"{res[2]*100:0.2f} %"})

    # 최종적으로 구성된 결과 리스트를 반환합니다.
    return result_list

model_loader.py

# TensorFlow 라이브러리를 tf라는 별칭으로 임포트합니다.
import tensorflow as tf

# 사전 훈련된 모델을 로드하는 함수를 정의합니다.
def load_model():
    # tf.keras.applications에서 제공하는 MobileNetV2 모델을 로드합니다.
    # weights='imagenet'은 ImageNet 데이터셋으로 사전 훈련된 가중치를 사용하겠다는 의미입니다.
    model = tf.keras.applications.MobileNetV2(weights="imagenet")
    
    # 모델이 성공적으로 로드되었음을 콘솔에 출력합니다.
    print("Model loaded")
    
    # 로드된 모델을 반환합니다.
    return model

# load_model 함수를 호출하여 모델을 로드하고, 이를 model 변수에 할당합니다.
model = load_model()

Using Page

요크셔테리어 이미지를 모델에 넣어서 확인해보자.

강아지 이미지를 웹페이지에 업로드하여 execute 해보면 결과를 얻을 수 있다.

[
  {
    "class": "Yorkshire_terrier",
    "confidence": "95.57 %"
  },
  {
    "class": "Australian_terrier",
    "confidence": "0.63 %"
  },
  {
    "class": "Norfolk_terrier",
    "confidence": "0.59 %"
  }
]

Yorkshire_terrier가 95.57%로 모델이 잘 작동해서 결과를 보여주는걸 확인할 수 있었다.