[python] FastAPI로 딥러닝 모델 API 구축하기 | 문자열 파일경로, Bytes 타입 입력받기

FastAPI를 이용한 API 생성

from typing import List
from pydantic import BaseModel

from fastapi import FastAPI, Depends, File, Request
import uvicorn

app = FastAPI()

class InputData(BaseModel):
    data: str

@app.on_event("startup")
async def load_ai_model():
    global model
    model = My_Model()
    print("AI Model Loaded")

async def parse_body(request: Request):
    data: bytes = await request.body()
    return data

@app.post("/model_inference_str_input/")
def AI_model_inference(input: InputData):
    output = model(input.data)

    return output

@app.post("/model_inference_bytes_input/")
def AI_model_inference_2(image_bytes: bytes = Depends(parse_body)):
    output = model(image_bytes)
    
    return output
 

if __name__ == "__main__":
    uvicorn.run(app, host="localhost", port=8000)

API를 만들 때 입력으로 파일 경로(문자열)를 받고 싶을 수도 있고 Bytes 포맷의 변수를 받고 싶을 수도 있다. 

 

• load_ai_model : 'on_event' 데코레이터를 사용하여 서버가 시작될 때 호출하여 AI 모델을 미리 불러온다
• AI_model_inference : 파일 경로(문자열)을 입력받는 함수
• AI_model_inference_2 : Bytes 변수를 입력받는 함수

 

uvicorn.run()을 사용하여 서버를 실행하며, 이를 실행하면 API가 http://localhost:8000 주소에서 실행된다.

 

이제 서버를 실행하고 http://localhost:8000/docs로 이동하여 자동 생성된 API 문서에서 두 엔드포인트에 대한 상세 정보를 확인할 수 있다. 또한 API 문서를 통해 입력 데이터의 형식과 요청 방법을 확인하고 테스트할 수 있다.

 

파일 경로(문자열)를 입력하는 API 호출

import requests
import numpy as np

url = 'http://localhost:8000/model_inference_str_input/'

input_image_path = 'image.jpg'

response = requests.post(url, json={'data': input_image_path})

if response.status_code == 200:
    result = response.json()
    print(result)
else:
    print("API 호출 실패:", response.text)

 

Bytes 변수를 입력하는 API 호출

import requests
import numpy as np

url = 'http://localhost:8000/model_inference_bytes_input/'

input_image_path = 'image.jpg'

# 이미지 파일을 바이트 형태로 읽어옵니다.
with open(input_image_path, "rb") as image_file:
    image_bytes = image_file.read()

response = requests.post(url, data = image_bytes)

if response.status_code == 200:
    result = response.json()
    print(result)
else:
    print("API 호출 실패:", response.text)