[캡스톤디자인프로젝트] YOLOv5 & Flask

YOLOv5 & Flask

이번에는 Flask를 이용하여 custom YOLOv5가 돌아가는 모델 서버를 만들어 볼 것이다. 
[참고한 자료]

Python 3.7.9
Flask 2.2.3

 

1. 실행 환경 설정

1-1. Flask 설치

설치하면서 Module Not Found가 뜨면 시키는대로 라이브러리를 install 해준다.

$ pip3 install Flask

 

1-2. 프로젝트 폴더 생성

원하는 위치에 프로젝트 폴더를 만든다.

$ mkdir 프로젝트명

 

1-3. app.py 만들기

생성된 프로젝트 안에 app.py 파일을 만든다. 

테스트용으로 "Hello World!"를 출력하는 코드를 작성해봤다.

from flask import Flask

# App INIT
app = Flask(__name__)

@app.route('/')
def index():
    return "Hello World!"

if __name__ == '__main__':
    app.run(debug=True)

 

1-4. app.py 실행

프로젝트 터미널에서 실행한다.

terminal >  flask run

 

등장한 주소로 접속하면 Hello World가 출력되어 있다.

2. Model 돌리기

2-1. requirement.txt 설치

프로젝트 제일 상단에 Yolov5를 돌리기 위해 필요한 조건들을 설정할 requirement.txt 파일을 만든 후, install 한다.

 

▷ requirement.txt

flask

# YOLOv5 requirements
# Base ------------------------------------------------------------------------
gitpython>=3.1.30
matplotlib>=3.2.2
numpy>=1.18.5
opencv-python>=4.1.1
Pillow>=7.1.2
psutil  # system resources
PyYAML>=5.3.1
requests>=2.23.0
scipy>=1.4.1
thop>=0.1.1  # FLOPs computation
torch>=1.7.0  # see https://pytorch.org/get-started/locally (recommended)
torchvision>=0.8.1
tqdm>=4.64.0
# protobuf<=3.20.1  # https://github.com/ultralytics/yolov5/issues/8012

# Logging ---------------------------------------------------------------------
tensorboard>=2.4.1
# clearml>=1.2.0
# comet

# Plotting --------------------------------------------------------------------
pandas>=1.1.4
seaborn>=0.11.0

# Export ----------------------------------------------------------------------
# coremltools>=6.0  # CoreML export
# onnx>=1.12.0  # ONNX export
# onnx-simplifier>=0.4.1  # ONNX simplifier
# nvidia-pyindex  # TensorRT export
# nvidia-tensorrt  # TensorRT export
# scikit-learn<=1.1.2  # CoreML quantization
# tensorflow>=2.4.1  # TF exports (-cpu, -aarch64, -macos)
# tensorflowjs>=3.9.0  # TF.js export
# openvino-dev  # OpenVINO export

# Deploy ----------------------------------------------------------------------
setuptools>=65.5.1 # Snyk vulnerability fix
# tritonclient[all]~=2.24.0

# Extras ----------------------------------------------------------------------
# ipython  # interactive notebook
# mss  # screenshots
# albumentations>=1.0.3
# pycocotools>=2.0.6  # COCO mAP
# ultralytics  # HUB https://hub.ultralytics.com

 

terminal >  pip3 install -r requirements.txt 

 

 

2-2. REST API src 받기

YOLOv5 공식 깃허브에 가면 Flask를 이용한 Yolov5 REST API 레퍼런스가 있다.

여기에서 일단 example_request.py, restapi.py 의 코드를 가져와서 가공해볼 것이다.

 

(1) example_request.py

import pprint
import requests

DETECTION_URL = 'http://localhost:5000/v1/object-detection/yolov5s'
IMAGE = 'zidane.jpg'

# Read image
with open(IMAGE, 'rb') as f:
    image_data = f.read()

response = requests.post(DETECTION_URL, files={'image': image_data}).json()

pprint.pprint(response)

 

(2) restapi.py

import argparse
import io
 
import torch
from flask import Flask, request
from PIL import Image
 
app = Flask(__name__)
models = {}
 
DETECTION_URL = '/v1/object-detection/<model>'
 
 
@app.route(DETECTION_URL, methods=['POST'])
def predict(model):
    if request.method != 'POST':
        return
 
    if request.files.get('image'):
        # Method 1
        # with request.files["image"] as f:
        #     im = Image.open(io.BytesIO(f.read()))
 
        # Method 2
        im_file = request.files['image']
        im_bytes = im_file.read()
        im = Image.open(io.BytesIO(im_bytes))
 
        if model in models:
            results = models[model](im, size=640)  # reduce size=320 for faster inference
            return results.pandas().xyxy[0].to_json(orient='records')
 
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Flask API exposing YOLOv5 model')
    parser.add_argument('--port', default=5000, type=int, help='port number')
    parser.add_argument('--model', nargs='+', default=['yolov5s'], help='model(s) to run, i.e. --model yolov5n yolov5s')
    opt = parser.parse_args()
 
    for m in opt.model:
        models[m] = torch.hub.load('ultralytics/yolov5', m, force_reload=True, skip_validation=True)
 
    app.run(host='0.0.0.0', port=opt.port)  # debug=True causes Restarting with stat

 

이때 (1)은 그대로 두고, (2)의 마지막 부분에서 model을 불러오는 for문을 내 모델의 경우로 변경해준다.

'custom'을 추가하고, './yolov5/best.pt' 자리에 restapi.py를 가준으로 내 모델이 있는 위치를 넣어주면 된다.

for m in opt.model:
        models[m] = torch.hub.load('ultralytics/yolov5', 'custom', './yolov5/best.pt', force_reload=True, skip_validation=True)

 

 

2-3. 모델 돌려보기

터미널에서 아래의 명령어로 모델을 실행시켜본다.

$  python restapi.py --port 5000 --model 모델명

 

※ 모델명은 모델 학습 시에 --name으로 설정해줬던 값을 사용하면 된다. 내 경우, Google Colab에서 설정한 model_v4가 name 값이었다.

!python /content/drive/MyDrive/Capstone_Model/yolov5/train.py --img 640 --batch 8 --epochs 150 --data /content/drive/MyDrive/Capstone_Model/yolov5/models/yolov5s.yaml --weights yolov5s.pt --name model_v4

 

실행 결과

 

2-4. 이미지 POST하는 요청 보내보기

cmd창에서 curl 명령어를 통해 테스트 이미지를 POST 하는 요청을 보내보자.

$  curl -X POST -F  "image=@이미지경로"  "http://localhost:5000/v1/object-detection/모델명"

실행 결과

 

flask & Yolov5 연결 성공 !