[FastAPI] 객체 탐지하기

FastAPI 서버 구축

 

GET 요청을 수락하고 localhost:8000에서 사용자가 이미지를 업로드한다. 그 다음 yolo 모델을 선택하고 POST를 통해 해당 데이터를 제출한다. POST 요청을 수락하고 입력 이미지에서 yolo 모델을 실행하고 JSON 출력을 반환한다.

 

from fastapi import FastAPI, Request, Form, File, UploadFile
from fastapi.responses import HTMLResponse
from PIL import Image
from io import BytesIO
import torch

app = FastAPI()

@app.get("/")
def read_root():
    return "FastAPI"

@app.post("/predict")
async def process_home_form(file: UploadFile = File(...), 
              model_name: str = Form(...)):

    model = torch.hub.load('ultralytics/yolov5', model_name, pretrained=True, force_reload = False)

    results = model(Image.open(BytesIO(await file.read())))

    json_results = results_to_json(results,model)
    
    return json_results

def results_to_json(results, model):
    return [
        [
          {
          "class": int(pred[5]),
          "class_name": model.model.names[int(pred[5])],
          "bbox": [int(x) for x in pred[:4].tolist()], # convert bbox results to int from float
          "confidence": float(pred[4]),
          }
        for pred in result
        ]
      for result in results.xyxy
      ]

if __name__ == '__main__':
    import uvicorn
    
    uvicorn.run(host='localhost', port=8000, reload=True, workers=1)

 

클라이언트에서 POST

 

import requests as r
import json
from pprint import pprint

def send_request(image = '../images/cup.jpg', model_name = 'yolov5s'):
    res = r.post("http://localhost:8000/predict", 
                    data={'model_name': model_name}, 
                    files = {'file': open(image , "rb")}
                    )

    pprint(json.loads(res.text))

if __name__ == '__main__':
    send_request()