YOLO Flask-API 배포

💪 들어가며..

해당 포스트에서는 사전에 학습한 YOLO ML 모델을 이용해서 결과를 반환해주는 Flask API에 대해서 작성해보겠습니다.
참고로 제가 생각하기에는 너무 단순무식한(?) 방법이라고 생각되기에 이보다 더 좋고 간단하게 모델을 serve하고 결과를 반환하는 방법이 있다면 꼭꼭꼭 알려주시면 감사하겠습니다.

😏 사전 작업

다른 Pytorch 모델은 단순 torch.load(pt경로명) 명령어를 통해 불러오는 것이 가능했습니다.
그래서 YOLOv5로 생성되는 pt도 위와 같은 명령어로 불러와서 사용할 수 있을 것이라고 생각되었는데, 잘 되지 않았습니다....
그래서 방법을 찾던 중 해당 과정을 통해 해결하기는 했습니다.
이를 위해서는 먼저 학습 결과와 관련된 파일을 저장해줍니다.
제 방법은 best.pt만을 사용하기는 하지만, 전부 필요한 듯 했습니다...

🦎 Flask-API

import io
from pickletools import read_uint1
from torchvision import models
import json
from flask import Flask, jsonify, request
from flask import make_response
import torchvision.transforms as transforms
import torch
from PIL import Image
import os

app = Flask(__name__)

# yolo model 불러오기
model = torch.hub.load('./yolov5/', 'custom', path='./yolov5/runs/train/mask_check_models5/weights/best.pt', source='local')

# POST 통신으로 들어오는 이미지를 저장하고 모델로 추론하는 과정
def save_image(file):
    file.save('./temp/'+ file.filename)

@app.route('/')
def web():
    return "Lungnaha's flask test page"





@app.route('/predict', methods=['POST'])
def predict():
    if request.method == 'POST':
        file = request.files['file']
        save_image(file) # 들어오는 이미지 저장
        train_img = './temp/' + file.filename
        temp = model(train_img)
        result = temp.pandas().xyxy[0]['name']
        check = True # 마스크 미착용 여부를 알려줄 변수
        answer = ""
        for i in range(len(result)):
            if result[i] == "no-mask":
                check = False
                answer = "Detect"
                break
        
        if check:
            answer = "Safe"

        name = file.filename

        res = {
            'answer' : answer,
            'name'   : name
        }
        os.remove('./temp/'+file.filename)

        return res

if __name__=="__main__":
    app.run(host="0.0.0.0",debug=True)

간단하게 위에 코드를 통해 동작하는 과정을 설명하면, POST 통신으로 이미지가 전달되면 해당 이미지를 YOLO 모델로 돌려서 결과를 JSON 형식으로 되돌려주는 API 입니다.

🕊 Send.py

import requests
import os
import json

def send_data(url):
    path_dir = '/root/pi'
    file_list = os.listdir(path_dir)
    for name in file_list:
        files = {
       
            'file':open('/root/pi/' + name, 'rb')
        }
        res = requests.post(url,files=files)
        #name = res['name']
        res_json = res.json()
        if res_json['answer'] == "Detect":
            os.system("mv /root/pi/" + name + " /root/web" )
        else:
            os.system("rm -f /root/pi/"+ name)
    return "Done"

url = "http://localhost/predict"
print(send_data(url))

위의 코드를 실행시켜서 /root/pi 에 위치한 이미지를 전부 전송할 수 있고, 결과를 받아 Detect 가 반환되면 이미지의 위치를 옮기고, Detect 가 반환되지 않으면 해당 이미지를 삭제하는 과정을 작성해보았습니다.

🧂 참고

참고로 저는 해당 Flask API를 k8s pods에서 실행시키고자 했기에 해당 컨테이너에 Flask-API 파일과 YOLO 폴더를 위치시켰습니다.

💡 추가(22.08.29)

현재 YOLOv5 관련 Serving을 진행하다가, 위의 파일처럼 검출된 객체에 대한 정보만 얻는 것이 아닌, Bounding Box 관련 정보가 필요한 상황이 생겨서 관련 내용을 추가합니다.

위와 동일한 방법으로 추가 정보를 얻을 수 있습니다.

temp = model(test_img)
info = temp.pandas().xyxy[0]['정보']

위의 코드에서 '정보' 부분에 name 외에도 아래와 같은 정보를 입력해서 값을 얻을 수 있습니다.

• xmin
• ymin
• xmax
• ymax
• confidence
• class
• name

저는 Bounding Box의 크기를 구하기 원했으므로 xmin, ymin, xmax, ymax 를 사용해서 Bounding Box의 크기를 구했습니다.
해당 코드는 아래와 같습니다.

row = temp.pandas().xyxy[0]['xmax'][cnt] - temp.pandas().xyxy[0]['xmin'][cnt]
col = temp.pandas().xyxy[0]['ymax'][cnt] - temp.pandas().xyxy[0]['ymin'][cnt]
size = int(row)*int(col)

자세한 코드는 제 Github를 참고하시면 됩니다.